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

smftools/informatics/h5ad_functions.py

@@ -0,0 +1,197 @@
+from pathlib import Path
+import pandas as pd
+import numpy as np
+import scipy.sparse as sp
+from typing import Optional, List, Dict, Union
+
+def add_demux_type_annotation(
+    adata,
+    double_demux_source,
+    sep: str = "\t",
+    read_id_col: str = "read_id",
+    barcode_col: str = "barcode",
+):
+    """
+    Add adata.obs["demux_type"]:
+        - "double" if read_id appears in the *double demux* TSV
+        - "single" otherwise
+
+    Rows where barcode == "unclassified" in the demux TSV are ignored.
+
+    Parameters
+    ----------
+    adata : AnnData
+        AnnData object whose obs_names are read_ids.
+    double_demux_source : str | Path | list[str]
+        Either:
+          - path to a TSV/TXT of dorado demux results
+          - a list of read_ids
+    """
+
+    # -----------------------------
+    # If it's a file → load TSV
+    # -----------------------------
+    if isinstance(double_demux_source, (str, Path)):
+        file_path = Path(double_demux_source)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File does not exist: {file_path}")
+
+        df = pd.read_csv(file_path, sep=sep, dtype=str)
+
+        # If the file has only one column → treat as a simple read list
+        if df.shape[1] == 1:
+            read_ids = df.iloc[:, 0].tolist()
+        else:
+            # Validate columns
+            if read_id_col not in df.columns:
+                raise ValueError(f"TSV must contain a '{read_id_col}' column.")
+            if barcode_col not in df.columns:
+                raise ValueError(f"TSV must contain a '{barcode_col}' column.")
+
+            # Drop unclassified reads
+            df = df[df[barcode_col].str.lower() != "unclassified"]
+
+            # Extract read_ids
+            read_ids = df[read_id_col].tolist()
+
+    # -----------------------------
+    # If user supplied list-of-ids
+    # -----------------------------
+    else:
+        read_ids = list(double_demux_source)
+
+    # Deduplicate for speed
+    double_set = set(read_ids)
+
+    # Boolean lookup in AnnData
+    is_double = adata.obs_names.isin(double_set)
+
+    adata.obs["demux_type"] = np.where(is_double, "double", "single")
+    adata.obs["demux_type"] = adata.obs["demux_type"].astype("category")
+
+    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 = "read_lenth_and_mapping_qc_performed",
+    extract_read_features_from_bam_callable = None,
+    bypass: bool = False,
+    force_redo: bool = True
+):
+    """
+    Populate adata.obs with read/mapping QC columns.
+
+    Parameters
+    ----------
+    adata
+        AnnData to annotate (modified in-place).
+    bam_files
+        Optional list of BAM files to extract metrics from. Ignored if read_metrics supplied.
+    read_metrics
+        Optional dict mapping obs_name -> [read_length, read_quality, reference_length, mapped_length, mapping_quality]
+        If provided, this will be used directly and bam_files will be ignored.
+    uns_flag
+        key in final_adata.uns used to record that QC was performed (kept the name with original misspelling).
+    extract_read_features_from_bam_callable
+        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)
+    """
+
+    # Only run if not already performed
+    already = bool(adata.uns.get(uns_flag, False))
+    if (already and not force_redo) or bypass:
+        # QC already performed; nothing to do
+        return
+
+    # Build read_metrics dict either from provided arg or by extracting from bam files
+    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")
+            if extractor is None:
+                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):
+                    raise ValueError(f"extract_read_features_from_bam returned non-dict for {bam}")
+                read_metrics.update(bam_read_metrics)
+        else:
+            # nothing to do
+            read_metrics = {}
+
+    # Convert read_metrics dict -> DataFrame (rows = read id)
+    # Values may be lists/tuples or scalars; prefer lists/tuples with 5 entries.
+    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)
+    else:
+        # Build DF robustly
+        # Convert values to lists where possible, else to [val, val, val...]
+        max_cols = 5
+        rows = {}
+        for k, v in read_metrics.items():
+            if isinstance(v, (list, tuple, np.ndarray)):
+                vals = list(v)
+            else:
+                # scalar -> replicate into 5 columns to preserve original behavior
+                vals = [v] * max_cols
+            # Ensure length >= 5
+            if len(vals) < max_cols:
+                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'
+        ])
+
+        # 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
+
+    # 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)
+
+    # safe divisions: use np.where to avoid warnings and replace inf with nan
+    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_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
+
+    # Add read level raw modification signal: sum over X rows
+    X = adata.X
+    if sp.issparse(X):
+        # sum returns (n_obs, 1) sparse matrix; convert to 1d array
+        raw_sig = np.asarray(X.sum(axis=1)).ravel()
+    else:
+        raw_sig = np.asarray(X.sum(axis=1)).ravel()
+
+    adata.obs['Raw_modification_signal'] = raw_sig
+
+    # mark as done
+    adata.uns[uns_flag] = True
+
+    return None

