smftools 0.1.3__py3-none-any.whl → 0.1.7__py3-none-any.whl

smftools/informatics/helpers/converted_BAM_to_adata_II.py

@@ -0,0 +1,369 @@
+import numpy as np
+import time
+import os
+import gc
+import pandas as pd
+import anndata as ad
+from tqdm import tqdm
+import multiprocessing
+from multiprocessing import Manager, Lock, current_process, Pool
+import traceback
+import gzip
+import torch
+
+from .. import readwrite
+from .binarize_converted_base_identities import binarize_converted_base_identities
+from .find_conversion_sites import find_conversion_sites
+from .count_aligned_reads import count_aligned_reads
+from .extract_base_identities import extract_base_identities
+from .make_dirs import make_dirs
+from .ohe_batching import ohe_batching
+
+if __name__ == "__main__":
+    multiprocessing.set_start_method("forkserver", force=True)
+
+def converted_BAM_to_adata_II(converted_FASTA, split_dir, mapping_threshold, experiment_name, conversion_types, bam_suffix, device='cpu', num_threads=8):
+    """
+    Converts BAM files into an AnnData object by binarizing modified base identities.
+
+    Parameters:
+        converted_FASTA (str): Path to the converted FASTA reference.
+        split_dir (str): Directory containing converted BAM files.
+        mapping_threshold (float): Minimum fraction of aligned reads required for inclusion.
+        experiment_name (str): Name for the output AnnData object.
+        conversion_types (list): List of modification types (e.g., ['unconverted', '5mC', '6mA']).
+        bam_suffix (str): File suffix for BAM files.
+        num_threads (int): Number of parallel processing threads.
+
+    Returns:
+        str: Path to the final AnnData object.
+    """
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    print(f"Using device: {device}")
+
+    ## Set Up Directories and File Paths
+    parent_dir = os.path.dirname(split_dir)
+    h5_dir = os.path.join(parent_dir, 'h5ads')
+    tmp_dir = os.path.join(parent_dir, 'tmp')
+    final_adata_path = os.path.join(h5_dir, f'{experiment_name}_{os.path.basename(split_dir)}.h5ad.gz')
+
+    if os.path.exists(final_adata_path):
+        print(f"{final_adata_path} already exists. Using existing AnnData object.")
+        return final_adata_path
+
+    make_dirs([h5_dir, tmp_dir])
+
+    ## Get BAM Files ##
+    bam_files = [f for f in os.listdir(split_dir) if f.endswith(bam_suffix) and not f.endswith('.bai') and 'unclassified' not in f]
+    bam_files.sort()
+    bam_path_list = [os.path.join(split_dir, f) for f in bam_files]
+    print(f"Found {len(bam_files)} BAM files: {bam_files}")
+
+    ## Process Conversion Sites
+    max_reference_length, record_FASTA_dict, chromosome_FASTA_dict = process_conversion_sites(converted_FASTA, conversion_types)
+
+    ## Filter BAM Files by Mapping Threshold
+    records_to_analyze = filter_bams_by_mapping_threshold(bam_path_list, bam_files, mapping_threshold)
+
+    ## Process BAMs in Parallel
+    final_adata = process_bams_parallel(bam_path_list, records_to_analyze, record_FASTA_dict, tmp_dir, h5_dir, num_threads, max_reference_length, device)
+
+    for chromosome, [seq, comp] in chromosome_FASTA_dict.items():
+        final_adata.var[f'{chromosome}_top_strand_FASTA_base'] = list(seq)
+        final_adata.var[f'{chromosome}_bottom_strand_FASTA_base'] = list(comp)
+        final_adata.uns[f'{chromosome}_FASTA_sequence'] = seq
+
+    ## Save Final AnnData
+    # print(f"Saving AnnData to {final_adata_path}")
+    # final_adata.write_h5ad(final_adata_path, compression='gzip')
+    return final_adata, final_adata_path
+
+
+def process_conversion_sites(converted_FASTA, conversion_types):
+    """
+    Extracts conversion sites and determines the max reference length.
+
+    Parameters:
+        converted_FASTA (str): Path to the converted reference FASTA.
+        conversion_types (list): List of modification types (e.g., ['unconverted', '5mC', '6mA']).
+
+    Returns:
+        max_reference_length (int): The length of the longest sequence.
+        record_FASTA_dict (dict): Dictionary of sequence information for **both converted & unconverted** records.
+    """
+    modification_dict = {}
+    record_FASTA_dict = {}
+    chromosome_FASTA_dict = {}
+    max_reference_length = 0
+    unconverted = conversion_types[0]
+    conversions = conversion_types[1:]
+
+    # Process the unconverted sequence once
+    modification_dict[unconverted] = find_conversion_sites(converted_FASTA, unconverted, conversion_types)
+    # Above points to record_dict[record.id] = [sequence_length, [], [], sequence, complement] with only unconverted record.id keys
+
+    # Get **max sequence length** from unconverted records
+    max_reference_length = max(values[0] for values in modification_dict[unconverted].values())
+
+    # Add **unconverted records** to `record_FASTA_dict`
+    for record, values in modification_dict[unconverted].items():
+        sequence_length, top_coords, bottom_coords, sequence, complement = values
+        chromosome = record.replace(f"_{unconverted}_top", "")
+
+        # Store **original sequence**
+        record_FASTA_dict[record] = [
+            sequence + "N" * (max_reference_length - sequence_length),
+            complement + "N" * (max_reference_length - sequence_length),
+            chromosome, record, sequence_length, max_reference_length - sequence_length, unconverted, "top"
+        ]
+
+        if chromosome not in chromosome_FASTA_dict:
+            chromosome_FASTA_dict[chromosome] = [sequence + "N" * (max_reference_length - sequence_length), complement + "N" * (max_reference_length - sequence_length)]
+
+    # Process converted records
+    for conversion in conversions:
+        modification_dict[conversion] = find_conversion_sites(converted_FASTA, conversion, conversion_types)
+        # Above points to record_dict[record.id] = [sequence_length, top_strand_coordinates, bottom_strand_coordinates, sequence, complement] with only unconverted record.id keys
+
+        for record, values in modification_dict[conversion].items():
+            sequence_length, top_coords, bottom_coords, sequence, complement = values
+            chromosome = record.split(f"_{unconverted}_")[0]  # Extract chromosome name
+
+            # Add **both strands** for converted records
+            for strand in ["top", "bottom"]:
+                converted_name = f"{chromosome}_{conversion}_{strand}"
+                unconverted_name = f"{chromosome}_{unconverted}_top"
+
+                record_FASTA_dict[converted_name] = [
+                    sequence + "N" * (max_reference_length - sequence_length),
+                    complement + "N" * (max_reference_length - sequence_length),
+                    chromosome, unconverted_name, sequence_length, 
+                    max_reference_length - sequence_length, conversion, strand
+                ]
+
+    print("Updated record_FASTA_dict Keys:", list(record_FASTA_dict.keys()))
+    return max_reference_length, record_FASTA_dict, chromosome_FASTA_dict
+
+
+def filter_bams_by_mapping_threshold(bam_path_list, bam_files, mapping_threshold):
+    """Filters BAM files based on mapping threshold."""
+    records_to_analyze = set()
+
+    for i, bam in enumerate(bam_path_list):
+        aligned_reads, unaligned_reads, record_counts = count_aligned_reads(bam)
+        aligned_percent = aligned_reads * 100 / (aligned_reads + unaligned_reads)
+        print(f"{aligned_percent:.2f}% of reads in {bam_files[i]} aligned successfully.")
+
+        for record, (count, percent) in record_counts.items():
+            if percent >= mapping_threshold:
+                records_to_analyze.add(record)
+
+    print(f"Analyzing the following FASTA records: {records_to_analyze}")
+    return records_to_analyze
+
+
+def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, tmp_dir, max_reference_length, device):
+    """Worker function to process a single BAM file (must be at top-level for multiprocessing)."""
+    adata_list = []
+
+    for record in records_to_analyze:
+        sample = os.path.basename(bam).split(sep=".bam")[0]
+        chromosome = record_FASTA_dict[record][2]
+        current_length = record_FASTA_dict[record][4]
+        mod_type, strand = record_FASTA_dict[record][6], record_FASTA_dict[record][7]
+
+        # Extract Base Identities
+        fwd_bases, rev_bases = extract_base_identities(bam, record, range(current_length), max_reference_length)
+
+        # Skip processing if both forward and reverse base identities are empty
+        if not fwd_bases and not rev_bases:
+            print(f"{timestamp()} [Worker {current_process().pid}] Skipping {sample} - No valid base identities for {record}.")
+            continue
+
+        merged_bin = {}
+
+        # Binarize the Base Identities if they exist
+        if fwd_bases:
+            fwd_bin = binarize_converted_base_identities(fwd_bases, strand, mod_type, bam, device)
+            merged_bin.update(fwd_bin)
+
+        if rev_bases:
+            rev_bin = binarize_converted_base_identities(rev_bases, strand, mod_type, bam, device)
+            merged_bin.update(rev_bin)
+
+        # Skip if merged_bin is empty (no valid binarized data)
+        if not merged_bin:
+            print(f"{timestamp()} [Worker {current_process().pid}] Skipping {sample} - No valid binarized data for {record}.")
+            continue
+
+        # Convert to DataFrame
+        # for key in merged_bin:
+        #     merged_bin[key] = merged_bin[key].cpu().numpy()  # Move to CPU & convert to NumPy
+        bin_df = pd.DataFrame.from_dict(merged_bin, orient='index')
+        sorted_index = sorted(bin_df.index)
+        bin_df = bin_df.reindex(sorted_index)
+
+        # One-Hot Encode Reads if there is valid data
+        one_hot_reads = {}
+
+        if fwd_bases:
+            fwd_ohe_files = ohe_batching(fwd_bases, tmp_dir, record, f"{bam_index}_fwd", batch_size=100000)
+            for ohe_file in fwd_ohe_files:
+                tmp_ohe_dict = ad.read_h5ad(ohe_file).uns
+                one_hot_reads.update(tmp_ohe_dict)
+                del tmp_ohe_dict
+
+        if rev_bases:
+            rev_ohe_files = ohe_batching(rev_bases, tmp_dir, record, f"{bam_index}_rev", batch_size=100000)
+            for ohe_file in rev_ohe_files:
+                tmp_ohe_dict = ad.read_h5ad(ohe_file).uns
+                one_hot_reads.update(tmp_ohe_dict)
+                del tmp_ohe_dict
+
+        # Skip if one_hot_reads is empty
+        if not one_hot_reads:
+            print(f"{timestamp()} [Worker {current_process().pid}] Skipping {sample} - No valid one-hot encoded data for {record}.")
+            continue
+
+        gc.collect()
+
+        # Convert One-Hot Encodings to Numpy Arrays
+        n_rows_OHE = 5
+        read_names = list(one_hot_reads.keys())
+
+        # Skip if no read names exist
+        if not read_names:
+            print(f"{timestamp()} [Worker {current_process().pid}] Skipping {sample} - No reads found in one-hot encoded data for {record}.")
+            continue
+
+        sequence_length = one_hot_reads[read_names[0]].reshape(n_rows_OHE, -1).shape[1]
+        df_A, df_C, df_G, df_T, df_N = [np.zeros((len(sorted_index), sequence_length), dtype=int) for _ in range(5)]
+
+        # Populate One-Hot Arrays
+        for j, read_name in enumerate(sorted_index):
+            if read_name in one_hot_reads:
+                one_hot_array = one_hot_reads[read_name].reshape(n_rows_OHE, -1)
+                df_A[j], df_C[j], df_G[j], df_T[j], df_N[j] = one_hot_array
+
+        # Convert to AnnData
+        X = bin_df.values.astype(np.float32)
+        adata = ad.AnnData(X)
+        adata.obs_names = bin_df.index.astype(str)
+        adata.var_names = bin_df.columns.astype(str)
+        adata.obs["Sample"] = [sample] * len(adata)
+        adata.obs["Reference"] = [chromosome] * len(adata)
+        adata.obs["Strand"] = [strand] * len(adata)
+        adata.obs["Dataset"] = [mod_type] * len(adata)
+        adata.obs["Reference_dataset_strand"] = [f"{chromosome}_{mod_type}_{strand}"] * len(adata)
+        adata.obs["Reference_strand"] = [f"{chromosome}_{strand}"] * len(adata)
+
+        # Attach One-Hot Encodings to Layers
+        adata.layers["A_binary_encoding"] = df_A
+        adata.layers["C_binary_encoding"] = df_C
+        adata.layers["G_binary_encoding"] = df_G
+        adata.layers["T_binary_encoding"] = df_T
+        adata.layers["N_binary_encoding"] = df_N
+
+        adata_list.append(adata)
+
+    return ad.concat(adata_list, join="outer") if adata_list else None
+
+def timestamp():
+    """Returns a formatted timestamp for logging."""
+    return time.strftime("[%Y-%m-%d %H:%M:%S]")
+
+
+def worker_function(bam_index, bam, records_to_analyze, shared_record_FASTA_dict, tmp_dir, h5_dir, max_reference_length, device, progress_queue):
+    """Worker function that processes a single BAM and writes the output to an H5AD file."""
+    worker_id = current_process().pid  # Get worker process ID
+    sample = os.path.basename(bam).split(sep=".bam")[0]
+
+    try:
+        print(f"{timestamp()} [Worker {worker_id}] Processing BAM: {sample}")
+
+        h5ad_path = os.path.join(h5_dir, f"{sample}.h5ad")
+        if os.path.exists(h5ad_path):
+            print(f"{timestamp()} [Worker {worker_id}] Skipping {sample}: Already processed.")
+            progress_queue.put(sample)
+            return
+
+        # Filter records specific to this BAM
+        bam_records_to_analyze = {record for record in records_to_analyze if record in shared_record_FASTA_dict}
+
+        if not bam_records_to_analyze:
+            print(f"{timestamp()} [Worker {worker_id}] No valid records to analyze for {sample}. Skipping.")
+            progress_queue.put(sample)
+            return
+
+        # Process BAM
+        adata = process_single_bam(bam_index, bam, bam_records_to_analyze, shared_record_FASTA_dict, tmp_dir, max_reference_length, device)
+
+        if adata is not None:
+            adata.write_h5ad(h5ad_path)
+            print(f"{timestamp()} [Worker {worker_id}] Completed processing for BAM: {sample}")
+
+            # Free memory
+            del adata
+            gc.collect()
+
+        progress_queue.put(sample)
+
+    except Exception as e:
+        print(f"{timestamp()} [Worker {worker_id}] ERROR while processing {sample}:\n{traceback.format_exc()}")
+        progress_queue.put(sample)  # Still signal completion to prevent deadlock
+
+def process_bams_parallel(bam_path_list, records_to_analyze, record_FASTA_dict, tmp_dir, h5_dir, num_threads, max_reference_length, device):
+    """Processes BAM files in parallel, writes each H5AD to disk, and concatenates them at the end."""
+    os.makedirs(h5_dir, exist_ok=True)  # Ensure h5_dir exists
+
+    print(f"{timestamp()} Starting parallel BAM processing with {num_threads} threads...")
+
+    # Ensure macOS uses forkserver to avoid spawning issues
+    try:
+        import multiprocessing
+        multiprocessing.set_start_method("forkserver", force=True)
+    except RuntimeError:
+        print(f"{timestamp()} [WARNING] Multiprocessing context already set. Skipping set_start_method.")
+
+    with Manager() as manager:
+        progress_queue = manager.Queue()
+        shared_record_FASTA_dict = manager.dict(record_FASTA_dict)
+
+        with Pool(processes=num_threads) as pool:
+            results = [
+                pool.apply_async(worker_function, (i, bam, records_to_analyze, shared_record_FASTA_dict, tmp_dir, h5_dir, max_reference_length, device, progress_queue))
+                for i, bam in enumerate(bam_path_list)
+            ]
+
+            print(f"{timestamp()} Submitted {len(bam_path_list)} BAMs for processing.")
+
+            # Track completed BAMs
+            completed_bams = set()
+            while len(completed_bams) < len(bam_path_list):
+                try:
+                    processed_bam = progress_queue.get(timeout=2400)  # Wait for a finished BAM
+                    completed_bams.add(processed_bam)
+                except Exception as e:
+                    print(f"{timestamp()} [ERROR] Timeout waiting for worker process. Possible crash? {e}")
+
+            pool.close()
+            pool.join()  # Ensure all workers finish
+
+    # Final Concatenation Step
+    h5ad_files = [os.path.join(h5_dir, f) for f in os.listdir(h5_dir) if f.endswith(".h5ad")]
+
+    if not h5ad_files:
+        print(f"{timestamp()} No valid H5AD files generated. Exiting.")
+        return None
+
+    print(f"{timestamp()} Concatenating {len(h5ad_files)} H5AD files into final output...")
+    final_adata = ad.concat([ad.read_h5ad(f) for f in h5ad_files], join="outer")
+
+    print(f"{timestamp()} Successfully generated final AnnData object.")
+    return final_adata

