smftools 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

smftools/informatics/fasta_functions.py

@@ -1,36 +1,57 @@
-from ..readwrite import make_dirs, time_string
+from __future__ import annotations
 
-import os
-import subprocess
+import gzip
+from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
-
-from typing import Union, List, Dict, Tuple
+from typing import Dict, Iterable, Tuple
 
 import numpy as np
-import gzip
-
+import pysam
 from Bio import SeqIO
-from Bio.SeqRecord import SeqRecord
 from Bio.Seq import Seq
+from Bio.SeqRecord import SeqRecord
 from pyfaidx import Fasta
-import pysam
 
-from concurrent.futures import ProcessPoolExecutor
-from itertools import chain
+from smftools.logging_utils import get_logger
+
+from ..readwrite import time_string
+
+logger = get_logger(__name__)
+
+
+def _convert_FASTA_record(
+    record: SeqRecord,
+    modification_type: str,
+    strand: str,
+    unconverted: str,
+) -> SeqRecord:
+    """Convert a FASTA record based on modification type and strand.
+
+    Args:
+        record: Input FASTA record.
+        modification_type: Modification type (e.g., ``5mC`` or ``6mA``).
+        strand: Strand label (``top`` or ``bottom``).
+        unconverted: Label for the unconverted record type.
+
+    Returns:
+        Bio.SeqRecord.SeqRecord: Converted FASTA record.
 
-def _convert_FASTA_record(record, modification_type, strand, unconverted):
-    """ Converts a FASTA record based on modification type and strand. """
+    Raises:
+        ValueError: If the modification type/strand combination is invalid.
+    """
     conversion_maps = {
-        ('5mC', 'top'): ('C', 'T'),
-        ('5mC', 'bottom'): ('G', 'A'),
-        ('6mA', 'top'): ('A', 'G'),
-        ('6mA', 'bottom'): ('T', 'C')
+        ("5mC", "top"): ("C", "T"),
+        ("5mC", "bottom"): ("G", "A"),
+        ("6mA", "top"): ("A", "G"),
+        ("6mA", "bottom"): ("T", "C"),
     }
 
     sequence = str(record.seq).upper()
 
     if modification_type == unconverted:
-        return SeqRecord(Seq(sequence), id=f"{record.id}_{modification_type}_top", description=record.description)
+        return SeqRecord(
+            Seq(sequence), id=f"{record.id}_{modification_type}_top", description=record.description
+        )
 
     if (modification_type, strand) not in conversion_maps:
         raise ValueError(f"Invalid combination: {modification_type}, {strand}")
@@ -38,62 +59,80 @@ def _convert_FASTA_record(record, modification_type, strand, unconverted):
     original_base, converted_base = conversion_maps[(modification_type, strand)]
     new_seq = sequence.replace(original_base, converted_base)
 
-    return SeqRecord(Seq(new_seq), id=f"{record.id}_{modification_type}_{strand}", description=record.description)
+    return SeqRecord(
+        Seq(new_seq), id=f"{record.id}_{modification_type}_{strand}", description=record.description
+    )
+
+
+def _process_fasta_record(
+    args: tuple[SeqRecord, Iterable[str], Iterable[str], str],
+) -> list[SeqRecord]:
+    """Process a single FASTA record for parallel conversion.
 
-def _process_fasta_record(args):
-    """
-    Processes a single FASTA record for parallel execution.
     Args:
-        args (tuple): (record, modification_types, strands, unconverted)
+        args: Tuple containing ``(record, modification_types, strands, unconverted)``.
+
     Returns:
-        list of modified SeqRecord objects.
+        list[Bio.SeqRecord.SeqRecord]: Converted FASTA records.
     """
     record, modification_types, strands, unconverted = args
     modified_records = []
-    
+
     for modification_type in modification_types:
         for i, strand in enumerate(strands):
             if i > 0 and modification_type == unconverted:
                 continue  # Ensure unconverted is added only once
 
-            modified_records.append(_convert_FASTA_record(record, modification_type, strand, unconverted))
+            modified_records.append(
+                _convert_FASTA_record(record, modification_type, strand, unconverted)
+            )
 
     return modified_records
 
-def generate_converted_FASTA(input_fasta, modification_types, strands, output_fasta, num_threads=4, chunk_size=500):
-    """
-    Converts an input FASTA file and writes a new converted FASTA file efficiently.
 
