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

smftools/informatics/converted_BAM_to_adata.py

@@ -1,61 +1,65 @@
-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 shutil
+import time
 import traceback
-import gzip
+from multiprocessing import Manager, Pool, current_process
+from pathlib import Path
+from typing import Iterable, Optional, Union
+
+import anndata as ad
+import numpy as np
+import pandas as pd
 import torch
 
-import shutil
-from pathlib import Path
-from typing import Union, Iterable, Optional
+from smftools.logging_utils import get_logger
 
-from ..readwrite import make_dirs, safe_write_h5ad
+from ..readwrite import make_dirs
+from .bam_functions import count_aligned_reads, extract_base_identities
 from .binarize_converted_base_identities import binarize_converted_base_identities
 from .fasta_functions import find_conversion_sites
-from .bam_functions import count_aligned_reads, extract_base_identities
 from .ohe import ohe_batching
 
+logger = get_logger(__name__)
+
 if __name__ == "__main__":
     multiprocessing.set_start_method("forkserver", force=True)
 
-def converted_BAM_to_adata(converted_FASTA, 
-                              split_dir,
-                              output_dir,
-                              input_already_demuxed,
-                              mapping_threshold, 
-                              experiment_name, 
-                              conversions, 
-                              bam_suffix, 
-                              device='cpu', 
-                              num_threads=8, 
-                              deaminase_footprinting=False,
-                              delete_intermediates=True,
-                              double_barcoded_path = None,
-):
-    """
-    Converts BAM files into an AnnData object by binarizing modified base identities.
 
-    Parameters:
-        converted_FASTA (Path): Path to the converted FASTA reference.
-        split_dir (Path): Directory containing converted BAM files.
-        output_dir (Path): Directory of the output dir
-        input_already_demuxed (bool): Whether input reads were originally demuxed
-        mapping_threshold (float): Minimum fraction of aligned reads required for inclusion.
-        experiment_name (str): Name for the output AnnData object.
-        conversions (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.
-        deaminase_footprinting (bool): Whether the footprinting was done with a direct deamination chemistry.
-        double_barcoded_path (Path): Path to dorado demux summary file of double ended barcodes
+def converted_BAM_to_adata(
+    converted_FASTA: str | Path,
+    split_dir: Path,
+    output_dir: Path,
+    input_already_demuxed: bool,
+    mapping_threshold: float,
+    experiment_name: str,
+    conversions: list[str],
+    bam_suffix: str,
+    device: str | torch.device = "cpu",
+    num_threads: int = 8,
+    deaminase_footprinting: bool = False,
+    delete_intermediates: bool = True,
+    double_barcoded_path: Path | None = None,
+) -> tuple[ad.AnnData | None, Path]:
+    """Convert BAM files into an AnnData object by binarizing modified base identities.
+
+    Args:
+        converted_FASTA: Path to the converted FASTA reference.
+        split_dir: Directory containing converted BAM files.
+        output_dir: Output directory for intermediate and final files.
+        input_already_demuxed: Whether input reads were originally demultiplexed.
+        mapping_threshold: Minimum fraction of aligned reads required for inclusion.
+        experiment_name: Name for the output AnnData object.
+        conversions: List of modification types (e.g., ``["unconverted", "5mC", "6mA"]``).
+        bam_suffix: File suffix for BAM files.
+        device: Torch device or device string.
+        num_threads: Number of parallel processing threads.
+        deaminase_footprinting: Whether the footprinting used direct deamination chemistry.
+        delete_intermediates: Whether to remove intermediate files after processing.
+        double_barcoded_path: Path to dorado demux summary file of double-ended barcodes.
 
     Returns:
-        str: Path to the final AnnData object.
+        tuple[anndata.AnnData | None, Path]: The AnnData object (if generated) and its path.
     """
     if torch.cuda.is_available():
         device = torch.device("cuda")
@@ -64,69 +68,88 @@ def converted_BAM_to_adata(converted_FASTA,
     else:
         device = torch.device("cpu")
 
-    print(f"Using device: {device}")
+    logger.debug(f"Using device: {device}")
 
     ## Set Up Directories and File Paths
-    h5_dir = output_dir / 'h5ads'
-    tmp_dir = output_dir / 'tmp'
+    h5_dir = output_dir / "h5ads"
+    tmp_dir = output_dir / "tmp"
     final_adata = None
-    final_adata_path = h5_dir / f'{experiment_name}.h5ad.gz'
+    final_adata_path = h5_dir / f"{experiment_name}.h5ad.gz"
 
     if final_adata_path.exists():
-        print(f"{final_adata_path} already exists. Using existing AnnData object.")
+        logger.debug(f"{final_adata_path} already exists. Using existing AnnData object.")
         return final_adata, final_adata_path
 
     make_dirs([h5_dir, tmp_dir])
 
     bam_files = sorted(
-        p for p in split_dir.iterdir()
-        if p.is_file()
-        and p.suffix == ".bam"
-        and "unclassified" not in p.name
+        p
+        for p in split_dir.iterdir()
+        if p.is_file() and p.suffix == ".bam" and "unclassified" not in p.name
     )
 
-    bam_path_list = [split_dir / f for f in bam_files]
-    print(f"Found {len(bam_files)} BAM files: {bam_files}")
+    bam_path_list = bam_files
+    logger.info(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, conversions, deaminase_footprinting)
+    max_reference_length, record_FASTA_dict, chromosome_FASTA_dict = process_conversion_sites(
+        converted_FASTA, conversions, deaminase_footprinting
+    )
 
     ## Filter BAM Files by Mapping Threshold
-    records_to_analyze = filter_bams_by_mapping_threshold(bam_path_list, bam_files, 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, chromosome_FASTA_dict, tmp_dir, h5_dir, num_threads, max_reference_length, device, deaminase_footprinting)
+    final_adata = process_bams_parallel(
+        bam_path_list,
+        records_to_analyze,
+        record_FASTA_dict,
+        chromosome_FASTA_dict,
+        tmp_dir,
+        h5_dir,
+        num_threads,
+        max_reference_length,
+        device,
+        deaminase_footprinting,
+    )
 
-    final_adata.uns['References'] = {}
+    final_adata.uns["References"] = {}
     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
-        final_adata.uns['References'][f'{chromosome}_FASTA_sequence'] = seq
+        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
+        final_adata.uns["References"][f"{chromosome}_FASTA_sequence"] = seq
 
     final_adata.obs_names_make_unique()
     cols = final_adata.obs.columns
 
     # Make obs cols categorical
     for col in cols:
-        final_adata.obs[col] = final_adata.obs[col].astype('category')
+        final_adata.obs[col] = final_adata.obs[col].astype("category")
 
     if input_already_demuxed:
         final_adata.obs["demux_type"] = ["already"] * final_adata.shape[0]
         final_adata.obs["demux_type"] = final_adata.obs["demux_type"].astype("category")
     else:
         from .h5ad_functions import add_demux_type_annotation
+
         double_barcoded_reads = double_barcoded_path / "barcoding_summary.txt"
+        logger.info("Adding demux type to each read")
         add_demux_type_annotation(final_adata, double_barcoded_reads)
 
     ## Delete intermediate h5ad files and temp directories
     if delete_intermediates:
+        logger.info("Deleting intermediate h5ad files")
         delete_intermediate_h5ads_and_tmpdir(h5_dir, tmp_dir)
-    
+
     return final_adata, final_adata_path
 
 
-def process_conversion_sites(converted_FASTA, conversions=['unconverted', '5mC'], deaminase_footprinting=False):
+def process_conversion_sites(
+    converted_FASTA, conversions=["unconverted", "5mC"], deaminase_footprinting=False
+):
     """
     Extracts conversion sites and determines the max reference length.
 
@@ -147,7 +170,9 @@ def process_conversion_sites(converted_FASTA, conversions=['unconverted', '5mC']
     conversion_types = conversions[1:]
 
     # Process the unconverted sequence once
-    modification_dict[unconverted] = find_conversion_sites(converted_FASTA, unconverted, conversions, deaminase_footprinting)
+    modification_dict[unconverted] = find_conversion_sites(
+        converted_FASTA, unconverted, conversions, deaminase_footprinting
+    )
     # Above points to record_dict[record.id] = [sequence_length, [], [], sequence, complement] with only unconverted record.id keys
 
     # Get **max sequence length** from unconverted records
@@ -166,15 +191,25 @@ def process_conversion_sites(converted_FASTA, conversions=['unconverted', '5mC']
         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"
+            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)]
+            chromosome_FASTA_dict[chromosome] = [
+                sequence + "N" * (max_reference_length - sequence_length),
+                complement + "N" * (max_reference_length - sequence_length),
+            ]
 
     # Process converted records
     for conversion in conversion_types:
-        modification_dict[conversion] = find_conversion_sites(converted_FASTA, conversion, conversions, deaminase_footprinting)
+        modification_dict[conversion] = find_conversion_sites(
+            converted_FASTA, conversion, conversions, deaminase_footprinting
+        )
         # 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():
@@ -193,11 +228,15 @@ def process_conversion_sites(converted_FASTA, conversions=['unconverted', '5mC']
                 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
+                    chromosome,
+                    unconverted_name,
+                    sequence_length,
+                    max_reference_length - sequence_length,
+                    conversion,
+                    strand,
                 ]
 
-    print("Updated record_FASTA_dict Keys:", list(record_FASTA_dict.keys()))
+    logger.debug("Updated record_FASTA_dict Keys:", list(record_FASTA_dict.keys()))
     return max_reference_length, record_FASTA_dict, chromosome_FASTA_dict
 
 
@@ -214,11 +253,21 @@ def filter_bams_by_mapping_threshold(bam_path_list, bam_files, mapping_threshold
             if percent >= mapping_threshold:
                 records_to_analyze.add(record)
 
-    print(f"Analyzing the following FASTA records: {records_to_analyze}")
+    logger.info(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, chromosome_FASTA_dict, tmp_dir, max_reference_length, device, deaminase_footprinting):
+def process_single_bam(
+    bam_index,
+    bam,
+    records_to_analyze,
+    record_FASTA_dict,
+    chromosome_FASTA_dict,
+    tmp_dir,
+    max_reference_length,
+    device,
+    deaminase_footprinting,
+):
     """Worker function to process a single BAM file (must be at top-level for multiprocessing)."""
     adata_list = []
 
@@ -230,34 +279,58 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
         sequence = chromosome_FASTA_dict[chromosome][0]
 
         # Extract Base Identities
-        fwd_bases, rev_bases, mismatch_counts_per_read, mismatch_trend_per_read = extract_base_identities(bam, record, range(current_length), max_reference_length, sequence)
+        fwd_bases, rev_bases, mismatch_counts_per_read, mismatch_trend_per_read = (
+            extract_base_identities(
+                bam, record, range(current_length), max_reference_length, sequence
+            )
+        )
         mismatch_trend_series = pd.Series(mismatch_trend_per_read)
 
         # 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}.")
+            logger.debug(
+                f"[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,deaminase_footprinting, mismatch_trend_per_read)
+            fwd_bin = binarize_converted_base_identities(
+                fwd_bases,
+                strand,
+                mod_type,
+                bam,
+                device,
+                deaminase_footprinting,
+                mismatch_trend_per_read,
+            )
             merged_bin.update(fwd_bin)
 
         if rev_bases:
-            rev_bin = binarize_converted_base_identities(rev_bases, strand, mod_type, bam, device, deaminase_footprinting, mismatch_trend_per_read)
+            rev_bin = binarize_converted_base_identities(
+                rev_bases,
+                strand,
+                mod_type,
+                bam,
+                device,
+                deaminase_footprinting,
+                mismatch_trend_per_read,
+            )
             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}.")
+            logger.debug(
+                f"[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')
+        bin_df = pd.DataFrame.from_dict(merged_bin, orient="index")
         sorted_index = sorted(bin_df.index)
         bin_df = bin_df.reindex(sorted_index)
 
@@ -265,14 +338,18 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
         one_hot_reads = {}
 
         if fwd_bases:
-            fwd_ohe_files = ohe_batching(fwd_bases, tmp_dir, record, f"{bam_index}_fwd", batch_size=100000)
+            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)
+            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)
@@ -280,7 +357,9 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
 
         # 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}.")
+            logger.debug(
+                f"[Worker {current_process().pid}] Skipping {sample} - No valid one-hot encoded data for {record}."
+            )
             continue
 
         gc.collect()
@@ -291,11 +370,15 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
 
         # 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}.")
+            logger.debug(
+                f"[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)]
+        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):
@@ -310,8 +393,8 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
         adata.var_names = bin_df.columns.astype(str)
         adata.obs["Sample"] = [sample] * len(adata)
         try:
-            barcode = sample.split('barcode')[1]
-        except:
+            barcode = sample.split("barcode")[1]
+        except Exception:
             barcode = np.nan
         adata.obs["Barcode"] = [int(barcode)] * len(adata)
         adata.obs["Barcode"] = adata.obs["Barcode"].astype(str)
@@ -323,49 +406,76 @@ def process_single_bam(bam_index, bam, records_to_analyze, record_FASTA_dict, ch
         adata.obs["Read_mismatch_trend"] = adata.obs_names.map(mismatch_trend_series)
 
         # 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.layers["A_binary_sequence_encoding"] = df_A
+        adata.layers["C_binary_sequence_encoding"] = df_C
+        adata.layers["G_binary_sequence_encoding"] = df_G
+        adata.layers["T_binary_sequence_encoding"] = df_T
+        adata.layers["N_binary_sequence_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, chromosome_FASTA_dict, tmp_dir, h5_dir, max_reference_length, device, deaminase_footprinting, progress_queue):
+def worker_function(
+    bam_index,
+    bam,
+    records_to_analyze,
+    shared_record_FASTA_dict,
+    chromosome_FASTA_dict,
+    tmp_dir,
+    h5_dir,
+    max_reference_length,
+    device,
+    deaminase_footprinting,
+    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 = bam.stem
 
     try:
-        print(f"{timestamp()} [Worker {worker_id}] Processing BAM: {sample}")
+        logger.info(f"[Worker {worker_id}] Processing BAM: {sample}")
 
         h5ad_path = h5_dir / bam.with_suffix(".h5ad").name
         if h5ad_path.exists():
-            print(f"{timestamp()} [Worker {worker_id}] Skipping {sample}: Already processed.")
+            logger.debug(f"[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}
+        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.")
+            logger.debug(
+                f"[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, chromosome_FASTA_dict, tmp_dir, max_reference_length, device, deaminase_footprinting)
+        adata = process_single_bam(
+            bam_index,
+            bam,
+            bam_records_to_analyze,
+            shared_record_FASTA_dict,
+            chromosome_FASTA_dict,
+            tmp_dir,
+            max_reference_length,
+            device,
+            deaminase_footprinting,
+        )
 
         if adata is not None:
             adata.write_h5ad(str(h5ad_path))
-            print(f"{timestamp()} [Worker {worker_id}] Completed processing for BAM: {sample}")
+            logger.info(f"[Worker {worker_id}] Completed processing for BAM: {sample}")
 
             # Free memory
             del adata
@@ -373,22 +483,37 @@ def worker_function(bam_index, bam, records_to_analyze, shared_record_FASTA_dict
 
         progress_queue.put(sample)
 
-    except Exception as e:
-        print(f"{timestamp()} [Worker {worker_id}] ERROR while processing {sample}:\n{traceback.format_exc()}")
+    except Exception:
+        logger.warning(
+            f"[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, chromosome_FASTA_dict, tmp_dir, h5_dir, num_threads, max_reference_length, device, deaminase_footprinting):
+
+def process_bams_parallel(
+    bam_path_list,
+    records_to_analyze,
+    record_FASTA_dict,
+    chromosome_FASTA_dict,
+    tmp_dir,
+    h5_dir,
+    num_threads,
+    max_reference_length,
+    device,
+    deaminase_footprinting,
+):
     """Processes BAM files in parallel, writes each H5AD to disk, and concatenates them at the end."""
     make_dirs(h5_dir)  # Ensure h5_dir exists
 
-    print(f"{timestamp()} Starting parallel BAM processing with {num_threads} threads...")
+    logger.info(f"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.")
+        logger.warning(f"Multiprocessing context already set. Skipping set_start_method.")
 
     with Manager() as manager:
         progress_queue = manager.Queue()
@@ -396,11 +521,26 @@ def process_bams_parallel(bam_path_list, records_to_analyze, 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, chromosome_FASTA_dict, tmp_dir, h5_dir, max_reference_length, device, deaminase_footprinting, progress_queue))
+                pool.apply_async(
+                    worker_function,
+                    (
+                        i,
+                        bam,
+                        records_to_analyze,
+                        shared_record_FASTA_dict,
+                        chromosome_FASTA_dict,
+                        tmp_dir,
+                        h5_dir,
+                        max_reference_length,
+                        device,
+                        deaminase_footprinting,
+                        progress_queue,
+                    ),
+                )
                 for i, bam in enumerate(bam_path_list)
             ]
 
-            print(f"{timestamp()} Submitted {len(bam_path_list)} BAMs for processing.")
+            logger.info(f"Submitted {len(bam_path_list)} BAMs for processing.")
 
             # Track completed BAMs
             completed_bams = set()
@@ -409,24 +549,25 @@ def process_bams_parallel(bam_path_list, records_to_analyze, record_FASTA_dict,
                     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}")
+                    logger.error(f"Timeout waiting for worker process. Possible crash? {e}")
 
             pool.close()
             pool.join()  # Ensure all workers finish
 
     # Final Concatenation Step
-    h5ad_files = [h5_dir / f for f in h5_dir.iterdir() if f.suffix == ".h5ad"]
+    h5ad_files = [f for f in h5_dir.iterdir() if f.suffix == ".h5ad"]
 
     if not h5ad_files:
-        print(f"{timestamp()} No valid H5AD files generated. Exiting.")
+        logger.debug(f"No valid H5AD files generated. Exiting.")
         return None
 
-    print(f"{timestamp()} Concatenating {len(h5ad_files)} H5AD files into final output...")
+    logger.info(f"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.")
+    logger.info(f"Successfully generated final AnnData object.")
     return final_adata
 
+
 def delete_intermediate_h5ads_and_tmpdir(
     h5_dir: Union[str, Path, Iterable[str], None],
     tmp_dir: Optional[Union[str, Path]] = None,
@@ -450,25 +591,27 @@ def delete_intermediate_h5ads_and_tmpdir(
     verbose : bool
         Print progress / warnings.
     """
+
     # Helper: remove a single file path (Path-like or string)
     def _maybe_unlink(p: Path):
+        """Remove a file path if it exists and is a file."""
         if not p.exists():
             if verbose:
-                print(f"[skip] not found: {p}")
+                logger.debug(f"[skip] not found: {p}")
             return
         if not p.is_file():
             if verbose:
-                print(f"[skip] not a file: {p}")
+                logger.debug(f"[skip] not a file: {p}")
             return
         if dry_run:
-            print(f"[dry-run] would remove file: {p}")
+            logger.debug(f"[dry-run] would remove file: {p}")
             return
         try:
             p.unlink()
             if verbose:
-                print(f"Removed file: {p}")
+                logger.info(f"Removed file: {p}")
         except Exception as e:
-            print(f"[error] failed to remove file {p}: {e}")
+            logger.warning(f"[error] failed to remove file {p}: {e}")
 
     # Handle h5_dir input (directory OR iterable of file paths)
     if h5_dir is not None:
@@ -483,7 +626,7 @@ def delete_intermediate_h5ads_and_tmpdir(
                 else:
                     if verbose:
                         # optional: comment this out if too noisy
-                        print(f"[skip] not matching pattern: {p.name}")
+                        logger.debug(f"[skip] not matching pattern: {p.name}")
         else:
             # treat as iterable of file paths
             for f in h5_dir:
@@ -493,25 +636,25 @@ def delete_intermediate_h5ads_and_tmpdir(
                     _maybe_unlink(p)
                 else:
                     if verbose:
-                        print(f"[skip] not matching pattern or not a file: {p}")
+                        logger.debug(f"[skip] not matching pattern or not a file: {p}")
 
     # Remove tmp_dir recursively (if provided)
     if tmp_dir is not None:
         td = Path(tmp_dir)
         if not td.exists():
             if verbose:
-                print(f"[skip] tmp_dir not found: {td}")
+                logger.debug(f"[skip] tmp_dir not found: {td}")
         else:
             if not td.is_dir():
                 if verbose:
-                    print(f"[skip] tmp_dir is not a directory: {td}")
+                    logger.debug(f"[skip] tmp_dir is not a directory: {td}")
             else:
                 if dry_run:
-                    print(f"[dry-run] would remove directory tree: {td}")
+                    logger.debug(f"[dry-run] would remove directory tree: {td}")
                 else:
                     try:
                         shutil.rmtree(td)
                         if verbose:
-                            print(f"Removed directory tree: {td}")
+                            logger.info(f"Removed directory tree: {td}")
                     except Exception as e:
-                        print(f"[error] failed to remove tmp dir {td}: {e}")
+                        logger.warning(f"[error] failed to remove tmp dir {td}: {e}")