smftools/informatics/helpers/demux_and_index_BAM.py

@@ -0,0 +1,52 @@
+## demux_and_index_BAM
+
+def demux_and_index_BAM(aligned_sorted_BAM, split_dir, bam_suffix, barcode_kit, barcode_both_ends, trim, fasta, make_bigwigs, threads):
+    """
+    A wrapper function for splitting BAMS and indexing them.
+    Parameters:
+        aligned_sorted_BAM (str): A string representing the file path of the aligned_sorted BAM file.
+        split_dir (str): A string representing the file path to the directory to split the BAMs into.
+        bam_suffix (str): A suffix to add to the bam file.
+        barcode_kit (str): Name of barcoding kit.
+        barcode_both_ends (bool): Whether to require both ends to be barcoded.
+        trim (bool): Whether to trim off barcodes after demultiplexing.
+        fasta (str): File path to the reference genome to align to.
+        make_bigwigs (bool): Whether to make bigwigs
+        threads (int): Number of threads to use.
+    
+    Returns:
+        bam_files (list): List of split BAM file path strings
+            Splits an input BAM file on barcode value and makes a BAM index file.
+    """
+    from .. import readwrite
+    import os
+    import subprocess
+    import glob
+    from .make_dirs import make_dirs
+    
+    input_bam = aligned_sorted_BAM + bam_suffix
+    command = ["dorado", "demux", "--kit-name", barcode_kit]
+    if barcode_both_ends:
+        command.append("--barcode-both-ends")
+    if not trim:
+        command.append("--no-trim")
+    if threads:
+        command += ["-t", str(threads)]
+    else:
+        pass
+    command += ["--emit-summary", "--sort-bam", "--output-dir", split_dir]
+    command.append(input_bam)
+    command_string = ' '.join(command)
+    print(f"Running: {command_string}")
+    subprocess.run(command)
+
+    # Make a BAM index file for the BAMs in that directory
+    bam_pattern = '*' + bam_suffix
+    bam_files = glob.glob(os.path.join(split_dir, bam_pattern))
+    bam_files = [bam for bam in bam_files if '.bai' not in bam and 'unclassified' not in bam]
+    bam_files.sort()
+
+    if not bam_files:
+        raise FileNotFoundError(f"No BAM files found in {split_dir} with suffix {bam_suffix}")
+    
+    return bam_files

