Dataset Module¶

def __init__(
    self,
    chunk_path: Path | None = None,
    signal_len: int = 400,
    kmer_len: int = 11,
    model_type: str = "ConvLSTMDwell",
    dwell_offset: int = 0,
    chunks: list[dict] | None = None,
    augmentation: dict | None = None,
    seq_encoding: str = "signal_kmer",
    signal_kmer_context: tuple[int, int] = (4, 4),
    allow_encoding_fallback: bool = True,
    left_context: int | None = None,
    right_context: int | None = None,
    confound_encoder: "ConfoundEncoder | None" = None,
    cl_regression: bool = False,
    signal_mode: str = "both",
    time_mask_bases: int = 0,
    time_mask_count: int = 1,
    shift_max_bases: float = 0.0,
    feature_noise_scale: float = 0.0,
    dwell_template_table: str | Path | None = None,
):
    """
    Initialize dataset.

    Args:
        chunk_path: Path to .npz file with training chunks
        signal_len: Expected signal length (will pad/truncate)
        kmer_len: Expected k-mer length
        model_type: Model architecture name (e.g., "ConvLSTMDwell", "TransformerDwell")
        dwell_offset: Shift dwell/feature window toward 3' end (bases).
            Compensates for physical offset between motor protein and
            sensing region. Requires feature_left >= kmer_context + offset.
        chunks: Pre-loaded list of chunk dicts. When provided, chunk_path is
            ignored and no disk I/O occurs (useful for grid search caching).
        augmentation: Signal augmentation config dict. Keys:
            - jitter_std (float): Gaussian noise std dev (0 = disabled)
            - scale_range (tuple[float, float]): Random scale range (1.0, 1.0 = disabled)
        allow_encoding_fallback: When ``seq_encoding="signal_kmer"`` and the
            corpus carries no base-to-signal maps at all, fall back to
            ``base_onehot`` (warning) rather than raising. A corpus that
            carries them for *some* chunks is damaged and raises either way.
            Read :attr:`effective_seq_encoding` for what the dataset
            actually yields — it is a different model input, not a tuning
            difference, so callers that persist a config must record it.
        left_context: Left signal context (samples before focus base).
            When both left_context and right_context are provided, crop
            asymmetrically around the focus base instead of center-cropping.
        right_context: Right signal context (samples after focus base).
        confound_encoder: Optional :class:`~leech.confounds.ConfoundEncoder`
            that maps each chunk to an integer confound class. When provided,
            each batch includes a ``confound_label`` tensor for adversarial
            training. Chunks whose confound value is unknown get ``-1``
            (ignored by the CE loss via ``ignore_index=-1``).
        cl_regression: When True, include ``cl_target`` in each batch
            (cl_value / 255.0 in [0,1]; sentinel -1.0 for missing).
        time_mask_bases: Max width in bases for time masking (0 = disabled).
            Zeros out contiguous blocks across signal, features, and sequence.
        time_mask_count: Number of time masks to apply per sample.
        shift_max_bases: Max cross-layer shift in bases (0 = disabled).
            Simulates motif anchor offset, applied consistently to all branches.
        feature_noise_scale: Per-channel Gaussian noise multiplier (0 = disabled).
            Noise std = feature_noise_scale * per-channel empirical std.
        dwell_template_table: Path to TSV with per-AA per-position expected
            dwell times. When provided, 20 dwell ratio channels are appended
            to features: ``dwell / expected_dwell[AA_i, pos]`` for each of
            20 AAs. The correct AA's channel has ratio closest to 1.0.
            TSV columns: aa, position, dwell_mean, ...
    """
    self.chunk_path = chunk_path
    self.signal_len = signal_len
    self.kmer_len = kmer_len
    self.model_type = model_type
    self.dwell_offset = dwell_offset
    self.augmentation = augmentation
    self.seq_encoding = seq_encoding
    self.signal_kmer_context = signal_kmer_context
    self.left_context = left_context
    self.right_context = right_context
    self._time_mask_bases = time_mask_bases
    self._time_mask_count = time_mask_count
    self._shift_max_bases = shift_max_bases
    self._feature_noise_scale = feature_noise_scale

    # Load dwell template table for 20-channel AA template features
    self._dwell_templates: np.ndarray | None = None
    self._template_aa_order: list[str] = []
    if dwell_template_table is not None:
        self._load_dwell_templates(Path(dwell_template_table))

    self._needs_features = model_type in FEATURE_MODELS

    # Use pre-loaded chunks or load from file. Loading from a path streams
    # the per-chunk arrays out of the npz a row block at a time instead of
    # holding a full numpy copy alongside the tensors built from it (#211);
    # pre-loaded chunks have already paid that cost.
    self._array_stream: _ArrayStream | None = None
    self._npz_members: set[str] = set()
    self._s2s_csr: tuple[np.ndarray, np.ndarray] | None = None
    self._s2s_rows: np.ndarray | None = None
    if chunks is not None:
        logger.info(f"Using {len(chunks)} pre-loaded chunks (skipping disk I/O)")
        self.chunks = chunks
    elif chunk_path is not None:
        self._load_from_path(
            Path(chunk_path), signal_mode=signal_mode, seq_encoding=seq_encoding
        )
    else:
        raise ValueError("Either chunk_path or chunks must be provided")

    # Filter chunks with valid numeric labels (label_int). The columnar
    # path applied this at load, recording the same mask so npz rows stay
    # aligned with self.chunks — a mismatch here would pair signals with
    # the wrong labels silently.
    if not isinstance(self.chunks, ChunkTable):
        self.chunks = [c for c in self.chunks if c["label_int"] is not None]

    if len(self.chunks) == 0:
        raise ValueError(f"No valid chunks found{f' in {chunk_path}' if chunk_path else ''}")

    # Pre-tensorize: encode sequences, labels, signals, and features once.
    # Each accumulator fills one preallocated contiguous tensor — see
    # _TensorFill for why stacking a list instead doubles the peak.
    n_chunks = len(self.chunks)
    fill_encoded_seqs = _TensorFill("Encoded-sequence", n_chunks)
    fill_labels = _TensorFill("Label", n_chunks)
    fill_signals = _TensorFill("Signal", n_chunks)
    fill_features = _TensorFill("Feature", n_chunks)
    fill_confounds = _TensorFill("Confound", n_chunks)
    fill_cl_targets = _TensorFill("CL target", n_chunks)
    self._encoded_seqs: list[torch.Tensor] = []
    self._labels: list[torch.Tensor] = []
    self._signals: list[torch.Tensor] = []
    self._features: list[torch.Tensor] = []
    self._confound_encoder = confound_encoder
    self._has_confound = confound_encoder is not None
    self._confound_labels: list[torch.Tensor] = []
    self._cl_regression = cl_regression
    self._cl_targets: list[torch.Tensor] = []

    # Determine effective encoding, from the whole corpus rather than from
    # chunk 0 (#230). A corpus carries the signal_kmer inputs for every
    # chunk or for none: none is the version-skew case the fallback exists
    # for, and anything in between is damage.
    self._effective_seq_encoding = seq_encoding
    if seq_encoding == "signal_kmer":
        n_covered, n_total = self._signal_kmer_coverage()
        if n_covered < n_total:
            n_missing = n_total - n_covered
            detail = (
                f"{n_missing} of {n_total} chunks ({n_missing / n_total:.1%}) lack "
                f"seq_to_sig_map/sequence_with_kmer_context"
                f"{f' in {chunk_path}' if chunk_path else ''}"
            )
            if n_covered:
                # Ragged, not legacy — and no automatic choice is right.
                # Encoding the uncovered rows anyway gives them all-zero
                # sequence channels; switching the whole corpus to
                # base_onehot on the strength of a few bad rows throws the
                # encoding away for every good one. Both are the silent
                # representation change this guard exists to prevent.
                raise ValueError(
                    f"seq_encoding='signal_kmer': {detail}. A corpus carries "
                    "base-to-signal maps for every chunk or for none; re-prepare "
                    "it, or ask for --seq-encoding base_onehot explicitly."
                )
            if not allow_encoding_fallback:
                raise ValueError(
                    f"seq_encoding='signal_kmer' was requested but {detail}. "
                    "Re-prepare the corpus with a leech version that writes "
                    "base-to-signal maps, or ask for --seq-encoding base_onehot "
                    "explicitly."
                )
            logger.warning("%s; falling back to base_onehot encoding", detail)
            self._effective_seq_encoding = "base_onehot"
            self._s2s_csr = None

    # Detect signal_residual channel and apply signal_mode
    self._signal_mode = signal_mode
    if self._array_stream is not None:
        # Arrays were deferred, so presence is a property of the file.
        self._has_signal_residual = bool(
            {"signal_residuals_flat", "signal_residuals"} & self._npz_members
        )
    else:
        self._has_signal_residual = self.chunks[0].get("signal_residual") is not None
    if signal_mode == "both" and self._has_signal_residual:
        self.signal_channels = 2
    else:
        self.signal_channels = 1

    # Detect multi-class: if any label_int > 1, use long dtype for
    # CrossEntropyLoss. Off the column when there is one — the generator
    # form builds a row view per chunk, 150 ms per 200k chunks for a max.
    if isinstance(self.chunks, ChunkTable):
        max_label = int(self.chunks.require_values("label_int").max())
    else:
        max_label = max(c["label_int"] for c in self.chunks)
    self._multiclass = max_label > 1

    # For signal_kmer encoding, the on-the-fly inputs are ~30 B/chunk
    # (seq_ints + seq_to_sig_map) vs ~77 KB/chunk for the encoded output
    # — a >2000x reduction. We stash the compact inputs here and run
    # `encode_signal_kmer` lazily in __getitem__, where DataLoader workers
    # parallelize it behind prefetch. base_onehot stays eagerly tensorized
    # (88 floats/chunk; not worth deferring).
    self._seq_ints: list[np.ndarray] = []
    self._seq_to_sig: list[np.ndarray] = []
    self._seq_ints_tensor: torch.Tensor | None = None

    # Arrays come either from the chunk dicts (pre-loaded / legacy corpus)
    # or a row-block stream over the npz. Both fill the same accumulators.
    # A streamed corpus is handled a whole block at a time: the arrays
    # arrive in blocks of ~1,500 rows, and taking them apart to prepare one
    # row at a time cost 58-67 us per chunk — six to seven minutes of
    # single-threaded startup on a 6.7M-chunk corpus, none of it I/O.
    fills = {
        "signals": fill_signals,
        "encoded_seqs": fill_encoded_seqs,
        "labels": fill_labels,
        "features": fill_features,
        "confounds": fill_confounds,
        "cl_targets": fill_cl_targets,
    }
    if self._block_fill_supported():
        self._fill_from_blocks(fills)
    else:
        self._fill_from_rows(fills)

    # Every per-chunk tensor now lives in one contiguous buffer. That isn't
    # just for cache friendliness — it's required for fork-safety. A
    # DataLoader with num_workers > 0 forks worker processes that COW-inherit
    # the parent's address space. CPython refcounts live inside each
    # PyObject header, so a worker iterating a list of N tensors writes to N
    # separate page-resident headers and faults every page into a private
    # copy, multiplying peak RSS by (1 + num_workers). A single contiguous
    # tensor keeps its data buffer outside Python's GC, so the buffer
    # pages are genuinely shared across the fork.
    self._signals_tensor, self._signals = fill_signals.finish()
    self._encoded_seqs_tensor, self._encoded_seqs = fill_encoded_seqs.finish()

    # Compact signal_kmer inputs — only populated when encoding == signal_kmer.
    # Stacked into fork-safe int tensors. encode_signal_kmer is then called
    # lazily per-sample in __getitem__ (workers parallelize behind prefetch).
    # Chunks have variable basecalled sequence lengths near the motif, so we
    # pad to the per-array max. Padding values are chosen so the encoder
    # gracefully ignores them: seq_ints=-1 hits the encoder's `base < 0`
    # skip; seq_to_sig=signal_len makes its `sig_start < signal_len` check
    # fail, writing nothing for the padded positions.
    #
    # The block-wise filler builds the padded int8 matrix straight off the
    # text column and leaves ``_seq_ints_tensor`` set; the per-chunk path
    # leaves a list to pad here. Either way the CSR expansion below runs.
    self._seq_to_sig_tensor: torch.Tensor | None = None
    if self._effective_seq_encoding == "signal_kmer" and (
        self._seq_ints or self._seq_ints_tensor is not None
    ):
        if self._seq_ints:
            max_seq_ints_len = max(s.shape[0] for s in self._seq_ints)
            n = len(self._seq_ints)
            padded_seq_ints = np.full((n, max_seq_ints_len), -1, dtype=np.int8)
            for i, si in enumerate(self._seq_ints):
                padded_seq_ints[i, : si.shape[0]] = si
            self._seq_ints_tensor = torch.from_numpy(padded_seq_ints)
            self._seq_ints = []
        n = int(self._seq_ints_tensor.shape[0])

        if self._s2s_csr is not None:
            # Streaming path: expand every map at once from the CSR pair,
            # instead of one astype/clip per chunk.
            values, offsets = self._s2s_csr
            crop_starts = None
            if self.left_context is not None and self.right_context is not None:
                crop_starts = self._crop_starts(values, offsets) - self.left_context
            padded_s2s = _expand_seq_to_sig_csr(
                values,
                offsets,
                self._s2s_rows,
                signal_len=signal_len,
                crop_starts=crop_starts,
            )
            self._s2s_csr = None
        else:
            max_s2s_len = max(s.shape[0] for s in self._seq_to_sig)
            padded_s2s = np.full((n, max_s2s_len), signal_len, dtype=np.int64)
            for i, s2s in enumerate(self._seq_to_sig):
                padded_s2s[i, : s2s.shape[0]] = s2s
        self._seq_to_sig_tensor = torch.from_numpy(padded_s2s)
        self._seq_to_sig = []

    self._labels_tensor, self._labels = fill_labels.finish()

    self._features_tensor: torch.Tensor | None = None
    if self._needs_features:
        self._features_tensor, self._features = fill_features.finish()

    self._confound_labels_tensor: torch.Tensor | None = None
    if self._has_confound:
        self._confound_labels_tensor, self._confound_labels = fill_confounds.finish()

    self._cl_targets_tensor: torch.Tensor | None = None
    if self._cl_regression:
        self._cl_targets_tensor, self._cl_targets = fill_cl_targets.finish()

    # Drop the raw numpy arrays from self.chunks now that everything has
    # been pre-tensorized. External code (samplers, label tally, feature
    # window introspection) still reads the small scalar/string fields,
    # so we keep self.chunks alive but null out the per-chunk arrays.
    # Without this, each chunk dict keeps a ~50 KB numpy view alive and
    # the same COW blowup hits during DataLoader fork.
    if not isinstance(self.chunks, ChunkTable):
        for chunk in self.chunks:
            for key in (
                "signal",
                "signal_residual",
                "dwell",
                "features",
                "seq_to_sig_map",
                "sequence_with_kmer_context",
            ):
                if key in chunk:
                    chunk[key] = None

    # Same reasoning for the streaming bookkeeping: one row per chunk each,
    # and a DataLoader that spawns workers pickles whatever is still here.
    self._s2s_csr = None
    self._s2s_rows = None
    if self._array_stream is not None:
        self._array_stream.keep = None

    # Precompute per-channel feature stds for feature noise augmentation.
    # Reuse the already-stacked features tensor when available.
    self._feature_stds: torch.Tensor | None = None
    if self._feature_noise_scale > 0 and self._needs_features:
        if self._features_tensor is not None:
            self._feature_stds = self._features_tensor.std(dim=0)  # (C, K)
        else:
            logger.warning("Feature shapes differ, feature noise disabled")
            self._feature_noise_scale = 0.0

    # Approx samples per base for cross-layer shift/mask
    self._samples_per_base = signal_len / max(kmer_len, 1)

    # Whether __getitems__ can gather a batch straight out of the
    # contiguous tensors. Every field has to be one: a field that degraded
    # to a list of per-chunk tensors has no batch to gather.
    self._batched_fetch = (
        self._signals_tensor is not None
        and self._labels_tensor is not None
        and (not self._needs_features or self._features_tensor is not None)
        and (not self._has_confound or self._confound_labels_tensor is not None)
        and (not self._cl_regression or self._cl_targets_tensor is not None)
        and (
            (self._seq_ints_tensor is not None and self._seq_to_sig_tensor is not None)
            if self._effective_seq_encoding == "signal_kmer"
            else self._encoded_seqs_tensor is not None
        )
    )

    _n_encoded = (
        self._encoded_seqs_tensor.shape[0]
        if self._encoded_seqs_tensor is not None
        else len(self._encoded_seqs)
    )
    logger.debug(
        f"Pre-tensorized {len(self.chunks)} chunks "
        f"({_n_encoded} sequences encoded, encoding={self._effective_seq_encoding})"
    )