smftools/informatics/{helpers/modkit_extract_to_adata.py → modkit_extract_to_adata.py}

@@ -1,11 +1,12 @@
-## modkit_extract_to_adata
-
 import concurrent.futures
 import gc
-from .count_aligned_reads import count_aligned_reads
+from .bam_functions import count_aligned_reads
 import pandas as pd
 from tqdm import tqdm
 import numpy as np
+from pathlib import Path
+from typing import Union, Iterable, Optional
+import shutil
 
 def filter_bam_records(bam, mapping_threshold):
     """Processes a single BAM file, counts reads, and determines records to analyze."""
@@ -336,29 +337,122 @@ def parallel_extract_stranded_methylation(dict_list, dict_to_skip, max_reference
             dict_list[dict_index][record][sample] = processed_data
     return dict_list
 
-def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name, mods, batch_size, mod_tsv_dir, delete_batch_hdfs=False, threads=None):
+def delete_intermediate_h5ads_and_tmpdir(
+    h5_dir: Union[str, Path, Iterable[str], None],
+    tmp_dir: Optional[Union[str, Path]] = None,
+    *,
+    dry_run: bool = False,
+    verbose: bool = True,
+):
+    """
+    Delete intermediate .h5ad files and a temporary directory.
+
+    Parameters
+    ----------
+    h5_dir : str | Path | iterable[str] | None
+        If a directory path is given, all files directly inside it will be considered.
+        If an iterable of file paths is given, those files will be considered.
+        Only files ending with '.h5ad' (and not ending with '.gz') are removed.
+    tmp_dir : str | Path | None
+        Path to a directory to remove recursively (e.g. a temp dir created earlier).
+    dry_run : bool
+        If True, print what *would* be removed but do not actually delete.
+    verbose : bool
+        Print progress / warnings.
+    """
+    # Helper: remove a single file path (Path-like or string)
+    def _maybe_unlink(p: Path):
+        if not p.exists():
+            if verbose:
+                print(f"[skip] not found: {p}")
+            return
+        if not p.is_file():
+            if verbose:
+                print(f"[skip] not a file: {p}")
+            return
+        if dry_run:
+            print(f"[dry-run] would remove file: {p}")
+            return
+        try:
+            p.unlink()
+            if verbose:
+                print(f"Removed file: {p}")
+        except Exception as e:
+            print(f"[error] failed to remove file {p}: {e}")
+
+    # Handle h5_dir input (directory OR iterable of file paths)
+    if h5_dir is not None:
+        # If it's a path to a directory, iterate its children
+        if isinstance(h5_dir, (str, Path)) and Path(h5_dir).is_dir():
+            dpath = Path(h5_dir)
+            for p in dpath.iterdir():
+                # only target top-level files (not recursing); require '.h5ad' suffix and exclude gz
+                name = p.name.lower()
+                if "h5ad" in name:
+                    _maybe_unlink(p)
+                else:
+                    if verbose:
+                        # optional: comment this out if too noisy
+                        print(f"[skip] not matching pattern: {p.name}")
+        else:
+            # treat as iterable of file paths
+            for f in h5_dir:
+                p = Path(f)
+                name = p.name.lower()
+                if name.endswith(".h5ad") and not name.endswith(".gz"):
+                    _maybe_unlink(p)
+                else:
+                    if verbose:
+                        print(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}")
+        else:
+            if not td.is_dir():
+                if verbose:
+                    print(f"[skip] tmp_dir is not a directory: {td}")
+            else:
+                if dry_run:
+                    print(f"[dry-run] would remove directory tree: {td}")
+                else:
+                    try:
+                        shutil.rmtree(td)
+                        if verbose:
+                            print(f"Removed directory tree: {td}")
+                    except Exception as e:
+                        print(f"[error] failed to remove tmp dir {td}: {e}")
+
+def modkit_extract_to_adata(fasta, bam_dir, out_dir, input_already_demuxed, mapping_threshold, experiment_name, mods, batch_size, mod_tsv_dir, delete_batch_hdfs=False, threads=None, double_barcoded_path = None):
     """
     Takes modkit extract outputs and organizes it into an adata object
 
     Parameters:
-        fasta (str): File path to the reference genome to align to.
-        bam_dir (str): File path to the directory containing the aligned_sorted split modified BAM files
+        fasta (Path): File path to the reference genome to align to.
+        bam_dir (Path): File path to the directory containing the aligned_sorted split modified BAM files
+        out_dir (Path): File path to output directory
+        input_already_demuxed (bool): Whether input reads were originally demuxed
         mapping_threshold (float): A value in between 0 and 1 to threshold the minimal fraction of aligned reads which map to the reference region. References with values above the threshold are included in the output adata.
         experiment_name (str): A string to provide an experiment name to the output adata file.
         mods (list): A list of strings of the modification types to use in the analysis.
         batch_size (int): An integer number of TSV files to analyze in memory at once while loading the final adata object.
-        mod_tsv_dir (str): String representing the path to the mod TSV directory
+        mod_tsv_dir (Path): path to the mod TSV directory
         delete_batch_hdfs (bool): Whether to delete the batch hdfs after writing out the final concatenated hdf. Default is False
+        double_barcoded_path (Path): Path to dorado demux summary file of double ended barcodes
 
     Returns:
-        final_adata_path (str): Path to the final adata
+        final_adata_path (Path): Path to the final adata
     """
     ###################################################
     # Package imports
     from .. import readwrite