-    Parameters:
-        input_fasta (str): Path to the unconverted FASTA file.
-        modification_types (list): List of modification types ('5mC', '6mA', or unconverted).
-        strands (list): List of strands ('top', 'bottom').
-        output_fasta (str): Path to the converted FASTA output file.
-        num_threads (int): Number of parallel threads to use.
-        chunk_size (int): Number of records to process per write batch.
+def generate_converted_FASTA(
+    input_fasta: str | Path,
+    modification_types: list[str],
+    strands: list[str],
+    output_fasta: str | Path,
+    num_threads: int = 4,
+    chunk_size: int = 500,
+) -> None:
+    """Convert a FASTA file and write converted records to disk.
 
-    Returns:
-        None (Writes the converted FASTA file).
+    Args:
+        input_fasta: Path to the unconverted FASTA file.
+        modification_types: List of modification types (``5mC``, ``6mA``, or unconverted).
+        strands: List of strands (``top``, ``bottom``).
+        output_fasta: Path to the converted FASTA output file.
+        num_threads: Number of parallel workers to use.
+        chunk_size: Number of records to process per write batch.
     """
     unconverted = modification_types[0]
     input_fasta = str(input_fasta)
     output_fasta = str(output_fasta)
 
     # Detect if input is gzipped
-    open_func = gzip.open if input_fasta.endswith('.gz') else open
-    file_mode = 'rt' if input_fasta.endswith('.gz') else 'r'
+    open_func = gzip.open if input_fasta.endswith(".gz") else open
+    file_mode = "rt" if input_fasta.endswith(".gz") else "r"
 
     def _fasta_record_generator():
-        """ Lazily yields FASTA records from file. """
+        """Lazily yields FASTA records from file."""
         with open_func(input_fasta, file_mode) as handle:
-            for record in SeqIO.parse(handle, 'fasta'):
+            for record in SeqIO.parse(handle, "fasta"):
                 yield record
 
-    with open(output_fasta, 'w') as output_handle, ProcessPoolExecutor(max_workers=num_threads) as executor:
+    with (
+        open(output_fasta, "w") as output_handle,
+        ProcessPoolExecutor(max_workers=num_threads) as executor,
+    ):
         # Process records in parallel using a named function (avoiding lambda)
         results = executor.map(
             _process_fasta_record,
-            ((record, modification_types, strands, unconverted) for record in _fasta_record_generator())
+            (
+                (record, modification_types, strands, unconverted)
+                for record in _fasta_record_generator()
+            ),
         )
 
         buffer = []
