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

smftools/cli/load_adata.py

@@ -1,8 +1,14 @@
+from __future__ import annotations
+
+import logging
 import shutil
 from pathlib import Path
 from typing import Iterable, Union
 
-from smftools.logging_utils import get_logger
+import numpy as np
+
+from smftools.constants import HMM_DIR, LOAD_DIR, LOGGING_DIR, PREPROCESS_DIR, SPATIAL_DIR
+from smftools.logging_utils import get_logger, setup_logging
 
 from .helpers import AdataPaths
 
@@ -76,6 +82,62 @@ def delete_tsvs(
                         logger.warning(f"[error] failed to remove tmp dir {td}: {e}")
 
 
+def load_adata(config_path: str):
+    """
+    CLI-facing wrapper for the load pipeline.
+
+    - Reads config CSV into ExperimentConfig
+    - Computes canonical paths for all downstream AnnData stages
+    - Registers those in the summary CSV
+    - Applies stage-skipping logic (hmm > spatial > pp_dedup > pp > raw)
+    - If needed, calls the core pipeline to actually build the raw AnnData
+
+    Returns
+    -------
+    adata : anndata.AnnData | None
+        Newly created AnnData object, or None if we skipped because a later-stage
+        AnnData already exists.
+    adata_path : pathlib.Path
+        Path to the "current" AnnData that should be used downstream.
+    cfg : ExperimentConfig
+        Config object for downstream steps.
+    """
+    from datetime import datetime
+    from importlib import resources
+
+    from ..readwrite import add_or_update_column_in_csv, make_dirs
+    from .helpers import get_adata_paths, load_experiment_config
+
+    # -----------------------------
+    # 1) Load config into cfg
+    # -----------------------------
+    cfg = load_experiment_config(config_path)
+
+    # Ensure base output dir
+    output_directory = Path(cfg.output_directory)
+    make_dirs([output_directory])
+
+    # -----------------------------
+    # 2) Compute and register paths
+    # -----------------------------
+    paths = get_adata_paths(cfg)
+
+    # -----------------------------
+    # 3) Stage skipping logic
+    # -----------------------------
+    if not getattr(cfg, "force_redo_load_adata", False):
+        if paths.raw.exists():
+            logger.info(
+                f"Raw AnnData from smftools load already exists: {paths.raw}\nSkipping smftools load"
+            )
+            return None, paths.raw, cfg
+
+    # If we get here, we actually want to run the full load pipeline
+    adata, adata_path, cfg = load_adata_core(cfg, paths, config_path=config_path)
+
+    return adata, adata_path, cfg
+
+
 def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     """
     Core load pipeline.
@@ -105,9 +167,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     cfg : ExperimentConfig
         (Same object, possibly with some fields updated, e.g. fasta path.)
     """
-    from pathlib import Path
-
-    import numpy as np
+    from datetime import datetime
 
     from ..informatics.bam_functions import (
         align_and_sort_BAM,
@@ -115,6 +175,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         concatenate_fastqs_to_bam,
         demux_and_index_BAM,
         extract_read_features_from_bam,
+        extract_read_tags_from_bam,
         split_and_index_BAM,
     )
     from ..informatics.basecalling import canoncall, modcall
@@ -125,7 +186,11 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         get_chromosome_lengths,
         subsample_fasta_from_bed,
     )
-    from ..informatics.h5ad_functions import add_read_length_and_mapping_qc
+    from ..informatics.h5ad_functions import (
+        add_read_length_and_mapping_qc,
+        add_read_tag_annotations,
+        add_secondary_supplementary_alignment_flags,
+    )
     from ..informatics.modkit_extract_to_adata import modkit_extract_to_adata
     from ..informatics.modkit_functions import extract_mods, make_modbed, modQC
     from ..informatics.pod5_functions import fast5_to_pod5
@@ -135,8 +200,25 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     from .helpers import write_gz_h5ad
 
     ################################### 1) General params and input organization ###################################
+    date_str = datetime.today().strftime("%y%m%d")
+    now = datetime.now()
+    time_str = now.strftime("%H%M%S")
+
+    log_level = getattr(logging, cfg.log_level.upper(), logging.INFO)
+
     output_directory = Path(cfg.output_directory)