-    from .get_native_references import get_native_references
-    from .extract_base_identities import extract_base_identities
-    from .ohe_batching import ohe_batching
+    from ..readwrite import safe_write_h5ad, make_dirs
+    from .fasta_functions import get_native_references
+    from .bam_functions import extract_base_identities
+    from .ohe import ohe_batching
     import pandas as pd
     import anndata as ad
     import os
@@ -368,41 +462,34 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
     from Bio.Seq import Seq
     from tqdm import tqdm
     import h5py
-    from .make_dirs import make_dirs
     ###################################################
 
     ################## Get input tsv and bam file names into a sorted list ################
-    # List all files in the directory
-    tsv_files = os.listdir(mod_tsv_dir)
-    bam_files = os.listdir(bam_dir)
-    # get current working directory
-    parent_dir = os.path.dirname(mod_tsv_dir)
-
     # Make output dirs
-    h5_dir = os.path.join(parent_dir, 'h5ads')
-    tmp_dir = os.path.join(parent_dir, 'tmp')
+    h5_dir = out_dir / 'h5ads'
+    tmp_dir = out_dir / 'tmp'
     make_dirs([h5_dir, tmp_dir])
-    existing_h5s =  os.listdir(h5_dir)
-    existing_h5s = [h5 for h5 in existing_h5s if '.h5ad.gz' in h5]
-    final_hdf = f'{experiment_name}_final_experiment_hdf5.h5ad'    
-    final_adata_path = os.path.join(h5_dir, final_hdf)
-
-    if os.path.exists(f"{final_adata_path}.gz"):
-        print(f'{final_adata_path}.gz already exists. Using existing adata')
-        return f"{final_adata_path}.gz"
+
+    existing_h5s =  h5_dir.iterdir()
+    existing_h5s = [h5 for h5 in existing_h5s if '.h5ad.gz' in str(h5)]
+    final_hdf = f'{experiment_name}.h5ad.gz'    
+    final_adata_path = h5_dir / final_hdf
+    final_adata = None
     
-    elif os.path.exists(f"{final_adata_path}"):
+    if final_adata_path.exists():
         print(f'{final_adata_path} already exists. Using existing adata')
-        return final_adata_path
+        return final_adata, final_adata_path
     