smftools/informatics/helpers/extract_base_identities.py

@@ -1,57 +1,44 @@
-## extract_base_identities
-
-# General
 def extract_base_identities(bam_file, chromosome, positions, max_reference_length):
     """
-    Extracts the base identities from every position within the mapped reads that have a reference coordinate
+    Efficiently extracts base identities from mapped reads with reference coordinates.
 
     Parameters:
-        bam (str): File path to the BAM file to align (excluding the file suffix).
-        chromosome (str): A string representing the name of the record within the reference FASTA.
-        positions (list): A list of position coordinates within the record to extract.
-        max_reference_length (int): The maximum length of a record in the reference set.
+        bam_file (str): Path to the BAM file.
+        chromosome (str): Name of the reference chromosome.
+        positions (list): Positions to extract (0-based).
+        max_reference_length (int): Maximum reference length for padding.
 
     Returns:
-        fwd_base_identities (dict): A dictionary, keyed by read name, that points to a list of base identities from forward mapped reads. If the read does not contain that position, fill the list at that index with a N value.
-        rev_base_identities (dict): A dictionary, keyed by read name, that points to a list of base identities from reverse mapped reads. If the read does not contain that position, fill the list at that index with a N value.
+        dict: Base identities from forward mapped reads.
+        dict: Base identities from reverse mapped reads.
     """
