yamcot 1.0.0__cp312-cp312-macosx_11_0_arm64.whl

yamcot/io.py

@@ -0,0 +1,522 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import Iterable, List, Tuple, Union
+
+import numpy as np
+
+from yamcot.functions import pfm_to_pwm
+from yamcot.ragged import RaggedData
+
+
+def read_fasta(path: str) -> RaggedData:
+    """Read a FASTA file and return integer-encoded sequences.
+
+    Parameters
+    ----------
+    path : str
+        Path to a FASTA formatted file.
+    return_ragged : bool, default True
+        If True, returns RaggedData object. If False, returns List[np.ndarray]
+        for backward compatibility.
+
+    Returns
+    -------
+    Union[RaggedData, List[np.ndarray]]
+        Integer-encoded sequences (dtype=np.int8).
+        Using RaggedData is more memory-efficient and faster for large files
+        as it avoids multiple small allocations.
+    """
+    # Translation table for fast conversion of strings to bytes (int8)
+    # 0=A, 1=C, 2=G, 3=T, 4=N (and others)
+    trans_table = bytearray([4] * 256)
+    for char, code in zip(b"ACGTacgt", [0, 1, 2, 3] * 2, strict=False):
+        trans_table[char] = code
+
+    # First collect lengths and data in one pass if possible,
+    # but FASTA requires parsing to determine lengths.
+    # For efficiency, we use a temporary list or intermediate buffer.
+
+    sequences: List[np.ndarray] = []
+
+    with open(path, "r") as handle:
+        current_seq_bytes = bytearray()
+        for line in handle:
+            line = line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if current_seq_bytes:
+                    encoded = np.frombuffer(current_seq_bytes.translate(trans_table), dtype=np.int8).copy()
+                    sequences.append(encoded)
+                    current_seq_bytes.clear()
+            else:
+                current_seq_bytes.extend(line.encode("ascii", errors="ignore"))
+
+        if current_seq_bytes:
+            encoded = np.frombuffer(current_seq_bytes.translate(trans_table), dtype=np.int8).copy()
+            sequences.append(encoded)
+
+    if not sequences:
+        return RaggedData(np.empty(0, dtype=np.int8), np.zeros(1, dtype=np.int64))
+
+    # Convert to RaggedData
+    n = len(sequences)
+    offsets = np.zeros(n + 1, dtype=np.int64)
+    for i, seq in enumerate(sequences):
+        offsets[i + 1] = offsets[i] + len(seq)
+
+    total_data = np.empty(offsets[-1], dtype=np.int8)
+    for i, seq in enumerate(sequences):
+        total_data[offsets[i] : offsets[i + 1]] = seq
+
+    return RaggedData(total_data, offsets)
+
+
+def write_fasta(sequences: Union[RaggedData, Iterable[np.ndarray]], path: str) -> None:
+    """Write integer-encoded sequences to a FASTA file.
+
+    Parameters
+    ----------
+    sequences : Union[RaggedData, Iterable[np.ndarray]]
+        RaggedData object or a collection of integer-encoded sequences (0=A, 1=C, 2=G, 3=T, 4=N).
+    path : str
+        Path to the output file.
+    """
+    # Array for converting indices back to symbols
+    decoder = np.array(["A", "C", "G", "T", "N"], dtype="U1")
+
+    with open(path, "w") as out:
+        if isinstance(sequences, RaggedData):
+            for i in range(sequences.num_sequences):
+                seq_int = sequences.get_slice(i)
+                safe_seq = np.clip(seq_int, 0, 4)
+                chars = decoder[safe_seq]
+                seq_str = "".join(chars)
+                out.write(f">{i}\n")
+                out.write(f"{seq_str}\n")
+        else:
+            for idx, seq_int in enumerate(sequences):
+                safe_seq = np.clip(seq_int, 0, 4)
+                chars = decoder[safe_seq]
+                seq_str = "".join(chars)
+                out.write(f">{idx}\n")
+                out.write(f"{seq_str}\n")
+
+
+def read_meme(path: str, index: int = 0) -> Tuple[np.ndarray, Tuple[str, int], int]:
+    """Read a specific motif from a MEME formatted file and return total count.
+
+    Parameters
+    ----------
+    path : str
+        Path to the MEME file.
+    index : int, default 0
+        The zero-based index of the motif to return.
+
+    Returns
+    -------
+    Tuple[np.ndarray, Tuple[str, int], int]
+        A tuple containing:
+        - The requested motif matrix (shape (4, L))
+        - A tuple with the motif's name and length
+        - Total number of motifs found in the file
+    """
+    target_motif: np.ndarray | None = None
+    target_info: Tuple[str, int] | None = None
+    motif_count = 0
+
+    with open(path) as handle:
+        line = handle.readline()
+        while line:
+            if line.startswith("MOTIF"):
+                # Check if this is the motif we are looking for
+                is_target = motif_count == index
+                motif_count += 1
+
+                parts = line.strip().split()
+                name = parts[1]
+
+                # Read header line containing motif length (w=)
+                header_line = handle.readline()
+                header = header_line.strip().split()
+
+                try:
+                    length_idx = header.index("w=") + 1
+                    length = int(header[length_idx])
+                except (ValueError, IndexError):
+                    length = 0
+
+                if is_target:
+                    matrix = []
+                    for _ in range(length):
+                        row_line = handle.readline()
+                        row = row_line.strip().split()
+                        if not row:
+                            continue
+                        matrix.append(list(map(float, row)))
+
+                    # Transpose into shape (4, length)
+                    target_motif = np.array(matrix, dtype=np.float32).T
+                    target_info = (name, length)
+                else:
+                    # Skip the matrix rows for other motifs to save time
+                    for _ in range(length):
+                        handle.readline()
+
+            line = handle.readline()
+
+    if target_motif is None:
+        if motif_count == 0:
+            raise ValueError(f"No motifs found in {path}")
+        else:
+            raise IndexError(f"Motif index {index} out of range. File contains {motif_count} motifs.")
+
+    # We know that if target_motif is not None, then target_info is also not None
+    # because they are set together in the same condition
+    assert target_info is not None
+
+    return target_motif, target_info, motif_count
+
+
+def write_meme(motifs: List[np.ndarray], info: List[Tuple[str, int]], path: str) -> None:
+    """Write a list of motifs to a MEME formatted file.
+
+    Parameters
+    ----------
+    motifs : List[np.ndarray]
+        List of motif matrices of shape (5, L).  Only the first four rows
+        (A, C, G, T) are written; the fifth row is ignored.
+    info : List[Tuple[str, int]]
+        A list of (name, length) tuples corresponding to the motifs.
+    path : str
+        Path of the output file.
+    """
+    with open(path, "w") as out:
+        out.write("MEME version 4\n\n")
+        out.write("ALPHABET= ACGT\n\n")
+        out.write("strands: + -\n\n")
+        out.write("Background letter frequencies\n")
+        out.write("A 0.25 C 0.25 G 0.25 T 0.25\n\n")
+        for motif, (name, length) in zip(motifs, info, strict=False):
+            out.write(f"MOTIF {name}\n")
+            out.write(f"letter-probability matrix: alength= 4 w= {length}\n")
+            for row in motif[:4].T:
+                out.write(" " + " ".join(f"{val:.6f}" for val in row) + "\n")
+            out.write("\n")
+
+
+def read_sitega(path: str) -> tuple[np.ndarray, int, float, float]:
+    """Parse SiteGA output file and return the motif matrix with metadata.
+
+    Parameters
+    ----------
+    path : str
+        Path to the SiteGA output file (typically ends with '.mat').
+
+    Returns
+    -------
+    tuple[np.ndarray, int, float, float]
+        A tuple containing:
+        - SiteGA matrix of shape (5, 5, length) representing dinucleotide dependencies
+        - Length of the motif
+        - Minimum score value
+        - Maximum score value
+    """
+    converter = {"A": 0, "C": 1, "G": 2, "T": 3}
+    with open(path) as file:
+        _name = file.readline().strip()
+        _number_of_lpd = int(file.readline().strip().split()[0])
+        length = int(file.readline().strip().split()[0])
+        minimum = float(file.readline().strip().split()[0])
+        maximum = float(file.readline().strip().split()[0])
+        sitega = np.zeros((5, 5, length), dtype=np.float32)
+        for line in file:
+            start, stop, value, _, dinucleotide = line.strip().split()
+            dinucleotide = dinucleotide.upper()
+            nuc_1, nuc_2 = converter[dinucleotide[0]], converter[dinucleotide[1]]
+            number_of_positions = int(stop) - int(start) + 1
+            for index in range(int(start), int(stop) + 1):
+                sitega[nuc_1][nuc_2][index] += float(value) / number_of_positions
+    return np.array(sitega, dtype=np.float32), length, minimum, maximum
+
+
+def parse_file_content(filepath: str) -> tuple[dict[int, list[np.ndarray]], int, int]:
+    """Parse BaMM file content, ignoring comments starting with '#'.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to the BaMM file to parse.
+
+    Returns
+    -------
+    tuple[dict[int, list[np.ndarray]], int, int]
+        A tuple containing:
+        - Dictionary mapping order indices to lists of coefficient arrays
+        - Maximum order found in the file
+        - Number of positions (length of motif)
+
+    Raises
+    ------
+    FileNotFoundError
+        If the specified file does not exist.
+    ValueError
+        If no valid data is found in the file or inconsistent orders are detected.
+    """
+    if not os.path.isfile(filepath):
+        raise FileNotFoundError(f"File {filepath} not found")
+
+    with open(filepath, "r") as f:
+        raw_text = f.read()
+
+    # Split blocks by double newline
+    raw_blocks = raw_text.strip().split("\n\n")
+    clean_blocks_data = []
+
+    for raw_block in raw_blocks:
+        lines = raw_block.strip().split("\n")
+
+        # Filter comments and empty lines
+        valid_lines = [line.strip() for line in lines if line.strip() and not line.strip().startswith("#")]
+
+        if not valid_lines:
+            continue
+
+        block_arrays = []
+        for line in valid_lines:
+            # Check for potential empty strings after split
+            parts = line.split()
+            if not parts:
+                continue
+            arr = np.array([float(x) for x in parts], dtype=np.float32)
+            block_arrays.append(arr)
+
+        clean_blocks_data.append(block_arrays)
+
+    if not clean_blocks_data:
+        raise ValueError(f"No valid data found in {filepath}")
+
+    num_positions = len(clean_blocks_data)
+    max_order = len(clean_blocks_data[0]) - 1
+
+    data_by_order = {}
+    for k in range(max_order + 1):
+        data_by_order[k] = []
+        for pos_idx in range(num_positions):
+            if len(clean_blocks_data[pos_idx]) <= k:
+                raise ValueError(f"Inconsistent orders in block {pos_idx}")
+            data_by_order[k].append(clean_blocks_data[pos_idx][k])
+
+    return data_by_order, max_order, num_positions
+
+
+def read_bamm(motif_path: str, bg_path: str, target_order: int) -> np.ndarray:
+    """Read BaMM files, apply ramp-up logic, and add padding for 'N' (index 4).
+
+    This function reads motif and background BaMM files, computes log-odds ratios,
+    applies a ramp-up strategy for lower-order coefficients, and adds padding for
+    ambiguous nucleotides ('N') by setting their scores to the minimum of each position.
+
+    Parameters
+    ----------
+    motif_path : str
+        Path to the motif BaMM file (.ihbcp format).
+    bg_path : str
+        Path to the background BaMM file (.hbcp format).
+    target_order : int
+        Target order for the BaMM model (determines tensor dimensions).
+
+    Returns
+    -------
+    np.ndarray
+        3D+ tensor of shape (5, 5, ..., 5, Length) where the number of dimensions
+        equals target_order + 2. The final dimension represents motif length,
+        and the first target_order+1 dimensions represent nucleotide dependencies
+        including 'N' padding at index 4.
+
+    Raises
+    ------
+    ValueError
+        If target order exceeds the maximum order in the file.
+    """
+    # 1. Parse Data
+    motif_raw, max_order_file, motif_length = parse_file_content(motif_path)
+    bg_raw, max_order_bg, _ = parse_file_content(bg_path)
+
+    if max_order_file > max_order_bg:
+        max_order_file = max_order_bg
+    if target_order > max_order_file:
+        target_order = max_order_file
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Target order {target_order} exceeds file max order {max_order_file}, target order set as max order"
+        )
+
+    # 2. Build 4x4...x4 Tensor slices (Standard Logic)
+    # We first build the pure ACGT tensor to compute minimums correctly
+    acgt_slices = []
+
+    for pos in range(motif_length):
+        current_k = min(pos, target_order)
+
+        p_motif = motif_raw[current_k][pos]
+        p_bg = bg_raw[current_k][0]
+
+        epsilon = 1e-10
+        log_odds = np.log2((p_motif + epsilon) / (p_bg + epsilon))
+
+        # Reshape & Broadcast (Ramp-up)
+        shape_k = [4] * (current_k + 1)
+        tensor_k = log_odds.reshape(shape_k)
+
+        if current_k < target_order:
+            missing_dims = target_order - current_k
+            expand_shape = [1] * missing_dims + shape_k
+            tensor_expanded = tensor_k.reshape(expand_shape)
+            target_shape_4 = [4] * (target_order + 1)
+            tensor_final = np.broadcast_to(tensor_expanded, target_shape_4).copy()
+        else:
+            tensor_final = tensor_k
+
+        acgt_slices.append(tensor_final)
+
+    # Stack to (4, 4, ..., 4, Length)
+    # This tensor contains only valid ACGT scores
+    acgt_tensor = np.stack(acgt_slices, axis=-1)
+
+    # 3. Create 5x5...x5 Tensor with N-padding
+
+    # Calculate global minimum per position (over all ACGT contexts)
+    # We want min over axes (0, 1, ..., target_order).
+    # acgt_tensor shape is (4, 4, ..., L). Last axis is Length.
+    reduce_axes = tuple(range(target_order + 1))
+    min_scores_per_pos = np.min(acgt_tensor, axis=reduce_axes)  # Shape (Length,)
+
+    # Define new shape: (5, 5, ..., 5, Length)
+    new_shape = [5] * (target_order + 1) + [motif_length]
+
+    # Initialize with min values broadcasted
+    # NumPy broadcasts from last dimension: (Length,) broadcasts to (5, 5, ..., 5, Length)
+    final_tensor = np.ones(new_shape, dtype=np.float32) * min_scores_per_pos
+
+    # 4. Copy ACGT data into the 5x5 structure
+    # We need to slice [0:4, 0:4, ..., :]
+    slice_objs = [slice(0, 4)] * (target_order + 1) + [slice(None)]
+    final_tensor[tuple(slice_objs)] = acgt_tensor
+
+    # Return as contiguous array for Numba
+    return np.array(final_tensor, dtype=np.float32)
+
+
+def write_sitega(motif, path: str) -> None:
+    """Write SiteGA motif to a file in the .mat format understood by mco_prc.exe.
+
+    Parameters
+    ----------
+    motif : SitegaMotif
+        The SiteGA motif to write.
+    path : str
+        Path to the output file.
+    """
+    sitega_matrix = motif.matrix
+    converter = {0: "A", 1: "C", 2: "G", 3: "T"}
+
+    # Список для хранения найденных сегментов (start, stop, value, dinucleotide)
+    segments = []
+
+    # Один проход для сбора всех данных
+    for nuc1 in range(4):
+        for nuc2 in range(4):
+            # Пропускаем, если все значения для динуклеотида нулевые
+            if np.all(np.abs(sitega_matrix[nuc1, nuc2, :]) <= 1e-9):
+                continue
+
+            dinucleotide = converter[nuc1] + converter[nuc2]
+            pos = 0
+
+            while pos < motif.length:
+                # Пропускаем нули
+                while pos < motif.length and abs(sitega_matrix[nuc1, nuc2, pos]) <= 1e-9:
+                    pos += 1
+
+                if pos >= motif.length:
+                    break
+
+                # Начало ненулевой последовательности
+                start_pos = pos
+                current_val = sitega_matrix[nuc1, nuc2, pos]
+
+                # Ищем конец последовательности с одинаковым значением
+                while pos + 1 < motif.length and abs(sitega_matrix[nuc1, nuc2, pos + 1] - current_val) < 1e-9:
+                    pos += 1
+
+                # Сохраняем сегмент
+                segments.append({"start": start_pos, "stop": pos, "val": current_val, "dinucl": dinucleotide})
+
+                pos += 1
+
+    # Количество строк данных теперь просто длина списка
+    lpd_count = len(segments)
+
+    with open(path, "w") as f:
+        f.write(f"{motif.name}\n")
+        f.write(f"{lpd_count}\tLPD count\n")
+        f.write(f"{motif.length}\tModel length\n")
+        f.write(f"{motif.minimum:.12f}\tMinimum\n")
+        f.write(f"{motif.maximum:.12f}\tRazmah\n")
+
+        # Записываем данные из собранного списка
+        for seg in segments:
+            range_length = seg["stop"] - seg["start"] + 1
+            total_value = seg["val"] * range_length
+
+            f.write(f"{seg['start']}\t{seg['stop']}\t{total_value:.12f}\t0\t{seg['dinucl'].lower()}\n")
+
+
+def write_pfm(pfm: np.ndarray, name: str, length: int, path: str) -> None:
+    """Write a Position Frequency Matrix to a file.
+
+    Parameters
+    ----------
+    pfm : np.ndarray
+        Position frequency matrix of shape (4, length).
+    name : str
+        Name of the motif.
+    length : int
+        Length of the motif.
+    path : str
+        Path to the output file.
+    """
+    with open(path, "w") as f:
+        f.write(f">{name}\n")
+        # Transpose the matrix to get the right format
+        np.savetxt(f, pfm.T, fmt="%.6f", delimiter="\t")
+
+
+def read_pfm(path: str) -> tuple[np.ndarray, int, float, float]:
+    """Read a Position Frequency Matrix (PFM) from a file and convert to PWM.
+
+    Parameters
+    ----------
+    path : str
+        Path to the PFM file.
+
+    Returns
+    -------
+    tuple[np.ndarray, int, float, float]
+        A tuple containing:
+        - PWM matrix with shape (5, L) where the 5th row contains column minima
+        - Length of the motif
+        - Minimum possible score
+        - Maximum possible score
+    """
+    pfm = np.loadtxt(path, comments=">").T
+    length = pfm.shape[1]
+    pwm = pfm_to_pwm(pfm)
+    minimum = np.sum(pwm.min(axis=0))
+    maximum = np.sum(pwm.max(axis=0))
+    pwm = np.concatenate((pwm, np.min(pwm, axis=0).reshape(1, pwm.shape[1])), axis=0)
+    pwm = np.array(pwm, dtype=np.float32)
+    return pwm, length, minimum, maximum