-    make_dirs([output_directory])
+    load_directory = output_directory / LOAD_DIR
+    logging_directory = load_directory / LOGGING_DIR
+
+    make_dirs([output_directory, load_directory])
+
+    if cfg.emit_log_file:
+        log_file = logging_directory / f"{date_str}_{time_str}_log.log"
+        make_dirs([logging_directory])
+    else:
+        log_file = None
+
+    setup_logging(level=log_level, log_file=log_file, reconfigure=log_file is not None)
 
     raw_adata_path = paths.raw
     pp_adata_path = paths.pp
@@ -150,11 +232,9 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
 
     # Direct methylation detection SMF specific parameters
     if cfg.smf_modality == "direct":
-        mod_bed_dir = cfg.output_directory / "mod_beds"
-        add_or_update_column_in_csv(cfg.summary_file, "mod_bed_dir", mod_bed_dir)
-        mod_tsv_dir = cfg.output_directory / "mod_tsvs"
-        add_or_update_column_in_csv(cfg.summary_file, "mod_tsv_dir", mod_tsv_dir)
-        bam_qc_dir = cfg.output_directory / "bam_qc"
+        mod_bed_dir = load_directory / "mod_beds"
+        mod_tsv_dir = load_directory / "mod_tsvs"
+        bam_qc_dir = load_directory / "bam_qc"
         mods = [cfg.mod_map[mod] for mod in cfg.mod_list]
 
         if not check_executable_exists("dorado"):
@@ -190,7 +270,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     # If the input files are fast5 files, convert the files to a pod5 file before proceeding.
     if cfg.input_type == "fast5":
         # take the input directory of fast5 files and write out a single pod5 file into the output directory.
-        output_pod5 = cfg.output_directory / "FAST5s_to_POD5.pod5"
+        output_pod5 = load_directory / "FAST5s_to_POD5.pod5"
         if output_pod5.exists():
             pass
         else:
@@ -204,7 +284,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     # If the input is a fastq or a directory of fastqs, concatenate them into an unaligned BAM and save the barcode
     elif cfg.input_type == "fastq":
         # Output file for FASTQ concatenation.
-        output_bam = cfg.output_directory / "canonical_basecalls.bam"
+        output_bam = load_directory / "canonical_basecalls.bam"
         if output_bam.exists():
             logger.debug("Output BAM already exists")
         else:
@@ -219,6 +299,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
                 rg_sample_field=None,
                 progress=False,
                 auto_pair=cfg.fastq_auto_pairing,
+                samtools_backend=cfg.samtools_backend,
             )
 
             logger.info(f"Found the following barcodes in FASTQ inputs: {summary['barcodes']}")
@@ -231,8 +312,6 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     else:
         pass
 
-    add_or_update_column_in_csv(cfg.summary_file, "input_data_path", cfg.input_data_path)
-
     # Determine if the input data needs to be basecalled
     if cfg.input_type == "pod5":
         logger.info(f"Detected pod5 inputs: {cfg.input_files}")
@@ -249,25 +328,24 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         model_basename = str(model_basename).replace(".", "_")
         if cfg.smf_modality == "direct":
             mod_string = "_".join(cfg.mod_list)
-            bam = cfg.output_directory / f"{model_basename}_{mod_string}_calls"
+            bam = load_directory / f"{model_basename}_{mod_string}_calls"
         else:
-            bam = cfg.output_directory / f"{model_basename}_canonical_basecalls"
+            bam = load_directory / f"{model_basename}_canonical_basecalls"
     else:
-        bam_base = cfg.input_data_path.name
-        bam = cfg.output_directory / bam_base
+        bam_base = cfg.input_data_path.stem
+        bam = cfg.input_data_path.parent / bam_base
 
     # Generate path names for the unaligned, aligned, as well as the aligned/sorted bam.
     unaligned_output = bam.with_suffix(cfg.bam_suffix)
+
     aligned_BAM = (
-        cfg.output_directory / (bam.stem + "_aligned")
+        load_directory / (bam.stem + "_aligned")
     )  # doing this allows specifying an input bam in a seperate directory as the aligned output bams
+
     aligned_output = aligned_BAM.with_suffix(cfg.bam_suffix)
     aligned_sorted_BAM = aligned_BAM.with_name(aligned_BAM.stem + "_sorted")
     aligned_sorted_output = aligned_sorted_BAM.with_suffix(cfg.bam_suffix)
 
