Skills › Scientific Agent Skills › scvelo › rna_velocity_workflow.py
rna_velocity_workflow.py
251 lines · skills/scvelo/scripts/rna_velocity_workflow.py at 5f2bc03
1"""RNA velocity from raw, aligned spliced/unspliced layers (scVelo 0.3.4).
3The workflow mutates AnnData in place; pass a copy to retain all input genes.
4Run --help for a local-file CLI. No datasets are downloaded automatically.
5"""
6from __future__ import annotations
8import argparse
9from importlib.metadata import version
10from pathlib import Path
12import anndata as ad
13import matplotlib.pyplot as plt
14import numpy as np
15import pandas as pd
16import scanpy as sc
17import scvelo as scv
18from packaging.version import Version
19from scipy import sparse
22def validate_layers(adata):
23 """Reject missing, nonfinite, negative, empty or ambiguously indexed counts."""
24 if not adata.obs_names.is_unique or not adata.var_names.is_unique:
25 raise ValueError("Cell and gene identifiers must be unique before alignment.")
26 if adata.n_obs < 5 or adata.n_vars < 3:
27 raise ValueError("This workflow needs at least 5 cells and 3 genes to run.")
28 for name in ("spliced", "unspliced"):
29 if name not in adata.layers:
30 raise ValueError(f"Missing '{name}' layer; supply raw quantifier counts.")
31 matrix = adata.layers[name]
32 values = matrix.data if sparse.issparse(matrix) else np.asarray(matrix)
33 if not np.isfinite(values).all() or (values < 0).any():
34 raise ValueError(f"'{name}' must contain finite nonnegative counts.")
35 if not np.any(values > 0):
36 raise ValueError(f"'{name}' contains no positive counts.")
37 if np.any(np.asarray(adata.layers["spliced"].sum(axis=1)).ravel() == 0):
38 raise ValueError("Remove cells with zero spliced library size before analysis.")
41def _check_versions(mode, make_plots):
42 """Reject confirmed upstream failures, without silently changing estimators."""
43 if Version(scv.__version__) == Version("0.3.4"):
44 if mode == "stochastic" and Version(np.__version__) >= Version("2"):
45 raise RuntimeError(
46 "scVelo 0.3.4 default stochastic GLS fails with NumPy >=2. "
47 "Use a separately validated legacy NumPy<2/Scanpy<1.12 environment, "
48 "or explicitly choose a different scientific model."
49 )
50 if (mode == "dynamical" or make_plots) and Version(pd.__version__) >= Version("3"):
51 raise RuntimeError("scVelo 0.3.4 dynamical fitting/plots require pandas<3.")
54def _save_plot(function, adata, destination, **kwargs):
55 """Own figure paths directly; scVelo's figdir/save prefix is not a path API."""
56 before = set(plt.get_fignums())
57 try:
58 function(adata, show=False, save=False, **kwargs)
59 figure = plt.gcf()
60 figure.savefig(destination, dpi=150, bbox_inches="tight")
61 finally:
62 for number in set(plt.get_fignums()) - before:
63 plt.close(number)
66def run_velocity_analysis(
67 adata,
68 groupby="leiden",
69 n_top_genes=2000,
70 n_neighbors=30,
71 mode="dynamical",
72 n_jobs=1,
73 output_dir="velocity_results",
74 *,
75 min_shared_counts=20,
76 n_pcs=30,
77 random_state=0,
78 make_plots=True,
79 recover_max_iter=10,
80):
81 """Analyze raw velocity layers in place, rebuilding X, PCA and neighbors.
83 Existing X (including integrated/scaled values) is replaced from spliced
84 counts; existing UMAP is recomputed for plotted runs. Cluster annotations
85 are preserved, but their biological validity is the caller's responsibility.
86 Input layers must be raw counts (fractional count estimates are permitted),
87 never logged/scaled/normalized. Numeric inspection cannot establish this
88 provenance. Save the original full-gene input before running.
89 """
90 if mode not in {"deterministic", "stochastic", "dynamical"}:
91 raise ValueError("mode must be deterministic, stochastic or dynamical")
92 validate_layers(adata)
93 if "velocity_workflow" in adata.uns or any(k in adata.layers for k in ("Ms", "Mu", "velocity")):
94 raise ValueError("Use fresh raw layers; existing velocity preprocessing must not be repeated.")
95 for name, value, minimum in (("n_top_genes", n_top_genes, 3), ("n_neighbors", n_neighbors, 2),
96 ("n_pcs", n_pcs, 1), ("n_jobs", n_jobs, 1),
97 ("recover_max_iter", recover_max_iter, 1)):
98 if not isinstance(value, int) or isinstance(value, bool) or value < minimum:
99 raise ValueError(f"{name} must be an integer >= {minimum}")
100 if not np.isfinite(min_shared_counts) or min_shared_counts < 0:
101 raise ValueError("min_shared_counts must be finite and nonnegative")
102 _check_versions(mode, make_plots)
103 if groupby is not None and groupby not in adata.obs:
104 raise ValueError(f"Grouping column {groupby!r} is missing; use groupby=None to omit ranking.")
105 if groupby is not None and adata.obs[groupby].isna().any():
106 raise ValueError("Grouping labels must not be missing.")
108 # Rebuild expression geometry from explicit count provenance. Store raw
109 # layers before normalization; filtered-out genes remain in the input file.
110 adata.X = adata.layers["spliced"].astype(np.float32).copy()
111 for key in ("spliced", "unspliced"):
112 adata.layers[f"{key}_counts"] = adata.layers[key].copy()
113 adata.layers[key] = adata.layers[key].astype(np.float32).copy()
114 for key in ("log1p", "neighbors", "pca"):
115 adata.uns.pop(key, None)
116 for key in ("initial_size", "initial_size_spliced", "initial_size_unspliced"):
117 if key in adata.obs:
118 del adata.obs[key]
119 print("[1/5] Filter counts, normalize layers, select genes and rebuild neighbors")
120 scv.pp.filter_and_normalize(adata, min_shared_counts=min_shared_counts,
121 layers_normalize=["spliced", "unspliced"])
122 if adata.n_vars < 3:
123 raise ValueError("Fewer than 3 genes survived count filtering; inspect coverage and thresholds.")
124 sc.pp.log1p(adata) # only X; spliced/unspliced stay on linear scale
125 sc.pp.highly_variable_genes(adata, n_top_genes=min(n_top_genes, adata.n_vars), subset=True)
126 if adata.n_vars < 3:
127 raise ValueError("Fewer than 3 variable genes remain; inspect the expression distribution.")
128 actual_pcs = min(n_pcs, adata.n_obs - 1, adata.n_vars - 1)
129 actual_neighbors = min(n_neighbors, adata.n_obs - 1)
130 sc.pp.pca(adata, n_comps=actual_pcs, random_state=random_state)
131 sc.pp.neighbors(adata, n_neighbors=actual_neighbors, n_pcs=actual_pcs,
132 use_rep="X_pca", random_state=random_state)
133 scv.pp.moments(adata, n_neighbors=None) # consume the explicit Scanpy graph
135 print(f"[2/5] Fit {mode} model")
136 if mode == "dynamical":
137 scv.tl.recover_dynamics(adata, var_names="all", n_jobs=n_jobs,
138 max_iter=recover_max_iter, show_progress_bar=False)
139 if not np.isfinite(adata.var["fit_likelihood"]).any():
140 raise ValueError("No finite dynamical fits; inspect phase portraits and input coverage.")
141 scv.tl.velocity(adata, mode=mode)
142 selected = adata.var["velocity_genes"].to_numpy(dtype=bool)
143 if selected.sum() < 2 or not np.isfinite(adata.layers["velocity"][:, selected]).all():
144 raise ValueError("Fewer than 2 usable velocity genes, or nonfinite fitted velocities.")
145 scv.tl.velocity_graph(adata, n_jobs=n_jobs, show_progress_bar=False)
146 if not adata.uns["velocity_graph"].nnz:
147 raise ValueError("Velocity graph is empty; inspect fitted genes and neighborhood geometry.")
149 print("[3/5] Compute coherence and relative time")
150 scv.tl.velocity_confidence(adata)
151 scv.tl.velocity_pseudotime(adata)
152 if mode == "dynamical":
153 scv.tl.latent_time(adata)
154 diagnostics = {}
155 for key in ("velocity_confidence", "velocity_pseudotime", "latent_time"):
156 if key not in adata.obs:
157 continue
158 values = adata.obs[key].to_numpy()
159 finite = values[np.isfinite(values)]
160 diagnostics[f"{key}_nonfinite"] = int(len(values) - len(finite))
161 diagnostics[f"{key}_constant"] = bool(len(finite) == 0 or np.ptp(finite) == 0)
162 if len(finite) != len(values) or diagnostics[f"{key}_constant"]:
163 print(f"[WARN] {key} is nonfinite or constant; it does not support an ordering/confidence claim.")
164 if groupby is not None:
165 adata.obs[groupby] = adata.obs[groupby].astype("category")
166 sizes = adata.obs[groupby].value_counts()
167 if len(sizes) >= 2 and (sizes >= 2).all():
168 scv.tl.rank_velocity_genes(adata, groupby=groupby, min_corr=0.3)
170 destination = Path(output_dir)
171 destination.mkdir(parents=True, exist_ok=True)
172 print("[4/5] Render figures" if make_plots else "[4/5] Figures disabled")
173 if make_plots:
174 sc.tl.umap(adata, random_state=random_state)
175 _save_plot(scv.pl.velocity_embedding_stream, adata, destination / "velocity_stream.png",
176 basis="umap", color=groupby, title="RNA velocity (projected)")
177 _save_plot(scv.pl.velocity_embedding, adata, destination / "velocity_arrows.png",
178 basis="umap", color=groupby, arrow_length=3, arrow_size=2)
179 _save_plot(scv.pl.scatter, adata, destination / "pseudotime.png", basis="umap",
180 color="velocity_pseudotime", cmap="gnuplot")
181 _save_plot(scv.pl.scatter, adata, destination / "velocity_quality.png", basis="umap",
182 color=["velocity_length", "velocity_confidence"], cmap="coolwarm", perc=[5, 95])
183 if mode == "dynamical":
184 _save_plot(scv.pl.scatter, adata, destination / "latent_time.png", basis="umap",
185 color="latent_time", cmap="gnuplot")
186 genes = adata.var["fit_likelihood"].dropna().nlargest(min(50, adata.n_vars)).index
187 _save_plot(scv.pl.heatmap, adata, destination / "dynamical_gene_heatmap.png",
188 var_names=genes, sortby="latent_time", col_color=groupby,
189 n_convolve=min(30, adata.n_obs))
191 adata.uns["velocity_workflow"] = {
192 "mode": mode, "n_neighbors": actual_neighbors, "n_pcs": actual_pcs,
193 "random_state": random_state, "input_layers": "raw spliced and unspliced counts",
194 "diagnostics": diagnostics,
195 "packages": {name: version(name) for name in ("scvelo", "scanpy", "anndata", "numpy", "pandas")},
196 }
197 print("[5/5] Save annotated data")
198 adata.write_h5ad(destination / "adata_velocity.h5ad")
199 print(f"[OK] Cells: {adata.n_obs}; selected genes: {adata.n_vars}; velocity genes: {selected.sum()}")
200 print("Velocity coherence is a descriptive score, not calibrated certainty.")
201 return adata
204def load_from_loom(loom_path, processed_h5ad=None):
205 """Read local loom counts and optionally align to processed metadata exactly.
207 No barcode rewriting or silent intersections. Resolve library prefixes and
208 gene IDs explicitly before combining mismatched files. Extra loom cells and
209 genes are permitted, but every processed cell/gene must have a match.
210 """
211 counts = ad.io.read_loom(loom_path, X_name="spliced", sparse=True)
212 validate_layers(counts)
213 if processed_h5ad is None:
214 return counts
215 processed = ad.read_h5ad(processed_h5ad)
216 if not processed.obs_names.is_unique or not processed.var_names.is_unique:
217 raise ValueError("Processed cell and gene identifiers must be unique.")
218 if not processed.obs_names.isin(counts.obs_names).all() or not processed.var_names.isin(counts.var_names).all():
219 raise ValueError("Processed IDs do not all match loom IDs; align barcodes and gene identifiers explicitly.")
220 aligned = counts[processed.obs_names, processed.var_names]
221 for key in ("spliced", "unspliced"):
222 processed.layers[key] = aligned.layers[key].copy()
223 validate_layers(processed)
224 return processed
227def main():
228 parser = argparse.ArgumentParser(description=__doc__)
229 parser.add_argument("input", help="Local .h5ad or .loom with raw spliced/unspliced layers")
230 parser.add_argument("--processed-h5ad", help="Metadata/embedding file to align with loom counts")
231 parser.add_argument("--groupby", default=None, help="Optional annotation column for coloring/ranking")
232 parser.add_argument("--mode", choices=("deterministic", "stochastic", "dynamical"), default="dynamical")
233 parser.add_argument("--n-top-genes", type=int, default=2000)
234 parser.add_argument("--n-neighbors", type=int, default=30)
235 parser.add_argument("--n-jobs", type=int, default=1)
236 parser.add_argument("--output-dir", default="velocity_results")
237 parser.add_argument("--no-plots", action="store_true")
238 args = parser.parse_args()
239 if Path(args.input).suffix.lower() == ".loom":
240 data = load_from_loom(args.input, args.processed_h5ad)
241 else:
242 if args.processed_h5ad:
243 parser.error("--processed-h5ad is only used with .loom input")
244 data = ad.read_h5ad(args.input)
245 run_velocity_analysis(data, groupby=args.groupby, n_top_genes=args.n_top_genes,
246 n_neighbors=args.n_neighbors, mode=args.mode, n_jobs=args.n_jobs,
247 output_dir=args.output_dir, make_plots=not args.no_plots)
250if __name__ == "__main__":
251 main()