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

smftools/cli/latent_adata.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Optional, Tuple
+
+import anndata as ad
+
+from smftools.constants import LATENT_DIR, LOGGING_DIR, SEQUENCE_INTEGER_ENCODING
+from smftools.logging_utils import get_logger, setup_logging
+
+logger = get_logger(__name__)
+
+
+def latent_adata(
+    config_path: str,
+) -> Tuple[Optional[ad.AnnData], Optional[Path]]:
+    """
+    CLI-facing wrapper for representation learning.
+
+    Called by: `smftools latent <config_path>`
+
+    Responsibilities:
+    - Determine which AnnData stages exist (pp, pp_dedup, spatial, hmm).
+    - Call `latent_adata_core(...)` when actual work is needed.
+
+    Returns
+    -------
+    latent_adata : AnnData | None
+        AnnData with latent analyses, or None if we skipped because a later-stage
+        AnnData already exists.
+    latent_adata_path : Path | None
+        Path to the “current” latent AnnData.
+    """
+    from ..readwrite import add_or_update_column_in_csv, safe_read_h5ad
+    from .helpers import get_adata_paths, load_experiment_config
+
+    # 1) Ensure config + basic paths via load_adata
+    cfg = load_experiment_config(config_path)
+
+    paths = get_adata_paths(cfg)
+
+    pp_path = paths.pp
+    pp_dedup_path = paths.pp_dedup
+    spatial_path = paths.spatial
+    hmm_path = paths.hmm
+    latent_path = paths.latent
+
+    # Stage-skipping logic for latent
+    if not getattr(cfg, "force_redo_latent_analyses", False):
+        # If latent exists, we consider latent analyses already done.
+        if latent_path.exists():
+            logger.info(f"Latent AnnData found: {latent_path}\nSkipping smftools latent")
+            return None, latent_path
+
+    # Helper to load from disk, reusing loaded_adata if it matches
+    def _load(path: Path):
+        adata, _ = safe_read_h5ad(path)
+        return adata
+
+    # 3) Decide which AnnData to use as the *starting point* for latent analyses
+    if latent_path.exists():
+        start_adata = _load(latent_path)
+        source_path = latent_path
+    elif hmm_path.exists():
+        start_adata = _load(hmm_path)
+        source_path = hmm_path
+    elif spatial_path.exists():
+        start_adata = _load(spatial_path)
+        source_path = spatial_path
+    elif pp_dedup_path.exists():
+        start_adata = _load(pp_dedup_path)
+        source_path = pp_dedup_path
+    elif pp_path.exists():
+        start_adata = _load(pp_path)
+        source_path = pp_path
+    else:
+        logger.warning(
+            "No suitable AnnData found for latent analyses (need at least preprocessed)."
+        )
+        return None, None
+
+    # 4) Run the latent core
+    adata_latent, latent_path = latent_adata_core(
+        adata=start_adata,
+        cfg=cfg,
+        paths=paths,
+        source_adata_path=source_path,
+        config_path=config_path,
+    )
+
+    return adata_latent, latent_path
+
+
+def latent_adata_core(
+    adata: ad.AnnData,
+    cfg,
+    paths: AdataPaths,
+    source_adata_path: Optional[Path] = None,
+    config_path: Optional[str] = None,
+) -> Tuple[ad.AnnData, Path]:
+    """
+    Core spatial analysis pipeline.
+
+    Assumes:
+    - `adata` is (typically) the preprocessed, duplicate-removed AnnData.
+    - `cfg` is the ExperimentConfig.
+
+    Does:
+    - Optional sample sheet load.
+    - Optional inversion & reindexing.
+    - PCA/UMAP/Leiden.
+    - Save latent AnnData to `latent_adata_path`.
+
+    Returns
+    -------
+    adata : AnnData
+        analyzed AnnData (same object, modified in-place).
+    adata_path : Path
+        Path where AnnData was written.
+    """
+    import os
+    import warnings
+    from datetime import datetime
+    from pathlib import Path
+
+    import numpy as np
+    import pandas as pd
+
+    from ..metadata import record_smftools_metadata
+    from ..plotting import (
+        plot_cp_sequence_components,
+        plot_embedding,
+        plot_nmf_components,
+        plot_pca,
+        plot_umap,
+    )
+    from ..preprocessing import (
+        invert_adata,
+        load_sample_sheet,
+        reindex_references_adata,
+    )
+    from ..readwrite import make_dirs, safe_read_h5ad
+    from ..tools import (
+        calculate_leiden,
+        calculate_nmf,
+        calculate_sequence_cp_decomposition,
+        calculate_umap,
+    )
+    from .helpers import write_gz_h5ad
+
+    # -----------------------------
+    # General setup
+    # -----------------------------
+    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)
+
+    latent_adata_path = paths.latent
+
+    output_directory = Path(cfg.output_directory)
+    latent_directory = output_directory / LATENT_DIR
+    logging_directory = latent_directory / LOGGING_DIR
+
+    make_dirs([output_directory, latent_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)
+
+    smf_modality = cfg.smf_modality
+    if smf_modality == "conversion":
+        deaminase = False
+    else:
+        deaminase = True
+
+    # -----------------------------
+    # Optional sample sheet metadata
+    # -----------------------------
+    if getattr(cfg, "sample_sheet_path", None):
+        load_sample_sheet(
+            adata,
+            cfg.sample_sheet_path,
+            mapping_key_column=cfg.sample_sheet_mapping_column,
+            as_category=True,
+            force_reload=cfg.force_reload_sample_sheet,
+        )
+
+    # -----------------------------
+    # Optional inversion along positions axis
+    # -----------------------------
+    if getattr(cfg, "invert_adata", False):
+        adata = invert_adata(adata)
+
+    # -----------------------------
+    # Optional reindexing by reference
+    # -----------------------------
+    reindex_references_adata(
+        adata,
+        reference_col=cfg.reference_column,
+        offsets=cfg.reindexing_offsets,
+        new_col=cfg.reindexed_var_suffix,
+    )
+
+    if adata.uns.get("reindex_references_adata_performed", False):
+        reindex_suffix = cfg.reindexed_var_suffix
+    else:
+        reindex_suffix = None
+
+    references = adata.obs[cfg.reference_column].cat.categories
+
+    # ============================================================
+    # 2) PCA/UMAP on *deduplicated* preprocessed AnnData
+    # ============================================================
+    latent_dir_dedup = latent_directory / "deduplicated"
+    umap_dir = latent_dir_dedup / "07_umaps"
+    nmf_dir = latent_dir_dedup / "07b_nmf"
+    nmf_sequence_dir = latent_dir_dedup / "07c_nmf_sequence"
+
+    var_filters = []
+    if smf_modality == "direct":
+        for ref in references:
+            for base in cfg.mod_target_bases:
+                var_filters.append(f"{ref}_{base}_site")
+    elif deaminase:
+        for ref in references:
+            var_filters.append(f"{ref}_C_site")
+    else:
+        for ref in references:
+            for base in cfg.mod_target_bases:
+                var_filters.append(f"{ref}_{base}_site")
+
+    # UMAP / Leiden
+    if umap_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
+        logger.debug(f"{umap_dir} already exists. Skipping UMAP plotting.")
+    else:
+        make_dirs([umap_dir])
+
+        adata = calculate_umap(
+            adata,
+            layer=cfg.layer_for_umap_plotting,
+            var_filters=var_filters,
+            n_pcs=10,
+            knn_neighbors=15,
+        )
+
+        calculate_leiden(adata, resolution=0.1)
+
+        umap_layers = ["leiden", cfg.sample_name_col_for_plotting, "Reference_strand"]
+        umap_layers += cfg.umap_layers_to_plot
+        plot_umap(adata, color=umap_layers, output_dir=umap_dir)
+        plot_pca(adata, color=umap_layers, output_dir=umap_dir)
+
+    # NMF
+    if nmf_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
+        logger.debug(f"{nmf_dir} already exists. Skipping NMF plotting.")
+    else:
+        make_dirs([nmf_dir])
+        adata = calculate_nmf(
+            adata,
+            layer=cfg.layer_for_umap_plotting,
+            var_filters=var_filters,
+            n_components=5,
+        )
+        nmf_layers = ["leiden", cfg.sample_name_col_for_plotting, "Reference_strand"]
+        nmf_layers += cfg.umap_layers_to_plot
+        plot_embedding(adata, basis="nmf", color=nmf_layers, output_dir=nmf_dir)
+        plot_nmf_components(adata, output_dir=nmf_dir)
+
+    # CP decomposition using sequence integer encoding (no var filters)
+    if nmf_sequence_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
+        logger.debug(f"{nmf_sequence_dir} already exists. Skipping sequence CP plotting.")
+    elif SEQUENCE_INTEGER_ENCODING not in adata.layers:
+        logger.warning(
+            "Layer %s not found; skipping sequence integer encoding CP.",
+            SEQUENCE_INTEGER_ENCODING,
+        )
+    else:
+        make_dirs([nmf_sequence_dir])
+        adata = calculate_sequence_cp_decomposition(
+            adata,
+            layer=SEQUENCE_INTEGER_ENCODING,
+            rank=5,
+            embedding_key="X_cp_sequence",
+            components_key="H_cp_sequence",
+            uns_key="cp_sequence",
+        )
+        nmf_layers = ["leiden", cfg.sample_name_col_for_plotting, "Reference_strand"]
+        nmf_layers += cfg.umap_layers_to_plot
+        plot_embedding(adata, basis="cp_sequence", color=nmf_layers, output_dir=nmf_sequence_dir)
+        plot_cp_sequence_components(
+            adata,
+            output_dir=nmf_sequence_dir,
+            components_key="H_cp_sequence",
+            uns_key="cp_sequence",
+        )
+
+    # ============================================================
+    # 5) Save latent AnnData
+    # ============================================================
+    if (not latent_adata_path.exists()) or getattr(cfg, "force_redo_latent_analyses", False):
+        logger.info("Saving latent analyzed AnnData (post preprocessing and duplicate removal).")
+        record_smftools_metadata(
+            adata,
+            step_name="latent",
+            cfg=cfg,
+            config_path=config_path,
+            input_paths=[source_adata_path] if source_adata_path else None,
+            output_path=latent_adata_path,
+        )
+        write_gz_h5ad(adata, latent_adata_path)
+
+    return adata, latent_adata_path

smftools/cli/load_adata.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import logging
 import shutil
 from pathlib import Path
 from typing import Iterable, Union
 
 import numpy as np
 
-from smftools.logging_utils import get_logger
+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
 
@@ -103,63 +105,29 @@ def load_adata(config_path: str):
     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")
+    from .helpers import get_adata_paths, load_experiment_config
 
     # -----------------------------
     # 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
-    )
+    cfg = load_experiment_config(config_path)
 
     # Ensure base output dir