-    add_or_update_column_in_csv(cfg.summary_file, "basecalled_bam", unaligned_output)
-    add_or_update_column_in_csv(cfg.summary_file, "aligned_bam", aligned_output)
-    add_or_update_column_in_csv(cfg.summary_file, "sorted_bam", aligned_sorted_output)
     ########################################################################################################################
 
     ################################### 2) FASTA Handling ###################################
@@ -281,11 +359,11 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     if cfg.fasta_regions_of_interest and ".bed" in cfg.fasta_regions_of_interest:
         fasta_stem = cfg.fasta.stem
         bed_stem = Path(cfg.fasta_regions_of_interest).stem
-        output_FASTA = cfg.output_directory / f"{fasta_stem}_subsampled_by_{bed_stem}.fasta"
+        output_FASTA = load_directory / f"{fasta_stem}_subsampled_by_{bed_stem}.fasta"
 
         logger.info("Subsampling FASTA records using the provided BED file")
         subsample_fasta_from_bed(
-            cfg.fasta, cfg.fasta_regions_of_interest, cfg.output_directory, output_FASTA
+            cfg.fasta, cfg.fasta_regions_of_interest, load_directory, output_FASTA
         )
         fasta = output_FASTA
     else:
@@ -296,7 +374,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     if cfg.smf_modality == "conversion":
         fasta_stem = fasta.stem
         converted_FASTA_basename = f"{fasta_stem}_converted.fasta"
-        converted_FASTA = cfg.output_directory / converted_FASTA_basename
+        converted_FASTA = load_directory / converted_FASTA_basename
 
         if "converted.fa" in fasta.name:
             logger.info(f"{fasta} is already converted. Using existing converted FASTA.")
@@ -308,8 +386,6 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
             generate_converted_FASTA(fasta, cfg.conversion_types, cfg.strands, converted_FASTA)
         fasta = converted_FASTA
 
-    add_or_update_column_in_csv(cfg.summary_file, "fasta", fasta)
-
     # Make a FAI and .chrom.names file for the fasta
     get_chromosome_lengths(fasta)
     ########################################################################################################################
@@ -370,13 +446,13 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         logger.debug(f"{aligned_sorted_output} already exists. Using existing aligned/sorted BAM.")
     else:
         logger.info(f"Aligning and sorting reads")
-        align_and_sort_BAM(fasta, unaligned_output, cfg)
+        align_and_sort_BAM(fasta, unaligned_output, aligned_output, cfg)
         # Deleted the unsorted aligned output
         aligned_output.unlink()
 
     if cfg.make_beds:
         # Make beds and provide basic histograms
