spatial-aim 1.0.0__py3-none-any.whl

aim/analysis/ksweep.py

@@ -0,0 +1,184 @@
+"""Cross-K comparison: gather each K's key analysis metrics from its per-K
+``analysis/data`` outputs into one table at the run root (``k_comparison.csv``),
+score them, and record which K each criterion selects (``k_selection.json``).
+
+``k_comparison.csv`` holds the measured metrics *and* what they reduce to: one
+combined score per criterion, the overall score across criteria, and a ``pareto``
+flag. Reads the files each K's post-mapping analysis already writes
+(``cossim_summary.csv``, ``topology_metrics.json``, ``modularity_metrics.json``),
+so it runs after the K-loop with no recomputation — including the label-shuffle
+nulls those steps measured, which the GUI needs alongside the observed values.
+Metrics missing for a given K come through as NaN. The scoring itself lives in
+:mod:`aim.metrics.kselection`, so a CLI run and the GUI's "Comparing K" card
+report the same numbers.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from aim.metrics import kselection
+
+logger = logging.getLogger(__name__)
+
+
+def _load_json(path: Path) -> dict:
+    """Load a JSON dict, or return an empty dict if the file is absent/unreadable."""
+    if not path.exists():
+        return {}
+    try:
+        with open(path, encoding="utf-8") as fh:
+            return json.load(fh)
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def _finite(x) -> float:
+    """Coerce a value to float, mapping non-finite (inf / NaN / None) to NaN so it
+    leaves a gap in the plot rather than distorting an autoscaled axis."""
+    try:
+        f = float(x)
+    except (TypeError, ValueError):
+        return float("nan")
+    return f if np.isfinite(f) else float("nan")
+
+
+def _cossim_cell(cs: pd.DataFrame, row: str, column: str) -> float:
+    """One cell of ``cossim_summary.csv`` as a float, NaN when absent."""
+    if row not in cs.index or column not in cs.columns:
+        return float("nan")
+    return _finite(cs.loc[row, column])
+
+
+def collect_k_metrics(output_folder: Path, ks: list[int]) -> pd.DataFrame:
+    """Gather the per-K comparison metrics into one DataFrame (one row per K, ascending).
+
+    Columns: ``k``, ``cossim_hard_norm_spot``, ``cossim_hard_norm_gene``,
+    ``cossim_hard_raw_spot``, ``cossim_hard_raw_gene``, ``nhood_mean_self_zscore``,
+    ``local_purity_zscore``, ``modularity_shared``, ``modularity_st_expression``,
+    ``mean_confidence``, plus the ``*_null`` label-shuffle baseline of every metric
+    that has one (the two z-scores are already measured against that shuffle, so
+    their null is 0 by construction and is not repeated here).
+    """
+    rows = []
+    for k in ks:
+        data_dir = output_folder / f"k_{k:03d}" / "analysis" / "data"
+
+        cossim_norm_spot = cossim_norm_gene = np.nan
+        cossim_raw_spot = cossim_raw_gene = np.nan
+        cossim_raw_spot_null = cossim_raw_gene_null = np.nan
+        cossim_path = data_dir / "cossim_summary.csv"
+        if cossim_path.exists():
+            cs = pd.read_csv(cossim_path, index_col=0)
+            cossim_norm_spot = _cossim_cell(cs, "hard-norm", "median_spot")
+            cossim_norm_gene = _cossim_cell(cs, "hard-norm", "median_gene")
+            cossim_raw_spot = _cossim_cell(cs, "hard-raw", "median_spot")
+            cossim_raw_gene = _cossim_cell(cs, "hard-raw", "median_gene")
+            cossim_raw_spot_null = _cossim_cell(cs, "hard-raw", "median_spot_null")
+            cossim_raw_gene_null = _cossim_cell(cs, "hard-raw", "median_gene_null")
+
+        topo = _load_json(data_dir / "topology_metrics.json")
+        # z-scores can be non-finite (inf) when a permutation null has zero variance
+        # for a degenerate/tiny state; coerce to NaN so they leave a gap instead of
+        # blowing up the plot's autoscaled y-axis.
+        # These sub-dicts are stored as null when a K has too few states for the
+        # test (e.g. K=1), so coalesce None -> {} before indexing.
+        nhood = topo.get("nhood_enrichment") or {}
+        purity = topo.get("local_purity") or {}
+        self_z = _finite(nhood.get("mean_self_zscore", np.nan))
+        purity_z = _finite(purity.get("z_score", np.nan))
+
+        modularity = _load_json(data_dir / "modularity_metrics.json")
+        mod_shared = _finite(modularity.get("modularity_shared", np.nan))
+        mod_st = _finite(modularity.get("modularity_st_expression", np.nan))
+
+        # Present only for mappers that define a confidence (nearest / wann);
+        # absent -> NaN -> the curve simply does not appear for this run.
+        confidence = _load_json(data_dir / "confidence_summary.json")
+        mean_conf = confidence.get("mean", np.nan)
+
+        rows.append(
+            {
+                "k": k,
+                "cossim_hard_norm_spot": cossim_norm_spot,
+                "cossim_hard_norm_gene": cossim_norm_gene,
+                "cossim_hard_raw_spot": cossim_raw_spot,
+                "cossim_hard_raw_gene": cossim_raw_gene,
+                "cossim_hard_raw_spot_null": cossim_raw_spot_null,
+                "cossim_hard_raw_gene_null": cossim_raw_gene_null,
+                "nhood_mean_self_zscore": self_z,
+                "local_purity_zscore": purity_z,
+                "modularity_shared": mod_shared,
+                "modularity_st_expression": mod_st,
+                "mean_confidence": mean_conf,
+            }
+        )
+
+    return pd.DataFrame(rows).sort_values("k").reset_index(drop=True)
+
+
+def score_k_metrics(df: pd.DataFrame) -> pd.DataFrame:
+    """``df`` with the K-selection columns appended: one combined score per
+    criterion, the ``overall`` score across them, and a ``pareto`` flag.
+
+    The combined scores' own label-shuffle nulls are not carried here — only
+    reconstruction has one, and the GUI recomputes it in memory for the chance
+    crosshair; the measured nulls it is derived from are already columns of
+    ``df``.
+    """
+    table = kselection.score_table(df)
+    out = df.copy()
+    for criterion in kselection.CRITERIA:
+        out[criterion.key] = table[criterion.key].to_numpy()
+    out["overall"] = kselection.overall(table)
+    out["pareto"] = kselection.pareto_mask(table)
+    return out
+
+
+def select_k(output_folder: Path, df: pd.DataFrame) -> dict:
+    """Write ``k_selection.json``: which K wins, per criterion and overall.
+
+    The proposal is the Pareto-optimal K plus the best K for each criterion and
+    for the harmonic mean across all of them, each with its score. Same
+    computation the GUI's "Comparing K" card performs, so the CLI no longer
+    hides it. Returns the proposal dict.
+    """
+    table = kselection.score_table(df)
+    best = kselection.best_ks(table)
+    best_json: dict[str, dict[str, float] | None] = {
+        name: (None if hit is None else {"k": hit[0], "score": hit[1]})
+        for name, hit in best.items()
+    }
+    proposal = {
+        "pareto_ks": [int(k) for k in table.loc[kselection.pareto_mask(table), "k"]],
+        "best": best_json,
+        "criteria": {c.key: c.label for c in kselection.CRITERIA},
+    }
+    json_path = output_folder / "k_selection.json"
+    with open(json_path, "w", encoding="utf-8") as fh:
+        json.dump(proposal, fh, indent=2)
+
+    chosen = ", ".join(
+        f"{name}={hit[0]}" for name, hit in best.items() if hit is not None
+    )
+    logger.info("K selection written to %s (%s)", json_path, chosen)
+    return proposal
+
+
+def compare_k_runs(output_folder: Path, ks: list[int]) -> None:
+    """Write ``k_comparison.csv`` at the run root: the swept K-runs' reconstruction
+    cosine similarity, spatial organisation and mapping modularity, together with
+    the criterion scores those reduce to, the overall score and the Pareto flag
+    (the GUI renders the comparison plots from it on demand). The winning K per
+    criterion and overall goes to ``k_selection.json`` (see :func:`select_k`)."""
+    output_folder = Path(output_folder)
+    df = collect_k_metrics(output_folder, ks)
+    csv_path = output_folder / "k_comparison.csv"
+    score_k_metrics(df).to_csv(csv_path, index=False)
+    logger.info("K-sweep comparison written to %s", csv_path)
+    select_k(output_folder, df)

aim/analysis/loading.py

@@ -0,0 +1,70 @@
+"""Load one K's saved sweep outputs for the post-mapping analysis."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+
+from aim.adata_schema import (
+    OBSM_MAPPING_SOFT,
+    OBS_MAPPING_HARD,
+    OBS_MAPPING_CONFIDENCE,
+    OBS_START_CLUSTER,
+    OBS_COMPUTED_STATE,
+)
+from aim.io import START_CLUSTER_TO_STATE_FILENAME
+
+logger = logging.getLogger(__name__)
+
+
+def load_spot_to_state_mapping_soft_and_hard(run_dir: Path, adata_st: AnnData) -> None:
+    """Load the spot->state matrix P from spot_to_state_mapping_soft.h5ad in run_dir.
+
+    Adds: adata_st.obsm[OBSM_MAPPING_SOFT] (S x K), assigned positionally, and,
+    when the mapper wrote one, adata_st.obs[OBS_MAPPING_CONFIDENCE] (S,).
+    Raises ValueError if the file's spot order does not match adata_st.obs_names.
+    """
+    run_dir = Path(run_dir)
+    mapping_path = run_dir / "spot_to_state_mapping_soft.h5ad"
+    if not mapping_path.exists():
+        raise FileNotFoundError(f"Required mapping output missing: {mapping_path}")
+
+    mapping_ad = ad.read_h5ad(mapping_path)
+    if not mapping_ad.obs_names.equals(adata_st.obs_names):
+        raise ValueError(
+            f"Spot order in {mapping_path} does not match adata_st.obs_names."
+        )
+    adata_st.obsm[OBSM_MAPPING_SOFT] = np.asarray(mapping_ad.X, dtype=np.float64)
+    adata_st.obs[OBS_MAPPING_HARD] = adata_st.obsm[OBSM_MAPPING_SOFT].argmax(axis=1)
+    if OBS_MAPPING_CONFIDENCE in mapping_ad.obs:
+        adata_st.obs[OBS_MAPPING_CONFIDENCE] = mapping_ad.obs[
+            OBS_MAPPING_CONFIDENCE
+        ].to_numpy()
+
+
+def load_start_cluster_to_state(run_dir: Path) -> np.ndarray:
+    """Load the start-cluster->state label array (L,), values 0..K-1, from
+    start_cluster_to_state.csv in run_dir.
+
+    Falls back to the pre-rename ``leiden_to_state.csv`` so run roots written by an
+    older sweep still load; only the ``state`` column is read either way.
+    """
+    run_dir = Path(run_dir)
+    path = run_dir / START_CLUSTER_TO_STATE_FILENAME
+    if not path.exists():
+        legacy = run_dir / "leiden_to_state.csv"
+        if not legacy.exists():
+            raise FileNotFoundError(f"Required mapping output missing: {path}")
+        path = legacy
+    return pd.read_csv(path)["state"].to_numpy()
+
+
+def infer_cell_to_state_cluster(adata_sc: AnnData, start_cluster_to_state: np.ndarray):
+    start_cluster_idx = adata_sc.obs[OBS_START_CLUSTER].astype(int).to_numpy()
+    cell_states = start_cluster_to_state[start_cluster_idx]
+    adata_sc.obs[OBS_COMPUTED_STATE] = pd.Categorical(cell_states.astype(str))

