finemo.postprocessing

  1"""Post-processing utilities for Fi-NeMo hit calling results.
  2
  3This module provides functions for:
  4- Collapsing overlapping hits based on similarity scores
  5- Intersecting hit sets across multiple runs
  6- Quality control and filtering operations
  7
  8The main operations are optimized using Numba for efficient processing
  9of large hit datasets.
 10"""
 11
 12import heapq
 13from typing import List, Union
 14
 15import numpy as np
 16from numpy import ndarray
 17import polars as pl
 18from numba import njit
 19from numba.types import Array, uint32, int32, float32  # type: ignore[attr-defined]
 20from jaxtyping import Float, Int
 21
 22
 23@njit(
 24    uint32[:](
 25        Array(uint32, 1, "C", readonly=True),
 26        Array(int32, 1, "C", readonly=True),
 27        Array(int32, 1, "C", readonly=True),
 28        Array(float32, 1, "C", readonly=True),
 29    ),
 30    cache=True,
 31)
 32def _collapse_hits(
 33    chrom_ids: Int[ndarray, " N"],
 34    starts: Int[ndarray, " N"],
 35    ends: Int[ndarray, " N"],
 36    similarities: Float[ndarray, " N"],
 37) -> Int[ndarray, " N"]:
 38    """Identify primary hits among overlapping hits using a sweep line algorithm.
 39
 40    This function uses a heap-based sweep line algorithm to efficiently identify
 41    the best hit (highest similarity) among sets of overlapping hits within each
 42    chromosome. Only one hit per overlapping group is marked as primary.
 43
 44    Parameters
 45    ----------
 46    chrom_ids : Int[ndarray, "N"]
 47        Chromosome identifiers for each hit, where N is the number of hits.
 48        Dtype should be uint32 for Numba compatibility.
 49    starts : Int[ndarray, "N"]
 50        Start positions of hits (adjusted for overlap computation).
 51        Dtype should be int32 for Numba compatibility.
 52    ends : Int[ndarray, "N"]
 53        End positions of hits (adjusted for overlap computation).
 54        Dtype should be int32 for Numba compatibility.
 55    similarities : Float[ndarray, "N"]
 56        Similarity scores used for selecting the best hit.
 57        Dtype should be float32 for Numba compatibility.
 58
 59    Returns
 60    -------
 61    Int[ndarray, "N"]
 62        Binary array where 1 indicates the hit is primary, 0 otherwise.
 63        Returns uint32 array for consistency with input types.
 64
 65    Notes
 66    -----
 67    This function is JIT-compiled with Numba for performance on large datasets.
 68    The algorithm maintains active intervals in a heap and resolves overlaps
 69    by keeping only the hit with the highest similarity score.
 70
 71    The sweep line algorithm processes hits in order and maintains a heap of
 72    currently active intervals. When a new interval is encountered, it is
 73    compared against all overlapping intervals in the heap, and only the
 74    interval with the highest similarity score remains marked as primary.
 75    """
 76    n = chrom_ids.shape[0]
 77    out = np.ones(n, dtype=np.uint32)
 78    heap = [(np.uint32(0), np.int32(0), -1) for _ in range(0)]
 79
 80    for i in range(n):
 81        chrom_new = chrom_ids[i]
 82        start_new = starts[i]
 83        end_new = ends[i]
 84        sim_new = similarities[i]
 85
 86        # Remove expired intervals from heap
 87        while heap and heap[0] < (chrom_new, start_new, -1):
 88            heapq.heappop(heap)
 89
 90        # Check overlaps with active intervals
 91        for _, _, idx in heap:
 92            cmp = sim_new > similarities[idx]
 93            out[idx] &= cmp
 94            out[i] &= not cmp
 95
 96        # Add current interval to heap
 97        heapq.heappush(heap, (chrom_new, end_new, i))
 98
 99    return out
