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

smftools/cli/spatial_adata.py

@@ -1,8 +1,16 @@
+from __future__ import annotations
+
 from pathlib import Path
 from typing import Optional, Tuple
 
 import anndata as ad
 
+from smftools.logging_utils import get_logger
+from smftools.optional_imports import require
+
+logger = get_logger(__name__)
+
+
 def spatial_adata(
     config_path: str,
 ) -> Tuple[Optional[ad.AnnData], Optional[Path]]:
@@ -26,10 +34,10 @@ def spatial_adata(
     spatial_adata_path : Path | None
         Path to the “current” spatial AnnData (or hmm AnnData if we skip to that).
     """
-    from ..readwrite import safe_read_h5ad, make_dirs, add_or_update_column_in_csv
+    from ..readwrite import add_or_update_column_in_csv, safe_read_h5ad
+    from .helpers import get_adata_paths
     from .load_adata import load_adata
     from .preprocess_adata import preprocess_adata
-    from .helpers import get_adata_paths
 
     # 1) Ensure config + basic paths via load_adata
     loaded_adata, loaded_path, cfg = load_adata(config_path)
@@ -45,21 +53,22 @@ def spatial_adata(
     if not getattr(cfg, "force_redo_spatial_analyses", False):
         # If HMM exists, it's the most processed stage — reuse it.
         if hmm_path.exists():
-            print(f"HMM AnnData found: {hmm_path}\nSkipping smftools spatial")
+            logger.info(f"HMM AnnData found: {hmm_path}\nSkipping smftools spatial")
             return None, hmm_path
 
         # If spatial exists, we consider spatial analyses already done.
         if spatial_path.exists():
-            print(f"Spatial AnnData found: {spatial_path}\nSkipping smftools spatial")
+            logger.info(f"Spatial AnnData found: {spatial_path}\nSkipping smftools spatial")
             return None, spatial_path
 
     # 2) Ensure preprocessing has been run
     #    This will create pp/pp_dedup as needed or return them if they already exist.
-    pp_adata, pp_adata_path_ret, pp_dedup_adata, pp_dedup_adata_path_ret = preprocess_adata(config_path)
+    pp_adata, pp_adata_path_ret, pp_dedup_adata, pp_dedup_adata_path_ret = preprocess_adata(
+        config_path
+    )
 
     # Helper to load from disk, reusing loaded_adata if it matches
     def _load(path: Path):
-        from ..readwrite import safe_read_h5ad
         if loaded_adata is not None and loaded_path == path:
             return loaded_adata
         adata, _ = safe_read_h5ad(path)
@@ -69,15 +78,19 @@ def spatial_adata(
     # Prefer in-memory pp_dedup_adata when preprocess_adata just ran.
     if pp_dedup_adata is not None:
         start_adata = pp_dedup_adata
+        source_path = pp_dedup_adata_path_ret
     else:
         if 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
         elif raw_path.exists():
             start_adata = _load(raw_path)
+            source_path = raw_path
         else:
-            print("No suitable AnnData found for spatial analyses (need at least raw).")
+            logger.warning("No suitable AnnData found for spatial analyses (need at least raw).")
             return None, None
 
     # 4) Run the spatial core
@@ -88,6 +101,8 @@ def spatial_adata(
         pp_adata_path=pp_path,
         pp_dup_rem_adata_path=pp_dedup_path,
         pp_adata_in_memory=pp_adata,
+        source_adata_path=source_path,
+        config_path=config_path,
     )
 
     # 5) Register spatial path in summary CSV
@@ -103,6 +118,8 @@ def spatial_adata_core(
     pp_adata_path: Path,
     pp_dup_rem_adata_path: Path,
     pp_adata_in_memory: Optional[ad.AnnData] = None,
+    source_adata_path: Optional[Path] = None,
+    config_path: Optional[str] = None,
 ) -> Tuple[ad.AnnData, Path]:
     """
     Core spatial analysis pipeline.
@@ -139,32 +156,33 @@ def spatial_adata_core(
 
     import numpy as np
     import pandas as pd
-    import scanpy as sc
 
-    from ..readwrite import make_dirs, safe_read_h5ad
-    from .helpers import write_gz_h5ad
+    sc = require("scanpy", extra="scanpy", purpose="spatial analyses")
 
-    from ..preprocessing import (
-        load_sample_sheet,
-        invert_adata,
-        reindex_references_adata,
-    )
+    from ..metadata import record_smftools_metadata
     from ..plotting import (
         combined_raw_clustermap,
         plot_rolling_grid,
         plot_spatial_autocorr_grid,
     )
+    from ..preprocessing import (
+        invert_adata,
+        load_sample_sheet,
+        reindex_references_adata,
+    )
+    from ..readwrite import make_dirs, safe_read_h5ad
     from ..tools import calculate_umap
+    from ..tools.position_stats import (
+        compute_positionwise_statistics,
+        plot_positionwise_matrices,
+    )
     from ..tools.spatial_autocorrelation import (
-        binary_autocorrelation_with_spacing,
         analyze_autocorr_matrix,
+        binary_autocorrelation_with_spacing,
         bootstrap_periodicity,
         rolling_autocorr_metrics,
     )
-    from ..tools.position_stats import (
-        compute_positionwise_statistics,
-        plot_positionwise_matrices,
-    )
+    from .helpers import write_gz_h5ad
 
     # -----------------------------
     # General setup
@@ -207,7 +225,12 @@ def spatial_adata_core(
         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
+
     pp_dir = output_directory / "preprocessed"
     references = adata.obs[cfg.reference_column].cat.categories
 
@@ -223,7 +246,9 @@ def spatial_adata_core(
             if pp_clustermap_dir.is_dir() and not getattr(
                 cfg, "force_redo_spatial_analyses", False
             ):
-                print(f"{pp_clustermap_dir} already exists. Skipping clustermap plotting for preprocessed AnnData.")
+                logger.debug(
+                    f"{pp_clustermap_dir} already exists. Skipping clustermap plotting for preprocessed AnnData."
+                )
             else:
                 make_dirs([pp_dir, pp_clustermap_dir])
 
@@ -232,6 +257,34 @@ def spatial_adata_core(
                 else:
                     pp_adata, _ = safe_read_h5ad(pp_adata_path)
 
+                # -----------------------------
+                # Optional sample sheet metadata
+                # -----------------------------
+                if getattr(cfg, "sample_sheet_path", None):
+                    load_sample_sheet(
+                        pp_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):
+                    pp_adata = invert_adata(pp_adata)
+
+                # -----------------------------
+                # Optional reindexing by reference
+                # -----------------------------
+                reindex_references_adata(
+                    pp_adata,
+                    reference_col=cfg.reference_column,
+                    offsets=cfg.reindexing_offsets,
+                    new_col=cfg.reindexed_var_suffix,
+                )
+
                 combined_raw_clustermap(
                     pp_adata,
                     sample_col=cfg.sample_name_col_for_plotting,
@@ -247,16 +300,19 @@ def spatial_adata_core(
                     cmap_a=cfg.clustermap_cmap_a,
                     min_quality=cfg.read_quality_filter_thresholds[0],
                     min_length=cfg.read_len_filter_thresholds[0],
-                    min_mapped_length_to_reference_length_ratio=cfg.read_len_to_ref_ratio_filter_thresholds[0],
+                    min_mapped_length_to_reference_length_ratio=cfg.read_len_to_ref_ratio_filter_thresholds[
+                        0
+                    ],
                     min_position_valid_fraction=cfg.min_valid_fraction_positions_in_read_vs_ref,
+                    demux_types=("double", "already"),
                     bins=None,
                     sample_mapping=None,
                     save_path=pp_clustermap_dir,
                     sort_by=cfg.spatial_clustermap_sortby,
                     deaminase=deaminase,
-                    index_col_suffix=cfg.reindexed_var_suffix,
+                    index_col_suffix=reindex_suffix,
                 )
-        
+
     # ============================================================
     # 2) Clustermaps + UMAP on *deduplicated* preprocessed AnnData
     # ============================================================
@@ -265,10 +321,10 @@ def spatial_adata_core(
     pp_umap_dir = pp_dir_dedup / "07_umaps"
 
     # Clustermaps on deduplicated adata
-    if pp_clustermap_dir_dedup.is_dir() and not getattr(
-        cfg, "force_redo_spatial_analyses", False
-    ):
-        print(f"{pp_clustermap_dir_dedup} already exists. Skipping clustermap plotting for deduplicated AnnData.")
+    if pp_clustermap_dir_dedup.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
+        logger.debug(
+            f"{pp_clustermap_dir_dedup} already exists. Skipping clustermap plotting for deduplicated AnnData."
+        )
     else:
         make_dirs([pp_dir_dedup, pp_clustermap_dir_dedup])
         combined_raw_clustermap(
@@ -286,19 +342,22 @@ def spatial_adata_core(
             cmap_a=cfg.clustermap_cmap_a,
             min_quality=cfg.read_quality_filter_thresholds[0],
             min_length=cfg.read_len_filter_thresholds[0],
-            min_mapped_length_to_reference_length_ratio=cfg.read_len_to_ref_ratio_filter_thresholds[0],
+            min_mapped_length_to_reference_length_ratio=cfg.read_len_to_ref_ratio_filter_thresholds[
+                0
+            ],
             min_position_valid_fraction=1 - cfg.position_max_nan_threshold,
+            demux_types=("double", "already"),
             bins=None,
             sample_mapping=None,
             save_path=pp_clustermap_dir_dedup,
             sort_by=cfg.spatial_clustermap_sortby,
             deaminase=deaminase,
-            index_col_suffix=cfg.reindexed_var_suffix,
+            index_col_suffix=reindex_suffix,
         )
 
     # UMAP / Leiden
     if pp_umap_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
-        print(f"{pp_umap_dir} already exists. Skipping UMAP plotting.")
+        logger.debug(f"{pp_umap_dir} already exists. Skipping UMAP plotting.")
     else:
         make_dirs([pp_umap_dir])
 
@@ -336,40 +395,48 @@ def spatial_adata_core(
     pp_autocorr_dir = pp_dir_dedup / "08_autocorrelations"
 
     if pp_autocorr_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
-        print(f"{pp_autocorr_dir} already exists. Skipping autocorrelation plotting.")
+        logger.debug(f"{pp_autocorr_dir} already exists. Skipping autocorrelation plotting.")
     else:
         positions = adata.var_names.astype(int).values
         lags = np.arange(cfg.autocorr_max_lag + 1)
 
         try:
             from joblib import Parallel, delayed
+
             _have_joblib = True
         except Exception:
             _have_joblib = False
 
-        samples = adata.obs[cfg.sample_name_col_for_plotting].astype("category").cat.categories.tolist()
+        samples = (
+            adata.obs[cfg.sample_name_col_for_plotting].astype("category").cat.categories.tolist()
+        )
         ref_col = getattr(cfg, "reference_strand_col", "Reference_strand")
         refs = adata.obs[ref_col].astype("category").cat.categories.tolist()
 
         for site_type in cfg.autocorr_site_types:
             layer_key = f"{site_type}_site_binary"
             if layer_key not in adata.layers:
-                print(f"Layer {layer_key} not found in adata.layers — skipping {site_type}.")
+                logger.debug(f"Layer {layer_key} not found in adata.layers — skipping {site_type}.")
                 continue
 
             X = adata.layers[layer_key]
             if getattr(X, "shape", (0,))[0] == 0:
-                print(f"Layer {layer_key} empty — skipping {site_type}.")
+                logger.debug(f"Layer {layer_key} empty — skipping {site_type}.")
                 continue
 
             rows = []
             counts = []
 
             if _have_joblib:
+
                 def _worker(row):
                     try:
                         ac, cnts = binary_autocorrelation_with_spacing(
-                            row, positions, max_lag=cfg.autocorr_max_lag, return_counts=True
+                            row,
+                            positions,
+                            max_lag=cfg.autocorr_max_lag,
+                            return_counts=True,
+                            normalize=cfg.autocorr_normalization_method,
                         )
                     except Exception:
                         ac = np.full(cfg.autocorr_max_lag + 1, np.nan, dtype=np.float32)
@@ -385,7 +452,11 @@ def spatial_adata_core(
             else:
                 for i in range(X.shape[0]):
                     ac, cnts = binary_autocorrelation_with_spacing(
-                        X[i], positions, max_lag=cfg.autocorr_max_lag, return_counts=True
+                        X[i],
+                        positions,
+                        max_lag=cfg.autocorr_max_lag,
+                        return_counts=True,
+                        normalize=cfg.autocorr_normalization_method,
                     )
                     rows.append(ac)
                     counts.append(cnts)
@@ -474,7 +545,9 @@ def spatial_adata_core(
                         try:
                             r = analyze_autocorr_matrix(
                                 ac_sel,
-                                cnt_sel if cnt_sel is not None else np.zeros_like(ac_sel, dtype=int),
+                                cnt_sel
+                                if cnt_sel is not None
+                                else np.zeros_like(ac_sel, dtype=int),
                                 lags,
                                 nrl_search_bp=(120, 260),
                                 pad_factor=4,
@@ -489,7 +562,9 @@ def spatial_adata_core(
 
             adata.uns[f"{site_type}_spatial_periodicity_metrics_by_group"] = metrics_by_group
 
-            global_nrl = adata.uns.get(f"{site_type}_spatial_periodicity_metrics", {}).get("nrl_bp", None)
+            global_nrl = adata.uns.get(f"{site_type}_spatial_periodicity_metrics", {}).get(
+                "nrl_bp", None
+            )
 
             rolling_cfg = {
                 "window_size": getattr(
@@ -554,27 +629,31 @@ def spatial_adata_core(
                             fixed_nrl_bp=global_nrl,
                         )
                     except Exception as e:
-                        warnings.warn(
+                        logger.warning(
                             f"rolling_autocorr_metrics failed for {site_type} "
                             f"{sample_name} {ref_label}: {e}"
                         )
                         continue
 
                     if "center" not in df_roll.columns:
-                        warnings.warn(
+                        logger.warning(
                             f"rolling_autocorr_metrics returned unexpected schema "
                             f"for {site_type} {sample_name} {ref_label}"
                         )
                         continue
 
-                    compact_df = df_roll[["center", "n_molecules", "nrl_bp", "snr", "xi", "fwhm_bp"]].copy()
+                    compact_df = df_roll[
+                        ["center", "n_molecules", "nrl_bp", "snr", "xi", "fwhm_bp"]
+                    ].copy()
                     compact_df["site"] = site_type
                     compact_df["sample"] = sample_name
                     compact_df["reference"] = ref_label if ref_label != "all" else "all"
 
                     if write_csvs:
                         safe_sample = str(sample_name).replace(os.sep, "_")
-                        safe_ref = str(ref_label if ref_label != "all" else "all").replace(os.sep, "_")
+                        safe_ref = str(ref_label if ref_label != "all" else "all").replace(
+                            os.sep, "_"
+                        )
                         out_csv = os.path.join(
                             site_out_dir,
                             f"{safe_sample}__{safe_ref}__rolling_metrics.csv",
@@ -582,7 +661,7 @@ def spatial_adata_core(
                         try:
                             compact_df.to_csv(out_csv, index=False)
                         except Exception as e:
-                            warnings.warn(f"Failed to write rolling CSV {out_csv}: {e}")
+                            logger.warning(f"Failed to write rolling CSV {out_csv}: {e}")
 
                     if write_plots:
                         try:
@@ -604,7 +683,7 @@ def spatial_adata_core(
                                     show=False,
                                 )
                             except Exception as e:
-                                warnings.warn(
+                                logger.warning(
                                     f"Failed to create rolling plot for {site_type} "
                                     f"{sample_name} {ref_label}: {e}"
                                 )
@@ -612,7 +691,9 @@ def spatial_adata_core(
                     combined_rows.append(
                         compact_df.assign(site=site_type, sample=sample_name, reference=ref_label)
                     )
-                    rolling_results_by_group[(sample_name, None if ref_label == "all" else ref_label)] = compact_df
+                    rolling_results_by_group[
+                        (sample_name, None if ref_label == "all" else ref_label)
+                    ] = compact_df
 
             adata.uns[f"{site_type}_rolling_metrics_by_group"] = rolling_results_by_group
 
@@ -624,9 +705,7 @@ def spatial_adata_core(
                 try:
                     combined_df_site.to_csv(combined_out_csv, index=False)
                 except Exception as e:
-                    warnings.warn(
-                        f"Failed to write combined rolling CSV for {site_type}: {e}"
-                    )
+                    logger.warning(f"Failed to write combined rolling CSV for {site_type}: {e}")
 
             rolling_dict = adata.uns[f"{site_type}_rolling_metrics_by_group"]
             plot_out_dir = os.path.join(pp_autocorr_dir, "rolling_plots")
@@ -650,6 +729,7 @@ def spatial_adata_core(
                 sample_col=cfg.sample_name_col_for_plotting,
                 window=cfg.autocorr_rolling_window_size,
                 rows_per_fig=cfg.rows_per_qc_autocorr_grid,
+                normalization_method=cfg.autocorr_normalization_method,
             )
 
     # ============================================================
@@ -658,7 +738,7 @@ def spatial_adata_core(
     pp_corr_dir = pp_dir_dedup / "09_correlation_matrices"
 
     if pp_corr_dir.is_dir() and not getattr(cfg, "force_redo_spatial_analyses", False):
-        print(f"{pp_corr_dir} already exists. Skipping correlation matrix plotting.")
+        logger.debug(f"{pp_corr_dir} already exists. Skipping correlation matrix plotting.")
     else:
         compute_positionwise_statistics(
             adata,
@@ -691,7 +771,15 @@ def spatial_adata_core(
     # 5) Save spatial AnnData
     # ============================================================
     if (not spatial_adata_path.exists()) or getattr(cfg, "force_redo_spatial_analyses", False):
-        print("Saving spatial analyzed AnnData (post preprocessing and duplicate removal).")
+        logger.info("Saving spatial analyzed AnnData (post preprocessing and duplicate removal).")
+        record_smftools_metadata(
+            adata,
+            step_name="spatial",
+            cfg=cfg,
+            config_path=config_path,
+            input_paths=[source_adata_path] if source_adata_path else None,
+            output_path=spatial_adata_path,
+        )
         write_gz_h5ad(adata, spatial_adata_path)
 
-    return adata, spatial_adata_path
+    return adata, spatial_adata_path