Source code for gunz_cm.preprocs.matrices.contacts

"""Module for selecting informative Hi-C contacts via a naive baseline.

This module provides two primitives that together identify the subset of
Hi-C contacts where a naive 3D baseline diverges from the observed
contact map. The naive baseline is a power-law of pairwise Euclidean
distance, ``C_pred(i,j) = d(i,j) ** -alpha``, so contacts whose residual
against the observed Hi-C is large are the "informative" ones (they
carry signal beyond what pure geometry can explain).

"""
__author__ = "Yeremia Gunawan Adhisantoso"

# =============================================================================
# METADATA
# =============================================================================
__email__ = "adhisant@tnt.uni-hannover.de"
__license__ = "Clear BSD"

# =============================================================================
# STANDARD LIBRARY IMPORTS
# =============================================================================
import typing as t

# =============================================================================
# THIRD-PARTY IMPORTS
# =============================================================================
from gunz_cm.exceptions import PreprocError
import numpy as np
import pandas as pd
from pydantic import ConfigDict, validate_call
from scipy.spatial.distance import pdist, squareform

# =============================================================================
# LOCAL APPLICATION IMPORTS
# =============================================================================

[docs]@validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def naive_predicted_contacts( p_naive: np.ndarray, alpha: float, k_band: int = 100, ) -> np.ndarray: """Compute naive baseline's predicted Hi-C contact matrix from structure. Uses a power-law of pairwise Euclidean distance: ``C_pred(i,j) = d(i,j) ** -alpha`` where ``d`` is the Euclidean distance between ``p_naive[i]`` and ``p_naive[j]``. Only pairs within ``k_band`` genomic-distance bins are populated; everything beyond the band is zeroed, and the main diagonal is zeroed as well. Parameters ---------- p_naive : np.ndarray Naive baseline structure of shape ``(N, 3)``. alpha : float Power-law exponent (typical: 0.6 to 0.9 for chromatin). k_band : int Max genomic distance in bins. Pairs beyond this band get value 0. Returns ------- np.ndarray Dense ``(N, N)`` symmetric array with the band populated; outside band = 0; diagonal = 0. """ if p_naive.ndim != 2: raise PreprocError("p_naive must be a 2D array of shape (N, 3).") if p_naive.shape[1] != 3: raise PreprocError( f"p_naive must have 3 columns (got {p_naive.shape[1]})." ) if k_band <= 0: raise PreprocError(f"k_band must be positive (got {k_band}).") n = p_naive.shape[0] # Pairwise Euclidean distances via scipy (already used elsewhere in preprocs). dist_sq = squareform(pdist(p_naive)) # Inverse power-law prediction. Clip to avoid division-by-zero on diagonal. np.fill_diagonal(dist_sq, 1.0) predicted = np.power(dist_sq, -alpha) np.fill_diagonal(predicted, 0.0) # Zero out everything beyond the band to mimic a localized contact signal. if k_band < n: rows, cols = np.indices((n, n)) predicted[cols - rows > k_band] = 0.0 # Mirror the upper-triangular cutoff across the diagonal for symmetry. predicted[rows - cols > k_band] = 0.0 return predicted
[docs]@validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def select_informative_contacts( hic_dense: np.ndarray, p_naive: np.ndarray, threshold_quantile: float = 0.5, k_band: int = 100, alpha: float | None = None, ) -> np.ndarray: """Return boolean mask of contacts where naive diverges from Hi-C. For each ``(i, j)`` pair in the upper-triangular band of width ``k_band``: * ``d_naive[i, j] = ||p_naive[i] - p_naive[j]||`` * ``hic_pred[i, j] = d_naive[i, j] ** -alpha`` if ``alpha`` is given, otherwise ``hic_pred[i, j] = hic_dense[i, j]`` (skip prediction). * ``residual[i, j] = |hic_dense[i, j] - hic_pred[i, j]|`` Returns a boolean mask where the residual exceeds the ``threshold_quantile``-quantile of all residuals in the band. Default ``0.5`` selects the upper half by residual. Parameters ---------- hic_dense : np.ndarray Square Hi-C matrix of shape ``(N, N)``. p_naive : np.ndarray Naive baseline structure of shape ``(N, 3)``. threshold_quantile : float Quantile in ``[0, 1]``. Higher = more selective. k_band : int Max genomic distance in bins. alpha : float | None If given, predict naive Hi-C via power law; else use ``hic_dense`` directly so the residual is zero and the mask is all-False. Returns ------- np.ndarray Boolean ``(N, N)`` array. ``True`` for "informative" contacts. """ if hic_dense.ndim != 2: raise PreprocError("hic_dense must be a 2D square array.") if hic_dense.shape[0] != hic_dense.shape[1]: raise PreprocError( f"hic_dense must be square (got shape {hic_dense.shape})." ) if p_naive.ndim != 2 or p_naive.shape[1] != 3: raise PreprocError( f"p_naive must have shape (N, 3) (got {p_naive.shape})." ) if hic_dense.shape[0] != p_naive.shape[0]: raise PreprocError( f"Shape mismatch: hic_dense has N={hic_dense.shape[0]} rows " f"but p_naive has N={p_naive.shape[0]} points." ) if not 0.0 <= threshold_quantile <= 1.0: raise PreprocError( f"threshold_quantile must be in [0, 1] (got {threshold_quantile})." ) if k_band <= 0: raise PreprocError(f"k_band must be positive (got {k_band}).") n = hic_dense.shape[0] # Pairwise Euclidean distances for the naive baseline. dist_sq = squareform(pdist(p_naive)) np.fill_diagonal(dist_sq, 1.0) if alpha is None: # No prediction -> residual is identically zero -> all-False mask. return np.zeros((n, n), dtype=bool) hic_pred = np.power(dist_sq, -alpha) np.fill_diagonal(hic_pred, 0.0) residual = np.abs(hic_dense.astype(np.float64) - hic_pred) # Upper-triangular band mask (col - row in [1, k_band]). rows, cols = np.triu_indices(n, k=1) band_mask = (cols - rows) <= k_band band_rows = rows[band_mask] band_cols = cols[band_mask] band_residuals = residual[band_rows, band_cols] if band_residuals.size == 0: return np.zeros((n, n), dtype=bool) threshold = float(np.quantile(band_residuals, threshold_quantile)) informative = np.zeros((n, n), dtype=bool) informative[band_rows, band_cols] = band_residuals > threshold return informative
[docs]@validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def filter_bimodal_outliers( df: pd.DataFrame, points: np.ndarray, raw_threshold: int = 2, dist_factor: float = 2.0, k_band: int = 50, mask: np.ndarray | None = None, pooled_points: np.ndarray | None = None, ) -> dict[str, t.Any]: """Identify bimodal outlier contacts: low raw count AND high predicted distance. A contact ``(i, j)`` is a "bimodal outlier" if **both** conditions hold: 1. ``hic_raw_count[i, j] < raw_threshold`` (likely PCR / sequencing noise) 2. ``predicted_distance[i, j] > dist_factor * local_median_distance`` (the structure says the two loci are far apart, contradicting the absence of a Hi-C signal — i.e., the contact is at once low-confidence in Hi-C and rejected by the structure). Such contacts carry no useful rank signal for Spearman correlation: the structure's "far apart" is not contradicting high-confidence Hi-C (because there is no high-confidence Hi-C), so the only thing they contribute is noise that deflates the global Spearman. Parameters ---------- df : pd.DataFrame Must contain columns ``['row_ids', 'col_ids', 'raw_counts']`` — typically the output of ``load_cm_data(return_raw_counts=True)``. ``row_ids`` and ``col_ids`` are bin indices in ``[0, N-1]``. points : np.ndarray ``(N, 3)`` 3D coordinates of the structure under evaluation. Pairwise Euclidean distance is computed only for these points and must align with ``df`` row/col ids. raw_threshold : int Contacts with ``raw_counts`` strictly less than this value are noise candidates. Default ``2`` drops 0-count and 1-count contacts. dist_factor : float Multiplier on the local median distance at each genomic distance ``s``. A contact at genomic distance ``s`` whose predicted 3D distance exceeds ``dist_factor * median_at_s`` is an outlier. Default ``2.0``. k_band : int Maximum genomic distance (in bins) considered. Contacts beyond ``k_band`` are not classified — they are returned as kept by default. Default ``50`` (covers roughly 100 kb to 2.5 Mb at 10 kb resolution). mask : np.ndarray | None Optional ``(N,)`` boolean mask of valid bins. Invalid bins are excluded from the distance matrix used for the per-structure median curve. pooled_points : np.ndarray | None Optional ``(N, 3)`` reference structure used to compute the *local median distance curve*. If ``None``, ``points`` is used. Useful when comparing structures of different complexity (e.g., smooth naive vs. wiggly real) so the median is derived from a common reference (e.g., 25 kb Golden). Returns ------- dict ``keep_mask`` : np.ndarray of bool, length ``n_total_in_band`` ``True`` for contacts that pass the filter (NOT outliers). ``is_outlier`` : np.ndarray of bool, same length as ``keep_mask``; ``True`` for contacts that ARE outliers. ``raw_counts`` : np.ndarray Raw counts of the contacts in the band. ``predicted_distances`` : np.ndarray Predicted 3D distances for the contacts in the band. ``genomic_distances`` : np.ndarray Genomic distances (in bins) for the contacts in the band. ``median_dist_curve`` : dict[int, float] Median predicted distance per genomic distance ``s`` from ``1`` to ``k_band``. Genomic distances with no pairs are omitted. ``n_total`` : int Total contacts in the band. ``n_outliers`` : int Number of outliers removed. ``n_kept`` : int Number of contacts kept. ``raw_threshold`` : int Echo of the parameter for downstream auditability. ``dist_factor`` : float Echo of the parameter for downstream auditability. Raises ------ PreprocError If ``raw_counts`` is missing from ``df``, or if any of ``raw_threshold < 1``, ``k_band < 1``, ``dist_factor < 1.0``, or ``points.ndim != 2 / shape[1] != 3``. """ # --- Input validation ------------------------------------------------ required_cols = {"row_ids", "col_ids", "raw_counts"} missing = required_cols - set(df.columns) if missing: raise PreprocError( f"DataFrame is missing required columns: {sorted(missing)}." ) if points.ndim != 2 or points.shape[1] != 3: raise PreprocError( f"points must have shape (N, 3) (got {points.shape})." ) if k_band < 1: raise PreprocError(f"k_band must be >= 1 (got {k_band}).") if raw_threshold < 1: raise PreprocError( f"raw_threshold must be >= 1 (got {raw_threshold})." ) if dist_factor < 1.0: raise PreprocError( f"dist_factor must be >= 1.0 (got {dist_factor})." ) if mask is not None and mask.shape[0] != points.shape[0]: raise PreprocError( f"mask length ({mask.shape[0]}) must match points.shape[0] " f"({points.shape[0]})." ) # --- Edge case: empty DataFrame -------------------------------------- if df.empty: return { "keep_mask": np.array([], dtype=bool), "is_outlier": np.array([], dtype=bool), "raw_counts": np.array([], dtype=np.int64), "predicted_distances": np.array([], dtype=np.float64), "genomic_distances": np.array([], dtype=np.int64), "median_dist_curve": {}, "n_total": 0, "n_outliers": 0, "n_kept": 0, "raw_threshold": int(raw_threshold), "dist_factor": float(dist_factor), } # --- Convert df to numpy arrays ------------------------------------- row_ids = df["row_ids"].to_numpy(dtype=np.int64) col_ids = df["col_ids"].to_numpy(dtype=np.int64) raw_counts_full = df["raw_counts"].to_numpy(dtype=np.int64) # Upper-triangular band mask: col > row (avoid self-ligation) AND # col - row <= k_band. band_mask = (col_ids > row_ids) & ((col_ids - row_ids) <= k_band) # --- Pairwise distance matrices -------------------------------------- # Per-contact predicted distance comes from `points`. Median curve # optionally comes from `pooled_points` (or `points`). dist_points = squareform(pdist(points)) dist_for_median = ( squareform(pdist(pooled_points)) if pooled_points is not None else dist_points ) # When a mask is supplied, exclude invalid-bin pairs from the # median curve so the median reflects only well-supported # genomic-distance bins. if mask is not None: n = points.shape[0] valid = np.asarray(mask, dtype=bool) if not valid.all(): # Build a reduced distance matrix over valid bins and a # mapping from full-bin index -> reduced-bin index. valid_idx = np.flatnonzero(valid) sub_dist = dist_for_median[np.ix_(valid_idx, valid_idx)] else: valid_idx = np.arange(n) sub_dist = dist_for_median else: valid_idx = np.arange(points.shape[0]) sub_dist = dist_for_median n_sub = sub_dist.shape[0] # --- Empty-band short circuit ---------------------------------------- if not np.any(band_mask): return { "keep_mask": np.zeros(int(band_mask.sum()), dtype=bool), "is_outlier": np.zeros(int(band_mask.sum()), dtype=bool), "raw_counts": np.array([], dtype=np.int64), "predicted_distances": np.array([], dtype=np.float64), "genomic_distances": np.array([], dtype=np.int64), "median_dist_curve": {}, "n_total": 0, "n_outliers": 0, "n_kept": 0, "raw_threshold": int(raw_threshold), "dist_factor": float(dist_factor), } # --- Per-contact arrays in the band ---------------------------------- band_rows = row_ids[band_mask] band_cols = col_ids[band_mask] raw_b = raw_counts_full[band_mask] gen_dist = (band_cols - band_rows).astype(np.int64) pred_dist = dist_points[band_rows, band_cols].astype(np.float64) # --- Build median distance curve ------------------------------------- # Compute median of sub_dist for each genomic distance s in [1, k_band] # over all valid (i, j) with j - i == s. median_dist_curve: dict[int, float] = {} if n_sub > 1: sub_rows, sub_cols = np.triu_indices(n_sub, k=1) sub_gen = sub_cols - sub_rows sub_d = sub_dist[sub_rows, sub_cols] for s in range(1, k_band + 1): sel = sub_gen == s if np.any(sel): median_dist_curve[int(s)] = float(np.median(sub_d[sel])) # --- Apply the outlier criterion ------------------------------------- # Map each band's genomic distance to the median at that s; if a # genomic distance has no median (e.g., beyond `k_band` for the # reduced structure), we conservatively do NOT flag it as an outlier. band_median = np.array( [median_dist_curve.get(int(s), np.nan) for s in gen_dist], dtype=np.float64, ) low_count = raw_b < int(raw_threshold) has_median = ~np.isnan(band_median) high_distance = has_median & ( pred_dist > float(dist_factor) * band_median ) is_outlier = low_count & high_distance keep_mask = ~is_outlier return { "keep_mask": keep_mask, "is_outlier": is_outlier, "raw_counts": raw_b, "predicted_distances": pred_dist, "genomic_distances": gen_dist, "median_dist_curve": median_dist_curve, "n_total": int(band_mask.sum()), "n_outliers": int(is_outlier.sum()), "n_kept": int(keep_mask.sum()), "raw_threshold": int(raw_threshold), "dist_factor": float(dist_factor), }