aim/analysis/onehot.py

@@ -0,0 +1,41 @@
+import json
+import logging
+from pathlib import Path
+import pandas as pd
+from anndata import AnnData
+
+from aim.adata_schema import (
+    OBSM_MAPPING_SOFT,
+)
+from aim.metrics.onehot import onehot_metrics
+
+logger = logging.getLogger(__name__)
+
+
+def analyse_spot_to_state_one_hotness(
+    adata_st: AnnData,
+    data_dir: Path,
+):
+    """One-hotness metrics for the spot->state mapping P.
+
+    Requires: adata_st.obsm[OBSM_MAPPING_SOFT].
+    Writes onehot_per_row_mapping.csv and onehot_summary_mapping.json under
+    data_dir.
+    """
+
+    m = onehot_metrics(adata_st.obsm[OBSM_MAPPING_SOFT])
+
+    pd.DataFrame(
+        {
+            "id": adata_st.obs_names,
+            "max_prob": m["max_prob"],
+            "gini_impurity": m["gini_impurity"],
+            "entropy": m["entropy"],
+        }
+    ).to_csv(data_dir / "onehot_per_row_mapping.csv", index=False)
+    with open(data_dir / "onehot_summary_mapping.json", "w") as f:
+        json.dump(
+            {"n_rows": m["n_rows"], "n_cols": m["n_cols"], "summary": m["summary"]},
+            f,
+            indent=2,
+        )

