diff --git a/.github/workflows/docs.yml.disabled b/.github/workflows/docs.yml.disabled deleted file mode 100644 index e645932..0000000 --- a/.github/workflows/docs.yml.disabled +++ /dev/null @@ -1,35 +0,0 @@ -name: Documentation - -on: - push: - branches: [main] - workflow_dispatch: - -permissions: - contents: write - -jobs: - build-and-deploy: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -r docs/requirements.txt - - - name: Build documentation - run: sphinx-build -b html docs docs/_build/html - - - name: Deploy to GitHub Pages - uses: peaceiris/actions-gh-pages@v4 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/_build/html - publish_branch: gh-pages diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..7a16aa9 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,28 @@ +name: Tests + +on: + push: + pull_request: + +permissions: + contents: read + +jobs: + smoke-test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + cache: pip + + - name: Install package and test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[dev]" + + - name: Run smoke test + run: python -m pytest tests/test_smoke.py diff --git a/.gitignore b/.gitignore index 288b3a5..9ad79b1 100644 --- a/.gitignore +++ b/.gitignore @@ -52,9 +52,6 @@ coverage.xml *.mo *.pot -# Sphinx documentation -docs/_build/ - # PyBuilder .pybuilder/ target/ diff --git a/README.md b/README.md index 039528a..211aef3 100644 --- a/README.md +++ b/README.md @@ -9,11 +9,16 @@ Notebooks for reproducing manuscript figures and analyses are available at [csglab/sclsd-manuscript](https://github.com/csglab/sclsd-manuscript). +This README is the maintained documentation for the `sclsd` package. It covers +installation, input data, configuration, training, inference, reproducibility, +and citation information. + ## Overview LSD reinterprets Waddington's epigenetic landscape as an energy landscape in a learned latent cell state space. Cell differentiation is modeled as a stochastic dynamical system governed by a gradient flow down this potential surface, combined with noise representing gene expression variability. The model jointly infers: + - **Cell state**: A latent representation of each cell's gene expression profile - **Differentiation state**: A 2D embedding capturing developmental progression - **Waddington potential**: An energy function whose gradient defines differentiation dynamics @@ -32,17 +37,13 @@ git clone https://github.com/csglab/sclsd.git cd sclsd pip install -e . ``` -> [!NOTE] -> As a lightweight Python package, installation time varies depending on pre-existing dependencies in your environment, starting from scratch, expect it to take less than 10 minutes. ### Dependencies -- Python ≥3.9 -- PyTorch ≥2.0.0 -- Pyro-PPL ≥1.8.0 -- torchdiffeq ≥0.2.0 -- scanpy ≥1.9.0 -- cellrank ≥2.0.0 +The package requires Python 3.9 or later. Runtime dependencies, including +PyTorch, Pyro, torchdiffeq, Scanpy, AnnData, and CellRank, are installed by +`pip`. The complete dependency specification is maintained in +[`pyproject.toml`](https://github.com/csglab/sclsd/blob/main/pyproject.toml). ## Quick Start @@ -138,7 +139,7 @@ cfg = LSDConfig() # Model architecture cfg.model.z_dim = 10 # Cell state dimensions -cfg.model.B_dim = 2 # Differentiation state dimensions (fixed at 2) +cfg.model.B_dim = 2 # Differentiation state dimensions (default: 2) cfg.model.V_coeff = 0.01 # Potential regularization # Training trajectories @@ -153,11 +154,12 @@ cfg.optimizer.adam.lr = 1e-3 # Learning rate ## Data Requirements Input AnnData should contain: + - Log-normalized expression in `adata.X` - Raw counts in `adata.layers["raw"]` - Library sizes in `adata.obs["librarysize"]` - Precomputed neighbor graph in `adata.obsp["connectivities"]` -- Pseudotime values (e.g., from diffusion pseudotime) for prior initialization +- Pseudotime values in `adata.obs` when initializing transitions from a pseudotime key; alternatively, users may supply a transition matrix directly ## Method @@ -169,11 +171,47 @@ where $V(z)$ is the Waddington potential parameterized by a neural network, and Training trajectories are generated by random walks on a k-nearest neighbor graph, biased by pseudotime to follow developmental progression. +## Reproducibility + +Dataset-specific training and postprocessing notebooks are available in the +[`sclsd-manuscript`](https://github.com/csglab/sclsd-manuscript) repository. +The preprocessed datasets used by those notebooks are available from +[Zenodo record 18331587](https://zenodo.org/records/18331587). + ## Citation -If you use sclsd, please cite: +If you use `sclsd` or the accompanying analyses, please cite: + +> Poursina, A., Hajhashemi, S., Mikaeili Namini, A., Saberi, A., Emad, A., & +> Najafabadi, H. S. (2026). A Latent Space Thermodynamic Model of Cell +> Differentiation. *bioRxiv*, 2026.03.04.709512. +> https://doi.org/10.64898/2026.03.04.709512 + +[View version 1 on bioRxiv](https://www.biorxiv.org/content/10.64898/2026.03.04.709512v1) + +### BibTeX + +```bibtex +@article{poursina2026latent, + title = {A Latent Space Thermodynamic Model of Cell Differentiation}, + author = {Poursina, Ali and Hajhashemi, Shayan and + {Mikaeili Namini}, Arsham and Saberi, Ali and + Emad, Amin and Najafabadi, Hamed S.}, + journal = {bioRxiv}, + pages = {2026.03.04.709512}, + year = {2026}, + publisher = {Cold Spring Harbor Laboratory}, + doi = {10.64898/2026.03.04.709512}, + url = {https://www.biorxiv.org/content/10.64898/2026.03.04.709512v1} +} +``` + +## Contact -> Poursina A, Hajhashemi S, Mikaeili Namini A, Saberi A, Emad A, Najafabadi HS. A Latent Space Thermodynamic Model of Cell Differentiation. 2026. +For questions about the sclsd software, contact Ali Poursina at +[ali.poursina@mail.mcgill.ca](mailto:ali.poursina@mail.mcgill.ca). Bug reports and +feature requests can also be submitted through the +[GitHub issue tracker](https://github.com/csglab/sclsd/issues). ## License diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..5b1828e --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Benchmark helpers that are not part of the public sclsd package.""" diff --git a/benchmarks/comment7_benchmark.py b/benchmarks/comment7_benchmark.py new file mode 100644 index 0000000..a468073 --- /dev/null +++ b/benchmarks/comment7_benchmark.py @@ -0,0 +1,244 @@ +"""Fresh-process old-versus-new benchmarks for comment 7.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +import platform +import resource +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np +import scanpy as sc +import scipy.sparse as sp +import torch + +from benchmarks.comment7_legacy import ( + calculate_transition_probs as legacy_transition_probs, +) +from benchmarks.comment7_legacy import ( + create_phylogeny_matrix as legacy_phylogeny_matrix, +) +from benchmarks.comment7_legacy import ( + random_walks_dense, +) +from benchmarks.freeze_comment7_baseline import DENTATE_PHYLOGENY + +BONE_MARROW_PHYLOGENY = { + "HSC_1": ["HSC_2", "Ery_1", "Mega"], + "Ery_1": ["Ery_2"], + "HSC_2": ["Precursors"], + "Precursors": ["Mono_1", "Mono_2", "DCs"], + "Ery_2": [], + "Mega": [], + "Mono_1": [], + "Mono_2": [], + "DCs": [], +} + +MOUSE_CORTEX_PHYLOGENY = { + "Apical progenitors": ["IPC", "IN nonMGE"], + "IPC": ["ULPN", "Migrating neurons"], + "IN nonMGE": ["IN MGE"], + "ULPN": ["DLPN"], + "Migrating neurons": ["DLPN"], + "DLPN": [], + "IN MGE": [], +} + +DATASETS = { + "dentate": (DENTATE_PHYLOGENY, "prior_pseudotime", 2048, 12), + "bone_marrow": (BONE_MARROW_PHYLOGENY, "palantir_pseudotime", 2048, 12), + "mouse_cortex": (MOUSE_CORTEX_PHYLOGENY, "prior_pseudotime", 8192, 30), +} + + +def _current_rss_bytes() -> int: + statm = Path("/proc/self/statm").read_text().split() + return int(statm[1]) * os.sysconf("SC_PAGE_SIZE") + + +def _peak_rss_bytes() -> int: + # Linux reports ru_maxrss in KiB. + return int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024) + + +def _csr_hash(matrix: sp.spmatrix) -> str: + matrix = matrix.tocsr() + digest = hashlib.sha256() + for array in (matrix.data, matrix.indices, matrix.indptr): + digest.update(np.ascontiguousarray(array).view(np.uint8)) + digest.update(np.asarray(matrix.shape, dtype=np.int64).view(np.uint8)) + return digest.hexdigest() + + +def _run_worker(args: argparse.Namespace) -> dict: + phylogeny, time_key, n_walks, path_len = DATASETS[args.dataset] + adata = sc.read(args.data) + connectivity = adata.obsp["connectivities"].tocsr() + pseudotime = adata.obs[time_key].to_numpy() + + if args.operation == "walks": + from sclsd.preprocessing.prior import prior_transition_matrix + + transition = prior_transition_matrix(adata, time_key, beta_t=50) + else: + transition = None + + current_phylogeny = None + current_transition = None + current_walks = None + if args.implementation == "current": + if args.operation == "phylogeny": + from sclsd.preprocessing.prior import _create_phylogeny_matrix + + current_phylogeny = _create_phylogeny_matrix + elif args.operation == "transition": + from sclsd.train.trainer import LSD + + current_transition = LSD.calculate_transition_probs + else: + from sclsd.train.walks import random_walks_sparse + + current_walks = random_walks_sparse + + gc.collect() + rss_before = _current_rss_bytes() + cpu_start = time.process_time() + wall_start = time.perf_counter() + + if args.operation == "phylogeny": + if args.implementation == "legacy": + mask = legacy_phylogeny_matrix(adata, phylogeny, "clusters") + else: + mask = current_phylogeny(adata, phylogeny, "clusters") + result = connectivity.multiply(mask).tocsr() + elif args.operation == "transition": + if args.implementation == "legacy": + binary = (connectivity.toarray() > 0).astype(float) + result = legacy_transition_probs(-pseudotime, binary, beta=50) + else: + binary = connectivity.copy() + binary.data = np.ones_like(binary.data, dtype=float) + result = current_transition(None, -pseudotime, binary, beta=50) + else: + if args.implementation == "legacy": + dense_transition = torch.from_numpy(transition.toarray()).float() + result = random_walks_dense( + dense_transition, n_walks, path_len, random_state=42 + ) + else: + result = current_walks( + transition, + n_steps=path_len, + n_trajectories=n_walks, + random_state=42, + ) + + wall_seconds = time.perf_counter() - wall_start + cpu_seconds = time.process_time() - cpu_start + peak_rss = _peak_rss_bytes() + + if sp.issparse(result): + output = { + "output_type": type(result).__name__, + "output_nnz": int(result.nnz), + "output_hash": _csr_hash(result), + } + elif isinstance(result, torch.Tensor): + array = result.cpu().numpy() + output = { + "output_type": str(result.dtype), + "output_nnz": None, + "output_hash": hashlib.sha256(array.tobytes()).hexdigest(), + } + else: + array = np.asarray(result) + output = { + "output_type": str(array.dtype), + "output_nnz": int(np.count_nonzero(array)), + "output_hash": hashlib.sha256(array.tobytes()).hexdigest(), + } + + return { + "dataset": args.dataset, + "implementation": args.implementation, + "operation": args.operation, + "n_obs": int(adata.n_obs), + "n_vars": int(adata.n_vars), + "connectivity_nnz": int(connectivity.nnz), + "wall_seconds": wall_seconds, + "cpu_seconds": cpu_seconds, + "rss_before_bytes": rss_before, + "peak_rss_bytes": peak_rss, + "peak_rss_delta_bytes": max(0, peak_rss - rss_before), + **output, + } + + +def _summarize(records: list[dict]) -> dict: + summary = {} + for key in ("wall_seconds", "cpu_seconds", "peak_rss_bytes", "peak_rss_delta_bytes"): + values = np.asarray([record[key] for record in records], dtype=float) + summary[key] = { + "median": float(np.median(values)), + "min": float(np.min(values)), + "max": float(np.max(values)), + } + return summary + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data", type=Path, required=True) + parser.add_argument("--dataset", choices=DATASETS, required=True) + parser.add_argument("--operation", choices=["phylogeny", "transition", "walks"], required=True) + parser.add_argument("--implementation", choices=["legacy", "current"], required=True) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--output", type=Path) + parser.add_argument("--worker", action="store_true") + args = parser.parse_args() + + if args.worker: + print(json.dumps(_run_worker(args))) + return + + records = [] + for _ in range(args.repeats): + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--data", + str(args.data), + "--dataset", + args.dataset, + "--operation", + args.operation, + "--implementation", + args.implementation, + "--worker", + ] + completed = subprocess.run(command, check=True, capture_output=True, text=True) + records.append(json.loads(completed.stdout.strip().splitlines()[-1])) + + payload = { + "python": platform.python_version(), + "repeats": args.repeats, + "records": records, + "summary": _summarize(records), + } + rendered = json.dumps(payload, indent=2) + "\n" + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered) + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/comment7_legacy.py b/benchmarks/comment7_legacy.py new file mode 100644 index 0000000..f2c86ea --- /dev/null +++ b/benchmarks/comment7_legacy.py @@ -0,0 +1,90 @@ +"""Frozen pre-comment-7 implementations used for equivalence benchmarks. + +These functions intentionally preserve the dense behavior that existed before +the sparse-memory revision. They must not be used by production code. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +import numpy as np +import scipy.sparse as sp +import torch + + +def get_all_descendants( + cluster: str, + phylogeny: Dict[str, List[str]], + descendants: Optional[set] = None, +) -> set: + """Frozen helper used by the legacy phylogeny implementation.""" + if descendants is None: + descendants = set() + + direct_children = phylogeny.get(cluster, []) + descendants.update(direct_children) + return descendants + + +def create_phylogeny_matrix(adata, phylogeny, cluster_key="clusters"): + """Exact dense implementation from preprocessing/prior.py.""" + clusters = adata.obs[cluster_key].unique().tolist() + cell_to_cluster = dict(zip(adata.obs_names, adata.obs[cluster_key])) + + all_descendants = {} + for cluster in clusters: + all_descendants[cluster] = get_all_descendants(cluster, phylogeny) + + n_cells = adata.shape[0] + phylo_matrix = np.zeros((n_cells, n_cells)) + + for i, cell_i in enumerate(adata.obs_names): + cluster_i = cell_to_cluster[cell_i] + for j, cell_j in enumerate(adata.obs_names): + cluster_j = cell_to_cluster[cell_j] + + if cluster_i == cluster_j: + phylo_matrix[i, j] = 1 + elif cluster_j in all_descendants.get(cluster_i, set()): + phylo_matrix[i, j] = 1 + + return sp.csr_matrix(phylo_matrix) + + +def calculate_transition_probs( + potential: np.ndarray, + connectivity_matrix: np.ndarray, + beta: float = 1.0, +) -> np.ndarray: + """Exact dense implementation from train/trainer.py.""" + potential = potential.astype(float) + energy_diff = potential[None, :] - potential[:, None] + boltzmann_weights = np.exp(-beta * energy_diff) + boltzmann_weights *= connectivity_matrix + row_sums = boltzmann_weights.sum(axis=1, keepdims=True) + 1e-12 + transition_matrix = boltzmann_weights / row_sums + return transition_matrix + + +def random_walks_dense( + transition: torch.Tensor, + n_trajectories: int, + path_len: int, + random_state: int, +) -> torch.Tensor: + """Exact dense CPU equivalent of LSD._random_walks.""" + torch.manual_seed(random_state) + n_cells = transition.shape[0] + current_states = torch.randint(0, n_cells, (n_trajectories,), dtype=torch.long) + walks = torch.empty((n_trajectories, path_len), dtype=torch.int) + walks[:, 0] = current_states + + for step in range(1, path_len): + next_states = torch.multinomial( + transition[current_states], num_samples=1 + ).squeeze(1) + walks[:, step] = next_states + current_states = next_states + + return walks diff --git a/benchmarks/freeze_comment7_baseline.py b/benchmarks/freeze_comment7_baseline.py new file mode 100644 index 0000000..2411f1f --- /dev/null +++ b/benchmarks/freeze_comment7_baseline.py @@ -0,0 +1,140 @@ +"""Freeze realistic legacy outputs before the comment-7 source update.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +from importlib import metadata +from pathlib import Path + +import numpy as np +import scanpy as sc +import scipy.sparse as sp +import torch + +from benchmarks.comment7_legacy import ( + calculate_transition_probs, + create_phylogeny_matrix, + random_walks_dense, +) + +DENTATE_PHYLOGENY = { + "nIPC": ["Neuroblast", "Radial Glia-like"], + "Neuroblast": ["Granule immature"], + "Granule immature": ["Granule mature"], + "Granule mature": [], + "Radial Glia-like": ["Astrocytes"], + "Astrocytes": [], +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _csr_payload(prefix: str, matrix: sp.spmatrix) -> dict: + matrix = matrix.tocsr() + return { + f"{prefix}_data": matrix.data, + f"{prefix}_indices": matrix.indices, + f"{prefix}_indptr": matrix.indptr, + f"{prefix}_shape": np.asarray(matrix.shape, dtype=np.int64), + } + + +def _csr_hash(matrix: sp.spmatrix) -> str: + matrix = matrix.tocsr() + digest = hashlib.sha256() + for array in (matrix.data, matrix.indices, matrix.indptr): + digest.update(np.ascontiguousarray(array).view(np.uint8)) + digest.update(np.asarray(matrix.shape, dtype=np.int64).view(np.uint8)) + return digest.hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--metadata", type=Path, required=True) + args = parser.parse_args() + + adata = sc.read(args.data) + connectivity = adata.obsp["connectivities"].tocsr() + pseudotime = adata.obs["prior_pseudotime"].to_numpy() + clusters = adata.obs["clusters"].astype(str).to_numpy(dtype="U") + + legacy_mask = create_phylogeny_matrix(adata, DENTATE_PHYLOGENY, "clusters") + effective_connectivity = connectivity.multiply(legacy_mask).tocsr() + + binary_connectivity = (connectivity.toarray() > 0).astype(float) + dense_transition = calculate_transition_probs( + -pseudotime, binary_connectivity, beta=50 + ) + sparse_transition = sp.csr_matrix(dense_transition) + walks = random_walks_dense( + torch.from_numpy(dense_transition).float(), + n_trajectories=2048, + path_len=12, + random_state=42, + ).numpy() + + payload = { + **_csr_payload("connectivity", connectivity), + **_csr_payload("effective_phylogeny", effective_connectivity), + **_csr_payload("transition", sparse_transition), + "clusters": clusters, + "pseudotime": pseudotime, + "legacy_walks": walks, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(args.output, **payload) + + repo = Path(__file__).resolve().parents[1] + package_names = ["numpy", "scipy", "pandas", "anndata", "scanpy", "torch", "pyro-ppl"] + package_versions = {name: metadata.version(name) for name in package_names} + metadata_payload = { + "created_before_source_update": True, + "python": platform.python_version(), + "packages": package_versions, + "source_sha256": { + "prior.py": _sha256(repo / "src/sclsd/preprocessing/prior.py"), + "trainer.py": _sha256(repo / "src/sclsd/train/trainer.py"), + "walks.py": _sha256(repo / "src/sclsd/train/walks.py"), + }, + "dataset": { + "path": str(args.data), + "sha256": _sha256(args.data), + "n_obs": int(adata.n_obs), + "n_vars": int(adata.n_vars), + "connectivity_nnz": int(connectivity.nnz), + "cluster_count": int(adata.obs["clusters"].nunique()), + }, + "configuration": { + "walk_count": 2048, + "path_len": 12, + "random_state": 42, + "transition_beta": 50, + }, + "legacy_outputs": { + "raw_phylogeny_mask_nnz": int(legacy_mask.nnz), + "effective_phylogeny_nnz": int(effective_connectivity.nnz), + "effective_phylogeny_sha256": _csr_hash(effective_connectivity), + "transition_nnz": int(sparse_transition.nnz), + "transition_sha256": _csr_hash(sparse_transition), + "walks_sha256": hashlib.sha256(walks.tobytes()).hexdigest(), + "walks_dtype": str(walks.dtype), + "walks_shape": list(walks.shape), + }, + } + args.metadata.parent.mkdir(parents=True, exist_ok=True) + args.metadata.write_text(json.dumps(metadata_payload, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results/bone_marrow_phylogeny_current.json b/benchmarks/results/bone_marrow_phylogeny_current.json new file mode 100644 index 0000000..60f94c5 --- /dev/null +++ b/benchmarks/results/bone_marrow_phylogeny_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.04323896486312151, + "cpu_seconds": 0.04325768200000013, + "rss_before_bytes": 631713792, + "peak_rss_bytes": 637632512, + "peak_rss_delta_bytes": 5918720, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.04605204053223133, + "cpu_seconds": 0.046071685999999445, + "rss_before_bytes": 631169024, + "peak_rss_bytes": 637386752, + "peak_rss_delta_bytes": 6217728, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.04360661655664444, + "cpu_seconds": 0.04363062300000031, + "rss_before_bytes": 631205888, + "peak_rss_bytes": 637427712, + "peak_rss_delta_bytes": 6221824, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.04263310041278601, + "cpu_seconds": 0.04265175399999954, + "rss_before_bytes": 631656448, + "peak_rss_bytes": 637726720, + "peak_rss_delta_bytes": 6070272, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.04396524466574192, + "cpu_seconds": 0.04398627700000013, + "rss_before_bytes": 631205888, + "peak_rss_bytes": 637485056, + "peak_rss_delta_bytes": 6279168, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + } + ], + "summary": { + "wall_seconds": { + "median": 0.04360661655664444, + "min": 0.04263310041278601, + "max": 0.04605204053223133 + }, + "cpu_seconds": { + "median": 0.04363062300000031, + "min": 0.04265175399999954, + "max": 0.046071685999999445 + }, + "peak_rss_bytes": { + "median": 637485056.0, + "min": 637386752.0, + "max": 637726720.0 + }, + "peak_rss_delta_bytes": { + "median": 6217728.0, + "min": 5918720.0, + "max": 6279168.0 + } + } +} diff --git a/benchmarks/results/bone_marrow_phylogeny_legacy.json b/benchmarks/results/bone_marrow_phylogeny_legacy.json new file mode 100644 index 0000000..c50148f --- /dev/null +++ b/benchmarks/results/bone_marrow_phylogeny_legacy.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 9.375272090546787, + "cpu_seconds": 9.375060449999998, + "rss_before_bytes": 600993792, + "peak_rss_bytes": 1078538240, + "peak_rss_delta_bytes": 477544448, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 9.337023792788386, + "cpu_seconds": 9.336520170999998, + "rss_before_bytes": 599674880, + "peak_rss_bytes": 1077223424, + "peak_rss_delta_bytes": 477548544, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 9.416737007908523, + "cpu_seconds": 9.416345341000001, + "rss_before_bytes": 599412736, + "peak_rss_bytes": 1076719616, + "peak_rss_delta_bytes": 477306880, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 9.680333655327559, + "cpu_seconds": 9.679520187000001, + "rss_before_bytes": 599539712, + "peak_rss_bytes": 1077002240, + "peak_rss_delta_bytes": 477462528, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 9.515539621002972, + "cpu_seconds": 9.515123518, + "rss_before_bytes": 599539712, + "peak_rss_bytes": 1077157888, + "peak_rss_delta_bytes": 477618176, + "output_type": "csr_matrix", + "output_nnz": 94899, + "output_hash": "1b70d2717e5c3a44a7e9eb32e5f6e8138a661a0f5c8ab6feb07166eb684bfadd" + } + ], + "summary": { + "wall_seconds": { + "median": 9.416737007908523, + "min": 9.337023792788386, + "max": 9.680333655327559 + }, + "cpu_seconds": { + "median": 9.416345341000001, + "min": 9.336520170999998, + "max": 9.679520187000001 + }, + "peak_rss_bytes": { + "median": 1077157888.0, + "min": 1076719616.0, + "max": 1078538240.0 + }, + "peak_rss_delta_bytes": { + "median": 477544448.0, + "min": 477306880.0, + "max": 477618176.0 + } + } +} diff --git a/benchmarks/results/bone_marrow_transition_current.json b/benchmarks/results/bone_marrow_transition_current.json new file mode 100644 index 0000000..27656a8 --- /dev/null +++ b/benchmarks/results/bone_marrow_transition_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.00529903918504715, + "cpu_seconds": 0.005321454000000614, + "rss_before_bytes": 631111680, + "peak_rss_bytes": 637964288, + "peak_rss_delta_bytes": 6852608, + "output_type": "csr_matrix", + "output_nnz": 114718, + "output_hash": "60d1e803a3eb53d35d023000dcea0d3a819d8802a37e159d6c27740d6baead03" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.005373898893594742, + "cpu_seconds": 0.005397865999999141, + "rss_before_bytes": 631820288, + "peak_rss_bytes": 638971904, + "peak_rss_delta_bytes": 7151616, + "output_type": "csr_matrix", + "output_nnz": 114718, + "output_hash": "60d1e803a3eb53d35d023000dcea0d3a819d8802a37e159d6c27740d6baead03" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.005222865380346775, + "cpu_seconds": 0.005247235999998878, + "rss_before_bytes": 631173120, + "peak_rss_bytes": 638214144, + "peak_rss_delta_bytes": 7041024, + "output_type": "csr_matrix", + "output_nnz": 114718, + "output_hash": "60d1e803a3eb53d35d023000dcea0d3a819d8802a37e159d6c27740d6baead03" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.005188260227441788, + "cpu_seconds": 0.005208514999999636, + "rss_before_bytes": 631054336, + "peak_rss_bytes": 638095360, + "peak_rss_delta_bytes": 7041024, + "output_type": "csr_matrix", + "output_nnz": 114718, + "output_hash": "60d1e803a3eb53d35d023000dcea0d3a819d8802a37e159d6c27740d6baead03" + }, + { + "dataset": "bone_marrow", + "implementation": "current", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.005155312828719616, + "cpu_seconds": 0.0051792279999993696, + "rss_before_bytes": 631468032, + "peak_rss_bytes": 638566400, + "peak_rss_delta_bytes": 7098368, + "output_type": "csr_matrix", + "output_nnz": 114718, + "output_hash": "60d1e803a3eb53d35d023000dcea0d3a819d8802a37e159d6c27740d6baead03" + } + ], + "summary": { + "wall_seconds": { + "median": 0.005222865380346775, + "min": 0.005155312828719616, + "max": 0.005373898893594742 + }, + "cpu_seconds": { + "median": 0.005247235999998878, + "min": 0.0051792279999993696, + "max": 0.005397865999999141 + }, + "peak_rss_bytes": { + "median": 638214144.0, + "min": 637964288.0, + "max": 638971904.0 + }, + "peak_rss_delta_bytes": { + "median": 7041024.0, + "min": 6852608.0, + "max": 7151616.0 + } + } +} diff --git a/benchmarks/results/bone_marrow_transition_legacy.json b/benchmarks/results/bone_marrow_transition_legacy.json new file mode 100644 index 0000000..95c8bd0 --- /dev/null +++ b/benchmarks/results/bone_marrow_transition_legacy.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.43377263750880957, + "cpu_seconds": 0.43372887999999854, + "rss_before_bytes": 599703552, + "peak_rss_bytes": 1496379392, + "peak_rss_delta_bytes": 896675840, + "output_type": "float64", + "output_nnz": 114718, + "output_hash": "7b4d90734c48f9e2877a6b8346ff160dc8852637de231600bae901dd1109d4db" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.42974887415766716, + "cpu_seconds": 0.42971178799999876, + "rss_before_bytes": 599662592, + "peak_rss_bytes": 1496141824, + "peak_rss_delta_bytes": 896479232, + "output_type": "float64", + "output_nnz": 114718, + "output_hash": "7b4d90734c48f9e2877a6b8346ff160dc8852637de231600bae901dd1109d4db" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.43343123886734247, + "cpu_seconds": 0.4334526479999994, + "rss_before_bytes": 599719936, + "peak_rss_bytes": 1496383488, + "peak_rss_delta_bytes": 896663552, + "output_type": "float64", + "output_nnz": 114718, + "output_hash": "7b4d90734c48f9e2877a6b8346ff160dc8852637de231600bae901dd1109d4db" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.4349979655817151, + "cpu_seconds": 0.43503339099999927, + "rss_before_bytes": 600170496, + "peak_rss_bytes": 1496846336, + "peak_rss_delta_bytes": 896675840, + "output_type": "float64", + "output_nnz": 114718, + "output_hash": "7b4d90734c48f9e2877a6b8346ff160dc8852637de231600bae901dd1109d4db" + }, + { + "dataset": "bone_marrow", + "implementation": "legacy", + "operation": "transition", + "n_obs": 5292, + "n_vars": 2000, + "connectivity_nnz": 114718, + "wall_seconds": 0.43536503333598375, + "cpu_seconds": 0.43536600599999886, + "rss_before_bytes": 599515136, + "peak_rss_bytes": 1496084480, + "peak_rss_delta_bytes": 896569344, + "output_type": "float64", + "output_nnz": 114718, + "output_hash": "7b4d90734c48f9e2877a6b8346ff160dc8852637de231600bae901dd1109d4db" + } + ], + "summary": { + "wall_seconds": { + "median": 0.43377263750880957, + "min": 0.42974887415766716, + "max": 0.43536503333598375 + }, + "cpu_seconds": { + "median": 0.43372887999999854, + "min": 0.42971178799999876, + "max": 0.43536600599999886 + }, + "peak_rss_bytes": { + "median": 1496379392.0, + "min": 1496084480.0, + "max": 1496846336.0 + }, + "peak_rss_delta_bytes": { + "median": 896663552.0, + "min": 896479232.0, + "max": 896675840.0 + } + } +} diff --git a/benchmarks/results/comment7_baseline_metadata.json b/benchmarks/results/comment7_baseline_metadata.json new file mode 100644 index 0000000..695ed36 --- /dev/null +++ b/benchmarks/results/comment7_baseline_metadata.json @@ -0,0 +1,45 @@ +{ + "created_before_source_update": true, + "python": "3.10.12", + "packages": { + "numpy": "1.26.4", + "scipy": "1.15.3", + "pandas": "2.3.3", + "anndata": "0.11.4", + "scanpy": "1.11.5", + "torch": "2.4.1", + "pyro-ppl": "1.9.1" + }, + "source_sha256": { + "prior.py": "ad4375a7e39f08a5e5afc2947e0a8b75234de4b841c01144b5afd2ccdadc0d52", + "trainer.py": "5f58825cfca289a923636aa4938b1f00f87b5026df585b6364d556701e0b5029", + "walks.py": "fdf18ffd82f98d57d7ff9dac4d59ae2b1e6b8901fd3e2aa09783e0cb0f6c3cb8" + }, + "dataset": { + "path": "../sclsd-manuscript/Zenodo/DentateGyrus/preprocessed_adata.h5ad", + "sha256": "6a8d2cce85e6e746d877af091b0175ec6b0a36440e9e1cf3ef55541efb6876c5", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "cluster_count": 6 + }, + "configuration": { + "walk_count": 2048, + "path_len": 12, + "random_state": 42, + "transition_beta": 50 + }, + "legacy_outputs": { + "raw_phylogeny_mask_nnz": 3128837, + "effective_phylogeny_nnz": 48504, + "effective_phylogeny_sha256": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19", + "transition_nnz": 53472, + "transition_sha256": "8a157a181ef8e36d911679b452b0694ccbd84093448b480b7265e7ca57e6025d", + "walks_sha256": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968", + "walks_dtype": "int32", + "walks_shape": [ + 2048, + 12 + ] + } +} diff --git a/benchmarks/results/comment7_benchmark_report.md b/benchmarks/results/comment7_benchmark_report.md new file mode 100644 index 0000000..da7c15e --- /dev/null +++ b/benchmarks/results/comment7_benchmark_report.md @@ -0,0 +1,80 @@ +# Comment 7 sparse-memory benchmark + +Date: 2026-07-17 + +## Method + +- Environment: Python 3.10.12, NumPy 1.26.4, SciPy 1.15.3, PyTorch 2.4.1+cu121. +- Each implementation ran five times in a fresh process. +- Imports and dataset loading occurred before the timed section. +- Times below are medians; detailed minima, maxima, input dimensions, output hashes, + and total process RSS are stored in the adjacent JSON files. +- Incremental peak RSS is Linux `ru_maxrss` minus RSS immediately before the measured + operation. It includes allocations retained by the operation but excludes dataset loading. +- CUDA was not available in the execution session, so no direct GPU measurement was made. + +## Frozen baseline + +The legacy implementations were copied and executed before production-source edits. +The source hashes, dataset hash, environment, configuration, and output hashes are in +[`comment7_baseline_metadata.json`](comment7_baseline_metadata.json). The frozen +Dentate Gyrus graph and effective outputs are stored in +[`tests/fixtures/dentate_gyrus_comment7_baseline.npz`](../../tests/fixtures/dentate_gyrus_comment7_baseline.npz). + +## Old-versus-new results + +| Dataset | Operation | Legacy wall (s) | Sparse wall (s) | Legacy CPU (s) | Sparse CPU (s) | Legacy incremental peak RSS (MiB) | Sparse incremental peak RSS (MiB) | +|---|---|---:|---:|---:|---:|---:|---:| +| Dentate Gyrus (2,460 cells; 53,472 edges) | Phylogeny | 2.0818 | 0.0197 | 2.0816 | 0.0197 | 141.86 | 3.13 | +| Dentate Gyrus | Transition probabilities | 0.0877 | 0.0029 | 0.0877 | 0.0029 | 184.97 | 3.62 | +| Dentate Gyrus (2,048 walks; length 12) | Random walks | 0.1571 | 0.2764 | 3.3273 | 0.2764 | 68.59 | 1.70 | +| Bone Marrow (5,292 cells; 114,718 edges) | Phylogeny | 9.4167 | 0.0436 | 9.4163 | 0.0436 | 455.42 | 5.93 | +| Bone Marrow | Transition probabilities | 0.4338 | 0.0052 | 0.4337 | 0.0052 | 855.12 | 6.71 | + +The sparse CPU walker reduced measured incremental peak RSS by approximately 40-fold on +Dentate Gyrus. Its wall time increased by approximately 0.12 seconds for the manuscript +configuration, while measured CPU time decreased because the legacy PyTorch operation used +multiple CPU threads. + +## New-only scaling check + +Mouse Cortex contains 12,814 cells and 250,236 connectivity edges. The manuscript walk +configuration uses 8,192 walks of length 30. + +| Operation | Median wall (s) | Median CPU (s) | Median incremental peak RSS (MiB) | +|---|---:|---:|---:| +| Phylogeny | 0.0768 | 0.0768 | 7.06 | +| Transition probabilities | 0.0096 | 0.0096 | 8.47 | +| Sparse CPU random walks | 2.8805 | 2.8805 | 19.10 | + +The corresponding dense float32 transition matrix would contain 164,198,596 values and +occupy approximately 626 MiB before accounting for the per-step indexed probability rows or +other model allocations. The sparse CPU path does not allocate this matrix on the GPU. + +## Equivalence and behavior + +- Dentate Gyrus effective phylogeny connectivity has the exact frozen legacy hash. +- Bone Marrow effective phylogeny connectivity has the exact legacy hash. +- Dentate Gyrus transition support is identical and values agree with the frozen dense + result at `rtol=1e-12`, `atol=1e-14`. +- Bone Marrow transition values agree at the same tolerance; the observed maximum absolute + difference was `6.66e-16`. +- Sparse walks are exactly reproducible for a fixed seed, use only nonzero transition edges, + match controlled dense inverse-CDF choices, and reproduce expected empirical frequencies. +- Walk indices intentionally differ from legacy `torch.multinomial` output because the CPU + and PyTorch samplers use different random-number streams. +- The existing public `random_walks_gpu` function remains unchanged. + +## Real workflow verification + +The complete Dentate Gyrus phylogeny, prior-transition, and walk-generation path passed with +the manuscript configuration. As in the legacy effective graph, three cells had no permitted +outgoing transition and were removed. The retained 2,457-cell CSR transition matrix contained +48,455 nonzero entries, and 2,048 length-12 int32 walks were generated successfully. + +## Limitations + +This change removes the dense cell-by-cell transition matrix from the main random-walk GPU +path and makes phylogeny/posterior transitions edge-sparse. It does not address other possible +memory hotspots, including dense expression materialization, fate projection, perturbation, +or arbitrary direct use of the retained dense GPU utility functions. diff --git a/benchmarks/results/dentate_phylogeny_current.json b/benchmarks/results/dentate_phylogeny_current.json new file mode 100644 index 0000000..5e0e134 --- /dev/null +++ b/benchmarks/results/dentate_phylogeny_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.01972712203860283, + "cpu_seconds": 0.019749906999999567, + "rss_before_bytes": 614023168, + "peak_rss_bytes": 617267200, + "peak_rss_delta_bytes": 3244032, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.019518598914146423, + "cpu_seconds": 0.019537773999999786, + "rss_before_bytes": 613101568, + "peak_rss_bytes": 616140800, + "peak_rss_delta_bytes": 3039232, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.020828742533922195, + "cpu_seconds": 0.02085322999999839, + "rss_before_bytes": 612646912, + "peak_rss_bytes": 616173568, + "peak_rss_delta_bytes": 3526656, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.019840167835354805, + "cpu_seconds": 0.019868595999998462, + "rss_before_bytes": 613777408, + "peak_rss_bytes": 617062400, + "peak_rss_delta_bytes": 3284992, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.019691831432282925, + "cpu_seconds": 0.019714757000000915, + "rss_before_bytes": 613777408, + "peak_rss_bytes": 617308160, + "peak_rss_delta_bytes": 3530752, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + } + ], + "summary": { + "wall_seconds": { + "median": 0.01972712203860283, + "min": 0.019518598914146423, + "max": 0.020828742533922195 + }, + "cpu_seconds": { + "median": 0.019749906999999567, + "min": 0.019537773999999786, + "max": 0.02085322999999839 + }, + "peak_rss_bytes": { + "median": 617062400.0, + "min": 616140800.0, + "max": 617308160.0 + }, + "peak_rss_delta_bytes": { + "median": 3284992.0, + "min": 3039232.0, + "max": 3530752.0 + } + } +} diff --git a/benchmarks/results/dentate_phylogeny_legacy.json b/benchmarks/results/dentate_phylogeny_legacy.json new file mode 100644 index 0000000..8a05a90 --- /dev/null +++ b/benchmarks/results/dentate_phylogeny_legacy.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 2.063926392234862, + "cpu_seconds": 2.0637603710000008, + "rss_before_bytes": 580685824, + "peak_rss_bytes": 729436160, + "peak_rss_delta_bytes": 148750336, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 2.110687952488661, + "cpu_seconds": 2.110709171, + "rss_before_bytes": 580800512, + "peak_rss_bytes": 729669632, + "peak_rss_delta_bytes": 148869120, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 2.059963595122099, + "cpu_seconds": 2.0599191359999995, + "rss_before_bytes": 580669440, + "peak_rss_bytes": 729681920, + "peak_rss_delta_bytes": 149012480, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 2.0818054266273975, + "cpu_seconds": 2.0815583580000006, + "rss_before_bytes": 581242880, + "peak_rss_bytes": 729980928, + "peak_rss_delta_bytes": 148738048, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "phylogeny", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 2.105445209890604, + "cpu_seconds": 2.105247917, + "rss_before_bytes": 580775936, + "peak_rss_bytes": 729497600, + "peak_rss_delta_bytes": 148721664, + "output_type": "csr_matrix", + "output_nnz": 48504, + "output_hash": "b5e83d6c5546004fc3ef83034052b93d332b4f1b0e1364de39cb9f7742e7aa19" + } + ], + "summary": { + "wall_seconds": { + "median": 2.0818054266273975, + "min": 2.059963595122099, + "max": 2.110687952488661 + }, + "cpu_seconds": { + "median": 2.0815583580000006, + "min": 2.0599191359999995, + "max": 2.110709171 + }, + "peak_rss_bytes": { + "median": 729669632.0, + "min": 729436160.0, + "max": 729980928.0 + }, + "peak_rss_delta_bytes": { + "median": 148750336.0, + "min": 148721664.0, + "max": 149012480.0 + } + } +} diff --git a/benchmarks/results/dentate_transition_current.json b/benchmarks/results/dentate_transition_current.json new file mode 100644 index 0000000..1015bcd --- /dev/null +++ b/benchmarks/results/dentate_transition_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "current", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.003134145401418209, + "cpu_seconds": 0.003159060999999852, + "rss_before_bytes": 614301696, + "peak_rss_bytes": 617816064, + "peak_rss_delta_bytes": 3514368, + "output_type": "csr_matrix", + "output_nnz": 53472, + "output_hash": "e78bda1585db946df9e36d081db774a3661e0fcf61c5564ec034bfd3d9ff99ac" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.0029108142480254173, + "cpu_seconds": 0.002933382000000151, + "rss_before_bytes": 613703680, + "peak_rss_bytes": 617390080, + "peak_rss_delta_bytes": 3686400, + "output_type": "csr_matrix", + "output_nnz": 53472, + "output_hash": "e78bda1585db946df9e36d081db774a3661e0fcf61c5564ec034bfd3d9ff99ac" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.0028299670666456223, + "cpu_seconds": 0.002850971000000868, + "rss_before_bytes": 612798464, + "peak_rss_bytes": 616595456, + "peak_rss_delta_bytes": 3796992, + "output_type": "csr_matrix", + "output_nnz": 53472, + "output_hash": "e78bda1585db946df9e36d081db774a3661e0fcf61c5564ec034bfd3d9ff99ac" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.0027580419555306435, + "cpu_seconds": 0.002777324000000192, + "rss_before_bytes": 613740544, + "peak_rss_bytes": 617631744, + "peak_rss_delta_bytes": 3891200, + "output_type": "csr_matrix", + "output_nnz": 53472, + "output_hash": "e78bda1585db946df9e36d081db774a3661e0fcf61c5564ec034bfd3d9ff99ac" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.0028610946610569954, + "cpu_seconds": 0.0028815540000000084, + "rss_before_bytes": 613236736, + "peak_rss_bytes": 617033728, + "peak_rss_delta_bytes": 3796992, + "output_type": "csr_matrix", + "output_nnz": 53472, + "output_hash": "e78bda1585db946df9e36d081db774a3661e0fcf61c5564ec034bfd3d9ff99ac" + } + ], + "summary": { + "wall_seconds": { + "median": 0.0028610946610569954, + "min": 0.0027580419555306435, + "max": 0.003134145401418209 + }, + "cpu_seconds": { + "median": 0.0028815540000000084, + "min": 0.002777324000000192, + "max": 0.003159060999999852 + }, + "peak_rss_bytes": { + "median": 617390080.0, + "min": 616595456.0, + "max": 617816064.0 + }, + "peak_rss_delta_bytes": { + "median": 3796992.0, + "min": 3514368.0, + "max": 3891200.0 + } + } +} diff --git a/benchmarks/results/dentate_transition_legacy.json b/benchmarks/results/dentate_transition_legacy.json new file mode 100644 index 0000000..3e3f477 --- /dev/null +++ b/benchmarks/results/dentate_transition_legacy.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.08998316433280706, + "cpu_seconds": 0.09000330199999951, + "rss_before_bytes": 581414912, + "peak_rss_bytes": 775143424, + "peak_rss_delta_bytes": 193728512, + "output_type": "float64", + "output_nnz": 53472, + "output_hash": "d3a13e20b0bcc61d7c6fecf02ac2ec7e3269122541948eb27f968f968d0be343" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.08766104187816381, + "cpu_seconds": 0.0876874179999998, + "rss_before_bytes": 579649536, + "peak_rss_bytes": 773632000, + "peak_rss_delta_bytes": 193982464, + "output_type": "float64", + "output_nnz": 53472, + "output_hash": "d3a13e20b0bcc61d7c6fecf02ac2ec7e3269122541948eb27f968f968d0be343" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.08808719646185637, + "cpu_seconds": 0.0881090380000007, + "rss_before_bytes": 579612672, + "peak_rss_bytes": 773570560, + "peak_rss_delta_bytes": 193957888, + "output_type": "float64", + "output_nnz": 53472, + "output_hash": "d3a13e20b0bcc61d7c6fecf02ac2ec7e3269122541948eb27f968f968d0be343" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.08749268297106028, + "cpu_seconds": 0.0875133090000002, + "rss_before_bytes": 580128768, + "peak_rss_bytes": 774180864, + "peak_rss_delta_bytes": 194052096, + "output_type": "float64", + "output_nnz": 53472, + "output_hash": "d3a13e20b0bcc61d7c6fecf02ac2ec7e3269122541948eb27f968f968d0be343" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "transition", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.08524742163717747, + "cpu_seconds": 0.08527022899999892, + "rss_before_bytes": 581160960, + "peak_rss_bytes": 775069696, + "peak_rss_delta_bytes": 193908736, + "output_type": "float64", + "output_nnz": 53472, + "output_hash": "d3a13e20b0bcc61d7c6fecf02ac2ec7e3269122541948eb27f968f968d0be343" + } + ], + "summary": { + "wall_seconds": { + "median": 0.08766104187816381, + "min": 0.08524742163717747, + "max": 0.08998316433280706 + }, + "cpu_seconds": { + "median": 0.0876874179999998, + "min": 0.08527022899999892, + "max": 0.09000330199999951 + }, + "peak_rss_bytes": { + "median": 774180864.0, + "min": 773570560.0, + "max": 775143424.0 + }, + "peak_rss_delta_bytes": { + "median": 193957888.0, + "min": 193728512.0, + "max": 194052096.0 + } + } +} diff --git a/benchmarks/results/dentate_walks_current.json b/benchmarks/results/dentate_walks_current.json new file mode 100644 index 0000000..79a7680 --- /dev/null +++ b/benchmarks/results/dentate_walks_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "current", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.2733072619885206, + "cpu_seconds": 0.27333117900000126, + "rss_before_bytes": 618790912, + "peak_rss_bytes": 620683264, + "peak_rss_delta_bytes": 1892352, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "09d1d09c6ac50dabb57e270d051b506904b166794a7c8fbf6338aade93612c20" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.27443698327988386, + "cpu_seconds": 0.27445300800000005, + "rss_before_bytes": 620027904, + "peak_rss_bytes": 620793856, + "peak_rss_delta_bytes": 765952, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "09d1d09c6ac50dabb57e270d051b506904b166794a7c8fbf6338aade93612c20" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.28869067784398794, + "cpu_seconds": 0.2887128050000012, + "rss_before_bytes": 619143168, + "peak_rss_bytes": 620933120, + "peak_rss_delta_bytes": 1789952, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "09d1d09c6ac50dabb57e270d051b506904b166794a7c8fbf6338aade93612c20" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.27729102596640587, + "cpu_seconds": 0.27730707100000096, + "rss_before_bytes": 619102208, + "peak_rss_bytes": 620888064, + "peak_rss_delta_bytes": 1785856, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "09d1d09c6ac50dabb57e270d051b506904b166794a7c8fbf6338aade93612c20" + }, + { + "dataset": "dentate", + "implementation": "current", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.2764015207067132, + "cpu_seconds": 0.2764254279999996, + "rss_before_bytes": 619253760, + "peak_rss_bytes": 620982272, + "peak_rss_delta_bytes": 1728512, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "09d1d09c6ac50dabb57e270d051b506904b166794a7c8fbf6338aade93612c20" + } + ], + "summary": { + "wall_seconds": { + "median": 0.2764015207067132, + "min": 0.2733072619885206, + "max": 0.28869067784398794 + }, + "cpu_seconds": { + "median": 0.2764254279999996, + "min": 0.27333117900000126, + "max": 0.2887128050000012 + }, + "peak_rss_bytes": { + "median": 620888064.0, + "min": 620683264.0, + "max": 620982272.0 + }, + "peak_rss_delta_bytes": { + "median": 1785856.0, + "min": 765952.0, + "max": 1892352.0 + } + } +} diff --git a/benchmarks/results/dentate_walks_legacy.json b/benchmarks/results/dentate_walks_legacy.json new file mode 100644 index 0000000..8a72fae --- /dev/null +++ b/benchmarks/results/dentate_walks_legacy.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.12815706990659237, + "cpu_seconds": 2.428806731, + "rss_before_bytes": 618860544, + "peak_rss_bytes": 690286592, + "peak_rss_delta_bytes": 71426048, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.17988937068730593, + "cpu_seconds": 3.9783924080000013, + "rss_before_bytes": 617476096, + "peak_rss_bytes": 689393664, + "peak_rss_delta_bytes": 71917568, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.15656772535294294, + "cpu_seconds": 3.3273125460000017, + "rss_before_bytes": 617881600, + "peak_rss_bytes": 689819648, + "peak_rss_delta_bytes": 71938048, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.15714468900114298, + "cpu_seconds": 3.0777459679999986, + "rss_before_bytes": 617635840, + "peak_rss_bytes": 689332224, + "peak_rss_delta_bytes": 71696384, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968" + }, + { + "dataset": "dentate", + "implementation": "legacy", + "operation": "walks", + "n_obs": 2460, + "n_vars": 1500, + "connectivity_nnz": 53472, + "wall_seconds": 0.30562551971524954, + "cpu_seconds": 6.710330459999998, + "rss_before_bytes": 618528768, + "peak_rss_bytes": 710868992, + "peak_rss_delta_bytes": 92340224, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "d9b966cff9677a9001fce384ec4240cbcd8c43d590fceb105e96d508684eb968" + } + ], + "summary": { + "wall_seconds": { + "median": 0.15714468900114298, + "min": 0.12815706990659237, + "max": 0.30562551971524954 + }, + "cpu_seconds": { + "median": 3.3273125460000017, + "min": 2.428806731, + "max": 6.710330459999998 + }, + "peak_rss_bytes": { + "median": 689819648.0, + "min": 689332224.0, + "max": 710868992.0 + }, + "peak_rss_delta_bytes": { + "median": 71917568.0, + "min": 71426048.0, + "max": 92340224.0 + } + } +} diff --git a/benchmarks/results/mouse_cortex_phylogeny_current.json b/benchmarks/results/mouse_cortex_phylogeny_current.json new file mode 100644 index 0000000..fe2ab2c --- /dev/null +++ b/benchmarks/results/mouse_cortex_phylogeny_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.07680433057248592, + "cpu_seconds": 0.07682317599999955, + "rss_before_bytes": 701759488, + "peak_rss_bytes": 709160960, + "peak_rss_delta_bytes": 7401472, + "output_type": "csr_matrix", + "output_nnz": 240930, + "output_hash": "90da8200b99af5fe41ac530946c6b45a4a0342df23d1454e9c78ebdde9a315c0" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.07630886416882277, + "cpu_seconds": 0.07633075400000067, + "rss_before_bytes": 701943808, + "peak_rss_bytes": 709578752, + "peak_rss_delta_bytes": 7634944, + "output_type": "csr_matrix", + "output_nnz": 240930, + "output_hash": "90da8200b99af5fe41ac530946c6b45a4a0342df23d1454e9c78ebdde9a315c0" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.07531259302049875, + "cpu_seconds": 0.0753304289999992, + "rss_before_bytes": 702316544, + "peak_rss_bytes": 706838528, + "peak_rss_delta_bytes": 4521984, + "output_type": "csr_matrix", + "output_nnz": 240930, + "output_hash": "90da8200b99af5fe41ac530946c6b45a4a0342df23d1454e9c78ebdde9a315c0" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.0774783305823803, + "cpu_seconds": 0.07750002000000045, + "rss_before_bytes": 702865408, + "peak_rss_bytes": 707145728, + "peak_rss_delta_bytes": 4280320, + "output_type": "csr_matrix", + "output_nnz": 240930, + "output_hash": "90da8200b99af5fe41ac530946c6b45a4a0342df23d1454e9c78ebdde9a315c0" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "phylogeny", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.07780855428427458, + "cpu_seconds": 0.07782924500000021, + "rss_before_bytes": 702197760, + "peak_rss_bytes": 711782400, + "peak_rss_delta_bytes": 9584640, + "output_type": "csr_matrix", + "output_nnz": 240930, + "output_hash": "90da8200b99af5fe41ac530946c6b45a4a0342df23d1454e9c78ebdde9a315c0" + } + ], + "summary": { + "wall_seconds": { + "median": 0.07680433057248592, + "min": 0.07531259302049875, + "max": 0.07780855428427458 + }, + "cpu_seconds": { + "median": 0.07682317599999955, + "min": 0.0753304289999992, + "max": 0.07782924500000021 + }, + "peak_rss_bytes": { + "median": 709160960.0, + "min": 706838528.0, + "max": 711782400.0 + }, + "peak_rss_delta_bytes": { + "median": 7401472.0, + "min": 4280320.0, + "max": 9584640.0 + } + } +} diff --git a/benchmarks/results/mouse_cortex_transition_current.json b/benchmarks/results/mouse_cortex_transition_current.json new file mode 100644 index 0000000..dde8300 --- /dev/null +++ b/benchmarks/results/mouse_cortex_transition_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "transition", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.008212902583181858, + "cpu_seconds": 0.008234408000001636, + "rss_before_bytes": 702537728, + "peak_rss_bytes": 709111808, + "peak_rss_delta_bytes": 6574080, + "output_type": "csr_matrix", + "output_nnz": 250236, + "output_hash": "faa56c17c15dd37d185c25ab785b8a82c1c7ba060569e4c1707412b7e891c918" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "transition", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.009271243587136269, + "cpu_seconds": 0.009294531000000106, + "rss_before_bytes": 702046208, + "peak_rss_bytes": 710545408, + "peak_rss_delta_bytes": 8499200, + "output_type": "csr_matrix", + "output_nnz": 250236, + "output_hash": "faa56c17c15dd37d185c25ab785b8a82c1c7ba060569e4c1707412b7e891c918" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "transition", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.010138964280486107, + "cpu_seconds": 0.010161818000000267, + "rss_before_bytes": 702042112, + "peak_rss_bytes": 710922240, + "peak_rss_delta_bytes": 8880128, + "output_type": "csr_matrix", + "output_nnz": 250236, + "output_hash": "faa56c17c15dd37d185c25ab785b8a82c1c7ba060569e4c1707412b7e891c918" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "transition", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.009996730834245682, + "cpu_seconds": 0.010016089000000505, + "rss_before_bytes": 702201856, + "peak_rss_bytes": 714915840, + "peak_rss_delta_bytes": 12713984, + "output_type": "csr_matrix", + "output_nnz": 250236, + "output_hash": "faa56c17c15dd37d185c25ab785b8a82c1c7ba060569e4c1707412b7e891c918" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "transition", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 0.009615094400942326, + "cpu_seconds": 0.009638960000000196, + "rss_before_bytes": 702152704, + "peak_rss_bytes": 712552448, + "peak_rss_delta_bytes": 10399744, + "output_type": "csr_matrix", + "output_nnz": 250236, + "output_hash": "faa56c17c15dd37d185c25ab785b8a82c1c7ba060569e4c1707412b7e891c918" + } + ], + "summary": { + "wall_seconds": { + "median": 0.009615094400942326, + "min": 0.008212902583181858, + "max": 0.010138964280486107 + }, + "cpu_seconds": { + "median": 0.009638960000000196, + "min": 0.008234408000001636, + "max": 0.010161818000000267 + }, + "peak_rss_bytes": { + "median": 710922240.0, + "min": 709111808.0, + "max": 714915840.0 + }, + "peak_rss_delta_bytes": { + "median": 8880128.0, + "min": 6574080.0, + "max": 12713984.0 + } + } +} diff --git a/benchmarks/results/mouse_cortex_walks_current.json b/benchmarks/results/mouse_cortex_walks_current.json new file mode 100644 index 0000000..d4c41ed --- /dev/null +++ b/benchmarks/results/mouse_cortex_walks_current.json @@ -0,0 +1,108 @@ +{ + "python": "3.10.12", + "repeats": 5, + "records": [ + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "walks", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 2.8961584316566586, + "cpu_seconds": 2.8959552239999997, + "rss_before_bytes": 713158656, + "peak_rss_bytes": 727347200, + "peak_rss_delta_bytes": 14188544, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "1abc28636c709f357c1b060ef908c4dccc10aad80da80bfea9f16cd713dc035b" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "walks", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 2.8579091550782323, + "cpu_seconds": 2.857932776, + "rss_before_bytes": 704888832, + "peak_rss_bytes": 724914176, + "peak_rss_delta_bytes": 20025344, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "1abc28636c709f357c1b060ef908c4dccc10aad80da80bfea9f16cd713dc035b" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "walks", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 2.8777672694996, + "cpu_seconds": 2.877525584999999, + "rss_before_bytes": 704634880, + "peak_rss_bytes": 728133632, + "peak_rss_delta_bytes": 23498752, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "1abc28636c709f357c1b060ef908c4dccc10aad80da80bfea9f16cd713dc035b" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "walks", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 2.8843463687226176, + "cpu_seconds": 2.8843807860000013, + "rss_before_bytes": 708177920, + "peak_rss_bytes": 728420352, + "peak_rss_delta_bytes": 20242432, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "1abc28636c709f357c1b060ef908c4dccc10aad80da80bfea9f16cd713dc035b" + }, + { + "dataset": "mouse_cortex", + "implementation": "current", + "operation": "walks", + "n_obs": 12814, + "n_vars": 1062, + "connectivity_nnz": 250236, + "wall_seconds": 2.8805352691560984, + "cpu_seconds": 2.8804612400000007, + "rss_before_bytes": 712642560, + "peak_rss_bytes": 726831104, + "peak_rss_delta_bytes": 14188544, + "output_type": "torch.int32", + "output_nnz": null, + "output_hash": "1abc28636c709f357c1b060ef908c4dccc10aad80da80bfea9f16cd713dc035b" + } + ], + "summary": { + "wall_seconds": { + "median": 2.8805352691560984, + "min": 2.8579091550782323, + "max": 2.8961584316566586 + }, + "cpu_seconds": { + "median": 2.8804612400000007, + "min": 2.857932776, + "max": 2.8959552239999997 + }, + "peak_rss_bytes": { + "median": 727347200.0, + "min": 724914176.0, + "max": 728420352.0 + }, + "peak_rss_delta_bytes": { + "median": 20025344.0, + "min": 14188544.0, + "max": 23498752.0 + } + } +} diff --git a/pyproject.toml b/pyproject.toml index 7a569a2..7739911 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,12 @@ readme = "README.md" license = "MIT" requires-python = ">=3.9" authors = [ - { name = "LSD Development Team" } + { name = "Ali Poursina", email = "ali.poursina@mail.mcgill.ca" }, + { name = "Ali Saberi" }, +] +maintainers = [ + { name = "Ali Poursina", email = "ali.poursina@mail.mcgill.ca" }, + { name = "Ali Saberi" }, ] keywords = [ "single-cell", @@ -62,19 +67,8 @@ dev = [ "mypy>=1.0.0", "pre-commit>=3.0.0", ] -docs = [ - "sphinx>=7.0.0", - "furo>=2024.0.0", - "sphinx-autodoc-typehints>=1.25.0", - "sphinx-copybutton>=0.5.0", - "nbsphinx>=0.9.0", - "myst-parser>=2.0.0", -] -all = ["sclsd[dev,docs]"] - [project.urls] Homepage = "https://github.com/csglab/sclsd" -Documentation = "https://csglab.github.io/sclsd" Repository = "https://github.com/csglab/sclsd" Issues = "https://github.com/csglab/sclsd/issues" diff --git a/src/sclsd/core/model.py b/src/sclsd/core/model.py index 09006ea..9bd474c 100644 --- a/src/sclsd/core/model.py +++ b/src/sclsd/core/model.py @@ -316,15 +316,21 @@ def model( z_loc, z_scale = self.B_decoder(B) # Timepoint lattice for variance parameter cell state - t = torch.exp(torch.linspace(0, 1, self.path_len).view(1, self.path_len, 1)) + t = torch.exp( + torch.linspace( + 0, + 1, + self.path_len, + device=z_scale.device, + dtype=z_scale.dtype, + ).view(1, self.path_len, 1) + ) t = t.expand( int(len(x) / self.path_len), self.path_len, self.z_dim ).contiguous() z_scale = z_scale.reshape( int(len(x) / self.path_len), self.path_len, self.z_dim ) - if torch.cuda.is_available(): - t = t.to(self.device) z_scale = (z_scale * t).reshape(-1, self.z_dim) # 4. Sample z @@ -399,9 +405,13 @@ def guide( pyro.factor("V_l2_reg", self.V_coeff * V.pow(2).max(), has_rsample=True) # Timepoint lattice for neural ODE - t = torch.linspace(0, 1, self.path_len) - if torch.cuda.is_available(): - t = t.to(self.device) + t = torch.linspace( + 0, + 1, + self.path_len, + device=z_loc.device, + dtype=z_loc.dtype, + ) z_t = z_loc.reshape( int(len(x) / self.path_len), self.path_len, self.z_dim ) @@ -424,12 +434,18 @@ def guide( B_scale = B_scale.reshape( int(len(x) / self.path_len), self.path_len, self.B_dim ) - t = torch.exp(torch.linspace(0, 1, self.path_len).view(1, self.path_len, 1)) + t = torch.exp( + torch.linspace( + 0, + 1, + self.path_len, + device=B_scale.device, + dtype=B_scale.dtype, + ).view(1, self.path_len, 1) + ) t = t.expand( int(len(x) / self.path_len), self.path_len, self.B_dim ).contiguous() - if torch.cuda.is_available(): - t = t.to(self.device) B_scale = (B_scale * t).reshape(-1, self.B_dim) # 4. Sample B with annealing factor diff --git a/src/sclsd/preprocessing/prior.py b/src/sclsd/preprocessing/prior.py index 23da93a..2824cc5 100644 --- a/src/sclsd/preprocessing/prior.py +++ b/src/sclsd/preprocessing/prior.py @@ -358,28 +358,37 @@ def _create_phylogeny_matrix( phylogeny: Dict[str, List[str]], cluster_key: str = "clusters", ) -> sp.csr_matrix: - """Create cell-cell matrix based on phylogeny relationships.""" + """Create a phylogeny mask over existing connectivity-graph edges.""" clusters = adata.obs[cluster_key].unique().tolist() - cell_to_cluster = dict(zip(adata.obs_names, adata.obs[cluster_key])) + cell_clusters = adata.obs[cluster_key].to_numpy() all_descendants = {} for cluster in clusters: all_descendants[cluster] = _get_all_descendants(cluster, phylogeny) - n_cells = adata.shape[0] - phylo_matrix = np.zeros((n_cells, n_cells)) - - for i, cell_i in enumerate(adata.obs_names): - cluster_i = cell_to_cluster[cell_i] - for j, cell_j in enumerate(adata.obs_names): - cluster_j = cell_to_cluster[cell_j] - - if cluster_i == cluster_j: - phylo_matrix[i, j] = 1 - elif cluster_j in all_descendants.get(cluster_i, set()): - phylo_matrix[i, j] = 1 - - return sp.csr_matrix(phylo_matrix) + connectivity = sp.csr_matrix(adata.obsp["connectivities"]) + connectivity.sum_duplicates() + connectivity.eliminate_zeros() + edges = connectivity.tocoo() + + keep = np.fromiter( + ( + cell_clusters[row] == cell_clusters[col] + or cell_clusters[col] in all_descendants.get(cell_clusters[row], set()) + for row, col in zip(edges.row, edges.col) + ), + dtype=bool, + count=edges.nnz, + ) + mask = sp.csr_matrix( + ( + np.ones(np.count_nonzero(keep), dtype=float), + (edges.row[keep], edges.col[keep]), + ), + shape=connectivity.shape, + ) + mask.sort_indices() + return mask def _get_tree_branches( diff --git a/src/sclsd/train/trainer.py b/src/sclsd/train/trainer.py index 13e75eb..4bc8337 100644 --- a/src/sclsd/train/trainer.py +++ b/src/sclsd/train/trainer.py @@ -26,6 +26,10 @@ from sclsd.core.config import LSDConfig, WalkConfig from sclsd.core.model import LSDModel +from sclsd.preprocessing.prior import ( + _create_phylogeny_matrix as create_sparse_phylogeny_matrix, +) +from sclsd.train.walks import random_walks_sparse from sclsd.utils.seed import set_all_seeds, clear_pyro_state, enable_pyro_validation try: @@ -63,8 +67,8 @@ class LSD: The single-cell data. walks : torch.Tensor Generated random walks for training. - P : torch.Tensor - Cell-cell transition probability matrix. + P : scipy.sparse.csr_matrix + Sparse cell-cell transition probability matrix used for random walks. Examples -------- @@ -87,6 +91,25 @@ def __init__( lib_size_key: str = "librarysize", raw_count_key: str = "raw", ): + missing_fields = [] + if adata.X is None: + missing_fields.append( + "adata.X is missing; provide log-normalized expression data." + ) + if raw_count_key not in adata.layers: + missing_fields.append( + f"adata.layers[{raw_count_key!r}] is missing; " + "provide the raw count matrix." + ) + if lib_size_key not in adata.obs: + missing_fields.append( + f"adata.obs[{lib_size_key!r}] is missing; " + "provide one library-size value per cell." + ) + if missing_fields: + details = "\n- ".join(missing_fields) + raise ValueError(f"Invalid AnnData input for LSD:\n- {details}") + self.config = config if isinstance(config, LSDConfig) else LSDConfig() model_cfg = self.config.model walk_cfg = replace(self.config.walks) @@ -500,32 +523,47 @@ def get_variables( def calculate_transition_probs( self, potential: np.ndarray, - connectivity_matrix: np.ndarray, + connectivity_matrix: Union[np.ndarray, sp.spmatrix], beta: float = 1.0, - ) -> np.ndarray: + ) -> sp.csr_matrix: """Compute cell-cell transition probabilities using Boltzmann weights. Parameters ---------- potential : np.ndarray Potential values of shape (n_cells,). - connectivity_matrix : np.ndarray - Binary connectivity matrix of shape (n_cells, n_cells). + connectivity_matrix : numpy.ndarray or scipy.sparse matrix + Connectivity matrix of shape (n_cells, n_cells). beta : float Boltzmann scaling factor. Returns ------- - np.ndarray - Transition probability matrix, row-normalized. + scipy.sparse.csr_matrix + Sparse transition probability matrix, row-normalized. """ - potential = potential.astype(float) - energy_diff = potential[None, :] - potential[:, None] - boltzmann_weights = np.exp(-beta * energy_diff) - boltzmann_weights *= connectivity_matrix - row_sums = boltzmann_weights.sum(axis=1, keepdims=True) + 1e-12 - transition_matrix = boltzmann_weights / row_sums - return transition_matrix + potential = np.asarray(potential, dtype=float).reshape(-1) + connectivity = sp.csr_matrix(connectivity_matrix, dtype=float) + connectivity.sum_duplicates() + connectivity.eliminate_zeros() + connectivity.sort_indices() + + rows = np.repeat( + np.arange(connectivity.shape[0]), np.diff(connectivity.indptr) + ) + energy_diff = potential[connectivity.indices] - potential[rows] + weights = np.exp(-beta * energy_diff) * connectivity.data + row_sums = np.bincount( + rows, weights=weights, minlength=connectivity.shape[0] + ) + normalized_weights = weights / (row_sums[rows] + 1e-12) + + transition = sp.csr_matrix( + (normalized_weights, connectivity.indices.copy(), connectivity.indptr.copy()), + shape=connectivity.shape, + ) + transition.eliminate_zeros() + return transition def set_adata(self, adata: "AnnData") -> None: """Set the AnnData object.""" @@ -557,9 +595,7 @@ def get_adata(self) -> "AnnData": "Run neighbors graph computation." ) connectivity = adata.obsp["connectivities"] - if not isinstance(connectivity, np.ndarray): - connectivity = connectivity.toarray() - binary_connectivity = (connectivity > 0).astype(float) + binary_connectivity = (sp.csr_matrix(connectivity) > 0).astype(float).tocsr() transition_matrix = self.calculate_transition_probs( potential=potential.squeeze(-1).cpu().numpy(), @@ -605,15 +641,19 @@ def prepare_walks(self, n_trajectories: Optional[int] = None) -> None: ) n_trajectories = self.walk_config.num_walks - self.P = self.P.to(self.device) - walks = self._random_walks(n_trajectories) - self.P = self.P.cpu() - self.walks = walks.cpu() + random_state = getattr(self.walk_config, "random_state", 42) + self.P = sp.csr_matrix(self.P) + self.walks = random_walks_sparse( + self.P, + n_steps=self.path_len, + n_trajectories=n_trajectories, + random_state=random_state, + ) def set_prior_transition( self, prior_time_key: Optional[str] = None, - prior_transition: Optional[np.ndarray] = None, + prior_transition: Optional[Union[np.ndarray, sp.spmatrix]] = None, random_state: int = 42, ) -> None: """Set the prior cell-cell transition matrix. @@ -622,7 +662,7 @@ def set_prior_transition( ---------- prior_time_key : str, optional Name of pseudotime key in adata.obs. - prior_transition : np.ndarray, optional + prior_transition : numpy.ndarray or scipy.sparse matrix, optional Precomputed prior transition matrix. random_state : int Random seed. @@ -635,52 +675,53 @@ def _get_connectivity_matrix(): "'connectivities' matrix not found in adata.obsp. " "Run neighbors graph computation (e.g. sc.pp.neighbors)." ) - mat = self.adata.obsp["connectivities"] - if not isinstance(mat, np.ndarray): - mat = mat.toarray() - return (mat > 0).astype(float) + mat = sp.csr_matrix(self.adata.obsp["connectivities"]) + return (mat > 0).astype(float).tocsr() if prior_transition is not None: - if not isinstance(prior_transition, np.ndarray): - prior_transition = prior_transition.toarray() - if prior_transition.shape != (n_cells, n_cells): + transition = sp.csr_matrix(prior_transition, dtype=np.float32) + if transition.shape != (n_cells, n_cells): raise ValueError( - f"Shape mismatch: prior_transition has shape {prior_transition.shape}, " + f"Shape mismatch: prior_transition has shape {transition.shape}, " f"but expected ({n_cells}, {n_cells}) from adata." ) - self.P = torch.from_numpy(prior_transition).float() + self.P = transition print("[LSD] Prior transition matrix set from user input.") return + if prior_time_key is not None and prior_time_key not in self.adata.obs: + raise KeyError( + f"prior_time_key={prior_time_key!r} was not found in adata.obs. " + "Add the pseudotime column or pass prior_transition instead." + ) + if self.phylogeny is not None: A = self._create_phylogeny_matrix() - if not isinstance(A, np.ndarray): - A = A.toarray() connectivity = _get_connectivity_matrix() - A *= connectivity + A = A.multiply(connectivity).tocsr() self.adata.obsp["phylogeny_matrix"] = A - row_sums = A.sum(axis=1) + row_sums = np.asarray(A.sum(axis=1)).ravel() valid_cells = row_sums > 0 if len(self.adata[~valid_cells]) != 0: print( f"[LSD] Removing {np.sum(~valid_cells)} cells with no transitions:" ) - self.adata = self.adata[valid_cells] + self.adata = self.adata[valid_cells].copy() if prior_time_key is not None: P = self._get_transition_from_pseudotime( - prior_time_key, self.adata.obsp["phylogeny_matrix"].toarray() + prior_time_key, self.adata.obsp["phylogeny_matrix"] ) print("[LSD] Prior transition matrix set from phylogeny and pseudotime.") else: raise KeyError("Run the function get_prior_transition first") - self.P = torch.from_numpy(P).float() + self.P = P.astype(np.float32) return if prior_time_key is not None: connectivity = _get_connectivity_matrix() P = self._get_transition_from_pseudotime(prior_time_key, connectivity) - self.P = torch.from_numpy(P).float() + self.P = P.astype(np.float32) print("[LSD] Prior transition matrix set from pseudotime and connectivities.") return @@ -692,8 +733,8 @@ def _get_connectivity_matrix(): def _get_transition_from_pseudotime( self, time_key: str, - connectivity: np.ndarray, - ) -> np.ndarray: + connectivity: Union[np.ndarray, sp.spmatrix], + ) -> sp.csr_matrix: """Compute transition matrix from pseudotime.""" potential = -self.adata.obs[time_key].values P = self.calculate_transition_probs(potential, connectivity, beta=50) @@ -726,29 +767,12 @@ def set_phylogeny(self, phylogeny: Dict[str, List[str]], cluster_key: str) -> No def _create_phylogeny_matrix(self) -> sp.csr_matrix: """Create phylogeny-based adjacency matrix.""" - clusters = self.adata.obs[self.cluster_key].unique().tolist() - cell_to_cluster = dict( - zip(self.adata.obs_names, self.adata.obs[self.cluster_key]) + return create_sparse_phylogeny_matrix( + self.adata, + self.phylogeny, + self.cluster_key, ) - all_descendants = {} - for cluster in clusters: - all_descendants[cluster] = self._get_all_descendants(cluster) - - n_cells = self.adata.shape[0] - phylo_matrix = np.zeros((n_cells, n_cells)) - - for i, cell_i in enumerate(self.adata.obs_names): - cluster_i = cell_to_cluster[cell_i] - for j, cell_j in enumerate(self.adata.obs_names): - cluster_j = cell_to_cluster[cell_j] - if cluster_i == cluster_j: - phylo_matrix[i, j] = 1 - elif cluster_j in all_descendants.get(cluster_i, set()): - phylo_matrix[i, j] = 1 - - return sp.csr_matrix(phylo_matrix) - def ode_solve( self, z0: torch.Tensor, diff --git a/src/sclsd/train/walks.py b/src/sclsd/train/walks.py index afad6de..9b2d973 100644 --- a/src/sclsd/train/walks.py +++ b/src/sclsd/train/walks.py @@ -6,11 +6,9 @@ from __future__ import annotations -from typing import Optional - import numpy as np -import torch import scipy.sparse as sp +import torch def prepare_transition_matrix_gpu( @@ -80,6 +78,71 @@ def random_walks_gpu( return walks +def random_walks_sparse( + transition_matrix: sp.spmatrix, + n_steps: int, + n_trajectories: int, + random_state: int = 42, +) -> torch.Tensor: + """Generate reproducible random walks directly from a sparse matrix. + + Sampling is performed on CPU over each cell's nonzero transition edges, so + the full cell-by-cell transition matrix is never copied to GPU memory. + + Parameters + ---------- + transition_matrix : scipy.sparse matrix + Row-wise transition probabilities. + n_steps : int + Number of states in each walk. + n_trajectories : int + Number of walks to generate. + random_state : int + Seed for the NumPy random-number generator. + + Returns + ------- + torch.Tensor + CPU tensor of shape ``(n_trajectories, n_steps)`` and dtype int32. + """ + transition = sp.csr_matrix(transition_matrix, dtype=float) + transition.sum_duplicates() + transition.eliminate_zeros() + transition.sort_indices() + + if transition.shape[0] != transition.shape[1]: + raise ValueError("Transition matrix must be square.") + if np.any(~np.isfinite(transition.data)) or np.any(transition.data < 0): + raise ValueError("Transition probabilities must be finite and non-negative.") + + row_sums = np.asarray(transition.sum(axis=1)).ravel() + empty_rows = np.flatnonzero(row_sums <= 0) + if len(empty_rows) > 0: + preview = ", ".join(map(str, empty_rows[:5])) + raise ValueError( + "Transition matrix contains rows with no outgoing transitions " + f"(for example: {preview})." + ) + + row_indices = np.repeat(np.arange(transition.shape[0]), np.diff(transition.indptr)) + transition.data = transition.data / row_sums[row_indices] + + rng = np.random.default_rng(random_state) + walks = np.empty((n_trajectories, n_steps), dtype=np.int32) + walks[:, 0] = rng.integers(0, transition.shape[0], size=n_trajectories) + + for step in range(1, n_steps): + for walk_idx in range(n_trajectories): + current = walks[walk_idx, step - 1] + start = transition.indptr[current] + end = transition.indptr[current + 1] + neighbors = transition.indices[start:end] + probabilities = transition.data[start:end] + walks[walk_idx, step] = rng.choice(neighbors, p=probabilities) + + return torch.from_numpy(walks) + + def prepare_walks( transition_matrix: torch.Tensor, n_trajectories: int, diff --git a/tests/fixtures/dentate_gyrus_comment7_baseline.npz b/tests/fixtures/dentate_gyrus_comment7_baseline.npz new file mode 100644 index 0000000..d7c5ba6 Binary files /dev/null and b/tests/fixtures/dentate_gyrus_comment7_baseline.npz differ diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..c5e3d2c --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,97 @@ +import numpy as np +import pandas as pd +import scipy.sparse as sp +import torch +from anndata import AnnData + +from sclsd import LSD, LSDConfig, clear_pyro_state, set_all_seeds + + +def test_toy_training_smoke(): + """Exercise the principal LSD training path on a tiny CPU dataset.""" + raw = np.array( + [ + [4, 1, 0, 0], + [3, 2, 0, 0], + [2, 3, 1, 0], + [1, 4, 1, 0], + [0, 3, 2, 1], + [0, 2, 3, 1], + [0, 1, 4, 2], + [0, 0, 4, 3], + ], + dtype=np.float32, + ) + library_size = raw.sum(axis=1) + normalized = np.log1p(raw / library_size[:, None] * 1e4) + + rows = [] + cols = [] + for cell in range(raw.shape[0]): + rows.append(cell) + cols.append(cell) + if cell > 0: + rows.append(cell) + cols.append(cell - 1) + if cell + 1 < raw.shape[0]: + rows.append(cell) + cols.append(cell + 1) + + connectivities = sp.csr_matrix( + (np.ones(len(rows), dtype=np.float32), (rows, cols)), + shape=(raw.shape[0], raw.shape[0]), + ) + adata = AnnData( + X=sp.csr_matrix(normalized), + obs=pd.DataFrame( + { + "librarysize": library_size, + "pseudotime": np.linspace(0, 1, raw.shape[0]), + }, + index=[f"cell-{i}" for i in range(raw.shape[0])], + ), + var=pd.DataFrame(index=[f"gene-{i}" for i in range(raw.shape[1])]), + ) + adata.layers["raw"] = sp.csr_matrix(raw) + adata.obsp["connectivities"] = connectivities + + cfg = LSDConfig() + cfg.model.z_dim = 2 + cfg.model.layer_dims.B_decoder = [4] + cfg.model.layer_dims.z_decoder = [4] + cfg.model.layer_dims.x_encoder = [4] + cfg.model.layer_dims.z_encoder = [4] + cfg.model.layer_dims.xl_encoder = [4] + cfg.model.layer_dims.potential = [4] + cfg.walks.batch_size = 2 + cfg.walks.path_len = 2 + cfg.walks.num_walks = 4 + + clear_pyro_state() + set_all_seeds(0) + lsd = LSD(adata, cfg, device=torch.device("cpu")) + lsd.set_prior_transition(prior_time_key="pseudotime", random_state=0) + lsd.prepare_walks() + lsd.train(num_epochs=1, plot_loss=False, random_state=0) + result = lsd.get_adata() + + assert result.obsm["X_cell_state"].shape == (adata.n_obs, cfg.model.z_dim) + assert result.obsm["X_diff_state"].shape == (adata.n_obs, cfg.model.B_dim) + assert result.obsp["transitions"].shape == (adata.n_obs, adata.n_obs) + assert sp.isspmatrix_csr(result.obsp["transitions"]) + + potential = result.obs["potential"].to_numpy().astype(float) + binary_connectivity = (adata.obsp["connectivities"].toarray() > 0).astype(float) + energy_diff = potential[None, :] - potential[:, None] + expected_weights = np.exp(-energy_diff) * binary_connectivity + expected_transitions = expected_weights / ( + expected_weights.sum(axis=1, keepdims=True) + 1e-12 + ) + np.testing.assert_allclose( + result.obsp["transitions"].toarray(), + expected_transitions, + rtol=1e-12, + atol=1e-14, + ) + assert np.isfinite(result.obs["potential"]).all() + assert np.isfinite(result.obs["entropy"]).all() diff --git a/tests/test_sparse_memory.py b/tests/test_sparse_memory.py new file mode 100644 index 0000000..9471602 --- /dev/null +++ b/tests/test_sparse_memory.py @@ -0,0 +1,203 @@ +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp +import torch +from anndata import AnnData + +from sclsd.preprocessing.prior import _create_phylogeny_matrix +from sclsd.train.trainer import LSD +from sclsd.train.walks import random_walks_sparse + +BASELINE_PATH = ( + __file__.replace( + "test_sparse_memory.py", "fixtures/dentate_gyrus_comment7_baseline.npz" + ) +) + +DENTATE_PHYLOGENY = { + "nIPC": ["Neuroblast", "Radial Glia-like"], + "Neuroblast": ["Granule immature"], + "Granule immature": ["Granule mature"], + "Granule mature": [], + "Radial Glia-like": ["Astrocytes"], + "Astrocytes": [], +} + + +def _load_csr(baseline, prefix): + return sp.csr_matrix( + ( + baseline[f"{prefix}_data"], + baseline[f"{prefix}_indices"], + baseline[f"{prefix}_indptr"], + ), + shape=tuple(baseline[f"{prefix}_shape"]), + ) + + +@pytest.fixture(scope="module") +def dentate_graph(): + with np.load(BASELINE_PATH, allow_pickle=False) as baseline: + connectivity = _load_csr(baseline, "connectivity") + expected_phylogeny = _load_csr(baseline, "effective_phylogeny") + expected_transition = _load_csr(baseline, "transition") + clusters = baseline["clusters"].copy() + pseudotime = baseline["pseudotime"].copy() + + obs = pd.DataFrame( + {"clusters": clusters, "prior_pseudotime": pseudotime}, + index=[f"dentate-cell-{i}" for i in range(len(clusters))], + ) + adata = AnnData( + X=sp.csr_matrix((len(clusters), 1), dtype=np.float32), + obs=obs, + var=pd.DataFrame(index=["placeholder-gene"]), + ) + adata.obsp["connectivities"] = connectivity + return adata, expected_phylogeny, expected_transition + + +def _assert_sparse_equal(actual, expected): + actual = actual.tocsr() + expected = expected.tocsr() + actual.sort_indices() + expected.sort_indices() + difference = actual - expected + difference.eliminate_zeros() + assert difference.nnz == 0 + + +def test_preprocessing_phylogeny_matches_frozen_effective_graph(dentate_graph): + adata, expected, _ = dentate_graph + mask = _create_phylogeny_matrix(adata, DENTATE_PHYLOGENY, "clusters") + actual = adata.obsp["connectivities"].multiply(mask).tocsr() + + _assert_sparse_equal(actual, expected) + assert sp.isspmatrix_csr(mask) + assert mask.nnz <= adata.obsp["connectivities"].nnz + + +def test_trainer_phylogeny_matches_frozen_effective_graph(dentate_graph): + adata, expected, _ = dentate_graph + lsd = object.__new__(LSD) + lsd.adata = adata + lsd.phylogeny = DENTATE_PHYLOGENY + lsd.cluster_key = "clusters" + + mask = lsd._create_phylogeny_matrix() + actual = adata.obsp["connectivities"].multiply(mask).tocsr() + + _assert_sparse_equal(actual, expected) + assert mask.nnz <= adata.obsp["connectivities"].nnz + + +def test_transition_probabilities_match_frozen_dense_result(dentate_graph): + adata, _, expected = dentate_graph + connectivity = adata.obsp["connectivities"].copy().tocsr() + connectivity.data = np.ones_like(connectivity.data, dtype=float) + potential = -adata.obs["prior_pseudotime"].to_numpy() + + actual = LSD.calculate_transition_probs(None, potential, connectivity, beta=50) + actual = actual.tocsr() + actual.sort_indices() + expected = expected.copy().tocsr() + expected.sort_indices() + + assert np.array_equal(actual.indptr, expected.indptr) + assert np.array_equal(actual.indices, expected.indices) + np.testing.assert_allclose(actual.data, expected.data, rtol=1e-12, atol=1e-14) + np.testing.assert_allclose( + np.asarray(actual.sum(axis=1)).ravel(), + np.asarray(expected.sum(axis=1)).ravel(), + rtol=1e-12, + atol=1e-14, + ) + + +def test_supplied_transition_remains_sparse(dentate_graph): + adata, _, transition = dentate_graph + lsd = object.__new__(LSD) + lsd.adata = adata + + lsd.set_prior_transition(prior_transition=transition) + + assert sp.isspmatrix_csr(lsd.P) + assert np.array_equal(lsd.P.data, transition.data.astype(np.float32)) + + +def test_sparse_walks_are_reproducible_and_follow_graph(dentate_graph): + _, _, transition = dentate_graph + first = random_walks_sparse(transition, 12, 2048, random_state=42) + second = random_walks_sparse(transition, 12, 2048, random_state=42) + + assert first.dtype == torch.int32 + assert first.shape == (2048, 12) + assert np.array_equal(first.numpy(), second.numpy()) + + walks = first.numpy() + for step in range(1, walks.shape[1]): + weights = np.asarray(transition[walks[:, step - 1], walks[:, step]]).ravel() + assert np.all(weights > 0) + + +def test_sparse_walk_empirical_frequencies_match_probabilities(): + probabilities = np.array([0.1, 0.2, 0.3, 0.4], dtype=float) + transition = sp.csr_matrix(np.tile(probabilities, (4, 1))) + + walks = random_walks_sparse( + transition, n_steps=2, n_trajectories=100_000, random_state=7 + ).numpy() + observed = np.bincount(walks[:, 1], minlength=4) / len(walks) + + np.testing.assert_allclose(observed, probabilities, atol=0.01, rtol=0) + + +def test_sparse_and_dense_inverse_cdf_agree_for_controlled_uniforms(): + transition = sp.csr_matrix( + [ + [0.0, 0.25, 0.0, 0.75], + [0.5, 0.0, 0.5, 0.0], + [0.0, 0.1, 0.2, 0.7], + [1.0, 0.0, 0.0, 0.0], + ] + ) + states = np.array([0, 1, 2, 3, 0, 2]) + uniforms = np.array([0.1, 0.6, 0.05, 0.9, 0.9, 0.25]) + + dense_choices = [] + sparse_choices = [] + dense = transition.toarray() + for state, uniform in zip(states, uniforms): + dense_choices.append(np.searchsorted(np.cumsum(dense[state]), uniform, side="right")) + + start = transition.indptr[state] + end = transition.indptr[state + 1] + neighbors = transition.indices[start:end] + probabilities = transition.data[start:end] + offset = np.searchsorted(np.cumsum(probabilities), uniform, side="right") + sparse_choices.append(neighbors[offset]) + + assert np.array_equal(dense_choices, sparse_choices) + + +def test_sparse_walks_reject_empty_rows(): + transition = sp.csr_matrix((3, 3), dtype=float) + + with pytest.raises(ValueError, match="no outgoing transitions"): + random_walks_sparse(transition, n_steps=2, n_trajectories=3, random_state=0) + + +def test_prepare_walks_keeps_transition_sparse(dentate_graph): + _, _, transition = dentate_graph + lsd = object.__new__(LSD) + lsd.P = transition.copy() + lsd.path_len = 3 + lsd.walk_config = SimpleNamespace(num_walks=16, random_state=11) + + lsd.prepare_walks() + + assert sp.isspmatrix_csr(lsd.P) + assert lsd.walks.shape == (16, 3)