from __future__ import annotations
import json
from collections.abc import Mapping as MappingABC
from collections.abc import Sequence as SequenceABC
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Mapping, Sequence
import numpy as np
from transcriptml.data.bundle import DatasetBundle, save_bundle_metadata
from transcriptml.data.region_edits import (
base_channel_indices as _shared_base_channel_indices,
base_symbols as _shared_base_symbols,
randomize_nucleotides_inplace as _shared_randomize_nucleotides,
region_bounds as _shared_region_bounds,
resolve_cds_channel as _shared_resolve_cds_channel,
shuffle_codons_inplace as _shared_shuffle_codons,
shuffle_nucleotides_inplace as _shared_shuffle_nucleotides,
valid_length_from_bases as _shared_valid_length,
write_base_symbols as _shared_write_base_symbols,
)
from transcriptml.data.schemas import SequenceSchema, get_schema
from transcriptml.interpret.codon_ism import CDSCodonStarts, find_cds_codon_starts
from transcriptml.progress import ProgressReporter, log_progress
RegionName = Literal["5utr", "cds", "3utr", "transcript"]
OperationName = Literal["shuffle_nucleotides", "shuffle_codons", "randomize_nucleotides", "cds_frameshift"]
_ANNOTATED_REGIONS: tuple[RegionName, ...] = ("5utr", "cds", "3utr")
_ALL_REGIONS_BY_OPERATION: dict[OperationName, tuple[RegionName, ...]] = {
"shuffle_nucleotides": _ANNOTATED_REGIONS,
"shuffle_codons": ("cds",),
"randomize_nucleotides": _ANNOTATED_REGIONS,
"cds_frameshift": ("cds",),
}
_OPERATION_ALIASES = {
"shuffle_nucleotides": "shuffle_nucleotides",
"shuffle_nt": "shuffle_nucleotides",
"nucleotide_shuffle": "shuffle_nucleotides",
"permute_nucleotides": "shuffle_nucleotides",
"permute_nt": "shuffle_nucleotides",
"true_scramble": "shuffle_nucleotides",
"shuffle_codons": "shuffle_codons",
"shuffle_codon": "shuffle_codons",
"codon_shuffle": "shuffle_codons",
"permute_codons": "shuffle_codons",
"permute_codon": "shuffle_codons",
"randomize_nucleotides": "randomize_nucleotides",
"randomize_nt": "randomize_nucleotides",
"random_nucleotides": "randomize_nucleotides",
"random_nt": "randomize_nucleotides",
"replace_random": "randomize_nucleotides",
"ablate": "randomize_nucleotides",
"cds_frameshift": "cds_frameshift",
"cds_frame_shift": "cds_frameshift",
"frameshift": "cds_frameshift",
"frame_shift": "cds_frameshift",
}
_REGION_ALIASES: dict[str, RegionName | tuple[RegionName, ...]] = {
"5utr": "5utr",
"utr5": "5utr",
"5putr": "5utr",
"fiveutr": "5utr",
"fiveprimeutr": "5utr",
"5primeutr": "5utr",
"cds": "cds",
"coding": "cds",
"codingsequence": "cds",
"orf": "cds",
"3utr": "3utr",
"utr3": "3utr",
"3putr": "3utr",
"threeutr": "3utr",
"threeprimeutr": "3utr",
"3primeutr": "3utr",
"transcript": "transcript",
"wholetranscript": "transcript",
"fulltranscript": "transcript",
"sequence": "transcript",
}
_LEGACY_REGION_KEYS = {
"5putrablation": "5utr",
"5utrablation": "5utr",
"utr5ablation": "5utr",
"cdsablation": "cds",
"3putrablation": "3utr",
"3utrablation": "3utr",
"utr3ablation": "3utr",
}
_REGION_SEED_OFFSETS = {
"5utr": 0x9E3779B97F4A7C15,
"cds": 0xC2B2AE3D27D4EB4F,
"3utr": 0x165667B19E3779F9,
"transcript": 0x85EBCA77C2B2AE63,
}
_OPERATION_SEED_OFFSETS = {
"shuffle_nucleotides": 0x27D4EB2F165667C5,
"shuffle_codons": 0xA24BAED4963EE407,
"randomize_nucleotides": 0x94D049BB133111EB,
}
[docs]
@dataclass(frozen=True)
class SequenceControlOperation:
"""One sequence-control operation applied to one or more regions."""
operation: OperationName
regions: tuple[RegionName, ...]
shift: int | None = None
def to_dict(self) -> dict[str, object]:
out: dict[str, object] = {"operation": self.operation, "regions": list(self.regions)}
if self.operation == "cds_frameshift":
out["shift"] = self.shift
return out
[docs]
@dataclass(frozen=True)
class SequenceControlConfig:
"""Normalized training-time RNA sequence-control configuration."""
operations: tuple[SequenceControlOperation, ...] = ()
seed: int = 0
save_dir: str | None = None
save: bool = False
cds_channel: str | int | None = None
def to_dict(self) -> dict[str, object]:
return {
"operations": [op.to_dict() for op in self.operations],
"seed": int(self.seed),
"save_dir": self.save_dir,
"save": bool(self.save),
"cds_channel": self.cds_channel,
}
def _control_key(value: str) -> str:
return (
str(value)
.strip()
.lower()
.replace("-", "")
.replace("_", "")
.replace("'", "")
.replace('"', "")
.replace(" ", "")
)
def _operation_name(value: str) -> OperationName:
key = str(value).strip().lower().replace("-", "_").replace(" ", "_")
try:
return _OPERATION_ALIASES[key] # type: ignore[return-value]
except KeyError as exc:
raise ValueError(
"sequence control operation must be one of: "
"shuffle_nucleotides, shuffle_codons, randomize_nucleotides, cds_frameshift"
) from exc
def _dedupe_regions(regions: Sequence[RegionName]) -> tuple[RegionName, ...]:
seen: set[RegionName] = set()
out: list[RegionName] = []
for region in regions:
if region not in seen:
out.append(region)
seen.add(region)
return tuple(out)
def _normalize_region_token(value: str, *, all_regions: Sequence[RegionName]) -> tuple[RegionName, ...]:
key = _control_key(value)
if key in {"all", "allregions", "annotated", "annotatedregions"}:
return tuple(all_regions)
if key in {"", "none", "false", "off", "no"}:
return ()
try:
region = _REGION_ALIASES[key]
except KeyError as exc:
raise ValueError("regions must be drawn from: 5utr, cds, 3utr, all, transcript") from exc
if isinstance(region, tuple):
return region
return (region,)
def _regions_from_value(
value: object,
*,
default_regions: Sequence[RegionName],
all_regions: Sequence[RegionName],
) -> tuple[RegionName, ...]:
if value is None or value is False:
return ()
if value is True:
return tuple(default_regions)
if isinstance(value, str):
regions: list[RegionName] = []
for token in value.replace(";", ",").replace("|", ",").split(","):
regions.extend(_normalize_region_token(token, all_regions=all_regions))
return _dedupe_regions(regions)
if isinstance(value, MappingABC):
enabled = value.get("enabled", True)
if enabled is False:
return ()
if "regions" in value:
return _regions_from_value(
value["regions"],
default_regions=default_regions,
all_regions=all_regions,
)
if "region" in value:
return _regions_from_value(
value["region"],
default_regions=default_regions,
all_regions=all_regions,
)
return tuple(default_regions)
if isinstance(value, SequenceABC) and not isinstance(value, (str, bytes, bytearray)):
regions = []
for item in value:
regions.extend(
_regions_from_value(
item,
default_regions=default_regions,
all_regions=all_regions,
)
)
return _dedupe_regions(regions)
raise TypeError(f"Unsupported regions value: {value!r}")
def _frameshift_from_value(value: object) -> int:
if isinstance(value, MappingABC):
for key in ("shift", "frameshift", "frame_shift", "amount"):
if key in value:
return _frameshift_from_value(value[key])
raise ValueError("cds_frameshift requires a shift of 1 or 2")
if value is None or isinstance(value, bool):
raise ValueError("cds_frameshift requires a shift of 1 or 2")
try:
shift = int(value)
except (TypeError, ValueError) as exc:
raise ValueError("cds_frameshift requires a shift of 1 or 2") from exc
if shift not in {1, 2}:
raise ValueError("cds_frameshift shift must be 1 or 2")
return shift
def _operation_from_entry(entry: Mapping[str, object]) -> SequenceControlOperation:
raw_operation = entry.get("operation", entry.get("type", entry.get("op")))
if raw_operation is None:
raise ValueError("Each sequence control operation needs an 'operation' or 'type' field")
operation = _operation_name(str(raw_operation))
shift = None
if operation == "cds_frameshift":
shift = _frameshift_from_value(
entry.get("shift", entry.get("frameshift", entry.get("frame_shift", entry.get("amount"))))
)
regions = _regions_from_value(
entry.get("regions", entry.get("region", True)),
default_regions=_ALL_REGIONS_BY_OPERATION[operation],
all_regions=_ALL_REGIONS_BY_OPERATION[operation],
)
return SequenceControlOperation(operation=operation, regions=regions, shift=shift)
def _legacy_operation(region: RegionName, mode: object) -> SequenceControlOperation | None:
if mode is None:
return None
key = _control_key(str(mode))
if key in {"", "none", "false", "off", "no"}:
return None
if key == "scramble":
operation: OperationName = "shuffle_codons" if region == "cds" else "shuffle_nucleotides"
elif key == "truescramble":
operation = "shuffle_nucleotides"
elif key == "ablate":
operation = "randomize_nucleotides"
elif key == "ablatesavemotif":
raise ValueError("Legacy ablate_savemotif is not supported by sequence_controls")
else:
raise ValueError(f"Unsupported legacy ablation mode: {mode!r}")
return SequenceControlOperation(operation=operation, regions=(region,))
def _merge_operations(operations: Sequence[SequenceControlOperation]) -> tuple[SequenceControlOperation, ...]:
merged: dict[OperationName, SequenceControlOperation] = {}
for operation in operations:
if not operation.regions:
continue
if operation.operation in {"shuffle_codons", "cds_frameshift"} and any(
region != "cds" for region in operation.regions
):
raise ValueError(f"{operation.operation} only supports the CDS region")
existing = merged.get(operation.operation)
if existing is not None and operation.operation == "cds_frameshift" and existing.shift != operation.shift:
raise ValueError("Duplicate cds_frameshift operations must use the same shift")
bucket = list(existing.regions) if existing is not None else []
for region in operation.regions:
if region not in bucket:
bucket.append(region)
merged[operation.operation] = SequenceControlOperation(
operation=operation.operation,
regions=tuple(bucket),
shift=operation.shift if operation.operation == "cds_frameshift" else None,
)
out = tuple(merged.values())
owners: dict[RegionName, OperationName] = {}
for operation in out:
if operation.operation == "cds_frameshift":
continue
for region in operation.regions:
if region == "transcript" and owners:
raise ValueError("The transcript region cannot be combined with other sequence-control regions")
if owners.get("transcript") is not None and region != "transcript":
raise ValueError("The transcript region cannot be combined with other sequence-control regions")
previous = owners.get(region)
if previous is not None and previous != operation.operation:
raise ValueError(
f"Region '{region}' is targeted by both {previous} and {operation.operation}; "
"use at most one sequence-control operation per region"
)
owners[region] = operation.operation
return out
[docs]
def normalize_sequence_control_config(config: object) -> SequenceControlConfig:
"""Normalize user-facing sequence-control config into explicit operations.
The preferred shape is::
{
"seed": 42,
"operations": [
{"operation": "shuffle_nucleotides", "regions": ["5utr", "3utr"]},
{"operation": "shuffle_codons", "regions": ["cds"]},
{"operation": "cds_frameshift", "shift": 1}
],
"save_dir": "data/saluki_control"
}
Top-level shortcuts such as ``"shuffle_nucleotides": ["5utr"]`` and the
older ``5pUTR_ablation``/``CDS_ablation``/``3pUTR_ablation`` names are also
accepted.
"""
if config is None or config is False:
return SequenceControlConfig()
if isinstance(config, SequenceControlConfig):
return config
entries: list[SequenceControlOperation] = []
seed = 0
save = False
save_dir: str | None = None
cds_channel: str | int | None = None
if isinstance(config, SequenceABC) and not isinstance(config, (str, bytes, bytearray)):
for entry in config:
if not isinstance(entry, MappingABC):
raise TypeError("sequence_controls list entries must be mappings")
entries.append(_operation_from_entry(entry))
return SequenceControlConfig(operations=_merge_operations(entries))
if not isinstance(config, MappingABC):
raise TypeError("sequence_controls must be a mapping, list of operations, or null")
if config.get("enabled", True) is False:
return SequenceControlConfig()
seed = int(config.get("seed", config.get("ablation_seed", 0)))
save = bool(config.get("save", False))
if config.get("save_dir") is not None:
save_dir = str(config["save_dir"])
save = True
cds_channel = config.get("cds_channel")
if isinstance(cds_channel, str) and cds_channel.isdigit():
cds_channel = int(cds_channel)
if any(key in config for key in ("operation", "type", "op")):
entries.append(_operation_from_entry(config))
raw_operations = config.get("operations")
if raw_operations is not None:
if not isinstance(raw_operations, SequenceABC) or isinstance(raw_operations, (str, bytes, bytearray)):
raise TypeError("sequence_controls.operations must be a list")
for entry in raw_operations:
if not isinstance(entry, MappingABC):
raise TypeError("sequence_controls.operations entries must be mappings")
entries.append(_operation_from_entry(entry))
shortcut_keys: tuple[tuple[str, OperationName], ...] = (
("shuffle_nucleotides", "shuffle_nucleotides"),
("shuffle_nt", "shuffle_nucleotides"),
("permute_nucleotides", "shuffle_nucleotides"),
("shuffle_codons", "shuffle_codons"),
("codon_shuffle", "shuffle_codons"),
("permute_codons", "shuffle_codons"),
("randomize_nucleotides", "randomize_nucleotides"),
("randomize_nt", "randomize_nucleotides"),
("random_nucleotides", "randomize_nucleotides"),
("replace_random", "randomize_nucleotides"),
("cds_frameshift", "cds_frameshift"),
("cds_frame_shift", "cds_frameshift"),
)
for key, operation in shortcut_keys:
if key not in config:
continue
if operation == "cds_frameshift":
value = config[key]
regions = _regions_from_value(
value.get("regions", value.get("region", True)) if isinstance(value, MappingABC) else True,
default_regions=_ALL_REGIONS_BY_OPERATION[operation],
all_regions=_ALL_REGIONS_BY_OPERATION[operation],
)
entries.append(
SequenceControlOperation(
operation=operation,
regions=regions,
shift=_frameshift_from_value(value),
)
)
continue
regions = _regions_from_value(
config[key],
default_regions=_ALL_REGIONS_BY_OPERATION[operation],
all_regions=_ALL_REGIONS_BY_OPERATION[operation],
)
entries.append(SequenceControlOperation(operation=operation, regions=regions))
for raw_key, value in config.items():
region = _LEGACY_REGION_KEYS.get(_control_key(str(raw_key)))
if region is None:
continue
legacy = _legacy_operation(region, value)
if legacy is not None:
entries.append(legacy)
return SequenceControlConfig(
operations=_merge_operations(entries),
seed=seed,
save_dir=save_dir,
save=save,
cds_channel=cds_channel,
)
def _base_channel_indices(schema: SequenceSchema) -> np.ndarray:
return _shared_base_channel_indices(schema)
def _resolve_cds_channel(schema: SequenceSchema, cds_channel: str | int | None) -> int:
return _shared_resolve_cds_channel(schema, cds_channel)
def _infer_valid_length(x: np.ndarray, base_channels: np.ndarray) -> int:
return _shared_valid_length(x, base_channels)
def _mixed_rng(seed: int, seq_index: int, operation: OperationName, region: RegionName) -> np.random.Generator:
value = (
(int(seed) & 0xFFFFFFFFFFFFFFFF)
^ (((int(seq_index) + 1) * 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF)
^ _OPERATION_SEED_OFFSETS[operation]
^ _REGION_SEED_OFFSETS[region]
)
value ^= (value >> 30) & 0xFFFFFFFFFFFFFFFF
value = (value * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
value ^= (value >> 27) & 0xFFFFFFFFFFFFFFFF
value = (value * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
value ^= (value >> 31) & 0xFFFFFFFFFFFFFFFF
return np.random.default_rng(int(value))
def _base_symbols(x: np.ndarray, start: int, end: int, base_channels: np.ndarray) -> np.ndarray:
return _shared_base_symbols(x, start, end, base_channels)
def _write_base_symbols(
x: np.ndarray,
start: int,
end: int,
symbols: np.ndarray,
base_channels: np.ndarray,
) -> None:
_shared_write_base_symbols(x, start, end, symbols, base_channels)
def _region_bounds(
region: RegionName,
*,
valid_length: int,
cds: CDSCodonStarts | None,
) -> tuple[int, int] | None:
return _shared_region_bounds(region, valid_length=valid_length, cds=cds)
def _shuffle_nucleotides(
x: np.ndarray,
*,
start: int,
end: int,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
_shared_shuffle_nucleotides(
x,
start=start,
end=end,
base_channels=base_channels,
rng=rng,
)
def _randomize_nucleotides(
x: np.ndarray,
*,
start: int,
end: int,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
_shared_randomize_nucleotides(
x,
start=start,
end=end,
base_channels=base_channels,
rng=rng,
)
def _frameshift_cds_channel(
x: np.ndarray,
*,
valid_length: int,
cds_channel: int,
shift: int,
) -> bool:
if valid_length <= 0:
return False
channel = int(cds_channel)
end = min(int(valid_length), int(x.shape[-1]))
locs = np.flatnonzero(np.asarray(x[channel, :end]) > 0)
if locs.size == 0:
return False
shifted = locs + int(shift)
shifted = shifted[shifted < end]
x[channel, :end] = 0
if shifted.size:
x[channel, shifted] = 1
return True
def _shuffle_codons(
x: np.ndarray,
*,
cds: CDSCodonStarts,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
_shared_shuffle_codons(x, cds=cds, base_channels=base_channels, rng=rng)
def sequence_control_base_channels(schema: str | SequenceSchema = "saluki6") -> np.ndarray:
"""Return base-channel indices using the sequence-control conventions.
This public helper lets interpretation analyses use the exact same channel
resolution and edit semantics as training-time sequence controls.
"""
return _base_channel_indices(get_schema(schema))
def resolve_sequence_control_cds_channel(
schema: str | SequenceSchema = "saluki6",
cds_channel: str | int | None = None,
) -> int:
"""Resolve the CDS channel using sequence-control conventions."""
return _resolve_cds_channel(get_schema(schema), cds_channel)
def sequence_control_valid_length(x: np.ndarray, base_channels: np.ndarray) -> int:
"""Infer represented length from the resolved base channels."""
return _infer_valid_length(x, np.asarray(base_channels, dtype=np.int64))
def sequence_control_region_bounds(
region: RegionName,
*,
valid_length: int,
cds: CDSCodonStarts | None,
) -> tuple[int, int] | None:
"""Return region bounds using the training-time sequence-control rules."""
return _region_bounds(region, valid_length=valid_length, cds=cds)
def shuffle_region_nucleotides_inplace(
x: np.ndarray,
*,
start: int,
end: int,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
"""Shuffle one region with the same semantics as sequence controls."""
_shuffle_nucleotides(
x,
start=start,
end=end,
base_channels=np.asarray(base_channels, dtype=np.int64),
rng=rng,
)
def randomize_region_nucleotides_inplace(
x: np.ndarray,
*,
start: int,
end: int,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
"""Replace one region with IID bases using sequence-control semantics."""
_randomize_nucleotides(
x,
start=start,
end=end,
base_channels=np.asarray(base_channels, dtype=np.int64),
rng=rng,
)
def shuffle_cds_codons_inplace(
x: np.ndarray,
*,
cds: CDSCodonStarts,
base_channels: np.ndarray,
rng: np.random.Generator,
) -> None:
"""Shuffle annotated CDS codons using sequence-control semantics."""
_shuffle_codons(
x,
cds=cds,
base_channels=np.asarray(base_channels, dtype=np.int64),
rng=rng,
)
def _empty_nested_counts() -> dict[str, dict[str, int]]:
return {
"shuffle_nucleotides": {region: 0 for region in _ANNOTATED_REGIONS + ("transcript",)},
"shuffle_codons": {region: 0 for region in _ANNOTATED_REGIONS + ("transcript",)},
"randomize_nucleotides": {region: 0 for region in _ANNOTATED_REGIONS + ("transcript",)},
"cds_frameshift": {region: 0 for region in _ANNOTATED_REGIONS + ("transcript",)},
}
def _stats(config: SequenceControlConfig, n_sequences: int) -> dict[str, object]:
return {
"n_sequences": int(n_sequences),
"config": config.to_dict(),
"edited": _empty_nested_counts(),
"skipped_empty_region": _empty_nested_counts(),
"skipped_no_cds": 0,
"skipped_empty_transcript": 0,
}
def _increment(stats: dict[str, object], group: str, operation: OperationName, region: RegionName) -> None:
nested = stats[group]
assert isinstance(nested, dict)
by_region = nested[operation]
assert isinstance(by_region, dict)
by_region[region] = int(by_region.get(region, 0)) + 1
[docs]
def apply_sequence_controls_array(
X: np.ndarray,
config: object,
*,
schema: str | SequenceSchema = "saluki6",
out_path: str | Path | None = None,
progress: bool = True,
) -> tuple[np.ndarray, dict[str, object]]:
"""Apply RNA sequence controls to an encoded array.
Base perturbations rewrite only base channels. ``cds_frameshift`` rewrites
only the CDS/codon-start channel. The splice-junction channel is copied
through unchanged.
"""
cfg = normalize_sequence_control_config(config)
arr = np.asarray(X)
if arr.ndim != 3:
raise ValueError(f"Expected X with shape (N, C, L), got {arr.shape}")
stats = _stats(cfg, int(arr.shape[0]))
if not cfg.operations:
return arr, stats
resolved = get_schema(schema)
if arr.shape[1] < resolved.n_channels:
raise ValueError(
f"X has {arr.shape[1]} channels, but schema '{resolved.name}' expects {resolved.n_channels}"
)
base_channels = _base_channel_indices(resolved)
needs_frameshift = any(op.operation == "cds_frameshift" for op in cfg.operations)
cds_channel_index = _resolve_cds_channel(resolved, cfg.cds_channel) if needs_frameshift else None
if out_path is None:
out = np.array(arr, copy=True)
else:
path = Path(out_path)
path.parent.mkdir(parents=True, exist_ok=True)
out = np.lib.format.open_memmap(path, mode="w+", dtype=arr.dtype, shape=arr.shape)
needs_cds = any(
op.operation != "cds_frameshift" and region != "transcript"
for op in cfg.operations
for region in op.regions
)
reporter = ProgressReporter(
"sequence-controls: edit transcripts",
total=int(arr.shape[0]),
unit="transcripts",
enabled=progress,
)
for seq_index in range(int(arr.shape[0])):
if out_path is not None:
out[seq_index] = arr[seq_index]
x_one = out[seq_index]
valid_length = _infer_valid_length(x_one, base_channels)
valid_length = min(valid_length, int(arr.shape[-1]))
if valid_length <= 0:
stats["skipped_empty_transcript"] = int(stats["skipped_empty_transcript"]) + 1
reporter.update()
continue
cds: CDSCodonStarts | None = None
if needs_cds:
try:
cds = find_cds_codon_starts(
x_one,
resolved,
valid_length=valid_length,
cds_channel=cfg.cds_channel,
)
except ValueError as exc:
raise ValueError(
"sequence_controls requires a resolvable CDS channel for 5utr/cds/3utr regions"
) from exc
if cds.starts.size == 0 or cds.cds_length < 3:
stats["skipped_no_cds"] = int(stats["skipped_no_cds"]) + 1
elif needs_frameshift and cds_channel_index is not None:
if not np.any(np.asarray(x_one[cds_channel_index, :valid_length]) > 0):
stats["skipped_no_cds"] = int(stats["skipped_no_cds"]) + 1
for op in cfg.operations:
if op.operation == "cds_frameshift":
if cds_channel_index is None:
continue
if op.shift is None:
raise ValueError("cds_frameshift requires a shift of 1 or 2")
shifted = _frameshift_cds_channel(
x_one,
valid_length=valid_length,
cds_channel=cds_channel_index,
shift=op.shift,
)
if shifted:
_increment(stats, "edited", op.operation, "cds")
continue
for region in op.regions:
bounds = _region_bounds(region, valid_length=valid_length, cds=cds)
if bounds is None:
continue
start, end = bounds
if end <= start:
_increment(stats, "skipped_empty_region", op.operation, region)
continue
rng = _mixed_rng(cfg.seed, seq_index, op.operation, region)
if op.operation == "shuffle_nucleotides":
_shuffle_nucleotides(x_one, start=start, end=end, base_channels=base_channels, rng=rng)
elif op.operation == "randomize_nucleotides":
_randomize_nucleotides(x_one, start=start, end=end, base_channels=base_channels, rng=rng)
elif op.operation == "shuffle_codons":
if cds is None:
continue
_shuffle_codons(x_one, cds=cds, base_channels=base_channels, rng=rng)
else: # pragma: no cover - guarded by normalization.
raise ValueError(f"Unsupported sequence-control operation: {op.operation}")
_increment(stats, "edited", op.operation, region)
reporter.update()
if hasattr(out, "flush"):
out.flush()
reporter.close()
return out, stats
[docs]
def apply_sequence_controls_to_bundle(
bundle: DatasetBundle,
config: object,
*,
default_save_dir: str | Path | None = None,
progress: bool = True,
) -> tuple[DatasetBundle, dict[str, object]]:
"""Apply sequence controls to a dataset bundle and optionally save it."""
cfg = normalize_sequence_control_config(config)
if not cfg.operations:
return bundle, _stats(cfg, int(bundle.X.shape[0]))
save_dir: Path | None = None
if cfg.save_dir is not None:
save_dir = Path(cfg.save_dir)
elif cfg.save and default_save_dir is not None:
save_dir = Path(default_save_dir)
log_progress(
"sequence-controls: "
+ ", ".join(f"{op.operation}({','.join(op.regions)})" for op in cfg.operations),
enabled=progress,
)
out_path = save_dir / "X.npy" if save_dir is not None else None
X_controlled, stats = apply_sequence_controls_array(
bundle.X,
cfg,
schema=bundle.schema,
out_path=out_path,
progress=progress,
)
if save_dir is not None:
save_dir.mkdir(parents=True, exist_ok=True)
if bundle.y is not None:
np.save(save_dir / "y.npy", np.asarray(bundle.y))
stats["save_dir"] = str(save_dir)
controlled_config = dict(bundle.config)
controlled_config["sequence_controls"] = cfg.to_dict()
controlled_config["sequence_control_stats"] = stats
controlled = DatasetBundle(
X=X_controlled,
y=bundle.y,
ids=bundle.ids,
schema=bundle.schema,
metadata=bundle.metadata,
splits=bundle.splits,
config=controlled_config,
)
if save_dir is not None:
save_bundle_metadata(controlled, save_dir)
(save_dir / "sequence_controls.json").write_text(
json.dumps({"config": cfg.to_dict(), "stats": stats}, indent=2),
encoding="utf-8",
)
log_progress(f"sequence-controls: saved controlled bundle to {save_dir}", enabled=progress)
return controlled, stats