Features Module¶

def compute_dwell_times(move_table: MoveTable) -> np.ndarray:
    """
    Compute per-base dwell times from move table.

    Dwell time = number of signal samples assigned to each base.

    Args:
        move_table: MoveTable object from extract_move_table()

    Returns:
        Array of shape (num_bases,) with dwell time for each base in signal samples

    Examples:
        >>> # Move array: [1,1,0,1,0,0,0,1,...] with stride=5
        >>> # Base 0: 1 move = 1 * 5 = 5 samples
        >>> # Base 1: 2 moves (1,1) = 2 * 5 = 10 samples
        >>> # Base 2: 4 moves (0,1,0,0,0) = 4 * 5 = 20 samples
    """
    seq_to_sig = move_table.to_seq_to_sig_map()
    dwells = np.diff(seq_to_sig)
    return dwells

def compute_signal_levels(
    signal: np.ndarray, seq_to_sig_map: np.ndarray, stat: str = "mean"
) -> np.ndarray:
    """
    Compute per-base signal level statistics.

    Args:
        signal: Normalized signal array
        seq_to_sig_map: Mapping from bases to signal indices (from MoveTable.to_seq_to_sig_map())
        stat: Statistic to compute ('mean', 'median', 'std', 'min', 'max')

    Returns:
        Array of shape (num_bases,) with signal level statistic for each base
    """
    # Rust fast path: compute all 4 stats at once, return requested one
    stat_to_idx = {"mean": 0, "median": 1, "std": 2}
    if HAS_RUST and stat in stat_to_idx:
        assert _rs_compute_signal_stats is not None
        sig = signal.astype(np.float32, copy=False)
        s2s = seq_to_sig_map.astype(np.int64, copy=False)
        means, medians, stds, ranges = _rs_compute_signal_stats(sig, s2s)
        return [np.asarray(means), np.asarray(medians), np.asarray(stds)][stat_to_idx[stat]]

    # Vectorized fast path (also the only place the mean is computed)
    if stat == "mean":
        return _per_base_means(signal, seq_to_sig_map)

    num_bases = len(seq_to_sig_map) - 1
    levels = np.zeros(num_bases, dtype=np.float32)
    lengths = np.diff(seq_to_sig_map)

    # Loop for median/std/min/max (no vectorized numpy equivalent)
    stat_funcs: dict[str, Callable[[np.ndarray], Any]] = {
        "mean": np.mean,
        "median": np.median,
        "std": np.std,
        "min": np.min,
        "max": np.max,
    }
    stat_func = stat_funcs[stat]

    for i in range(num_bases):
        if lengths[i] > 0:
            start_idx = seq_to_sig_map[i]
            end_idx = seq_to_sig_map[i + 1]
            base_signal = signal[start_idx:end_idx]
            if len(base_signal) > 0:
                levels[i] = float(stat_func(base_signal))

    return levels

def normalize_read_signal(
    raw_signal: np.ndarray,
    method: str = "median_mad",
    pa_mean: float | None = None,
    pa_stdev: float | None = None,
    cal_offset: float | None = None,
    cal_scale: float | None = None,
) -> tuple[np.ndarray, dict[str, float]]:
    """
    Normalize raw signal data.

    Named to avoid colliding with :func:`escapepod.normalize_signal`, which is a
    *different* transform: it takes int16 DAC input and divides by the bare MAD,
    with no 1.4826 factor. The ``median_mad`` method here delegates to
    :func:`escapepod.mad_normalize` instead, which is the matching one.

    Args:
        raw_signal: Raw DAC signal values
        method: Normalization method ('median_mad', 'zscore', 'quantile', 'pa_scaling')
        pa_mean: Global shift for pa_scaling normalization (from basecaller model)
        pa_stdev: Global scale for pa_scaling normalization (from basecaller model)
        cal_offset: POD5 calibration offset (DAC -> pA conversion)
        cal_scale: POD5 calibration scale (DAC -> pA conversion)

    Returns:
        Tuple of (normalized_signal, normalization_params)
    """
    if method == "median_mad":
        # Median Absolute Deviation normalization (robust to outliers).
        # escapepod's mad_normalize is the same 1.4826-scaled transform that
        # leech_core already uses (via escapepod_signal::mad_normalize_robust),
        # so the Python and Rust paths agree by construction. It also degrades
        # gracefully on a constant signal (dead pore / flat read), where the
        # naive form divides by a zero MAD and yields all-NaN.
        median = np.median(raw_signal)
        mad = np.median(np.abs(raw_signal - median))
        scale_factor = 1.4826
        normalized = mad_normalize(np.ascontiguousarray(raw_signal, dtype=np.float32))
        params = {"median": float(median), "mad": float(mad), "scale_factor": scale_factor}

    elif method == "zscore":
        # Standard z-score normalization
        mean = np.mean(raw_signal)
        std = np.std(raw_signal)
        normalized = (raw_signal - mean) / std
        params = {"mean": float(mean), "std": float(std)}

    elif method == "quantile":
        # Quantile normalization (winsorize extreme values)
        q01 = np.quantile(raw_signal, 0.01)
        q99 = np.quantile(raw_signal, 0.99)
        clipped = np.clip(raw_signal, q01, q99)
        median = np.median(clipped)
        mad = np.median(np.abs(clipped - median))
        normalized = (raw_signal - median) / (mad * 1.4826)
        params = {"median": float(median), "mad": float(mad), "q01": float(q01), "q99": float(q99)}

    elif method == "pa_scaling":
        # Remora-style global normalization using basecaller model parameters.
        # DAC -> pA (via POD5 calibration) -> global normalization.
        if pa_mean is None or pa_stdev is None:
            raise ValueError("pa_scaling requires pa_mean and pa_stdev parameters")
        if cal_offset is None or cal_scale is None:
            raise ValueError("pa_scaling requires cal_offset and cal_scale from POD5 calibration")

        # Convert DAC to pA: pA = (raw + offset) * scale (pod5 convention)
        pa_signal = (raw_signal + cal_offset) * cal_scale
        # Apply global normalization
        normalized = (pa_signal - pa_mean) / pa_stdev
        params = {
            "pa_mean": pa_mean,
            "pa_stdev": pa_stdev,
            "cal_offset": cal_offset,
            "cal_scale": cal_scale,
        }

    else:
        raise ValueError(f"Unknown normalization method: {method}")

    return normalized.astype(np.float32), params