-        bed_dir = cfg.output_directory / "beds"
+        bed_dir = load_directory / "beds"
         if bed_dir.is_dir():
             logger.debug(
                 f"{bed_dir} already exists. Skipping BAM -> BED conversion for {aligned_sorted_output}"
@@ -384,7 +460,14 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         else:
             logger.info("Making bed files from the aligned and sorted BAM file")
             aligned_BAM_to_bed(
-                aligned_sorted_output, cfg.output_directory, fasta, cfg.make_bigwigs, cfg.threads
+                aligned_sorted_output,
+                load_directory,
+                fasta,
+                cfg.make_bigwigs,
+                cfg.threads,
+                samtools_backend=cfg.samtools_backend,
+                bedtools_backend=cfg.bedtools_backend,
+                bigwig_backend=cfg.bigwig_backend,
             )
     ########################################################################################################################
 
@@ -404,13 +487,19 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         else:
             make_dirs([cfg.split_path])
             logger.info("Demultiplexing samples into individual aligned/sorted BAM files")
-            all_bam_files = split_and_index_BAM(aligned_sorted_BAM, cfg.split_path, cfg.bam_suffix)
+            all_bam_files = split_and_index_BAM(
+                aligned_sorted_BAM,
+                cfg.split_path,
+                cfg.bam_suffix,
+                samtools_backend=cfg.samtools_backend,
+            )
 
             unclassified_bams = [p for p in all_bam_files if "unclassified" in p.name]
             bam_files = sorted(p for p in all_bam_files if "unclassified" not in p.name)
 
         se_bam_files = bam_files
         bam_dir = cfg.split_path
+        double_barcoded_path = None
 
     else:
         if single_barcoded_path.is_dir():
@@ -489,19 +578,34 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         else:
             logger.info("Making BED files from BAM files for each sample")
             for bam in bam_files:
-                aligned_BAM_to_bed(bam, cfg.split_path, fasta, cfg.make_bigwigs, cfg.threads)
+                aligned_BAM_to_bed(
+                    bam,
+                    cfg.split_path,
+                    fasta,
+                    cfg.make_bigwigs,
+                    cfg.threads,
+                    samtools_backend=cfg.samtools_backend,
+                    bedtools_backend=cfg.bedtools_backend,
+                    bigwig_backend=cfg.bigwig_backend,
+                )
     ########################################################################################################################
 
     ################################### 6) SAMTools based BAM QC ######################################################################
 
     # 5) Samtools QC metrics on split BAM files
-    bam_qc_dir = cfg.split_path / "bam_qc"
+    bam_qc_dir = load_directory / "bam_qc"
     if bam_qc_dir.is_dir():
         logger.debug(f"{bam_qc_dir} already exists. Using existing BAM QC calculations.")
     else:
         make_dirs([bam_qc_dir])
         logger.info("Performing BAM QC")
-        bam_qc(bam_files, bam_qc_dir, cfg.threads, modality=cfg.smf_modality)
+        bam_qc(
+            bam_files,
+            bam_qc_dir,
+            cfg.threads,
+            modality=cfg.smf_modality,
+            samtools_backend=cfg.samtools_backend,
+        )
     ########################################################################################################################
 
     ################################### 7) AnnData loading ######################################################################
@@ -518,7 +622,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         raw_adata, raw_adata_path = converted_BAM_to_adata(
             fasta,
             bam_dir,
-            cfg.output_directory,
+            load_directory,
             cfg.input_already_demuxed,
             cfg.mapping_threshold,
             cfg.experiment_name,
@@ -529,6 +633,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
             deaminase_footprinting,
             delete_intermediates=cfg.delete_intermediate_hdfs,
             double_barcoded_path=double_barcoded_path,
+            samtools_backend=cfg.samtools_backend,
         )
     else:
         if mod_bed_dir.is_dir():
@@ -574,7 +679,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         raw_adata, raw_adata_path = modkit_extract_to_adata(
             fasta,
             bam_dir,
-            cfg.output_directory,
+            load_directory,
             cfg.input_already_demuxed,
             cfg.mapping_threshold,
             cfg.experiment_name,
@@ -584,6 +689,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
             cfg.delete_batch_hdfs,
             cfg.threads,
             double_barcoded_path,
+            cfg.samtools_backend,
         )
         if cfg.delete_intermediate_tsvs:
             delete_tsvs(mod_tsv_dir)
@@ -604,8 +710,28 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         extract_read_features_from_bam_callable=extract_read_features_from_bam,
         bypass=cfg.bypass_add_read_length_and_mapping_qc,
         force_redo=cfg.force_redo_add_read_length_and_mapping_qc,
+        samtools_backend=cfg.samtools_backend,
+    )
+
+    logger.info("Adding BAM tags and BAM flags to adata.obs")
+    add_read_tag_annotations(
+        raw_adata,
+        se_bam_files,
+        tag_names=getattr(cfg, "bam_tag_names", ["NM", "MD", "MM", "ML"]),
+        include_flags=True,
+        include_cigar=True,
+        extract_read_tags_from_bam_callable=extract_read_tags_from_bam,
+        samtools_backend=cfg.samtools_backend,
     )
 
+    if getattr(cfg, "annotate_secondary_supplementary", False):
+        logger.info("Annotating secondary/supplementary alignments from aligned BAM")
+        add_secondary_supplementary_alignment_flags(
+            raw_adata,
+            aligned_sorted_output,
+            samtools_backend=cfg.samtools_backend,
+        )
+
     raw_adata.obs["Raw_modification_signal"] = np.nansum(raw_adata.X, axis=1)
     ########################################################################################################################
 
@@ -618,7 +744,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
             raw_adata,
             cfg.input_data_path,
             n_jobs=cfg.threads,
-            csv_path=output_directory / "read_to_pod5_origin_mapping.csv",
+            csv_path=load_directory / "read_to_pod5_origin_mapping.csv",
         )
     ########################################################################################################################
 
@@ -637,12 +763,12 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     ############################################### MultiQC HTML Report ###############################################
 
     # multiqc ###
-    mqc_dir = cfg.split_path / "multiqc"
+    mqc_dir = load_directory / "multiqc"
     if mqc_dir.is_dir():
-        logger.debug(f"{mqc_dir} already exists, skipping multiqc")
+        logger.info(f"{mqc_dir} already exists, skipping multiqc")
     else:
         logger.info("Running multiqc")