@@ -102,14 +141,24 @@ def generate_converted_FASTA(input_fasta, modification_types, strands, output_fa
 
             # Write out in chunks to save memory
             if len(buffer) >= chunk_size:
-                SeqIO.write(buffer, output_handle, 'fasta')
+                SeqIO.write(buffer, output_handle, "fasta")
                 buffer.clear()
 
         # Write any remaining records
         if buffer:
-            SeqIO.write(buffer, output_handle, 'fasta')
+            SeqIO.write(buffer, output_handle, "fasta")
+
 
 def index_fasta(fasta: str | Path, write_chrom_sizes: bool = True) -> Path:
+    """Index a FASTA file and optionally write chromosome sizes.
+
+    Args:
+        fasta: Path to the FASTA file.
+        write_chrom_sizes: Whether to write a ``.chrom.sizes`` file.
+
+    Returns:
+        Path: Path to the index file or chromosome sizes file.
+    """
     fasta = Path(fasta)
     pysam.faidx(str(fasta))  # creates <fasta>.fai
 
@@ -123,9 +172,15 @@ def index_fasta(fasta: str | Path, write_chrom_sizes: bool = True) -> Path:
         return chrom_sizes
     return fai
 
+
 def get_chromosome_lengths(fasta: str | Path) -> Path:
-    """
-    Create (or reuse) <fasta>.chrom.sizes, derived from the FASTA index.
+    """Create or reuse ``<fasta>.chrom.sizes`` derived from the FASTA index.
+
+    Args:
+        fasta: Path to the FASTA file.
+
+    Returns:
+        Path: Path to the chromosome sizes file.
     """
     fasta = Path(fasta)
     fai = fasta.with_suffix(fasta.suffix + ".fai")
@@ -133,7 +188,7 @@ def get_chromosome_lengths(fasta: str | Path) -> Path:
         index_fasta(fasta, write_chrom_sizes=True)  # will also create .chrom.sizes
     chrom_sizes = fasta.with_suffix(".chrom.sizes")
     if chrom_sizes.exists():
-        print(f"Using existing chrom length file: {chrom_sizes}")
+        logger.debug(f"Using existing chrom length file: {chrom_sizes}")
         return chrom_sizes
 
     # Build chrom.sizes from .fai
@@ -143,10 +198,15 @@ def get_chromosome_lengths(fasta: str | Path) -> Path:
             out.write(f"{chrom}\t{size}\n")
     return chrom_sizes
 
+
 def get_native_references(fasta_file: str | Path) -> Dict[str, Tuple[int, str]]:
-    """
-    Return {record_id: (length, sequence)} from a FASTA.
-    Direct methylation specific
+    """Return record lengths and sequences from a FASTA file.
+
+    Args:
+        fasta_file: Path to the FASTA file.
+
+    Returns:
+        dict[str, tuple[int, str]]: Mapping of record ID to ``(length, sequence)``.
     """
     fasta_file = Path(fasta_file)
     print(f"{time_string()}: Opening FASTA file {fasta_file}")
@@ -157,28 +217,35 @@ def get_native_references(fasta_file: str | Path) -> Dict[str, Tuple[int, str]]:
             record_dict[rec.id] = (len(seq), seq)
     return record_dict
 
-def find_conversion_sites(fasta_file, modification_type, conversions, deaminase_footprinting=False):
-    """
-    Finds genomic coordinates of modified bases (5mC or 6mA) in a reference FASTA file.
-
-    Parameters:
-        fasta_file (str): Path to the converted reference FASTA.
-        modification_type (str): Modification type ('5mC' or '6mA') or 'unconverted'.
-        conversions (list): List of conversion types. The first element is the unconverted record type.
-        deaminase_footprinting (bool): Whether the footprinting was done with a direct deamination chemistry.
-
-    Returns: 
-        dict: Dictionary where keys are **both unconverted & converted record names**.
-              Values contain:
-              [sequence length, top strand coordinates, bottom strand coordinates, sequence, complement sequence].
+
+def find_conversion_sites(
+    fasta_file: str | Path,
+    modification_type: str,
+    conversions: list[str],
+    deaminase_footprinting: bool = False,
+) -> dict[str, list]:
+    """Find genomic coordinates of modified bases in a reference FASTA.
+
+    Args:
+        fasta_file: Path to the converted reference FASTA.
+        modification_type: Modification type (``5mC``, ``6mA``, or ``unconverted``).
+        conversions: List of conversion types (first entry is the unconverted record type).
+        deaminase_footprinting: Whether the footprinting used direct deamination chemistry.
+
+    Returns:
+        dict[str, list]: Mapping of record name to
+        ``[sequence length, top strand coordinates, bottom strand coordinates, sequence, complement]``.
+
+    Raises:
+        ValueError: If the modification type is invalid.
     """
     unconverted = conversions[0]
     record_dict = {}
 
     # Define base mapping based on modification type
     base_mappings = {
-        '5mC': ('C', 'G'),  # Cytosine and Guanine
-        '6mA': ('A', 'T')   # Adenine and Thymine
+        "5mC": ("C", "G"),  # Cytosine and Guanine
+        "6mA": ("A", "T"),  # Adenine and Thymine
     }
 
     # Read FASTA file and process records
@@ -200,22 +267,35 @@ def find_conversion_sites(fasta_file, modification_type, conversions, deaminase_
                     top_strand_coordinates = np.where(seq_array == top_base)[0].tolist()
                     bottom_strand_coordinates = np.where(seq_array == bottom_base)[0].tolist()
 
-                    record_dict[record.id] = [sequence_length, top_strand_coordinates, bottom_strand_coordinates, sequence, complement]
+                    record_dict[record.id] = [
+                        sequence_length,
+                        top_strand_coordinates,
+                        bottom_strand_coordinates,
+                        sequence,
+                        complement,
+                    ]
 
                 else:
-                    raise ValueError(f"Invalid modification_type: {modification_type}. Choose '5mC', '6mA', or 'unconverted'.")
+                    raise ValueError(
+                        f"Invalid modification_type: {modification_type}. Choose '5mC', '6mA', or 'unconverted'."
+                    )
 
     return record_dict
 
+
 def subsample_fasta_from_bed(
     input_FASTA: str | Path,
     input_bed: str | Path,
     output_directory: str | Path,
-    output_FASTA: str | Path
+    output_FASTA: str | Path,
 ) -> None:
-    """
-    Take a genome-wide FASTA file and a BED file containing
-    coordinate windows of interest. Outputs a subsampled FASTA.
+    """Subsample a FASTA using BED coordinates.
+
+    Args:
+        input_FASTA: Genome-wide FASTA path.
+        input_bed: BED file path containing coordinate windows of interest.
+        output_directory: Directory to write the subsampled FASTA.
+        output_FASTA: Output FASTA path.
     """
 
     # Normalize everything to Path
@@ -227,22 +307,20 @@ def subsample_fasta_from_bed(
     # Ensure output directory exists
     output_directory.mkdir(parents=True, exist_ok=True)
 
-    output_FASTA_path = output_directory / output_FASTA
-
     # Load the FASTA file using pyfaidx
-    fasta = Fasta(str(input_FASTA))   # pyfaidx requires string paths
+    fasta = Fasta(str(input_FASTA))  # pyfaidx requires string paths
 
     # Open BED + output FASTA
-    with input_bed.open("r") as bed, output_FASTA_path.open("w") as out_fasta:
+    with input_bed.open("r") as bed, output_FASTA.open("w") as out_fasta:
         for line in bed:
             fields = line.strip().split()
             chrom = fields[0]
-            start = int(fields[1]) # BED is 0-based
-            end   = int(fields[2]) # BED is 0-based and end is exclusive
-            desc  = " ".join(fields[3:]) if len(fields) > 3 else ""
+            start = int(fields[1])  # BED is 0-based
+            end = int(fields[2])  # BED is 0-based and end is exclusive
+            desc = " ".join(fields[3:]) if len(fields) > 3 else ""
 
             if chrom not in fasta:
-                print(f"Warning: {chrom} not found in FASTA")
+                logger.warning(f"{chrom} not found in FASTA")
                 continue
 
             # pyfaidx is 1-based indexing internally, but [start:end] works with BED coords
@@ -252,4 +330,4 @@ def subsample_fasta_from_bed(
             if desc:
                 header += f"    {desc}"
 
-            out_fasta.write(f"{header}\n{sequence}\n")
+            out_fasta.write(f"{header}\n{sequence}\n")

smftools/informatics/h5ad_functions.py

@@ -1,8 +1,18 @@
+import glob
+import os
+from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
-import pandas as pd
+from typing import Dict, List, Optional, Union
+
 import numpy as np
+import pandas as pd
 import scipy.sparse as sp
-from typing import Optional, List, Dict, Union
+from pod5 import Reader
+
+from smftools.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
 
 def add_demux_type_annotation(
     adata,
@@ -71,14 +81,15 @@ def add_demux_type_annotation(
 
     return adata
 
+
 def add_read_length_and_mapping_qc(
     adata,
     bam_files: Optional[List[str]] = None,
     read_metrics: Optional[Dict[str, Union[list, tuple]]] = None,
     uns_flag: str = "add_read_length_and_mapping_qc_performed",
-    extract_read_features_from_bam_callable = None,
+    extract_read_features_from_bam_callable=None,
     bypass: bool = False,
-    force_redo: bool = True
+    force_redo: bool = True,
 ):
     """
     Populate adata.obs with read/mapping QC columns.
@@ -98,6 +109,7 @@ def add_read_length_and_mapping_qc(
         Optional callable(bam_path) -> dict mapping read_name -> list/tuple of metrics.
         If not provided and bam_files is given, function will attempt to call `extract_read_features_from_bam`
         from the global namespace (your existing helper).
+
     Returns
     -------
     None (mutates final_adata in-place)
@@ -113,9 +125,13 @@ def add_read_length_and_mapping_qc(
     if read_metrics is None:
         read_metrics = {}
         if bam_files:
-            extractor = extract_read_features_from_bam_callable or globals().get("extract_read_features_from_bam")
+            extractor = extract_read_features_from_bam_callable or globals().get(
+                "extract_read_features_from_bam"
+            )
             if extractor is None:
-                raise ValueError("No `read_metrics` provided and `extract_read_features_from_bam` not found.")
+                raise ValueError(
+                    "No `read_metrics` provided and `extract_read_features_from_bam` not found."
+                )
             for bam in bam_files:
                 bam_read_metrics = extractor(bam)
                 if not isinstance(bam_read_metrics, dict):
@@ -130,11 +146,11 @@ def add_read_length_and_mapping_qc(
     if len(read_metrics) == 0:
         # fill with NaNs
         n = adata.n_obs
-        adata.obs['read_length'] = np.full(n, np.nan)
-        adata.obs['mapped_length'] = np.full(n, np.nan)
-        adata.obs['reference_length'] = np.full(n, np.nan)
-        adata.obs['read_quality'] = np.full(n, np.nan)
-        adata.obs['mapping_quality'] = np.full(n, np.nan)
+        adata.obs["read_length"] = np.full(n, np.nan)
+        adata.obs["mapped_length"] = np.full(n, np.nan)
+        adata.obs["reference_length"] = np.full(n, np.nan)
+        adata.obs["read_quality"] = np.full(n, np.nan)
+        adata.obs["mapping_quality"] = np.full(n, np.nan)
     else:
         # Build DF robustly
         # Convert values to lists where possible, else to [val, val, val...]
@@ -151,35 +167,45 @@ def add_read_length_and_mapping_qc(
                 vals = vals + [np.nan] * (max_cols - len(vals))
             rows[k] = vals[:max_cols]
 
-        df = pd.DataFrame.from_dict(rows, orient='index', columns=[
-            'read_length', 'read_quality', 'reference_length', 'mapped_length', 'mapping_quality'
-        ])
+        df = pd.DataFrame.from_dict(
+            rows,
+            orient="index",
+            columns=[
+                "read_length",
+                "read_quality",
+                "reference_length",
+                "mapped_length",
+                "mapping_quality",
+            ],
+        )
 
         # Reindex to final_adata.obs_names so order matches adata
         # If obs_names are not present as keys in df, the results will be NaN
         df_reindexed = df.reindex(adata.obs_names).astype(float)
 
-        adata.obs['read_length'] = df_reindexed['read_length'].values
-        adata.obs['mapped_length'] = df_reindexed['mapped_length'].values
-        adata.obs['reference_length'] = df_reindexed['reference_length'].values
-        adata.obs['read_quality'] = df_reindexed['read_quality'].values
-        adata.obs['mapping_quality'] = df_reindexed['mapping_quality'].values
+        adata.obs["read_length"] = df_reindexed["read_length"].values
+        adata.obs["mapped_length"] = df_reindexed["mapped_length"].values
+        adata.obs["reference_length"] = df_reindexed["reference_length"].values
+        adata.obs["read_quality"] = df_reindexed["read_quality"].values
+        adata.obs["mapping_quality"] = df_reindexed["mapping_quality"].values
 
     # Compute ratio columns safely (avoid divide-by-zero and preserve NaN)
     # read_length_to_reference_length_ratio
-    rl = pd.to_numeric(adata.obs['read_length'], errors='coerce').to_numpy(dtype=float)
-    ref_len = pd.to_numeric(adata.obs['reference_length'], errors='coerce').to_numpy(dtype=float)
-    mapped_len = pd.to_numeric(adata.obs['mapped_length'], errors='coerce').to_numpy(dtype=float)
+    rl = pd.to_numeric(adata.obs["read_length"], errors="coerce").to_numpy(dtype=float)
+    ref_len = pd.to_numeric(adata.obs["reference_length"], errors="coerce").to_numpy(dtype=float)
+    mapped_len = pd.to_numeric(adata.obs["mapped_length"], errors="coerce").to_numpy(dtype=float)
 
     # safe divisions: use np.where to avoid warnings and replace inf with nan
-    with np.errstate(divide='ignore', invalid='ignore'):
+    with np.errstate(divide="ignore", invalid="ignore"):
         rl_to_ref = np.where((ref_len != 0) & np.isfinite(ref_len), rl / ref_len, np.nan)
-        mapped_to_ref = np.where((ref_len != 0) & np.isfinite(ref_len), mapped_len / ref_len, np.nan)
+        mapped_to_ref = np.where(
+            (ref_len != 0) & np.isfinite(ref_len), mapped_len / ref_len, np.nan
+        )
         mapped_to_read = np.where((rl != 0) & np.isfinite(rl), mapped_len / rl, np.nan)
 
-    adata.obs['read_length_to_reference_length_ratio'] = rl_to_ref
-    adata.obs['mapped_length_to_reference_length_ratio'] = mapped_to_ref
-    adata.obs['mapped_length_to_read_length_ratio'] = mapped_to_read
+    adata.obs["read_length_to_reference_length_ratio"] = rl_to_ref
+    adata.obs["mapped_length_to_reference_length_ratio"] = mapped_to_ref
+    adata.obs["mapped_length_to_read_length_ratio"] = mapped_to_read
 
     # Add read level raw modification signal: sum over X rows
     X = adata.X
@@ -189,9 +215,149 @@ def add_read_length_and_mapping_qc(
     else:
         raw_sig = np.asarray(X.sum(axis=1)).ravel()
 
-    adata.obs['Raw_modification_signal'] = raw_sig
+    adata.obs["Raw_modification_signal"] = raw_sig
 
     # mark as done
     adata.uns[uns_flag] = True
 
-    return None
+    return None
+
+
+def _collect_read_origins_from_pod5(pod5_path: str, target_ids: set[str]) -> dict[str, str]:
+    """
+    Worker function: scan one POD5 file and return a mapping
+    {read_id: pod5_basename} only for read_ids in `target_ids`.
+    """
+    basename = os.path.basename(pod5_path)
+    mapping: dict[str, str] = {}
+
+    with Reader(pod5_path) as reader:
+        for read in reader.reads():
+            # Cast read id to string
+            rid = str(read.read_id)
+            if rid in target_ids:
+                mapping[rid] = basename
+
+    return mapping
+
+
+def annotate_pod5_origin(
+    adata,
+    pod5_path_or_dir: str | Path,
+    pattern: str = "*.pod5",
+    n_jobs: int | None = None,
+    fill_value: str | None = "unknown",
+    verbose: bool = True,
+    csv_path: str | None = None,
+):
+    """
+    Add `pod5_origin` column to `adata.obs`, containing the POD5 basename
+    each read came from.
+
+    Parameters
+    ----------
+    adata
+        AnnData with obs_names == read_ids (as strings).
+    pod5_path_or_dir
+        Directory containing POD5 files or path to a single POD5 file.
+    pattern
+        Glob pattern for POD5 files inside `pod5_dir`.
+    n_jobs
+        Number of worker processes. If None or <=1, runs serially.
+    fill_value
+        Value to use when a read_id is not found in any POD5 file.
+        If None, leaves missing as NaN.
+    verbose
+        Print progress info.
+    csv_path
+        Path to a csv of the read to pod5 origin mapping
+
+    Returns
+    -------
+    None (modifies `adata` in-place).
+    """
+    pod5_path_or_dir = Path(pod5_path_or_dir)
+
+    # --- Resolve input into a list of pod5 files ---
+    if pod5_path_or_dir.is_dir():
+        pod5_files = sorted(str(p) for p in pod5_path_or_dir.glob(pattern))
+        if not pod5_files:
+            raise FileNotFoundError(
+                f"No POD5 files matching {pattern!r} in {str(pod5_path_or_dir)!r}"
+            )
+    elif pod5_path_or_dir.is_file():
+        if pod5_path_or_dir.suffix.lower() != ".pod5":
+            raise ValueError(f"Expected a .pod5 file, got: {pod5_path_or_dir}")
+        pod5_files = [str(pod5_path_or_dir)]
+    else:
+        raise FileNotFoundError(f"Path does not exist: {pod5_path_or_dir}")
+
+    # Make sure obs_names are strings
+    obs_names = adata.obs_names.astype(str)
+    target_ids = set(obs_names)  # only these are interesting
+
+    if verbose:
+        logger.info(f"Found {len(pod5_files)} POD5 files.")
+        logger.info(f"Tracking {len(target_ids)} read IDs from AnnData.")
+
+    # --- Collect mappings (possibly multiprocessed) ---
+    global_mapping: dict[str, str] = {}
+
+    if n_jobs is None or n_jobs <= 1:
+        # Serial version (less overhead, useful for debugging)
+        if verbose:
+            logger.debug("Running in SERIAL mode.")
+        for f in pod5_files:
+            if verbose:
+                logger.debug(f"  Scanning {os.path.basename(f)} ...")
+            part = _collect_read_origins_from_pod5(f, target_ids)
+            global_mapping.update(part)
+    else:
+        if verbose:
+            logger.debug(f"Running in PARALLEL mode with {n_jobs} workers.")
+        with ProcessPoolExecutor(max_workers=n_jobs) as ex:
+            futures = {
+                ex.submit(_collect_read_origins_from_pod5, f, target_ids): f for f in pod5_files
+            }
+            for fut in as_completed(futures):
+                f = futures[fut]
+                try:
+                    part = fut.result()
+                except Exception as e:
+                    logger.warning(f"Error while processing {f}: {e}")
+                    continue
+                global_mapping.update(part)
+                if verbose:
+                    logger.info(f"  Finished {os.path.basename(f)} ({len(part)} matching reads)")
+
+    if verbose:
+        logger.info(f"Total reads matched: {len(global_mapping)}")
+
+    # --- Populate obs['pod5_origin'] in AnnData order, memory-efficiently ---
+    origin = np.empty(adata.n_obs, dtype=object)
+    default = None if fill_value is None else fill_value
+    for i, rid in enumerate(obs_names):
+        origin[i] = global_mapping.get(rid, default)
+
+    adata.obs["pod5_origin"] = origin
+    if verbose:
+        logger.info("Assigned `pod5_origin` to adata.obs.")
+
+    # --- Optionally write a CSV ---
+    if csv_path is not None:
+        if verbose:
+            logger.info(f"Writing CSV mapping to: {csv_path}")
+
+        # Create DataFrame in AnnData order for easier cross-referencing
+        df = pd.DataFrame(
+            {
+                "read_id": obs_names,
+                "pod5_origin": origin,
+            }
+        )
+        df.to_csv(csv_path, index=False)
+
+        if verbose:
+            logger.info("CSV saved.")
+
+    return global_mapping