def compute_dwell_features(dwells: np.ndarray, window: int = 5) -> dict[str, np.ndarray]:
    """
    Compute windowed dwell time features.

    Args:
        dwells: Per-base dwell times
        window: Size of sliding window for context features

    Returns:
        Dictionary with feature arrays:
            - 'dwell': raw dwell times
            - 'dwell_log': log-transformed dwell times
            - 'dwell_mean': local mean in window
            - 'dwell_std': local std in window
            - 'dwell_ratio': ratio to local mean
    """
    # Avoid log(0) by adding small epsilon
    eps = 1e-6
    dwell_log = np.log(dwells + eps)

    # Compute local statistics with padding using vectorized sliding window
    pad_width = window // 2
    padded = np.pad(dwells, pad_width, mode="edge")

    windows = sliding_window_view(padded, window)
    dwell_mean = windows.mean(axis=1).astype(np.float32)
    dwell_std = windows.std(axis=1).astype(np.float32)

    # Ratio of dwell to local mean (normalized dwell)
    dwell_ratio = dwells / (dwell_mean + eps)

    return {
        "dwell": dwells.astype(np.float32),
        "dwell_log": dwell_log.astype(np.float32),
        "dwell_mean": dwell_mean,
        "dwell_std": dwell_std,
        "dwell_ratio": dwell_ratio,
    }

def compute_signal_features(
    signal: np.ndarray, seq_to_sig_map: np.ndarray
) -> dict[str, np.ndarray]:
    """
    Compute comprehensive per-base signal features.

    Args:
        signal: Normalized signal array
        seq_to_sig_map: Base to signal mapping

    Returns:
        Dictionary with per-base features:
            - 'level_mean': mean signal level
            - 'level_median': median signal level
            - 'level_std': signal standard deviation
            - 'level_range': max - min signal
    """
    # Rust fast path: compute all 4 stats in one call (avoids Python dispatch overhead)
    if HAS_RUST:
        assert _rs_compute_signal_stats is not None
        sig = signal.astype(np.float32, copy=False)
        s2s = seq_to_sig_map.astype(np.int64, copy=False)
        means, medians, stds, ranges = _rs_compute_signal_stats(sig, s2s)
        return {
            "level_mean": np.asarray(means),
            "level_median": np.asarray(medians),
            "level_std": np.asarray(stds),
            "level_range": np.asarray(ranges),
        }

    # Python fallback
    num_bases = len(seq_to_sig_map) - 1

    features = {
        # Vectorized; the loop below covers the three that have no numpy
        # equivalent. Channel order here is the feature-row order -- see
        # `LeechRead.feature_channels`.
        "level_mean": _per_base_means(signal, seq_to_sig_map),
        "level_median": np.zeros(num_bases, dtype=np.float32),
        "level_std": np.zeros(num_bases, dtype=np.float32),
        "level_range": np.zeros(num_bases, dtype=np.float32),
    }

    lengths = np.diff(seq_to_sig_map)

    # Loop only for median/std/range (no vectorized numpy equivalent)
    for i in range(num_bases):
        if lengths[i] > 0:
            start = seq_to_sig_map[i]
            end = seq_to_sig_map[i + 1]
            base_sig = signal[start:end]
            if len(base_sig) > 0:
                features["level_median"][i] = np.median(base_sig)
                features["level_std"][i] = np.std(base_sig)
                features["level_range"][i] = np.max(base_sig) - np.min(base_sig)

    return features

def extract_move_table(alignment: pysam.AlignedSegment) -> MoveTable:
    """
    Extract move table from BAM alignment record.

    Args:
        alignment: pysam AlignedSegment with mv, ns, and ts tags

    Returns:
        MoveTable object

    Raises:
        ValueError: If required tags are missing
    """
    if not alignment.has_tag("mv"):
        raise ValueError(f"Read {alignment.query_name} missing 'mv' tag")
    if not alignment.has_tag("ns"):
        raise ValueError(f"Read {alignment.query_name} missing 'ns' tag")

    # Parse move table: first element is stride, rest is the move array
    mv_tag: Any = alignment.get_tag("mv")
    stride = int(mv_tag[0])
    moves = np.array(mv_tag[1:], dtype=np.int8)

    # Get signal metadata
    num_samples = int(alignment.get_tag("ns"))
    trim_offset = int(alignment.get_tag("ts")) if alignment.has_tag("ts") else 0

    read_id = alignment.query_name
    if read_id is None:
        raise ValueError("Alignment has no query_name")

    return MoveTable(
        stride=stride,
        moves=moves,
        read_id=read_id,
        num_samples=num_samples,
        trim_offset=trim_offset,
    )