-    make_dirs([cfg.output_directory])
+    output_directory = Path(cfg.output_directory)
+    make_dirs([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(
+            logger.info(
                 f"Raw AnnData from smftools load already exists: {paths.raw}\nSkipping smftools load"
             )
             return None, paths.raw, cfg
@@ -199,6 +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 datetime import datetime
 
     from ..informatics.bam_functions import (
         align_and_sort_BAM,
@@ -206,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
@@ -216,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
@@ -226,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
@@ -241,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"):
@@ -281,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:
@@ -295,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:
@@ -323,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}")
@@ -341,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 ###################################
@@ -373,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:
@@ -388,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.")
@@ -400,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)
     ########################################################################################################################
@@ -462,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}"
@@ -477,7 +461,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
             logger.info("Making bed files from the aligned and sorted BAM file")
             aligned_BAM_to_bed(
                 aligned_sorted_output,
-                cfg.output_directory,
+                load_directory,
                 fasta,
                 cfg.make_bigwigs,
                 cfg.threads,
@@ -515,6 +499,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
 
         se_bam_files = bam_files
         bam_dir = cfg.split_path
+        double_barcoded_path = None
 
     else:
         if single_barcoded_path.is_dir():
@@ -608,7 +593,7 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
     ################################### 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:
@@ -637,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,
@@ -694,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,
@@ -728,6 +713,25 @@ def load_adata_core(cfg, paths: AdataPaths, config_path: str | None = None):
         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)
     ########################################################################################################################
 
@@ -740,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",
         )
     ########################################################################################################################
 
@@ -759,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.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 ###############################################