100
101
102def collapse_hits(
103    hits_df: Union[pl.DataFrame, pl.LazyFrame], overlap_frac: float
104) -> pl.DataFrame:
105    """Collapse overlapping hits by selecting the best hit per overlapping group.
106
107    This function identifies overlapping hits and marks only the highest-similarity
108    hit as primary in each overlapping group. Overlap is determined by a fractional
109    threshold based on the average length of the two hits being compared.
110
111    Parameters
112    ----------
113    hits_df : Union[pl.DataFrame, pl.LazyFrame]
114        Hit data containing required columns: chr (or peak_id if no chr), start, end,
115        hit_similarity. Will be collected to DataFrame if passed as LazyFrame.
116    overlap_frac : float
117        Overlap fraction threshold for considering hits as overlapping.
118        For two hits with lengths x and y, minimum overlap = overlap_frac * (x + y) / 2.
119        Must be between 0 and 1, where 0 means any overlap and 1 means complete overlap.
120
121    Returns
122    -------
123    pl.DataFrame
124        Original hit data with an additional 'is_primary' column (1 for primary hits, 0 otherwise).
125        All original columns are preserved, with the new column added at the end.
126
127    Raises
128    ------
129    KeyError
130        If required columns (chr/peak_id, start, end, hit_similarity) are missing.
131
132    Notes
133    -----
134    The algorithm transforms coordinates by scaling by 2 and adjusting by the overlap
135    fraction to create effective overlap regions for efficient processing. This allows
136    using a sweep line algorithm to identify overlaps in a single pass.
137
138    The transformation works as follows:
139    - Original coordinates: [start, end]
140    - Length = end - start
141    - Adjusted start = start * 2 + length * overlap_frac
142    - Adjusted end = end * 2 - length * overlap_frac
143
144    This creates regions that overlap only when the original regions have sufficient
145    overlap according to the specified fraction.
146
147    Examples
148    --------
149    >>> hits_collapsed = collapse_hits(hits_df, overlap_frac=0.2)
150    >>> primary_hits = hits_collapsed.filter(pl.col("is_primary") == 1)
151    >>> print(f"Kept {primary_hits.height}/{hits_df.height} hits as primary")
152    """
153    # Ensure we're working with a DataFrame
154    if isinstance(hits_df, pl.LazyFrame):
155        hits_df = hits_df.collect()
156
157    chroms = hits_df["chr"].unique(maintain_order=True)
158
159    if not chroms.is_empty():
160        chrom_to_id = {chrom: i for i, chrom in enumerate(chroms)}
161        # Transform coordinates for overlap computation
162        # Scale by 2 and adjust by overlap fraction to create effective overlap regions
163        df = hits_df.select(
164            chrom_id=pl.col("chr").replace_strict(chrom_to_id, return_dtype=pl.UInt32),
165            start_trim=pl.col("start") * 2
166            + ((pl.col("end") - pl.col("start")) * overlap_frac).cast(pl.Int32),
167            end_trim=pl.col("end") * 2
168            - ((pl.col("end") - pl.col("start")) * overlap_frac).cast(pl.Int32),
169            similarity=pl.col("hit_similarity"),
170        )
171    else:
172        # Fall back to peak_id when chr column is not available
173        df = hits_df.select(
174            chrom_id=pl.col("peak_id"),
175            start_trim=pl.col("start") * 2
176            + ((pl.col("end") - pl.col("start")) * overlap_frac).cast(pl.Int32),
177            end_trim=pl.col("end") * 2
178            - ((pl.col("end") - pl.col("start")) * overlap_frac).cast(pl.Int32),
179            similarity=pl.col("hit_similarity"),
180        )
181
182    # Rechunk for efficient array access
183    df = df.rechunk()
184    chrom_ids = df["chrom_id"].to_numpy(allow_copy=False)
185    starts = df["start_trim"].to_numpy(allow_copy=False)
186    ends = df["end_trim"].to_numpy(allow_copy=False)
187    similarities = df["similarity"].to_numpy(allow_copy=False)
188
189    # Run the collapse algorithm
190    is_primary = _collapse_hits(chrom_ids, starts, ends, similarities)
191
192    # Add primary indicator column to original DataFrame
193    df_out = hits_df.with_columns(is_primary=pl.Series(is_primary, dtype=pl.UInt32))
194
195    return df_out
196
197
198def intersect_hits(
199    hits_dfs: List[Union[pl.DataFrame, pl.LazyFrame]], relaxed: bool
200) -> pl.DataFrame:
201    """Intersect hit datasets across multiple runs to find common hits.
202
203    This function finds hits that appear consistently across multiple Fi-NeMo
204    runs, which can be useful for identifying robust motif instances that are
205    not sensitive to parameter variations or random initialization.
206
207    Parameters
208    ----------
209    hits_dfs : List[Union[pl.DataFrame, pl.LazyFrame]]
210        List of hit DataFrames from different Fi-NeMo runs. Each DataFrame must
211        contain the columns specified by the intersection criteria. LazyFrames
212        will be collected before processing.
213    relaxed : bool
214        If True, uses relaxed intersection criteria with only motif names and
215        untrimmed coordinates. If False, uses strict criteria including all
216        coordinate and metadata columns.
217
218    Returns
219    -------
220    pl.DataFrame
221        DataFrame containing hits that appear in all input datasets.
222        Columns from later datasets are suffixed with their index (e.g., '_1', '_2').
223        The first dataset's columns retain their original names.
224
225    Raises
226    ------
227    ValueError
228        If fewer than one hits DataFrame is provided.
229    KeyError
230        If required columns for the specified intersection criteria are missing
231        from any of the input DataFrames.
232
233    Notes
234    -----
235    Relaxed intersection is useful when comparing results across different
236    region definitions or motif trimming parameters, but may produce less
237    precise matches. Strict intersection requires identical region definitions
238    and is recommended for most use cases.
239
240    The intersection columns used are:
241    - Relaxed: ["chr", "start_untrimmed", "end_untrimmed", "motif_name", "strand"]
242    - Strict: ["chr", "start", "end", "start_untrimmed", "end_untrimmed",
243               "motif_name", "strand", "peak_name", "peak_id"]
244
245    The function performs successive inner joins starting with the first DataFrame,
246    so the final result contains only hits present in all input datasets.
247
248    Examples
249    --------
250    >>> common_hits = intersect_hits([hits_df1, hits_df2], relaxed=False)
251    >>> print(f"Found {common_hits.height} hits common to both runs")
252    >>>
253    >>> # Compare relaxed vs strict intersection
254    >>> relaxed_hits = intersect_hits([hits_df1, hits_df2], relaxed=True)
255    >>> strict_hits = intersect_hits([hits_df1, hits_df2], relaxed=False)
256    >>> print(f"Relaxed: {relaxed_hits.height}, Strict: {strict_hits.height}")
257    """
258    if relaxed:
259        # Relaxed criteria: only motif identity and untrimmed positions
260        join_cols = ["chr", "start_untrimmed", "end_untrimmed", "motif_name", "strand"]
261    else:
262        # Strict criteria: all coordinate and metadata columns
263        join_cols = [
264            "chr",
265            "start",
266            "end",
267            "start_untrimmed",
268            "end_untrimmed",
269            "motif_name",
270            "strand",
271            "peak_name",
272            "peak_id",
273        ]
274
275    if len(hits_dfs) < 1:
276        raise ValueError("At least one hits dataframe required")
277
278    # Ensure all DataFrames are collected
279    collected_dfs = []
280    for df in hits_dfs:
281        if isinstance(df, pl.LazyFrame):
282            collected_dfs.append(df.collect())
283        else:
284            collected_dfs.append(df)
285
286    # Start with first DataFrame and successively intersect with others
287    hits_df = collected_dfs[0]
288    for i in range(1, len(collected_dfs)):
289        hits_df = hits_df.join(
290            collected_dfs[i],
291            on=join_cols,
292            how="inner",
293            suffix=f"_{i}",
294            join_nulls=True,
295            coalesce=True,
296        )
297
298    return hits_df