diff --git a/src/ops_model/data/paths.py b/src/ops_model/data/paths.py index d795adc4..332c83a0 100644 --- a/src/ops_model/data/paths.py +++ b/src/ops_model/data/paths.py @@ -18,10 +18,26 @@ def _resolve_base() -> Path: ) ) + @staticmethod + def _resolve_model_base() -> Path: + """Base dir for read-only model assets (checkpoints). + + Anchored to the canonical data root (OPS_BASE_PATH), NOT the output dir: + models are inputs and must not follow OPS_OUTPUT_BASE_DIR, which + redirects to rerun/research trees where the checkpoints don't exist. + Override with OPS_MODELS_BASE_DIR if models live elsewhere. + """ + return Path( + os.environ.get( + "OPS_MODELS_BASE_DIR", + os.environ.get("OPS_BASE_PATH", "/hpc/projects/icd.fast.ops"), + ) + ) + @classmethod def model_checkpoints_dir(cls) -> Path: - """Root directory holding all model checkpoints.""" - return cls._resolve_base() / "models" / "model_checkpoints" + """Root directory holding all model checkpoints (read-only inputs).""" + return cls._resolve_model_base() / "models" / "model_checkpoints" @classmethod def checkpoint(cls, *parts: str) -> Path: diff --git a/src/ops_model/post_process/combination/pca_optimization/phase2.py b/src/ops_model/post_process/combination/pca_optimization/phase2.py index 6d094df2..7176afd2 100644 --- a/src/ops_model/post_process/combination/pca_optimization/phase2.py +++ b/src/ops_model/post_process/combination/pca_optimization/phase2.py @@ -84,6 +84,7 @@ def aggregate_channels( agg_method: str = "mean", chromosome_csv: Optional[str] = None, umap_type: str = "max", + leiden_resolutions: Optional[Tuple[float, ...]] = None, ) -> str: """Load per-channel (or per-signal) h5ads, concatenate, normalize, score, save. @@ -367,6 +368,9 @@ def aggregate_channels( chad_map=chad_map, chad_path_override=CHAD_ANNOTATION_PATH, _logger=_logger, + # None -> save_extra_overlays keeps its full DEFAULT_LEIDEN_RESOLUTIONS + **({"leiden_resolutions": leiden_resolutions} + if leiden_resolutions is not None else {}), ) # Re-save h5ads now that leiden_r* columns + neighbors graph have been # added to the in-memory adata objects by save_extra_overlays diff --git a/src/ops_model/post_process/per_exp_embedding.py b/src/ops_model/post_process/per_exp_embedding.py new file mode 100644 index 00000000..83a06d02 --- /dev/null +++ b/src/ops_model/post_process/per_exp_embedding.py @@ -0,0 +1,326 @@ +"""Per-experiment embedding post-processing. + +Runs the pca_optimization aggregation machinery for a SINGLE experiment, one +marker (reporter) at a time, reusing ``aggregate_channels`` so all of its rich +outputs are produced (UMAP / PHATE overlays + interactive HTMLs, mAP +consistency / distinctiveness bars, sweep plots, coord CSVs, gene/guide h5ads). +On top of that it adds a gene x gene correlation heatmap (PNG / SVG / interactive +HTML + a downloadable CSV of the correlation values) and records every +post-processing decision in ``decisions.yaml``. + +There is no cross-experiment correction and no second-pass PCA — those only +matter when combining multiple experiments / markers; here each marker of the +single experiment is embedded on its own. + +All outputs land in the existing CellDINO dir: +``/3-assembly/cell_dino_features_v2/embeddings//``. +""" + +from __future__ import annotations + +import glob +from dataclasses import asdict, dataclass, field +from pathlib import Path + +import numpy as np +import yaml + + +# --- Post-processing decisions ------------------------------------------------ +# Captured in embeddings//decisions.yaml instead of being encoded in +# directory names (the multi-exp pipeline uses a nested path-per-decision layout). +@dataclass +class EmbeddingDecisions: + distance: str = "cosine" # cosine | euclidean (mAP / consistency scoring) + pca_variance: float = 0.80 # fraction of variance kept by the correlation-heatmap PCA + norm_method: str = "ntc" # ntc | global + zscore_per_experiment: bool = True + agg_method: str = "mean" # cells->guides / guides->genes reduction: mean | median + umap_type: str = "max" + # Leiden clustering resolutions for the overlays + GO enrichment. The full + # multi-exp default is ~13 resolutions; per-exp we only need a few (GO + # enrichment is the long pole and scales with resolution count). + leiden_resolutions: list = field(default_factory=lambda: [4.0, 10.0, 30.0]) + random_seed: int = 42 + # populated per marker at runtime + marker: str | None = None + n_genes: int | None = None + n_pcs: int | None = None + source_gene_h5ad: str | None = None + + @classmethod + def from_yaml(cls, path: str | Path) -> "EmbeddingDecisions": + data = yaml.safe_load(Path(path).read_text()) or {} + known = {f for f in cls.__dataclass_fields__} + return cls(**{k: v for k, v in data.items() if k in known}) + + +def _pca_to_variance(X: np.ndarray, variance: float, seed: int) -> tuple[np.ndarray, int]: + """PCA keeping enough components to explain ``variance`` fraction. + + sklearn accepts a float n_components as a variance target directly. + """ + from sklearn.decomposition import PCA + + pca = PCA(n_components=variance, svd_solver="full", random_state=seed) + X_pca = pca.fit_transform(X) + return X_pca, int(pca.n_components_) + + +def _correlation_heatmap(X_ops: np.ndarray, labels: list[str], out_stem: Path, marker: str): + """gene x gene correlation of mean-centered embeddings -> PNG + SVG + HTML. + + This is the "image half" of the paper joint heatmap: + X_ops_c = X_ops - X_ops.mean(axis=0, keepdims=True) + corr_ops = np.corrcoef(X_ops_c) + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + plt.rcParams["pdf.fonttype"] = 42 # editable text in vector output + + Path(out_stem).parent.mkdir(parents=True, exist_ok=True) + + X_ops_c = X_ops - X_ops.mean(axis=0, keepdims=True) + corr_ops = np.corrcoef(X_ops_c) + + # Downloadable values: gene x gene correlation matrix as CSV (labelled). + import pandas as pd + + pd.DataFrame(corr_ops, index=labels, columns=labels).to_csv(f"{out_stem}.csv") + + # static PNG + SVG + n = corr_ops.shape[0] + fig, ax = plt.subplots(figsize=(10, 10)) + im = ax.imshow(corr_ops, cmap="RdBu_r", vmin=-1, vmax=1, interpolation="nearest") + ax.set_title(f"{marker}: gene x gene embedding correlation (n={n})") + ax.set_xticks([]) + ax.set_yticks([]) + fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Pearson r") + fig.tight_layout() + fig.savefig(f"{out_stem}.png", dpi=200) + fig.savefig(f"{out_stem}.svg") + plt.close(fig) + + # interactive HTML (hover shows the gene pair + r) + import plotly.graph_objects as go + + fig = go.Figure( + go.Heatmap( + z=corr_ops, x=labels, y=labels, zmin=-1, zmax=1, colorscale="RdBu_r", + colorbar=dict(title="Pearson r"), + ) + ) + fig.update_layout( + title=f"{marker}: gene x gene embedding correlation (n={n})", + width=900, height=900, xaxis_showticklabels=False, yaxis_showticklabels=False, + ) + fig.write_html(f"{out_stem}.html", include_plotlyjs="cdn") + + return corr_ops + + +def _organize_plots(plots_dir: Path) -> None: + """Group the flat aggregate_channels plot files into category subdirs so the + output isn't a sprawling list of PNGs. Files are matched by name; anything + unmatched and any existing subdir (leiden/, canonical_leiden/, + marker_overlay/, ...) is left in place. + """ + if not plots_dir.is_dir(): + return + # (subdir, predicate) in priority order — first match wins. + groups = [ + ("umap", lambda n: "umap" in n), + ("phate", lambda n: "phate" in n), + # EBI complex + binary overlays share one subdir; checked before + # map_metrics (they contain "ebi"). + ("ebi_overlay", lambda n: "ebi_complex_overlay" in n or "ebi_binary_overlay" in n), + ("map_metrics", lambda n: n.startswith("map_") or "violin" in n or "consistency" in n + or "distinctiveness" in n or "activity" in n or "ebi" in n), + ("sweep", lambda n: "sweep" in n), + ("channel_qc", lambda n: "peak" in n or "per_channel" in n), + ] + for f in list(plots_dir.iterdir()): + if not f.is_file(): + continue # leave existing subdirs untouched + n = f.name.lower() + for sub, match in groups: + if match(n): + dest = plots_dir / sub + dest.mkdir(exist_ok=True) + f.rename(dest / f.name) + break + + +def run_marker( + gene_h5ad: Path, + guide_h5ad: Path, + out_dir: Path, + decisions: EmbeddingDecisions, +) -> dict: + """Post-process a single marker: rich aggregate_channels outputs + corr heatmap. + + ``gene_h5ad`` / ``guide_h5ad`` are the CellDINO combine's gene_bulked / + guide_bulked h5ads for this marker. + """ + import anndata as ad + + from ops_model.post_process.combination.pca_optimization.phase2 import ( + aggregate_channels, + ) + + marker = decisions.marker + out_dir.mkdir(parents=True, exist_ok=True) + + # aggregate_channels reads /per_channel/_{guide,gene}.h5ad. + # CellDINO writes {guide,gene}_bulked_.h5ad, so bridge the filename with + # a hardlink (same inode: no data copy, no symlink / broken-link risk); fall + # back to a copy only across filesystems. + import os + import shutil + + per_channel = out_dir / "per_channel" + per_channel.mkdir(parents=True, exist_ok=True) + for src, name in ((guide_h5ad, f"{marker}_guide.h5ad"), (gene_h5ad, f"{marker}_gene.h5ad")): + dst = per_channel / name + if dst.exists() or dst.is_symlink(): + dst.unlink() + try: + os.link(Path(src).resolve(), dst) + except OSError: + shutil.copy2(Path(src).resolve(), dst) + + # Reuse the full aggregation + plotting machinery (UMAP/PHATE/mAP/sweep/HTML). + agg_result = aggregate_channels( + output_dir=str(out_dir), + norm_method=decisions.norm_method, + per_unit_subdir="per_channel", + distance=decisions.distance, + random_seed=decisions.random_seed, + agg_method=decisions.agg_method, + umap_type=decisions.umap_type, + leiden_resolutions=tuple(decisions.leiden_resolutions), + ) + + # Group the flat aggregate_channels plots into category subdirs (de-sprawl). + _organize_plots(out_dir / "plots") + # canonical_leiden holds the GO-term-annotated cluster embeddings (top GO + # term labelled at each cluster centroid). Move it under leiden/ as + # go_annotated/ (clearer than the upstream "canonical" name). + _plots = out_dir / "plots" + _cl = _plots / "canonical_leiden" + if _cl.is_dir(): + (_plots / "leiden").mkdir(exist_ok=True) + _dest = _plots / "leiden" / "go_annotated" + if _dest.exists(): + shutil.rmtree(_dest) + shutil.move(str(_cl), str(_dest)) + + # Correlation heatmap on the gene-level embedding (mean-centered), PCA-reduced + # to the configured variance fraction. + gene = ad.read_h5ad(gene_h5ad) + X = np.asarray(gene.X) + if decisions.zscore_per_experiment: + X = (X - X.mean(axis=0, keepdims=True)) / (X.std(axis=0, keepdims=True) + 1e-8) + X_ops, n_pcs = _pca_to_variance(X, decisions.pca_variance, decisions.random_seed) + pert_col = "perturbation" if "perturbation" in gene.obs.columns else gene.obs.columns[0] + labels = gene.obs[pert_col].astype(str).tolist() + + decisions.n_genes = int(X_ops.shape[0]) + decisions.n_pcs = n_pcs + decisions.source_gene_h5ad = str(gene_h5ad) + + _correlation_heatmap(X_ops, labels, out_dir / "correlation_heatmap" / "corr_heatmap", marker) + + (out_dir / "decisions.yaml").write_text( + yaml.safe_dump(asdict(decisions), sort_keys=False) + ) + return {"marker": marker, "aggregate": agg_result, "n_genes": decisions.n_genes, "n_pcs": n_pcs} + + +def run_per_exp_embeddings( + experiment: str, + feature_dir: str | Path | None = None, + embeddings_dir: str | Path | None = None, + decisions_yaml: str | Path | None = None, + slurm: bool = True, + slurm_params: dict | None = None, + **decision_overrides, +) -> list[dict]: + """Post-process every marker of one experiment's CellDINO embeddings. + + Discovers gene_bulked_.h5ad in the CellDINO anndata_objects dir and + fans one :func:`run_marker` job per marker out to SLURM (each marker is + independent), writing to //. Only discovery and job + submission run locally; all compute (aggregate_channels + heatmap) runs on + SLURM. Pass ``slurm=False`` to run in-process instead. + """ + from ops_utils.data.experiment import OpsDataset + + ds = OpsDataset(experiment) + celldino_dir = ds.results / "cell_dino_features_v2" + if feature_dir is None: + feature_dir = celldino_dir / "anndata_objects" + feature_dir = Path(feature_dir) + if embeddings_dir is None: + # Keep outputs inside the existing CellDINO dir (no separate tree). + embeddings_dir = celldino_dir / "embeddings" + embeddings_dir = Path(embeddings_dir) + + base = EmbeddingDecisions.from_yaml(decisions_yaml) if decisions_yaml else EmbeddingDecisions() + for k, v in decision_overrides.items(): + if hasattr(base, k) and v is not None: + setattr(base, k, v) + + gene_files = sorted(glob.glob(str(feature_dir / "gene_bulked_*.h5ad"))) + if not gene_files: + raise FileNotFoundError( + f"No gene_bulked_*.h5ad in {feature_dir}. Run celldino_inference " + f"(extraction + combine) before embedding post-processing." + ) + + # Build one independent job per marker (each marker's aggregate_channels + + # heatmap is self-contained), then fan out to SLURM — matches cell_dino_main. + jobs = [] + for gf in gene_files: + marker = Path(gf).stem.replace("gene_bulked_", "") + guide = feature_dir / f"guide_bulked_{marker}.h5ad" + if not guide.exists(): + raise FileNotFoundError(f"Missing guide_bulked for {marker}: {guide}") + dec = EmbeddingDecisions(**{**asdict(base), "marker": marker}) + jobs.append({ + "name": f"embed_{experiment}_{marker}", + "func": run_marker, + "kwargs": { + "gene_h5ad": Path(gf), + "guide_h5ad": guide, + "out_dir": embeddings_dir / marker, + "decisions": dec, + }, + "metadata": {"experiment": experiment, "marker": marker}, + }) + + if not slurm: + return [run_marker(**j["kwargs"]) for j in jobs] + + from ops_utils.hpc.slurm_batch_utils import submit_parallel_jobs + + params = { + "timeout_min": 720, + "mem": "64G", + "cpus_per_task": 8, + "slurm_partition": "cpu", + } + if slurm_params: + params.update(slurm_params) + return submit_parallel_jobs( + jobs_to_submit=jobs, + experiment=f"{experiment}_embeddings", + slurm_params=params, + log_dir=f"slurm_embeddings_postprocess/{experiment}", + manifest_prefix="embedding_postprocess", + wait_for_completion=True, + verbose=True, + )