-    from .. import readwrite
     import pysam
-    from tqdm import tqdm
-    
+    import numpy as np
+    from collections import defaultdict
+    import time
+
+    timestamp = time.strftime("[%Y-%m-%d %H:%M:%S]")
+
     positions = set(positions)
-    # Initialize a base identity dictionary that will hold key-value pairs that are: key (read-name) and value (list of base identities at positions of interest)
-    fwd_base_identities = {}
-    rev_base_identities = {}
-    # Open the postion sorted BAM file
-    print('{0}: Reading BAM file: {1}'.format(readwrite.time_string(), bam_file))
+    fwd_base_identities = defaultdict(lambda: np.full(max_reference_length, 'N', dtype='<U1'))
+    rev_base_identities = defaultdict(lambda: np.full(max_reference_length, 'N', dtype='<U1'))
+
+    #print(f"{timestamp} Reading reads from {chromosome} BAM file: {bam_file}")
     with pysam.AlignmentFile(bam_file, "rb") as bam:
-        # Iterate over every read in the bam that comes from the chromosome of interest
-        print('{0}: Iterating over reads in bam'.format(readwrite.time_string()))
         total_reads = bam.mapped
-        for read in tqdm(bam.fetch(chromosome), desc='Extracting base identities from reads in BAM', total=total_reads):  
-            # Only iterate over mapped reads
-            if read.is_mapped:
-                # Get sequence of read. PySam reports fwd mapped reads as the true read sequence. Pysam reports rev mapped reads as the reverse complement of the read.
-                query_sequence = read.query_sequence
-                # If the read aligned as a reverse complement, mark that the read is reversed
-                if read.is_reverse:
-                    # Initialize the read key in a temp base_identities dictionary by pointing to a N filled list of length reference_length.
-                    rev_base_identities[read.query_name] = ['N'] * max_reference_length
-                    # Iterate over a list of tuples for the given read. The tuples contain the 0-indexed position relative to the read.query_sequence start, as well the 0-based index relative to the reference.
-                    for read_position, reference_position in read.get_aligned_pairs(matches_only=True):
-                        # If the aligned read's reference coordinate is in the positions set and if the read position was successfully mapped
-                        if reference_position in positions and read_position:
-                            # get the base_identity in the read corresponding to that position
-                            rev_base_identities[read.query_name][reference_position] = query_sequence[read_position]
-                else:
-                    # Initialize the read key in a temp base_identities dictionary by pointing to a N filled list of length reference_length.
-                    fwd_base_identities[read.query_name] = ['N'] * max_reference_length
-                    # Iterate over a list of tuples for the given read. The tuples contain the 0-indexed position relative to the read.query_sequence start, as well the 0-based index relative to the reference.
-                    for read_position, reference_position in read.get_aligned_pairs(matches_only=True):
-                        # If the aligned read's reference coordinate is in the positions set and if the read position was successfully mapped
-                        if reference_position in positions and read_position:
-                            # get the base_identity in the read corresponding to that position
-                            fwd_base_identities[read.query_name][reference_position] = query_sequence[read_position]
-
-    return fwd_base_identities, rev_base_identities
+        for read in bam.fetch(chromosome):
+            if not read.is_mapped:
+                continue  # Skip unmapped reads
+
+            read_name = read.query_name
+            query_sequence = read.query_sequence
+            base_dict = rev_base_identities if read.is_reverse else fwd_base_identities
+
+            # Use get_aligned_pairs directly with positions filtering
+            aligned_pairs = read.get_aligned_pairs(matches_only=True)
+
+            for read_position, reference_position in aligned_pairs:
+                if reference_position in positions:
+                    base_dict[read_name][reference_position] = query_sequence[read_position]
+
+    return dict(fwd_base_identities), dict(rev_base_identities)