-    # Filter file names that contain the search string in their filename and keep them in a list
-    tsvs = [tsv for tsv in tsv_files if 'extract.tsv' in tsv and 'unclassified' not in tsv]
-    bams = [bam for bam in bam_files if '.bam' in bam and '.bai' not in bam and 'unclassified' not in bam]
-    # Sort file list by names and print the list of file names
-    tsvs.sort()
-    tsv_path_list = [os.path.join(mod_tsv_dir, tsv) for tsv in tsvs]
-    bams.sort()
-    bam_path_list = [os.path.join(bam_dir, bam) for bam in bams]
+    # List all files in the directory
+    tsvs = sorted(
+        p for p in mod_tsv_dir.iterdir()
+        if p.is_file() and 'unclassified' not in p.name and 'extract.tsv' in p.name)
+    bams = sorted(
+        p for p in bam_dir.iterdir()
+        if p.is_file() and p.suffix == '.bam' and 'unclassified' not in p.name and '.bai' not in p.name)
+
+    tsv_path_list = [mod_tsv_dir / tsv for tsv in tsvs]
+    bam_path_list = [bam_dir / bam for bam in bams]
     print(f'{len(tsvs)} sample tsv files found: {tsvs}')
     print(f'{len(bams)} sample bams found: {bams}')
     ##########################################################################################
@@ -416,7 +503,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
     ########### Determine the maximum record length to analyze in the dataset ################
     # Get all references within the FASTA and indicate the length and identity of the record sequence
     max_reference_length = 0
-    reference_dict = get_native_references(fasta) # returns a dict keyed by record name. Points to a tuple of (reference length, reference sequence)
+    reference_dict = get_native_references(str(fasta)) # returns a dict keyed by record name. Points to a tuple of (reference length, reference sequence)
     # Get the max record length in the dataset.
     for record in records_to_analyze:
         if reference_dict[record][0] > max_reference_length:
@@ -430,11 +517,11 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
     # One hot encode read sequences and write them out into the tmp_dir as h5ad files. 
     # Save the file paths in the bam_record_ohe_files dict.
     bam_record_ohe_files = {}
-    bam_record_save = os.path.join(tmp_dir, 'tmp_file_dict.h5ad')
+    bam_record_save = tmp_dir / 'tmp_file_dict.h5ad'
     fwd_mapped_reads = set()
     rev_mapped_reads = set()
     # If this step has already been performed, read in the tmp_dile_dict
-    if os.path.exists(bam_record_save):
+    if bam_record_save.exists():
         bam_record_ohe_files = ad.read_h5ad(bam_record_save).uns
         print('Found existing OHE reads, using these')
     else:
@@ -444,8 +531,9 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
             for record in records_to_analyze:
                 current_reference_length = reference_dict[record][0]
                 positions = range(current_reference_length)
+                ref_seq = reference_dict[record][1]
                 # Extract the base identities of reads aligned to the record
-                fwd_base_identities, rev_base_identities = extract_base_identities(bam, record, positions, max_reference_length)
+                fwd_base_identities, rev_base_identities, mismatch_counts_per_read, mismatch_trend_per_read = extract_base_identities(bam, record, positions, max_reference_length, ref_seq)
                 # Store read names of fwd and rev mapped reads
                 fwd_mapped_reads.update(fwd_base_identities.keys())
                 rev_mapped_reads.update(rev_base_identities.keys())
@@ -487,7 +575,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
             bam_path_list = bam_path_list[batch_size:]
         print('{0}: tsvs in batch {1} '.format(readwrite.time_string(), tsv_batch))
 
-        batch_already_processed = sum([1 for h5 in existing_h5s if f'_{batch}_' in h5])
+        batch_already_processed = sum([1 for h5 in existing_h5s if f'_{batch}_' in h5.name])
     ###################################################
         if batch_already_processed:
             print(f'Batch {batch} has already been processed into h5ads. Skipping batch and using existing files')
@@ -675,7 +763,6 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
 
 
             # Save the sample files in the batch as gzipped hdf5 files
-            os.chdir(h5_dir)
             print('{0}: Converting batch {1} dictionaries to anndata objects'.format(readwrite.time_string(), batch))
             for dict_index, dict_type in enumerate(dict_list):
                 if dict_index not in dict_to_skip:
@@ -708,6 +795,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
                                 temp_adata.var_names = temp_adata.var_names.astype(str)
                                 print('{0}: Adding {1} anndata for sample {2}'.format(readwrite.time_string(), sample_types[dict_index], final_sample_index))
                                 temp_adata.obs['Sample'] = [str(final_sample_index)] * len(temp_adata)
