smftools 0.2.5__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