smftools/informatics/helpers/extract_mods.py

@@ -1,6 +1,6 @@
 ## extract_mods
 
-def extract_mods(thresholds, mod_tsv_dir, split_dir, bam_suffix):
+def extract_mods(thresholds, mod_tsv_dir, split_dir, bam_suffix, skip_unclassified=True, modkit_summary=False, threads=None):
     """
     Takes all of the aligned, sorted, split modified BAM files and runs Nanopore Modkit Extract to load the modification data into zipped TSV files
 
@@ -9,6 +9,9 @@ def extract_mods(thresholds, mod_tsv_dir, split_dir, bam_suffix):
         mod_tsv_dir (str): A string representing the file path to the directory to hold the modkit extract outputs.
         split_dit (str): A string representing the file path to the directory containing the converted aligned_sorted_split BAM files.
         bam_suffix (str): The suffix to use for the BAM file.
+        skip_unclassified (bool): Whether to skip unclassified bam file for modkit extract command
+        modkit_summary (bool): Whether to run and display modkit summary
+        threads (int): Number of threads to use
 
     Returns:
         None
@@ -23,29 +26,58 @@ def extract_mods(thresholds, mod_tsv_dir, split_dir, bam_suffix):
     os.chdir(mod_tsv_dir)
     filter_threshold, m6A_threshold, m5C_threshold, hm5C_threshold = thresholds
     bam_files = glob.glob(os.path.join(split_dir, f"*{bam_suffix}"))
+
+    if threads:
+        threads = str(threads)
+    else:
+        pass
+
     for input_file in bam_files:
         print(input_file)
         # Extract the file basename
         file_name = os.path.basename(input_file)
-        # Construct the output TSV file path
-        output_tsv_temp = os.path.join(mod_tsv_dir, file_name)
-        output_tsv = output_tsv_temp.replace(bam_suffix, "") + "_extract.tsv"
-        # Run modkit summary
-        subprocess.run(["modkit", "summary", input_file])
-        # Run modkit extract
-        subprocess.run([
-            "modkit", "extract",
-            "--filter-threshold", f'{filter_threshold}',
-            "--mod-thresholds", f"m:{m5C_threshold}",
-            "--mod-thresholds", f"a:{m6A_threshold}",
-            "--mod-thresholds", f"h:{hm5C_threshold}",
-            input_file, "null",
-            "--read-calls", output_tsv
-        ])
-        # Zip the output TSV
-        print(f'zipping {output_tsv}')
-        with zipfile.ZipFile(f"{output_tsv}.zip", 'w', zipfile.ZIP_DEFLATED) as zipf:
-            zipf.write(output_tsv, os.path.basename(output_tsv))
-        # Remove the non-zipped TSV
-        print(f'removing {output_tsv}')
-        os.remove(output_tsv)
+        if skip_unclassified and "unclassified" in file_name:
+            print("Skipping modkit extract on unclassified reads")
+        else:
+            # Construct the output TSV file path
+            output_tsv_temp = os.path.join(mod_tsv_dir, file_name)
+            output_tsv = output_tsv_temp.replace(bam_suffix, "") + "_extract.tsv"
+            if os.path.exists(f"{output_tsv}.gz"):
+                print(f"{output_tsv}.gz already exists, skipping modkit extract")
+            else:
+                print(f"Extracting modification data from {input_file}")
+                if modkit_summary:
+                    # Run modkit summary
+                    subprocess.run(["modkit", "summary", input_file])
+                else:
+                    pass
+                # Run modkit extract
+                if threads:
+                    extract_command = [
+                        "modkit", "extract",
+                        "calls", "--mapped-only",
+                        "--filter-threshold", f'{filter_threshold}',
+                        "--mod-thresholds", f"m:{m5C_threshold}",
+                        "--mod-thresholds", f"a:{m6A_threshold}",
+                        "--mod-thresholds", f"h:{hm5C_threshold}",
+                        "-t", threads,
+                        input_file, output_tsv
+                        ]
+                else:
+                    extract_command = [
+                        "modkit", "extract",
+                        "calls", "--mapped-only",
+                        "--filter-threshold", f'{filter_threshold}',
+                        "--mod-thresholds", f"m:{m5C_threshold}",
+                        "--mod-thresholds", f"a:{m6A_threshold}",
+                        "--mod-thresholds", f"h:{hm5C_threshold}",
+                        input_file, output_tsv
+                        ]                    
+                subprocess.run(extract_command)
+                # Zip the output TSV
+                print(f'zipping {output_tsv}')
+                if threads:
+                    zip_command = ["pigz", "-f", "-p", threads, output_tsv]
+                else:
+                    zip_command = ["pigz", "-f", output_tsv]
+                subprocess.run(zip_command, check=True)

smftools/informatics/helpers/extract_read_features_from_bam.py

@@ -0,0 +1,31 @@
+# extract_read_features_from_bam
+
+def extract_read_features_from_bam(bam_file_path):
+    """
+    Make a dict of reads from a bam that points to a list of read metrics: read length, read median Q-score, reference length.
+    Params:
+        bam_file_path (str):
+    Returns:
+        read_metrics (dict)
+    """
+    import pysam
+    import numpy as np
+    # Open the BAM file
+    print(f'Extracting read features from BAM: {bam_file_path}')
+    with pysam.AlignmentFile(bam_file_path, "rb") as bam_file:
+        read_metrics = {}
+        reference_lengths = bam_file.lengths  # List of lengths for each reference (chromosome)
+        for read in bam_file:
+            # Skip unmapped reads
+            if read.is_unmapped:
+                continue
+            # Extract the read metrics
+            read_quality = read.query_qualities
+            median_read_quality = np.median(read_quality)
+            # Extract the reference (chromosome) name and its length
+            reference_name = read.reference_name
+            reference_index = bam_file.references.index(reference_name)
+            reference_length = reference_lengths[reference_index]
+            read_metrics[read.query_name] = [read.query_length, median_read_quality, reference_length]
+
+    return read_metrics