+                                temp_adata.obs['Barcode'] = [str(final_sample_index)] * len(temp_adata)
                                 temp_adata.obs['Reference'] = [f'{record}'] * len(temp_adata)
                                 temp_adata.obs['Strand'] = [strand] * len(temp_adata)
                                 temp_adata.obs['Dataset'] = [dataset] * len(temp_adata)
@@ -804,7 +892,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
 
                     try:
                         print('{0}: Writing {1} anndata out as a hdf5 file'.format(readwrite.time_string(), sample_types[dict_index]))
-                        adata.write_h5ad('{0}_{1}_{2}_SMF_binarized_sample_hdf5.h5ad.gz'.format(readwrite.date_string(), batch, sample_types[dict_index]), compression='gzip')
+                        adata.write_h5ad(h5_dir / '{0}_{1}_{2}_SMF_binarized_sample_hdf5.h5ad.gz'.format(readwrite.date_string(), batch, sample_types[dict_index]), compression='gzip')
                     except:
                         print(f"Skipping writing anndata for sample")
 
@@ -813,11 +901,10 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
             gc.collect()
 
     # Iterate over all of the batched hdf5 files and concatenate them.
-    os.chdir(h5_dir)
-    files = os.listdir(h5_dir)        
+    files = h5_dir.iterdir()       
     # Filter file names that contain the search string in their filename and keep them in a list
-    hdfs = [hdf for hdf in files if 'hdf5.h5ad' in hdf and hdf != final_hdf]
-    combined_hdfs = [hdf for hdf in hdfs if "combined" in hdf]
+    hdfs = [hdf for hdf in files if 'hdf5.h5ad' in hdf.name and hdf != final_hdf]
+    combined_hdfs = [hdf for hdf in hdfs if "combined" in hdf.name]
     if len(combined_hdfs) > 0:
         hdfs = combined_hdfs
     else:
@@ -825,7 +912,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
     # Sort file list by names and print the list of file names
     hdfs.sort()
     print('{0} sample files found: {1}'.format(len(hdfs), hdfs))
-    hdf_paths = [os.path.join(h5_dir, hd5) for hd5 in hdfs]
+    hdf_paths = [h5_dir / hd5 for hd5 in hdfs]
     final_adata = None
     for hdf_index, hdf in enumerate(hdf_paths):
         print('{0}: Reading in {1} hdf5 file'.format(readwrite.time_string(), hdfs[hdf_index]))
@@ -844,6 +931,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
 
     ohe_bases = ['A', 'C', 'G', 'T'] # ignore N bases for consensus
     ohe_layers = [f"{ohe_base}_binary_encoding" for ohe_base in ohe_bases]
+    final_adata.uns['References'] = {}
     for record in records_to_analyze:
         # Add FASTA sequence to the object
         sequence = record_seq_dict[record][0]
@@ -851,6 +939,7 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
         final_adata.var[f'{record}_top_strand_FASTA_base'] = list(sequence)
         final_adata.var[f'{record}_bottom_strand_FASTA_base'] = list(complement)
         final_adata.uns[f'{record}_FASTA_sequence'] = sequence
+        final_adata.uns['References'][f'{record}_FASTA_sequence'] = sequence
         # Add consensus sequence of samples mapped to the record to the object
         record_subset = final_adata[final_adata.obs['Reference'] == record]
         for strand in record_subset.obs['Strand'].cat.categories:
@@ -866,19 +955,16 @@ def modkit_extract_to_adata(fasta, bam_dir, mapping_threshold, experiment_name,
                 consensus_sequence_list = [layer_map[i] for i in nucleotide_indexes]
                 final_adata.var[f'{record}_{strand}_{mapping_dir}_consensus_sequence_from_all_samples'] = consensus_sequence_list
 
-    #final_adata.write_h5ad(final_adata_path)
+    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"
+        add_demux_type_annotation(final_adata, double_barcoded_reads)
 
     # Delete the individual h5ad files and only keep the final concatenated file
     if delete_batch_hdfs:
-        files = os.listdir(h5_dir)
-        hdfs_to_delete = [hdf for hdf in files if 'hdf5.h5ad' in hdf and hdf != final_hdf]
-        hdf_paths_to_delete = [os.path.join(h5_dir, hdf) for hdf in hdfs_to_delete]
-        # Iterate over the files and delete them
-        for hdf in hdf_paths_to_delete:
-            try:
-                os.remove(hdf)
-                print(f"Deleted file: {hdf}")
-            except OSError as e:
-                print(f"Error deleting file {hdf}: {e}")
+        delete_intermediate_h5ads_and_tmpdir(h5_dir, tmp_dir)
 
     return final_adata, final_adata_path

smftools/informatics/modkit_functions.py

@@ -0,0 +1,129 @@
+import os
+import subprocess
+import glob
+import zipfile
+from pathlib import Path
+
+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
+
+    Parameters:
+        thresholds (list): A list of thresholds to use for marking each basecalled base as passing or failing on canonical and modification call status.
+        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
+        Runs modkit extract on input aligned_sorted_split modified BAM files to output zipped TSVs containing modification calls.
+
+    """
+    filter_threshold, m6A_threshold, m5C_threshold, hm5C_threshold = thresholds
+    bam_files = sorted(p for p in split_dir.iterdir() if bam_suffix in p.name and '.bai' not in p.name)
+    if skip_unclassified:
+        bam_files = [p for p in bam_files if "unclassified" not in p.name]
+    print(f"Running modkit extract for the following bam files: {bam_files}")
+
+    if threads:
+        threads = str(threads)
+    else:
+        pass
+
+    for input_file in bam_files:
+        print(input_file)
+        # Construct the output TSV file path
+        output_tsv = mod_tsv_dir / (input_file.stem + "_extract.tsv")
+        output_tsv_gz = output_tsv.parent / (output_tsv.name + '.gz')
+        if output_tsv_gz.exists():
+            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", str(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,
+                    str(input_file), str(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}",
+                    str(input_file), str(output_tsv)
+                    ]                    
+            subprocess.run(extract_command)
+            # Zip the output TSV
+            print(f'zipping {output_tsv}')
+            if threads:
+                zip_command = ["pigz", "-f", "-p", threads, str(output_tsv)]
+            else:
+                zip_command = ["pigz", "-f", str(output_tsv)]
+            subprocess.run(zip_command, check=True)
+    return
+
+def make_modbed(aligned_sorted_output, thresholds, mod_bed_dir):
+    """
+    Generating position methylation summaries for each barcoded sample starting from the overall BAM file that was direct output of dorado aligner.
+    Parameters:
+        aligned_sorted_output (str): A string representing the file path to the aligned_sorted non-split BAM file.
+    
+    Returns:
+        None
+    """
+    import os
+    import subprocess
+    
+    filter_threshold, m6A_threshold, m5C_threshold, hm5C_threshold = thresholds
+    command = [
+        "modkit", "pileup", str(aligned_sorted_output), str(mod_bed_dir),
+        "--partition-tag", "BC",
+        "--only-tabs",
+        "--filter-threshold", f'{filter_threshold}',
+        "--mod-thresholds", f"m:{m5C_threshold}",
+        "--mod-thresholds", f"a:{m6A_threshold}",
+        "--mod-thresholds", f"h:{hm5C_threshold}"
+    ]
+    subprocess.run(command)
+
+def modQC(aligned_sorted_output, thresholds):
+    """
+    Output the percentile of bases falling at a call threshold (threshold is a probability between 0-1) for the overall BAM file.
+    It is generally good to look at these parameters on positive and negative controls.
+
+    Parameters:
+        aligned_sorted_output (str): A string representing the file path of the aligned_sorted non-split BAM file output by the dorado aligned.
+        thresholds (list): A list of floats to pass for call thresholds.
+
+    Returns:
+        None
+    """
+    import subprocess
+
+    filter_threshold, m6A_threshold, m5C_threshold, hm5C_threshold = thresholds
+    subprocess.run(["modkit", "sample-probs", str(aligned_sorted_output)])
+    command = [
+        "modkit", "summary", str(aligned_sorted_output),
+        "--filter-threshold", f"{filter_threshold}",
+        "--mod-thresholds", f"m:{m5C_threshold}",
+        "--mod-thresholds", f"a:{m6A_threshold}",
+        "--mod-thresholds", f"h:{hm5C_threshold}"
+    ]
+    subprocess.run(command)