aim/analysis/reconstruction.py

@@ -0,0 +1,131 @@
+"""Reconstruction cosine similarity of the mapped spots vs. measured ST expression."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import anndata as ad
+import pandas as pd
+
+from aim.adata_schema import (
+    OBS_MAPPING_HARD,
+    OBSM_LOGNORM_SHARED_GENES,
+    OBSM_MAPPING_SOFT,
+    UNS_SHARED_GENES,
+    UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES,
+    UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES_NORM,
+    UNS_START_CLUSTER_SIZES,
+)
+from aim.metrics import (
+    CossimResult,
+    assemble_state_centroids,
+    compute_and_save_cossim,
+    cossim_null_medians,
+    predict_expression,
+)
+
+from .utils import to_dense
+
+logger = logging.getLogger(__name__)
+
+
+N_PERM_RECONSTRUCTION = 20  # label shuffles for the reconstruction null
+
+
+def analyse_reconstruction(
+    adata_sc: ad.AnnData,
+    adata_st: ad.AnnData,
+    start_cluster_to_state: np.ndarray,
+    output_data_dir: Path,
+    n_perm: int = N_PERM_RECONSTRUCTION,
+    seed: int = 0,
+) -> None:
+    """Cosine similarity of predicted (``mapping @ state_centroids``) vs. measured
+    spot expression, per gene and per spot, for the four soft/hard x raw/norm combos.
+
+    Each *hard* combo also gets a label-shuffling null (``median_*_null``): the same
+    medians with the spot->state labels permuted, averaged over ``n_perm`` shuffles.
+    That is the chance level of the reconstruction — with only positive expression
+    profiles in play it is far above zero, so it is the reference the K-sweep view
+    compares against. The soft combos have no null (their assignment is a
+    distribution, not a label) and come through as NaN.
+
+    ``start_cluster_to_state`` is this K's start-cluster->state cut (L,).
+    Requires: adata_st.obsm[OBSM_MAPPING_SOFT], adata_st.obs[OBS_MAPPING_HARD],
+        adata_st.obsm[OBSM_LOGNORM_SHARED_GENES], adata_sc.uns[UNS_SHARED_GENES],
+        adata_sc.uns[UNS_START_CLUSTER_SIZES],
+        adata_sc.uns[UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES],
+        adata_sc.uns[UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES_NORM].
+    Writes cossim_summary.csv and per-combo cossim JSONs under output_data_dir.
+    """
+
+    P = adata_st.obsm[OBSM_MAPPING_SOFT]
+    spot_states_hard = adata_st.obs[OBS_MAPPING_HARD].to_numpy()
+    k = P.shape[1]
+
+    shared_genes = list(adata_sc.uns[UNS_SHARED_GENES])
+    sizes = adata_sc.uns[UNS_START_CLUSTER_SIZES]
+
+    # Measured ST expression on the shared genes: raw counts from a view, and the
+    # normalize+log1p variant precomputed over the shared genes only.
+    adata_st_shared = adata_st[:, shared_genes]
+    adata_st_norm = ad.AnnData(
+        X=to_dense(adata_st.obsm[OBSM_LOGNORM_SHARED_GENES]),
+        obs=pd.DataFrame(index=adata_st.obs_names),
+        var=pd.DataFrame(index=shared_genes),
+    )
+
+    centroids_raw = assemble_state_centroids(
+        start_cluster_to_state,
+        k,
+        adata_sc.uns[UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES],
+        sizes,
+    )
+    centroids_norm = assemble_state_centroids(
+        start_cluster_to_state,
+        k,
+        adata_sc.uns[UNS_START_CLUSTER_EXPR_SUMS_SHARED_GENES_NORM],
+        sizes,
+    )
+
+    # "hard" combos index each spot's winning-state centroid row directly;
+    # "soft" combos are a genuine mapping @ centroids product.
+    spot_names = adata_st.obs_names.tolist()
+    combos = {
+        "soft-raw": (P, centroids_raw, adata_st_shared, False),
+        "hard-raw": (spot_states_hard, centroids_raw, adata_st_shared, True),
+        "soft-norm": (P, centroids_norm, adata_st_norm, False),
+        "hard-norm": (spot_states_hard, centroids_norm, adata_st_norm, True),
+    }
+    cossim_summary: dict[str, dict] = {}
+    cossim_dir = output_data_dir / "cossim"
+    cossim_results: dict[str, CossimResult] = {}
+    rng = np.random.default_rng(seed)
+    for label, (m, centroids, st_ref, is_hard) in combos.items():
+        pred = centroids[m] if is_hard else predict_expression(m, centroids)
+        pred_adata = ad.AnnData(
+            X=pred.T.astype(np.float32),
+            obs=pd.DataFrame(index=shared_genes),
+            var=pd.DataFrame(index=spot_names),
+        )
+        result = compute_and_save_cossim(
+            st_ref, pred_adata, cossim_dir, suffix=f"-{label}"
+        )
+        cossim_results[label] = result
+        # ``st_ref`` is already restricted to the shared genes in centroid column
+        # order, so the null reduces the very same matrices as the observed value.
+        null = (
+            cossim_null_medians(to_dense(st_ref), centroids, m, n_perm=n_perm, rng=rng)
+            if label == "hard-raw"
+            else {"median_spot": float("nan"), "median_gene": float("nan")}
+        )
+        cossim_summary[label] = {
+            "median_gene": result.median_gene,
+            "median_spot": result.median_spot,
+            "median_gene_null": null["median_gene"],
+            "median_spot_null": null["median_spot"],
+        }
+
+    pd.DataFrame(cossim_summary).T.to_csv(output_data_dir / "cossim_summary.csv")

aim/analysis/topology.py

@@ -0,0 +1,142 @@
+"""Spatial organisation of the mapped spots for the post-mapping analysis."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+import anndata as ad
+import squidpy as sq
+
+from aim.adata_schema import (
+    OBS_MAPPING_HARD,
+    OBSM_SPATIAL,
+    OBSP_SPATIAL_CONNECTIVITIES,
+)
+from aim.metrics.topology import local_spatial_purity, permutation_test
+
+logger = logging.getLogger(__name__)
+
+
+K_SPATIAL = 6  # neighbours for the squidpy spatial KNN graph
+N_PERM_SPATIAL = 100  # permutations for the spatial-organisation null
+
+
+def spatial_connectivities(adata_st: ad.AnnData, k: int):
+    """Squidpy spatial KNN connectivity matrix, built once and cached on adata_st.
+
+    Depends only on the spot coordinates, which are constant across the K-sweep,
+    so it is computed on the first K and reused for every later K. Serves both the
+    local spatial purity metric and the neighbourhood enrichment. Returns the
+    (n_spots x n_spots) sparse adjacency in adata_st.obsp[OBSP_SPATIAL_CONNECTIVITIES].
+    """
+    if OBSP_SPATIAL_CONNECTIVITIES not in adata_st.obsp:
+        sq.gr.spatial_neighbors_knn(adata_st, n_neighs=k)
+    return adata_st.obsp[OBSP_SPATIAL_CONNECTIVITIES]
+
+
+def analyse_spatial_organization(
+    adata_st: ad.AnnData,
+    output_data_dir: Path,
+    k: int = K_SPATIAL,
+    n_perm: int = N_PERM_SPATIAL,
+    seed: int = 0,
+):
+    """Local spatial purity + self neighbourhood enrichment of the mapped spot states.
+
+    Local spatial purity comes with a permutation-null z-score; the self
+    neighbourhood-enrichment z-score summarises how much each state preferentially
+    borders itself (see ``_nhood_enrichment``). Individually undefined metrics come
+    back as NaN/None; asserts len(coords) == len(spot_states) and
+    len(spot_states) > k.
+
+    Requires: adata_st.obs[OBS_MAPPING_HARD], adata_st.obsm[OBSM_SPATIAL].
+    Writes topology_metrics.json under output_data_dir.
+    """
+    spot_states = adata_st.obs[OBS_MAPPING_HARD].to_numpy()
+    coords = np.asarray(adata_st.obsm[OBSM_SPATIAL])
+
+    result: dict = {
+        "n_spots": int(len(spot_states)),
+        "n_mapped_states": int(len(np.unique(spot_states))),
+        "k": int(k),
+        "n_perm": int(n_perm),
+        "local_purity": None,
+        "nhood_enrichment": None,
+    }
+    assert len(coords) == len(
+        spot_states
+    ), "Spatial organisation skipped: coordinates length mismatch "
+    assert len(spot_states) > k, "too few spots"
+
+    rng = np.random.default_rng(seed)
+
+    conn = spatial_connectivities(adata_st, k)
+    obs_lsp = local_spatial_purity(spot_states, conn)
+    result["local_purity"] = permutation_test(
+        obs_lsp, spot_states, lambda l: local_spatial_purity(l, conn), n_perm, rng
+    )
+
+    result["nhood_enrichment"] = _nhood_enrichment(
+        adata_st, spot_states, k, n_perm, seed
+    )
+
+    with open(output_data_dir / "topology_metrics.json", "w") as f:
+        json.dump(result, f, indent=4)
+
+
+def _nhood_enrichment(
+    adata_st: ad.AnnData,
+    spot_states: np.ndarray,
+    k: int,
+    n_perm: int,
+    seed: int,
+) -> dict | None:
+    """Self neighbourhood-enrichment z-score of the mapped states, averaged over
+    states. Needs >= 2 mapped states (returns None otherwise).
+
+    Only the *self* term is computed -- the diagonal of squidpy's neighbourhood
+    enrichment -- since ``mean_self_zscore`` is the sole value consumed downstream.
+    For each state ``s`` we count the spatial-graph edges whose two endpoints are
+    both in ``s`` and z-score that against a label-shuffling null (the same null
+    ``sq.gr.nhood_enrichment`` uses), then average the per-state z-scores. This
+    skips squidpy's full K x K enrichment matrix and its per-K call overhead, so it
+    is markedly cheaper than a full ``nhood_enrichment`` run. A high value means
+    states preferentially border their own kind.
+
+    Values are comparable to (not bit-identical with) squidpy's diagonal: same
+    graph and null, different permutation RNG.
+    """
+    if len(np.unique(spot_states)) < 2:
+        return None
+
+    A = spatial_connectivities(adata_st, k).tocoo()
+    row, col = A.row, A.col
+    n_states = int(spot_states.max()) + 1
+
+    def within_state_edges(labels: np.ndarray) -> np.ndarray:
+        """Per-state count of spatial edges whose endpoints share that state."""
+        same = labels[row] == labels[col]
+        return np.bincount(labels[row][same], minlength=n_states).astype(np.float64)
+
+    observed = within_state_edges(spot_states)
+    rng = np.random.default_rng(seed)
+    null = np.stack(
+        [within_state_edges(rng.permutation(spot_states)) for _ in range(n_perm)]
+    )
+    mean = null.mean(axis=0)
+    std = null.std(axis=0)
+
+    present = np.unique(spot_states)
+    z = np.full(n_states, np.nan)
+    ok = std > 0
+    z[ok] = (observed[ok] - mean[ok]) / std[ok]
+    self_z = z[present]
+    return {
+        "n_states": int(len(present)),
+        "mean_self_zscore": (
+            float(np.nanmean(self_z)) if np.isfinite(self_z).any() else float("nan")
+        ),
+    }

aim/analysis/utils.py

@@ -0,0 +1,16 @@
+"""Shared array/AnnData helpers for the post-mapping analysis."""
+
+from __future__ import annotations
+import numpy as np
+
+__all__ = ["to_dense"]
+
+
+def to_dense(x, dtype=np.float32) -> np.ndarray:
+    """Return ``x`` as a dense ndarray, unwrapping AnnData ``.X`` and densifying
+    scipy-sparse inputs."""
+    if hasattr(x, "X"):
+        x = x.X
+    if hasattr(x, "toarray"):
+        x = x.toarray()
+    return np.asarray(x, dtype=dtype)