finemo.data_io

   1"""Data input/output module for the Fi-NeMo motif instance calling pipeline.
   2
   3This module handles loading and processing of various genomic data formats including:
   4- Peak region files (ENCODE NarrowPeak format)
   5- Genome sequences (FASTA format)
   6- Contribution scores (bigWig, HDF5 formats)
   7- Neural network model outputs
   8- Motif data from TF-MoDISco
   9- Hit calling results
  10
  11The module supports multiple input formats used for contribution scores
  12and provides utilities for data conversion and quality control.
  13"""
  14
  15import json
  16import os
  17import warnings
  18from contextlib import ExitStack
  19from typing import List, Dict, Tuple, Optional, Any, Union, Callable
  20
  21import numpy as np
  22from numpy import ndarray
  23import h5py
  24import hdf5plugin  # noqa: F401, imported for side effects (HDF5 plugin registration)
  25import polars as pl
  26import pyBigWig
  27import pyfaidx
  28from jaxtyping import Float, Int
  29
  30from tqdm import tqdm
  31
  32
  33def load_txt(path: str) -> List[str]:
  34    """Load a text file containing one item per line.
  35
  36    Parameters
  37    ----------
  38    path : str
  39        Path to the text file.
  40
  41    Returns
  42    -------
  43    List[str]
  44        List of strings, one per line (first column if tab-delimited).
  45    """
  46    entries = []
  47    with open(path) as f:
  48        for line in f:
  49            item = line.rstrip("\n").split("\t")[0]
  50            entries.append(item)
  51
  52    return entries
  53
  54
  55def load_mapping(path: str, value_type: Callable[[str], Any]) -> Dict[str, Any]:
  56    """Load a two-column tab-delimited mapping file.
  57
  58    Parameters
  59    ----------
  60    path : str
  61        Path to the mapping file. Must be tab-delimited with exactly two columns.
  62    value_type : Callable[[str], Any]
  63        Type constructor to apply to values (e.g., int, float, str).
  64        Must accept a string and return the converted value.
  65
  66    Returns
  67    -------
  68    Dict[str, Any]
  69        Dictionary mapping keys to values of the specified type.
  70
  71    Raises
  72    ------
  73    ValueError
  74        If lines don't contain exactly two tab-separated values.
  75    FileNotFoundError
  76        If the specified file does not exist.
  77    """
  78    mapping = {}
  79    with open(path) as f:
  80        for line in f:
  81            key, val = line.rstrip("\n").split("\t")
  82            mapping[key] = value_type(val)
  83
  84    return mapping
  85
  86
  87def load_mapping_tuple(
  88    path: str, value_type: Callable[[str], Any]
  89) -> Dict[str, Tuple[Any, ...]]:
  90    """Load a mapping file where values are tuples from multiple columns.
  91
  92    Parameters
  93    ----------
  94    path : str
  95        Path to the mapping file. Must be tab-delimited with multiple columns.
  96    value_type : Callable[[str], Any]
  97        Type constructor to apply to each value element.
  98        Must accept a string and return the converted value.
  99
 100    Returns
 101    -------
 102    Dict[str, Tuple[Any, ...]]
 103        Dictionary mapping keys to tuples of values of the specified type.
 104        The first column is used as the key, remaining columns as tuple values.
 105
 106    Raises
 107    ------
 108    ValueError
 109        If lines don't contain at least two tab-separated values.
 110    FileNotFoundError
 111        If the specified file does not exist.
 112    """
 113    mapping = {}
 114    with open(path) as f:
 115        for line in f:
 116            entries = line.rstrip("\n").split("\t")
 117            key = entries[0]
 118            val = entries[1:]
 119            mapping[key] = tuple(value_type(i) for i in val)
 120
 121    return mapping
 122
 123
 124# ENCODE NarrowPeak format column definitions
 125NARROWPEAK_SCHEMA: List[str] = [
 126    "chr",
 127    "peak_start",
 128    "peak_end",
 129    "peak_name",
 130    "peak_score",
 131    "peak_strand",
 132    "peak_signal",
 133    "peak_pval",
 134    "peak_qval",
 135    "peak_summit",
 136]
 137NARROWPEAK_DTYPES: List[Any] = [
 138    pl.String,
 139    pl.Int32,
 140    pl.Int32,
 141    pl.String,
 142    pl.UInt32,
 143    pl.String,
 144    pl.Float32,
 145    pl.Float32,
 146    pl.Float32,
 147    pl.Int32,
 148]
 149
 150
 151def load_peaks(
 152    peaks_path: str, chrom_order_path: Optional[str], half_width: int
 153) -> pl.DataFrame:
 154    """Load peak region data from ENCODE NarrowPeak format file.
 155
 156    Parameters
 157    ----------
 158    peaks_path : str
 159        Path to the NarrowPeak format file.
 160    chrom_order_path : str, optional
 161        Path to file defining chromosome ordering. If None, uses order from peaks file.
 162    half_width : int
 163        Half-width of regions around peak summits.
 164
 165    Returns
 166    -------
 167    pl.DataFrame
 168        DataFrame containing peak information with columns:
 169        - chr: Chromosome name
 170        - peak_region_start: Start coordinate of centered region
 171        - peak_name: Peak identifier
 172        - peak_id: Sequential peak index
 173        - chr_id: Numeric chromosome identifier
 174    """
 175    peaks = (
 176        pl.scan_csv(
 177            peaks_path,
 178            has_header=False,
 179            new_columns=NARROWPEAK_SCHEMA,
 180            separator="\t",
 181            quote_char=None,
 182            schema_overrides=NARROWPEAK_DTYPES,
 183            null_values=[".", "NA", "null", "NaN"],
 184        )
 185        .select(
 186            chr=pl.col("chr"),
 187            peak_region_start=pl.col("peak_start") + pl.col("peak_summit") - half_width,
 188            peak_name=pl.col("peak_name"),
 189        )
 190        .with_row_index(name="peak_id")
 191        .collect()
 192    )
 193
 194    if chrom_order_path is not None:
 195        chrom_order = load_txt(chrom_order_path)
 196    else:
 197        chrom_order = []
 198
 199    chrom_order_set = set(chrom_order)
 200    chrom_order_peaks = [
 201        i
 202        for i in peaks.get_column("chr").unique(maintain_order=True)
 203        if i not in chrom_order_set
 204    ]
 205    chrom_order.extend(chrom_order_peaks)
 206    chrom_ind_map = {val: ind for ind, val in enumerate(chrom_order)}
 207
 208    peaks = peaks.with_columns(
 209        pl.col("chr").replace_strict(chrom_ind_map).alias("chr_id")
 210    )
 211
 212    return peaks
 213
 214
 215# DNA sequence alphabet for one-hot encoding
 216SEQ_ALPHABET: np.ndarray = np.array(["A", "C", "G", "T"], dtype="S1")
 217
 218
 219def one_hot_encode(sequence: str, dtype: Any = np.int8) -> Int[ndarray, "4 L"]:
 220    """Convert DNA sequence string to one-hot encoded matrix.
 221
 222    Parameters
 223    ----------
 224    sequence : str
 225        DNA sequence string containing A, C, G, T characters.
 226    dtype : np.dtype, default np.int8
 227        Data type for the output array.
 228
 229    Returns
 230    -------
 231    Int[ndarray, "4 L"]
 232        One-hot encoded sequence where rows correspond to A, C, G, T and
 233        L is the sequence length.
 234
 235    Notes
 236    -----
 237    The output array has shape (4, len(sequence)) with rows corresponding to
 238    nucleotides A, C, G, T in that order. Non-standard nucleotides (N, etc.)
 239    result in all-zero columns.
 240    """
 241    sequence = sequence.upper()
 242
 243    seq_chararray = np.frombuffer(sequence.encode("UTF-8"), dtype="S1")
 244    one_hot = (seq_chararray[None, :] == SEQ_ALPHABET[:, None]).astype(dtype)
 245
 246    return one_hot
 247
 248
 249def load_regions_from_bw(
 250    peaks: pl.DataFrame, fa_path: str, bw_paths: List[str], half_width: int
 251) -> Tuple[Int[ndarray, "N 4 L"], Float[ndarray, "N L"]]:
 252    """Load genomic sequences and contribution scores from FASTA and bigWig files.
 253
 254    Parameters
 255    ----------
 256    peaks : pl.DataFrame
 257        Peak regions DataFrame from load_peaks() containing columns:
 258        'chr', 'peak_region_start'.
 259    fa_path : str
 260        Path to genome FASTA file (.fa or .fasta format).
 261    bw_paths : List[str]
 262        List of paths to bigWig files containing contribution scores.
 263        Must be non-empty.
 264    half_width : int
 265        Half-width of regions to extract around peak centers.
 266        Total region width will be 2 * half_width.
 267
 268    Returns
 269    -------
 270    sequences : Int[ndarray, "N 4 L"]
 271        One-hot encoded DNA sequences where N is the number of peaks,
 272        4 represents A,C,G,T nucleotides, and L is the region length (2 * half_width).
 273    contribs : Float[ndarray, "N L"]
 274        Contribution scores averaged across input bigWig files.
 275        Shape is (N peaks, L region_length).
 276
 277    Notes
 278    -----
 279    BigWig files only provide projected contribution scores, not hypothetical scores.
 280    Regions extending beyond chromosome boundaries are zero-padded.
 281    Missing values in bigWig files are converted to zero.
 282    """
 283    num_peaks = peaks.height
 284    region_width = half_width * 2
 285
 286    sequences = np.zeros((num_peaks, 4, region_width), dtype=np.int8)
 287    contribs = np.zeros((num_peaks, region_width), dtype=np.float16)
 288
 289    # Load genome reference
 290    genome = pyfaidx.Fasta(fa_path, one_based_attributes=False)
 291
 292    bws = [pyBigWig.open(i) for i in bw_paths]
 293    contrib_buffer = np.zeros((len(bw_paths), half_width * 2), dtype=np.float16)
 294
 295    try:
 296        for ind, row in tqdm(
 297            enumerate(peaks.iter_rows(named=True)),
 298            disable=None,
 299            unit="regions",
 300            total=num_peaks,
 301        ):
 302            chrom = row["chr"]
 303            start = row["peak_region_start"]
 304            end = start + 2 * half_width
 305
 306            sequence_data: pyfaidx.FastaRecord = genome[chrom][start:end]  # type: ignore
 307            sequence: str = sequence_data.seq  # type: ignore
 308            start_adj: int = sequence_data.start  # type: ignore
 309            end_adj: int = sequence_data.end  # type: ignore
 310            a = start_adj - start
 311            b = end_adj - start
 312
 313            if b > a:
 314                sequences[ind, :, a:b] = one_hot_encode(sequence)
 315
 316                for j, bw in enumerate(bws):
 317                    contrib_buffer[j, :] = np.nan_to_num(
 318                        bw.values(chrom, start_adj, end_adj, numpy=True)
 319                    )
 320
 321                contribs[ind, a:b] = np.mean(contrib_buffer, axis=0)
 322
 323    finally:
 324        for bw in bws:
 325            bw.close()
 326
 327    return sequences, contribs
 328
 329
 330def load_regions_from_chrombpnet_h5(
 331    h5_paths: List[str], half_width: int
 332) -> Tuple[Int[ndarray, "N 4 L"], Float[ndarray, "N 4 L"]]:
 333    """Load genomic sequences and contribution scores from ChromBPNet HDF5 files.
 334
 335    Parameters
 336    ----------
 337    h5_paths : List[str]
 338        List of paths to ChromBPNet HDF5 files containing sequences and SHAP scores.
 339        Must be non-empty and contain compatible data shapes.
 340    half_width : int
 341        Half-width of regions to extract around the center.
 342        Total region width will be 2 * half_width.
 343
 344    Returns
 345    -------
 346    sequences : Int[ndarray, "N 4 L"]
 347        One-hot encoded DNA sequences where N is the number of regions,
 348        4 represents A,C,G,T nucleotides, and L is the region length (2 * half_width).
 349    contribs : Float[ndarray, "N 4 L"]
 350        SHAP contribution scores averaged across input files.
 351        Shape is (N regions, 4 nucleotides, L region_length).
 352
 353    Notes
 354    -----
 355    ChromBPNet files store sequences in 'raw/seq' and SHAP scores in 'shap/seq'.
 356    All input files must have the same dimensions and number of regions.
 357    Missing values in contribution scores are converted to zero.
 358    """
 359    with ExitStack() as stack:
 360        h5s = [stack.enter_context(h5py.File(i)) for i in h5_paths]
 361
 362        start = h5s[0]["raw/seq"].shape[-1] // 2 - half_width  # type: ignore  # HDF5 array access
 363        end = start + 2 * half_width
 364
 365        sequences = h5s[0]["raw/seq"][:, :, start:end].astype(np.int8)  # type: ignore  # HDF5 array access
 366        contribs = np.mean(
 367            [np.nan_to_num(f["shap/seq"][:, :, start:end]) for f in h5s],  # type: ignore  # HDF5 array access
 368            axis=0,
 369            dtype=np.float16,
 370        )
 371
 372    return sequences, contribs  # type: ignore  # HDF5 arrays converted to NumPy
 373
 374
 375def load_regions_from_bpnet_h5(
 376    h5_paths: List[str], half_width: int
 377) -> Tuple[Int[ndarray, "N 4 L"], Float[ndarray, "N 4 L"]]:
 378    """Load genomic sequences and contribution scores from BPNet HDF5 files.
 379
 380    Parameters
 381    ----------
 382    h5_paths : List[str]
 383        List of paths to BPNet HDF5 files containing sequences and contribution scores.
 384        Must be non-empty and contain compatible data shapes.
 385    half_width : int
 386        Half-width of regions to extract around the center.
 387        Total region width will be 2 * half_width.
 388
 389    Returns
 390    -------
 391    sequences : Int[ndarray, "N 4 L"]
 392        One-hot encoded DNA sequences where N is the number of regions,
 393        4 represents A,C,G,T nucleotides, and L is the region length (2 * half_width).
 394    contribs : Float[ndarray, "N 4 L"]
 395        Hypothetical contribution scores averaged across input files.
 396        Shape is (N regions, 4 nucleotides, L region_length).
 397
 398    Notes
 399    -----
 400    BPNet files store sequences in 'input_seqs' and hypothetical scores in 'hyp_scores'.
 401    The data requires axis swapping to convert from (n, length, 4) to (n, 4, length) format.
 402    All input files must have the same dimensions and number of regions.
 403    Missing values in contribution scores are converted to zero.
 404    """
 405    with ExitStack() as stack:
 406        h5s = [stack.enter_context(h5py.File(i)) for i in h5_paths]
 407
 408        start = h5s[0]["input_seqs"].shape[-2] // 2 - half_width  # type: ignore  # HDF5 array access
 409        end = start + 2 * half_width
 410
 411        sequences = h5s[0]["input_seqs"][:, start:end, :].swapaxes(1, 2).astype(np.int8)  # type: ignore  # HDF5 array access with axis swap
 412        contribs = np.mean(
 413            [
 414                np.nan_to_num(f["hyp_scores"][:, start:end, :].swapaxes(1, 2))  # type: ignore  # HDF5 array access
 415                for f in h5s
 416            ],
 417            axis=0,
 418            dtype=np.float16,
 419        )
 420
 421    return sequences, contribs
 422
 423
 424def load_npy_or_npz(path: str) -> ndarray:
 425    """Load array data from .npy or .npz file.
 426
 427    Parameters
 428    ----------
 429    path : str
 430        Path to .npy or .npz file. File must exist and contain valid NumPy data.
 431
 432    Returns
 433    -------
 434    ndarray
 435        Loaded array data. For .npz files, returns the first array ('arr_0').
 436        For .npy files, returns the array directly.
 437
 438    Raises
 439    ------
 440    FileNotFoundError
 441        If the specified file does not exist.
 442    KeyError
 443        If .npz file does not contain 'arr_0' key.
 444    """
 445    f = np.load(path)
 446    if isinstance(f, np.ndarray):
 447        arr = f
 448    else:
 449        arr = f["arr_0"]
 450
 451    return arr
 452
 453
 454def load_regions_from_modisco_fmt(
 455    shaps_paths: List[str], ohe_path: str, half_width: int
 456) -> Tuple[Int[ndarray, "N 4 L"], Float[ndarray, "N 4 L"]]:
 457    """Load genomic sequences and contribution scores from TF-MoDISco format files.
 458
 459    Parameters
 460    ----------
 461    shaps_paths : List[str]
 462        List of paths to .npy/.npz files containing SHAP/attribution scores.
 463        Must be non-empty and all files must have compatible shapes.
 464    ohe_path : str
 465        Path to .npy/.npz file containing one-hot encoded sequences.
 466        Must have shape (n_regions, 4, sequence_length).
 467    half_width : int
 468        Half-width of regions to extract around the center.
 469        Total region width will be 2 * half_width.
 470
 471    Returns
 472    -------
 473    sequences : Int[ndarray, "N 4 L"]
 474        One-hot encoded DNA sequences where N is the number of regions,
 475        4 represents A,C,G,T nucleotides, and L is the region length (2 * half_width).
 476    contribs : Float[ndarray, "N 4 L"]
 477        SHAP contribution scores averaged across input files.
 478        Shape is (N regions, 4 nucleotides, L region_length).
 479
 480    Notes
 481    -----
 482    All SHAP files must have the same shape as the sequence file.
 483    Missing values in contribution scores are converted to zero.
 484    The center of the input sequences is used as the reference point for extraction.
 485    """
 486    sequences_raw = load_npy_or_npz(ohe_path)
 487
 488    start = sequences_raw.shape[-1] // 2 - half_width
 489    end = start + 2 * half_width
 490
 491    sequences = sequences_raw[:, :, start:end].astype(np.int8)
 492
 493    shaps = [np.nan_to_num(load_npy_or_npz(p)[:, :, start:end]) for p in shaps_paths]
 494    contribs = np.mean(shaps, axis=0, dtype=np.float16)
 495
 496    return sequences, contribs
 497
 498
 499def load_regions_npz(
 500    npz_path: str,
 501) -> Tuple[
 502    Int[ndarray, "N 4 L"],
 503    Union[Float[ndarray, "N 4 L"], Float[ndarray, "N L"]],
 504    pl.DataFrame,
 505    bool,
 506]:
 507    """Load preprocessed genomic regions from NPZ file.
 508
 509    Parameters
 510    ----------
 511    npz_path : str
 512        Path to NPZ file containing sequences, contributions, and optional coordinates.
 513        Must contain 'sequences' and 'contributions' arrays at minimum.
 514
 515    Returns
 516    -------
 517    sequences : Int[ndarray, "N 4 L"]
 518        One-hot encoded DNA sequences where N is the number of regions,
 519        4 represents A,C,G,T nucleotides, and L is the region length.
 520    contributions : Union[Float[ndarray, "N 4 L"], Float[ndarray, "N L"]]
 521        Contribution scores in either hypothetical format (N, 4, L) or
 522        projected format (N, L). Shape depends on the input data format.
 523    peaks_df : pl.DataFrame
 524        DataFrame containing peak region information with columns:
 525        'chr', 'chr_id', 'peak_region_start', 'peak_id', 'peak_name'.
 526    has_peaks : bool
 527        Whether the file contains genomic coordinate information.
 528        If False, placeholder coordinate data is used.
 529
 530    Notes
 531    -----
 532    If genomic coordinates are not present in the NPZ file, creates placeholder
 533    coordinate data and issues a warning. The placeholder data uses 'NA' for
 534    chromosome names and sequential indices for peak IDs.
 535
 536    Raises
 537    ------
 538    KeyError
 539        If required arrays 'sequences' or 'contributions' are missing from the file.
 540    """
 541    data = np.load(npz_path)
 542
 543    if "chr" not in data.keys():
 544        warnings.warn(
 545            "No genome coordinates present in the input .npz file. Returning sequences and contributions only."
 546        )
 547        has_peaks = False
 548        num_regions = data["sequences"].shape[0]
 549        peak_data = {
 550            "chr": np.array(["NA"] * num_regions, dtype="U"),
 551            "chr_id": np.arange(num_regions, dtype=np.uint32),
 552            "peak_region_start": np.zeros(num_regions, dtype=np.int32),
 553            "peak_id": np.arange(num_regions, dtype=np.uint32),
 554            "peak_name": np.array(["NA"] * num_regions, dtype="U"),
 555        }
 556
 557    else:
 558        has_peaks = True
 559        peak_data = {
 560            "chr": data["chr"],
 561            "chr_id": data["chr_id"],
 562            "peak_region_start": data["start"],
 563            "peak_id": data["peak_id"],
 564            "peak_name": data["peak_name"],
 565        }
 566
 567    peaks_df = pl.DataFrame(peak_data)
 568
 569    return data["sequences"], data["contributions"], peaks_df, has_peaks
 570
 571
 572def write_regions_npz(
 573    sequences: Int[ndarray, "N 4 L"],
 574    contributions: Union[Float[ndarray, "N 4 L"], Float[ndarray, "N L"]],
 575    out_path: str,
 576    peaks_df: Optional[pl.DataFrame] = None,
 577) -> None:
 578    """Write genomic regions and contribution scores to compressed NPZ file.
 579
 580    Parameters
 581    ----------
 582    sequences : Int[ndarray, "N 4 L"]
 583        One-hot encoded DNA sequences where N is the number of regions,
 584        4 represents A,C,G,T nucleotides, and L is the region length.
 585    contributions : Union[Float[ndarray, "N 4 L"], Float[ndarray, "N L"]]
 586        Contribution scores in either hypothetical format (N, 4, L) or
 587        projected format (N, L).
 588    out_path : str
 589        Output path for the NPZ file. Parent directory must exist.
 590    peaks_df : Optional[pl.DataFrame]
 591        DataFrame containing peak region information with columns:
 592        'chr', 'chr_id', 'peak_region_start', 'peak_id', 'peak_name'.
 593        If None, only sequences and contributions are saved.
 594
 595    Raises
 596    ------
 597    ValueError
 598        If the number of regions in sequences/contributions doesn't match peaks_df.
 599    FileNotFoundError
 600        If the parent directory of out_path does not exist.
 601
 602    Notes
 603    -----
 604    The output file is compressed using NumPy's savez_compressed format.
 605    If peaks_df is provided, genomic coordinate information is included
 606    in the output file for downstream analysis.
 607    """
 608    if peaks_df is None:
 609        warnings.warn(
 610            "No genome coordinates provided. Writing sequences and contributions only."
 611        )
 612        np.savez_compressed(out_path, sequences=sequences, contributions=contributions)
 613
 614    else:
 615        num_regions = peaks_df.height
 616        if (num_regions != sequences.shape[0]) or (
 617            num_regions != contributions.shape[0]
 618        ):
 619            raise ValueError(
 620                f"Input sequences of shape {sequences.shape} and/or "
 621                f"input contributions of shape {contributions.shape} "
 622                f"are not compatible with peak region count of {num_regions}"
 623            )
 624
 625        chr_arr = peaks_df.get_column("chr").to_numpy().astype("U")
 626        chr_id_arr = peaks_df.get_column("chr_id").to_numpy()
 627        start_arr = peaks_df.get_column("peak_region_start").to_numpy()
 628        peak_id_arr = peaks_df.get_column("peak_id").to_numpy()
 629        peak_name_arr = peaks_df.get_column("peak_name").to_numpy().astype("U")
 630        np.savez_compressed(
 631            out_path,
 632            sequences=sequences,
 633            contributions=contributions,
 634            chr=chr_arr,
 635            chr_id=chr_id_arr,
 636            start=start_arr,
 637            peak_id=peak_id_arr,
 638            peak_name=peak_name_arr,
 639        )
 640
 641
 642def trim_motif(cwm: Float[ndarray, "4 W"], trim_threshold: float) -> Tuple[int, int]:
 643    """Determine trimmed start and end positions for a motif based on contribution magnitude.
 644
 645    This function identifies the core region of a motif by finding positions where
 646    the total absolute contribution exceeds a threshold relative to the maximum.
 647
 648    Parameters
 649    ----------
 650    cwm : Float[ndarray, "4 W"]
 651        Contribution weight matrix for the motif where 4 represents A,C,G,T
 652        nucleotides and W is the motif width.
 653    trim_threshold : float
 654        Fraction of maximum score to use as trimming threshold (0.0 to 1.0).
 655        Higher values result in more aggressive trimming.
 656
 657    Returns
 658    -------
 659    start : int
 660        Start position of the trimmed motif (inclusive).
 661    end : int
 662        End position of the trimmed motif (exclusive).
 663
 664    Notes
 665    -----
 666    The trimming is based on the sum of absolute contributions across all nucleotides
 667    at each position. Positions with contributions below trim_threshold * max_score
 668    are removed from the motif edges.
 669
 670    Adapted from https://github.com/jmschrei/tfmodisco-lite/blob/570535ee5ccf43d670e898d92d63af43d68c38c5/modiscolite/report.py#L213-L236
 671    """
 672    score = np.sum(np.abs(cwm), axis=0)
 673    trim_thresh = np.max(score) * trim_threshold
 674    pass_inds = np.nonzero(score >= trim_thresh)
 675    start = max(int(np.min(pass_inds)), 0)  # type: ignore  # nonzero returns tuple of arrays
 676    end = min(int(np.max(pass_inds)) + 1, len(score))  # type: ignore  # nonzero returns tuple of arrays
 677
 678    return start, end
 679
 680
 681def softmax(x: Float[ndarray, "4 W"], temp: float = 100) -> Float[ndarray, "4 W"]:
 682    """Apply softmax transformation with temperature scaling.
 683
 684    Parameters
 685    ----------
 686    x : Float[ndarray, "4 W"]
 687        Input array to transform where 4 represents A,C,G,T nucleotides
 688        and W is the motif width.
 689    temp : float, default 100
 690        Temperature parameter for softmax scaling. Higher values create
 691        sharper probability distributions.
 692
 693    Returns
 694    -------
 695    Float[ndarray, "4 W"]
 696        Softmax-transformed array with same shape as input. Each column
 697        sums to 1.0, representing nucleotide probabilities at each position.
 698
 699    Notes
 700    -----
 701    The softmax is applied along the nucleotide axis (axis=0), normalizing
 702    each position to have probabilities that sum to 1. The temperature
 703    parameter controls the sharpness of the distribution.
 704    """
 705    norm_x = x - np.mean(x, axis=1, keepdims=True)
 706    exp = np.exp(temp * norm_x)
 707    return exp / np.sum(exp, axis=0, keepdims=True)
 708
 709
 710def _motif_name_sort_key(data: Tuple[str, Any]) -> Union[Tuple[int, int], Tuple[int, str]]:
 711    """Generate sort key for TF-MoDISco motif names.
 712
 713    This function creates a sort key that orders motifs by pattern number,
 714    with non-standard patterns sorted to the end.
 715
 716    Parameters
 717    ----------
 718    data : Tuple[str, Any]
 719        Tuple containing motif name as first element and additional data.
 720        The motif name should follow the format 'pattern_N' or 'pattern#N' where N is an integer.
 721
 722    Returns
 723    -------
 724    Union[Tuple[int, int], Tuple[int, str]]
 725        Sort key tuple for ordering motifs. Standard pattern names return
 726        (0, pattern_number) while non-standard names return (1, name).
 727
 728    Notes
 729    -----
 730    This function is used internally by load_modisco_motifs to ensure
 731    consistent motif ordering across runs.
 732    """
 733    pattern_name = data[0]
 734    try:
 735        return (0, int(pattern_name.split("_")[-1]))
 736    except (ValueError, IndexError):
 737        try:
 738            return (0, int(pattern_name.split("#")[-1]))
 739        except (ValueError, IndexError):
 740            return (1, pattern_name)
 741
 742
 743MODISCO_PATTERN_GROUPS = ["pos_patterns", "neg_patterns"]
 744
 745
 746def load_modisco_motifs(
 747    modisco_h5_path: str,
 748    trim_coords: Optional[Dict[str, Tuple[int, int]]],
 749    trim_thresholds: Optional[Dict[str, float]],
 750    trim_threshold_default: float,
 751    motif_type: str,
 752    motifs_include: Optional[List[str]],
 753    motif_name_map: Optional[Dict[str, str]],
 754    motif_lambdas: Optional[Dict[str, float]],
 755    motif_lambda_default: float,
 756    include_rc: bool,
 757) -> Tuple[pl.DataFrame, Float[ndarray, "M 4 W"], Int[ndarray, "M W"], ndarray]:
 758    """Load motif data from TF-MoDISco HDF5 file with customizable processing options.
 759
 760    This function extracts contribution weight matrices and associated metadata from
 761    TF-MoDISco results, with support for custom naming, trimming, and regularization
 762    parameters.
 763
 764    Parameters
 765    ----------
 766    modisco_h5_path : str
 767        Path to TF-MoDISco HDF5 results file containing pattern groups.
 768    trim_coords : Optional[Dict[str, Tuple[int, int]]]
 769        Manual trim coordinates for specific motifs {motif_name: (start, end)}.
 770        Takes precedence over automatic trimming based on thresholds.
 771    trim_thresholds : Optional[Dict[str, float]]
 772        Custom trim thresholds for specific motifs {motif_name: threshold}.
 773        Values should be between 0.0 and 1.0.
 774    trim_threshold_default : float
 775        Default trim threshold for motifs not in trim_thresholds.
 776        Fraction of maximum contribution used for trimming.
 777    motif_type : str
 778        Type of motif to extract. Must be one of:
 779        - 'cwm': Contribution weight matrix (normalized)
 780        - 'hcwm': Hypothetical contribution weight matrix
 781        - 'pfm': Position frequency matrix
 782        - 'pfm_softmax': Softmax-transformed position frequency matrix
 783    motifs_include : Optional[List[str]]
 784        List of motif names to include. If None, includes all motifs found.
 785        Names should follow format 'pos_patterns.pattern_N' or 'neg_patterns.pattern_N'.
 786    motif_name_map : Optional[Dict[str, str]]
 787        Mapping from original to custom motif names {orig_name: new_name}.
 788        New names must be unique across all motifs.
 789    motif_lambdas : Optional[Dict[str, float]]
 790        Custom lambda regularization values for specific motifs {motif_name: lambda}.
 791        Higher values increase sparsity penalty for the corresponding motif.
 792    motif_lambda_default : float
 793        Default lambda value for motifs not specified in motif_lambdas.
 794    include_rc : bool
 795        Whether to include reverse complement motifs in addition to forward motifs.
 796        If True, doubles the number of motifs returned.
 797
 798    Returns
 799    -------
 800    motifs_df : pl.DataFrame
 801        DataFrame containing motif metadata with columns: motif_id, motif_name,
 802        motif_name_orig, strand, motif_start, motif_end, motif_scale, lambda.
 803    cwms : Float[ndarray, "M 4 W"]
 804        Contribution weight matrices for all motifs where M is the number of motifs,
 805        4 represents A,C,G,T nucleotides, and W is the motif width.
 806    trim_masks : Int[ndarray, "M W"]
 807        Binary masks indicating core motif regions (1) vs trimmed regions (0).
 808        Shape is (M motifs, W motif_width).
 809    names : ndarray
 810        Array of unique motif names (forward strand only).
 811
 812    Raises
 813    ------
 814    ValueError
 815        If motif_type is not one of the supported types, or if motif names
 816        in motif_name_map are not unique.
 817    FileNotFoundError
 818        If the specified HDF5 file does not exist.
 819    KeyError
 820        If required datasets are missing from the HDF5 file.
 821
 822    Notes
 823    -----
 824    Motif trimming removes low-contribution positions from the edges based on
 825    the position-wise sum of absolute contributions across nucleotides. The trimming
 826    helps focus on the core binding site.
 827
 828    Adapted from https://github.com/jmschrei/tfmodisco-lite/blob/570535ee5ccf43d670e898d92d63af43d68c38c5/modiscolite/report.py#L252-L272
 829    """
 830    motif_data_lsts = {
 831        "motif_name": [],
 832        "motif_name_orig": [],
 833        "strand": [],
 834        "motif_start": [],
 835        "motif_end": [],
 836        "motif_scale": [],
 837        "lambda": [],
 838    }
 839    motif_lst = []
 840    trim_mask_lst = []
 841
 842    if motifs_include is not None:
 843        motifs_include_set = set(motifs_include)
 844    else:
 845        motifs_include_set = None
 846
 847    if motif_name_map is None:
 848        motif_name_map = {}
 849
 850    if motif_lambdas is None:
 851        motif_lambdas = {}
 852
 853    if trim_coords is None:
 854        trim_coords = {}
 855    if trim_thresholds is None:
 856        trim_thresholds = {}
 857
 858    if len(motif_name_map.values()) != len(set(motif_name_map.values())):
 859        raise ValueError("Specified motif names are not unique")
 860
 861    with h5py.File(modisco_h5_path, "r") as modisco_results:
 862        for name in MODISCO_PATTERN_GROUPS:
 863            if name not in modisco_results.keys():
 864                continue
 865
 866            metacluster = modisco_results[name]
 867            for _, (pattern_name, pattern) in enumerate(
 868                sorted(metacluster.items(), key=_motif_name_sort_key)  # type: ignore  # HDF5 access
 869            ):
 870                pattern_tag = f"{name}.{pattern_name}"
 871
 872                if (
 873                    motifs_include_set is not None
 874                    and pattern_tag not in motifs_include_set
 875                ):
 876                    continue
 877
 878                motif_lambda = motif_lambdas.get(pattern_tag, motif_lambda_default)
 879                pattern_tag_orig = pattern_tag
 880                pattern_tag = motif_name_map.get(pattern_tag, pattern_tag)
 881
 882                cwm_raw = pattern["contrib_scores"][:].T  # type: ignore
 883                cwm_norm = np.sqrt((cwm_raw**2).sum())
 884
 885                cwm_fwd = cwm_raw / cwm_norm
 886                cwm_rev = cwm_fwd[::-1, ::-1]
 887
 888                if pattern_tag in trim_coords:
 889                    start_fwd, end_fwd = trim_coords[pattern_tag]
 890                else:
 891                    trim_threshold = trim_thresholds.get(
 892                        pattern_tag, trim_threshold_default
 893                    )
 894                    start_fwd, end_fwd = trim_motif(cwm_fwd, trim_threshold)
 895
 896                cwm_len = cwm_fwd.shape[1]
 897                start_rev, end_rev = cwm_len - end_fwd, cwm_len - start_fwd
 898
 899                trim_mask_fwd = np.zeros(cwm_fwd.shape[1], dtype=np.int8)
 900                trim_mask_fwd[start_fwd:end_fwd] = 1
 901                trim_mask_rev = np.zeros(cwm_rev.shape[1], dtype=np.int8)
 902                trim_mask_rev[start_rev:end_rev] = 1
 903
 904                if motif_type == "cwm":
 905                    motif_fwd = cwm_fwd
 906                    motif_rev = cwm_rev
 907                    motif_norm = cwm_norm
 908
 909                elif motif_type == "hcwm":
 910                    motif_raw = pattern["hypothetical_contribs"][:].T  # type: ignore
 911                    motif_norm = np.sqrt((motif_raw**2).sum())
 912
 913                    motif_fwd = motif_raw / motif_norm
 914                    motif_rev = motif_fwd[::-1, ::-1]
 915
 916                elif motif_type == "pfm":
 917                    motif_raw = pattern["sequence"][:].T  # type: ignore
 918                    motif_norm = 1
 919
 920                    motif_fwd = motif_raw / np.sum(motif_raw, axis=0, keepdims=True)
 921                    motif_rev = motif_fwd[::-1, ::-1]
 922
 923                elif motif_type == "pfm_softmax":
 924                    motif_raw = pattern["sequence"][:].T  # type: ignore
 925                    motif_norm = 1
 926
 927                    motif_fwd = softmax(motif_raw)
 928                    motif_rev = motif_fwd[::-1, ::-1]
 929
 930                else:
 931                    raise ValueError(
 932                        f"Invalid motif_type: {motif_type}. Must be one of 'cwm', 'hcwm', 'pfm', 'pfm_softmax'."
 933                    )
 934
 935                motif_data_lsts["motif_name"].append(pattern_tag)
 936                motif_data_lsts["motif_name_orig"].append(pattern_tag_orig)
 937                motif_data_lsts["strand"].append("+")
 938                motif_data_lsts["motif_start"].append(start_fwd)
 939                motif_data_lsts["motif_end"].append(end_fwd)
 940                motif_data_lsts["motif_scale"].append(motif_norm)
 941                motif_data_lsts["lambda"].append(motif_lambda)
 942
 943                if include_rc:
 944                    motif_data_lsts["motif_name"].append(pattern_tag)
 945                    motif_data_lsts["motif_name_orig"].append(pattern_tag_orig)
 946                    motif_data_lsts["strand"].append("-")
 947                    motif_data_lsts["motif_start"].append(start_rev)
 948                    motif_data_lsts["motif_end"].append(end_rev)
 949                    motif_data_lsts["motif_scale"].append(motif_norm)
 950                    motif_data_lsts["lambda"].append(motif_lambda)
 951
 952                    motif_lst.extend([motif_fwd, motif_rev])
 953                    trim_mask_lst.extend([trim_mask_fwd, trim_mask_rev])
 954
 955                else:
 956                    motif_lst.append(motif_fwd)
 957                    trim_mask_lst.append(trim_mask_fwd)
 958
 959    motifs_df = pl.DataFrame(motif_data_lsts).with_row_index(name="motif_id")
 960    cwms = np.stack(motif_lst, dtype=np.float16, axis=0)
 961    trim_masks = np.stack(trim_mask_lst, dtype=np.int8, axis=0)
 962    names = (
 963        motifs_df.filter(pl.col("strand") == "+").get_column("motif_name").to_numpy()
 964    )
 965
 966    return motifs_df, cwms, trim_masks, names
 967
 968
 969def load_modisco_seqlets(
 970    modisco_h5_path: str,
 971    peaks_df: pl.DataFrame,
 972    motifs_df: pl.DataFrame,
 973    half_width: int,
 974    modisco_half_width: int,
 975    lazy: bool = False,
 976) -> Union[pl.DataFrame, pl.LazyFrame]:
 977    """Load seqlet data from TF-MoDISco HDF5 file and convert to genomic coordinates.
 978
 979    This function extracts seqlet instances from TF-MoDISco results and converts
 980    their relative positions to absolute genomic coordinates using peak region
 981    information.
 982
 983    Parameters
 984    ----------
 985    modisco_h5_path : str
 986        Path to TF-MoDISco HDF5 results file containing seqlet data.
 987    peaks_df : pl.DataFrame
 988        DataFrame containing peak region information with columns:
 989        'peak_id', 'chr', 'chr_id', 'peak_region_start'.
 990    motifs_df : pl.DataFrame
 991        DataFrame containing motif metadata with columns:
 992        'motif_name_orig', 'strand', 'motif_name', 'motif_start', 'motif_end'.
 993    half_width : int
 994        Half-width of the current analysis regions.
 995    modisco_half_width : int
 996        Half-width of the regions used in the original TF-MoDISco analysis.
 997        Used to calculate coordinate offsets.
 998    lazy : bool, default False
 999        If True, returns a LazyFrame for efficient chaining of operations.
1000        If False, collects the result into a DataFrame.
1001
1002    Returns
1003    -------
1004    Union[pl.DataFrame, pl.LazyFrame]
1005        Seqlets with genomic coordinates containing columns:
1006        - chr: Chromosome name
1007        - chr_id: Numeric chromosome identifier
1008        - start: Start coordinate of trimmed motif instance
1009        - end: End coordinate of trimmed motif instance
1010        - start_untrimmed: Start coordinate of full motif instance
1011        - end_untrimmed: End coordinate of full motif instance
1012        - is_revcomp: Whether the motif is reverse complemented
1013        - strand: Motif strand ('+' or '-')
1014        - motif_name: Motif name (may be remapped)
1015        - peak_id: Peak identifier
1016        - peak_region_start: Peak region start coordinate
1017
1018    Notes
1019    -----
1020    Seqlets are deduplicated based on chromosome ID, start position (untrimmed),
1021    motif name, and reverse complement status to avoid redundant instances.
1022
1023    The coordinate transformation accounts for differences in region sizes
1024    between the original TF-MoDISco analysis and the current analysis.
1025    """
1026
1027    start_lst = []
1028    end_lst = []
1029    is_revcomp_lst = []
1030    strand_lst = []
1031    peak_id_lst = []
1032    pattern_tags = []
1033
1034    with h5py.File(modisco_h5_path, "r") as modisco_results:
1035        for name in MODISCO_PATTERN_GROUPS:
1036            if name not in modisco_results.keys():
1037                continue
1038
1039            metacluster = modisco_results[name]
1040
1041            key = _motif_name_sort_key
1042            for _, (pattern_name, pattern) in enumerate(
1043                sorted(metacluster.items(), key=key)  # type: ignore  # HDF5 access
1044            ):
1045                pattern_tag = f"{name}.{pattern_name}"
1046
1047                starts = pattern["seqlets/start"][:].astype(np.int32)  # type: ignore
1048                ends = pattern["seqlets/end"][:].astype(np.int32)  # type: ignore
1049                is_revcomps = pattern["seqlets/is_revcomp"][:].astype(bool)  # type: ignore
1050                strands = ["+" if not i else "-" for i in is_revcomps]
1051                peak_ids = pattern["seqlets/example_idx"][:].astype(np.uint32)  # type: ignore
1052
1053                n_seqlets = int(pattern["seqlets/n_seqlets"][0])  # type: ignore
1054
1055                start_lst.append(starts)
1056                end_lst.append(ends)
1057                is_revcomp_lst.append(is_revcomps)
1058                strand_lst.extend(strands)
1059                peak_id_lst.append(peak_ids)
1060                pattern_tags.extend([pattern_tag for _ in range(n_seqlets)])
1061
1062    df_data = {
1063        "seqlet_start": np.concatenate(start_lst),
1064        "seqlet_end": np.concatenate(end_lst),
1065        "is_revcomp": np.concatenate(is_revcomp_lst),
1066        "strand": strand_lst,
1067        "peak_id": np.concatenate(peak_id_lst),
1068        "motif_name_orig": pattern_tags,
1069    }
1070
1071    offset = half_width - modisco_half_width
1072
1073    seqlets_df = (
1074        pl.LazyFrame(df_data)
1075        .join(motifs_df.lazy(), on=("motif_name_orig", "strand"), how="inner")
1076        .join(peaks_df.lazy(), on="peak_id", how="inner")
1077        .select(
1078            chr=pl.col("chr"),
1079            chr_id=pl.col("chr_id"),
1080            start=pl.col("peak_region_start")
1081            + pl.col("seqlet_start")
1082            + pl.col("motif_start")
1083            + offset,
1084            end=pl.col("peak_region_start")
1085            + pl.col("seqlet_start")
1086            + pl.col("motif_end")
1087            + offset,
1088            start_untrimmed=pl.col("peak_region_start")
1089            + pl.col("seqlet_start")
1090            + offset,
1091            end_untrimmed=pl.col("peak_region_start") + pl.col("seqlet_end") + offset,
1092            is_revcomp=pl.col("is_revcomp"),
1093            strand=pl.col("strand"),
1094            motif_name=pl.col("motif_name"),
1095            peak_id=pl.col("peak_id"),
1096            peak_region_start=pl.col("peak_region_start"),
1097        )
1098        .unique(subset=["chr_id", "start_untrimmed", "motif_name", "is_revcomp"])
1099    )
1100
1101    seqlets_df = seqlets_df if lazy else seqlets_df.collect()
1102
1103    return seqlets_df
1104
1105
1106def write_modisco_seqlets(
1107    seqlets_df: Union[pl.DataFrame, pl.LazyFrame], out_path: str
1108) -> None:
1109    """Write TF-MoDISco seqlets to TSV file.
1110
1111    Parameters
1112    ----------
1113    seqlets_df : Union[pl.DataFrame, pl.LazyFrame]
1114        Seqlets DataFrame with genomic coordinates. Must contain columns
1115        that are safe to drop: 'chr_id', 'is_revcomp'.
1116    out_path : str
1117        Output TSV file path.
1118
1119    Notes
1120    -----
1121    Removes internal columns 'chr_id' and 'is_revcomp' before writing
1122    to create a clean output format suitable for downstream analysis.
1123    """
1124    seqlets_df = seqlets_df.drop(["chr_id", "is_revcomp"])
1125    if isinstance(seqlets_df, pl.LazyFrame):
1126        seqlets_df = seqlets_df.collect()
1127    seqlets_df.write_csv(out_path, separator="\t")
1128
1129
1130HITS_DTYPES = {
1131    "chr": pl.String,
1132    "start": pl.Int32,
1133    "end": pl.Int32,
1134    "start_untrimmed": pl.Int32,
1135    "end_untrimmed": pl.Int32,
1136    "motif_name": pl.String,
1137    "hit_coefficient": pl.Float32,
1138    "hit_coefficient_global": pl.Float32,
1139    "hit_similarity": pl.Float32,
1140    "hit_correlation": pl.Float32,
1141    "hit_importance": pl.Float32,
1142    "hit_importance_sq": pl.Float32,
1143    "strand": pl.String,
1144    "peak_name": pl.String,
1145    "peak_id": pl.UInt32,
1146}
1147HITS_COLLAPSED_DTYPES = HITS_DTYPES | {"is_primary": pl.UInt32}
1148
1149
1150def load_hits(
1151    hits_path: str, lazy: bool = False, schema: Dict[str, Any] = HITS_DTYPES
1152) -> Union[pl.DataFrame, pl.LazyFrame]:
1153    """Load motif hit data from TSV file.
1154
1155    Parameters
1156    ----------
1157    hits_path : str
1158        Path to TSV file containing motif hit results.
1159    lazy : bool, default False
1160        If True, returns a LazyFrame for efficient chaining operations.
1161        If False, collects the result into a DataFrame.
1162    schema : Dict[str, Any], default HITS_DTYPES
1163        Schema defining column names and data types for the hit data.
1164
1165    Returns
1166    -------
1167    Union[pl.DataFrame, pl.LazyFrame]
1168        Hit data with an additional 'count' column set to 1 for aggregation.
1169    """
1170    hits_df = pl.scan_csv(
1171        hits_path, separator="\t", quote_char=None, schema=schema
1172    ).with_columns(pl.lit(1).alias("count"))
1173
1174    return hits_df if lazy else hits_df.collect()
1175
1176
1177def write_hits_processed(
1178    hits_df: Union[pl.DataFrame, pl.LazyFrame],
1179    out_path: str,
1180    schema: Optional[Dict[str, Any]] = HITS_DTYPES,
1181) -> None:
1182    """Write processed hit data to TSV file with optional column filtering.
1183
1184    Parameters
1185    ----------
1186    hits_df : Union[pl.DataFrame, pl.LazyFrame]
1187        Hit data to write to file.
1188    out_path : str
1189        Output path for the TSV file.
1190    schema : Optional[Dict[str, Any]], default HITS_DTYPES
1191        Schema defining which columns to include in output.
1192        If None, all columns are written.
1193    """
1194    if schema is not None:
1195        hits_df = hits_df.select(schema.keys())
1196
1197    if isinstance(hits_df, pl.LazyFrame):
1198        hits_df = hits_df.collect()
1199
1200    hits_df.write_csv(out_path, separator="\t")
1201
1202
1203def write_hits(
1204    hits_df: Union[pl.DataFrame, pl.LazyFrame],
1205    peaks_df: pl.DataFrame,
1206    motifs_df: pl.DataFrame,
1207    qc_df: pl.DataFrame,
1208    out_dir: str,
1209    motif_width: int,
1210) -> None:
1211    """Write comprehensive hit results to multiple output files.
1212
1213    This function combines hit data with peak, motif, and quality control information
1214    to generate complete output files including genomic coordinates and scores.
1215
1216    Parameters
1217    ----------
1218    hits_df : Union[pl.DataFrame, pl.LazyFrame]
1219        Hit data containing motif instance information.
1220    peaks_df : pl.DataFrame
1221        Peak region information for coordinate conversion.
1222    motifs_df : pl.DataFrame
1223        Motif metadata for annotation and trimming information.
1224    qc_df : pl.DataFrame
1225        Quality control data for normalization factors.
1226    out_dir : str
1227        Output directory for results files. Will be created if it doesn't exist.
1228    motif_width : int
1229        Width of motif instances for coordinate calculations.
1230
1231    Notes
1232    -----
1233    Creates three output files:
1234    - hits.tsv: Complete hit data with all instances
1235    - hits_unique.tsv: Deduplicated hits by genomic position and motif (excludes rows with NA chromosome coordinates)
1236    - hits.bed: BED format file for genome browser visualization
1237    
1238    Rows where the chromosome field is NA are filtered out during deduplication
1239    to ensure that data_unique only contains well-defined genomic coordinates.
1240    """
1241    os.makedirs(out_dir, exist_ok=True)
1242    out_path_tsv = os.path.join(out_dir, "hits.tsv")
1243    out_path_tsv_unique = os.path.join(out_dir, "hits_unique.tsv")
1244    out_path_bed = os.path.join(out_dir, "hits.bed")
1245
1246    data_all = (
1247        hits_df.lazy()
1248        .join(peaks_df.lazy(), on="peak_id", how="inner")
1249        .join(qc_df.lazy(), on="peak_id", how="inner")
1250        .join(motifs_df.lazy(), on="motif_id", how="inner")
1251        .select(
1252            chr_id=pl.col("chr_id"),
1253            chr=pl.col("chr"),
1254            start=pl.col("peak_region_start")
1255            + pl.col("hit_start")
1256            + pl.col("motif_start"),
1257            end=pl.col("peak_region_start") + pl.col("hit_start") + pl.col("motif_end"),
1258            start_untrimmed=pl.col("peak_region_start") + pl.col("hit_start"),
1259            end_untrimmed=pl.col("peak_region_start")
1260            + pl.col("hit_start")
1261            + motif_width,
1262            motif_name=pl.col("motif_name"),
1263            hit_coefficient=pl.col("hit_coefficient"),
1264            hit_coefficient_global=pl.col("hit_coefficient")
1265            * (pl.col("global_scale") ** 2),
1266            hit_similarity=pl.col("hit_similarity"),
1267            hit_correlation=pl.col("hit_similarity"),
1268            hit_importance=pl.col("hit_importance") * pl.col("global_scale"),
1269            hit_importance_sq=pl.col("hit_importance_sq")
1270            * (pl.col("global_scale") ** 2),
1271            strand=pl.col("strand"),
1272            peak_name=pl.col("peak_name"),
1273            peak_id=pl.col("peak_id"),
1274            motif_lambda=pl.col("lambda"),
1275        )
1276        .sort(["chr_id", "start"])
1277        .select(HITS_DTYPES.keys())
1278    )
1279
1280    data_unique = data_all.filter(pl.col("chr").is_not_null()).unique(
1281        subset=["chr", "start", "motif_name", "strand"], maintain_order=True
1282    )
1283
1284    data_bed = data_unique.select(
1285        chr=pl.col("chr"),
1286        start=pl.col("start"),
1287        end=pl.col("end"),
1288        motif_name=pl.col("motif_name"),
1289        score=pl.lit(0),
1290        strand=pl.col("strand"),
1291    )
1292
1293    data_all.collect().write_csv(out_path_tsv, separator="\t")
1294    data_unique.collect().write_csv(out_path_tsv_unique, separator="\t")
1295    data_bed.collect().write_csv(out_path_bed, include_header=False, separator="\t")
1296
1297
1298def write_qc(qc_df: pl.DataFrame, peaks_df: pl.DataFrame, out_path: str) -> None:
1299    """Write quality control data with peak information to TSV file.
1300
1301    Parameters
1302    ----------
1303    qc_df : pl.DataFrame
1304        Quality control metrics for each peak region.
1305    peaks_df : pl.DataFrame
1306        Peak region information for coordinate annotation.
1307    out_path : str
1308        Output path for the TSV file.
1309    """
1310    df = (
1311        qc_df.lazy()
1312        .join(peaks_df.lazy(), on="peak_id", how="inner")
1313        .sort(["chr_id", "peak_region_start"])
1314        .drop("chr_id")
1315        .collect()
1316    )
1317    df.write_csv(out_path, separator="\t")
1318
1319
1320def write_motifs_df(motifs_df: pl.DataFrame, out_path: str) -> None:
1321    """Write motif metadata to TSV file.
1322
1323    Parameters
1324    ----------
1325    motifs_df : pl.DataFrame
1326        Motif metadata DataFrame.
1327    out_path : str
1328        Output path for the TSV file.
1329    """
1330    motifs_df.write_csv(out_path, separator="\t")
1331
1332
1333MOTIF_DTYPES = {
1334    "motif_id": pl.UInt32,
1335    "motif_name": pl.String,
1336    "motif_name_orig": pl.String,
1337    "strand": pl.String,
1338    "motif_start": pl.UInt32,
1339    "motif_end": pl.UInt32,
1340    "motif_scale": pl.Float32,
1341    "lambda": pl.Float32,
1342}
1343
1344
1345def load_motifs_df(motifs_path: str) -> Tuple[pl.DataFrame, ndarray]:
1346    """Load motif metadata from TSV file.
1347
1348    Parameters
1349    ----------
1350    motifs_path : str
1351        Path to motif metadata TSV file.
1352
1353    Returns
1354    -------
1355    motifs_df : pl.DataFrame
1356        Motif metadata with predefined schema.
1357    motif_names : ndarray
1358        Array of unique forward-strand motif names.
1359    """
1360    motifs_df = pl.read_csv(motifs_path, separator="\t", schema=MOTIF_DTYPES)
1361    motif_names = (
1362        motifs_df.filter(pl.col("strand") == "+").get_column("motif_name").to_numpy()
1363    )
1364
1365    return motifs_df, motif_names
1366
1367
1368def write_motif_cwms(cwms: Float[ndarray, "M 4 W"], out_path: str) -> None:
1369    """Write motif contribution weight matrices to .npy file.
1370
1371    Parameters
1372    ----------
1373    cwms : Float[ndarray, "M 4 W"]
1374        Contribution weight matrices for M motifs, 4 nucleotides, W width.
1375    out_path : str
1376        Output path for the .npy file.
1377    """
1378    np.save(out_path, cwms)
1379
1380
1381def load_motif_cwms(cwms_path: str) -> Float[ndarray, "M 4 W"]:
1382    """Load motif contribution weight matrices from .npy file.
1383
1384    Parameters
1385    ----------
1386    cwms_path : str
1387        Path to .npy file containing CWMs.
1388
1389    Returns
1390    -------
1391    Float[ndarray, "M 4 W"]
1392        Loaded contribution weight matrices.
1393    """
1394    return np.load(cwms_path)
1395
1396
1397def write_params(params: Dict[str, Any], out_path: str) -> None:
1398    """Write parameter dictionary to JSON file.
1399
1400    Parameters
1401    ----------
1402    params : Dict[str, Any]
1403        Parameter dictionary to serialize.
1404    out_path : str
1405        Output path for the JSON file.
1406    """
1407    with open(out_path, "w") as f:
1408        json.dump(params, f, indent=4)
1409
1410
1411def load_params(params_path: str) -> Dict[str, Any]:
1412    """Load parameter dictionary from JSON file.
1413
1414    Parameters
1415    ----------
1416    params_path : str
1417        Path to JSON file containing parameters.
1418
1419    Returns
1420    -------
1421    Dict[str, Any]
1422        Loaded parameter dictionary.
1423    """
1424    with open(params_path) as f:
1425        params = json.load(f)
1426
1427    return params
1428
1429
1430def write_occ_df(occ_df: pl.DataFrame, out_path: str) -> None:
1431    """Write occurrence data to TSV file.
1432
1433    Parameters
1434    ----------
1435    occ_df : pl.DataFrame
1436        Occurrence data DataFrame.
1437    out_path : str
1438        Output path for the TSV file.
1439    """
1440    occ_df.write_csv(out_path, separator="\t")
1441
1442
1443def write_seqlet_confusion_df(seqlet_confusion_df: pl.DataFrame, out_path: str) -> None:
1444    """Write seqlet confusion matrix data to TSV file.
1445
1446    Parameters
1447    ----------
1448    seqlet_confusion_df : pl.DataFrame
1449        Seqlet confusion matrix DataFrame.
1450    out_path : str
1451        Output path for the TSV file.
1452    """
1453    seqlet_confusion_df.write_csv(out_path, separator="\t")
1454
1455
1456def write_report_data(
1457    report_df: pl.DataFrame, cwms: Dict[str, Dict[str, ndarray]], out_dir: str
1458) -> None:
1459    """Write comprehensive motif report data including CWMs and metadata.
1460
1461    Parameters
1462    ----------
1463    report_df : pl.DataFrame
1464        Report metadata DataFrame.
1465    cwms : Dict[str, Dict[str, ndarray]]
1466        Nested dictionary of motif names to CWM types to arrays.
1467    out_dir : str
1468        Output directory for report files.
1469    """
1470    cwms_dir = os.path.join(out_dir, "CWMs")
1471    os.makedirs(cwms_dir, exist_ok=True)
1472
1473    for m, v in cwms.items():
1474        motif_dir = os.path.join(cwms_dir, m)
1475        os.makedirs(motif_dir, exist_ok=True)
1476        for cwm_type, cwm in v.items():
1477            np.savetxt(os.path.join(motif_dir, f"{cwm_type}.txt"), cwm)
1478
1479    report_df.write_csv(os.path.join(out_dir, "motif_report.tsv"), separator="\t")