-        run_multiqc(cfg.split_path, mqc_dir)
+        run_multiqc(bam_qc_dir, mqc_dir)
     ########################################################################################################################
 
     ############################################### delete intermediate BAM files ###############################################
@@ -665,93 +791,3 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     ########################################################################################################################
 
     return raw_adata, raw_adata_path, cfg
-
-
-def load_adata(config_path: str):
-    """
-    CLI-facing wrapper for the load pipeline.
-
-    - Reads config CSV into ExperimentConfig
-    - Computes canonical paths for all downstream AnnData stages
-    - Registers those in the summary CSV
-    - Applies stage-skipping logic (hmm > spatial > pp_dedup > pp > raw)
-    - If needed, calls the core pipeline to actually build the raw AnnData
-
-    Returns
-    -------
-    adata : anndata.AnnData | None
-        Newly created AnnData object, or None if we skipped because a later-stage
-        AnnData already exists.
-    adata_path : pathlib.Path
-        Path to the "current" AnnData that should be used downstream.
-    cfg : ExperimentConfig
-        Config object for downstream steps.
-    """
-    from datetime import datetime
-    from importlib import resources
-
-    from ..config import ExperimentConfig, LoadExperimentConfig
-    from ..readwrite import add_or_update_column_in_csv, make_dirs
-    from .helpers import get_adata_paths
-
-    date_str = datetime.today().strftime("%y%m%d")
-
-    # -----------------------------
-    # 1) Load config into cfg
-    # -----------------------------
-    loader = LoadExperimentConfig(config_path)
-    defaults_dir = resources.files("smftools").joinpath("config")
-    cfg, report = ExperimentConfig.from_var_dict(
-        loader.var_dict, date_str=date_str, defaults_dir=defaults_dir
-    )
-
-    # Ensure base output dir
-    make_dirs([cfg.output_directory])
-
-    # -----------------------------
-    # 2) Compute and register paths
-    # -----------------------------
-    paths = get_adata_paths(cfg)
-
-    # experiment-level metadata in summary CSV
-    add_or_update_column_in_csv(cfg.summary_file, "experiment_name", cfg.experiment_name)
-    add_or_update_column_in_csv(cfg.summary_file, "config_path", config_path)
-    add_or_update_column_in_csv(cfg.summary_file, "input_data_path", cfg.input_data_path)
-    add_or_update_column_in_csv(cfg.summary_file, "input_files", [cfg.input_files])
-
-    # AnnData stage paths
-    add_or_update_column_in_csv(cfg.summary_file, "load_adata", paths.raw)
-    add_or_update_column_in_csv(cfg.summary_file, "pp_adata", paths.pp)
-    add_or_update_column_in_csv(cfg.summary_file, "pp_dedup_adata", paths.pp_dedup)
-    add_or_update_column_in_csv(cfg.summary_file, "spatial_adata", paths.spatial)
-    add_or_update_column_in_csv(cfg.summary_file, "hmm_adata", paths.hmm)
-
-    # -----------------------------
-    # 3) Stage skipping logic
-    # -----------------------------
-    if not getattr(cfg, "force_redo_load_adata", False):
-        if paths.hmm.exists():
-            logger.debug(f"HMM AnnData already exists: {paths.hmm}\nSkipping smftools load")
-            return None, paths.hmm, cfg
-        if paths.spatial.exists():
-            logger.debug(f"Spatial AnnData already exists: {paths.spatial}\nSkipping smftools load")
-            return None, paths.spatial, cfg
-        if paths.pp_dedup.exists():
-            logger.debug(
-                f"Preprocessed deduplicated AnnData already exists: {paths.pp_dedup}\n"
-                f"Skipping smftools load"
-            )
-            return None, paths.pp_dedup, cfg
-        if paths.pp.exists():
-            logger.debug(f"Preprocessed AnnData already exists: {paths.pp}\nSkipping smftools load")
-            return None, paths.pp, cfg
-        if paths.raw.exists():
-            logger.debug(
-                f"Raw AnnData from smftools load already exists: {paths.raw}\nSkipping smftools load"
-            )
-            return None, paths.raw, cfg
-
-    # If we get here, we actually want to run the full load pipeline
-    adata, adata_path, cfg = load_adata_core(cfg, paths, config_path=config_path)
-
-    return adata, adata_path, cfg