spatialrna 0.3.0.dev2__tar.gz → 0.3.0.dev4__tar.gz

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: spatialrna
-Version: 0.3.0.dev2
+Version: 0.3.0.dev4
 Summary: A package for spatial RNA graph data processing and analysis.
-Author: Ruqian Lyu
-Author-email: Anne Pan <anne.pan@svi.edu.au>
+Author: Yan Pan, Ruqian Lyu
+Maintainer-email: Yan Pan <anne.pan@svi.edu.au>
 License: MIT
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/pyproject.toml

@@ -6,8 +6,9 @@ build-backend = "setuptools.build_meta"
 name="spatialrna"
 description = "A package for spatial RNA graph data processing and analysis."
 readme = "README.md"
-version = "0.3.0.dev2"
-authors = [{ name = "Anne Pan", email = "anne.pan@svi.edu.au" },{ name = "Ruqian Lyu" }]
+version = "0.3.0.dev4"
+authors = [{ name = "Yan Pan" }, { name = "Ruqian Lyu" }]
+maintainers = [{ name = "Yan Pan", email = "anne.pan@svi.edu.au" }]
 license = { text = "MIT" }
 requires-python = ">=3.10"
 classifiers = [

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/__init__.py

@@ -6,9 +6,6 @@ any GNN (train / validate / embed) -> clusters -> spatial maps.
     df = sr.readers.load_transcripts("/path/to/xenium_or_atera_output")
     from spatialrna import SpatialRNA                 # tile graphs / subgraphs; needs torch + PyG
 
-The ISTIST pipeline builds on this with its own model (CrossSampleGAT), transcript
-annotation, statistics and report.
-
 Submodules (imported on first use, so the light ones need neither torch nor a GPU):
 
     readers         Xenium / Atera transcripts.zarr.zip -> transcripts

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/cluster.py

@@ -6,19 +6,6 @@
 
 GaussianMixtureGPU -- a Gaussian mixture model (EM) that runs on the GPU with PyTorch.
 
-Replaces PyCave (torch < 2 only, unmaintained) for niche clustering. Same method:
-greedy k-means++ seeding and a few k-means iterations, then expectation-maximisation, with diagonal ("diag") or full
-covariance matrices. Data stay on the GPU in float32 and are processed in chunks;
-sufficient statistics (sums) are accumulated in float64 for numerical stability.
-
-The algorithm is written once against a few array operations (`_TorchOps` for the
-pipeline, `_NumpyOps` for testing it against scikit-learn without a GPU), so both
-run the identical code path.
-
-    gmm = GaussianMixtureGPU(n_components=12, covariance_type="diag", device="cuda")
-    gmm.fit(X)                 # X: (n, d) numpy float32
-    labels = gmm.predict(X)    # numpy int64
-    gmm.save("gmm.pt")
 """
 from __future__ import annotations
 

spatialrna-0.3.0.dev4/src/spatialrna/gene_encoding.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import json
+import os
+import pickle
+import sys
+from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Union
+
+import numpy as np
+
+PathLike = Union[str, os.PathLike]
+Encoding = Dict[object, object]
+
+
+# ---------------------------------------------------------------------------
+# Building
+# ---------------------------------------------------------------------------
+def build_gene_encoding(gene_names: Sequence[str], drop_controls: bool = False) -> Dict[str, int]:
+    """{gene: index} in the given order (index = position). Duplicates are an error,
+    since two genes sharing an index would be indistinguishable to the model.
+    drop_controls: leave out control-probe / blank-codeword names (readers.CONTROL_NAME)."""
+    names = [str(g) for g in gene_names]
+    if drop_controls:
+        from .readers import CONTROL_NAME
+        names = [g for g in names if not CONTROL_NAME.match(g)]
+    dup = sorted({g for g in names if names.count(g) > 1}) if len(set(names)) != len(names) else []
+    if dup:
+        raise ValueError(f"duplicate gene names: {dup[:20]}")
+    return {g: i for i, g in enumerate(names)}
+
+
+def cohort_gene_names(cohort_zarr_path: PathLike) -> List[str]:
+    """Root attr `gene_names` of a cohort store (what its gene_identity codes index)."""
+    import zarr
+    root = zarr.open_group(str(cohort_zarr_path), mode="r")
+    names = root.attrs.get("gene_names")
+    if not names:
+        raise ValueError(f"{cohort_zarr_path} has no root attribute 'gene_names'")
+    return [str(g) for g in names]
+
+
+def encoding_from_cohort(cohort_zarr_path: PathLike, drop_controls: bool = False) -> Dict[str, int]:
+    """{gene: index} from a cohort store's gene_names. Index = the gene_identity code,
+    so stored codes and encoding agree."""
+    return build_gene_encoding(cohort_gene_names(cohort_zarr_path), drop_controls=drop_controls)
+
+
+# ---------------------------------------------------------------------------
+# Saving / loading
+# ---------------------------------------------------------------------------
+def save_gene_encoding(encoding: Mapping, path: PathLike) -> str:
+    """Write the encoding atomically (temp file + rename). `.json` -> JSON (index maps
+    only; readable without Python), anything else -> pickle (the earlier .pkl format)."""
+    path = os.fspath(path)
+    enc = dict(encoding)
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    tmp = f"{path}.tmp{os.getpid()}"
+    if path.endswith(".json"):
+        if not is_index_encoding(enc):
+            raise ValueError("a .json gene encoding must map genes to integers (use .pkl, or to_index_encoding)")
+        with open(tmp, "w") as fh:
+            json.dump({str(k): int(v) for k, v in enc.items()}, fh, indent=0)
+    else:
+        with open(tmp, "wb") as fh:
+            pickle.dump(enc, fh, protocol=4)
+    os.replace(tmp, path)
+    return path
+
+
+def load_gene_encoding(path: PathLike) -> Encoding:
+    """Read a saved encoding (.json, or a pickle -- only open pickles you trust)."""
+    path = os.fspath(path)
+    if path.endswith(".json"):
+        with open(path) as fh:
+            enc = {k: int(v) for k, v in json.load(fh).items()}
+    else:
+        with open(path, "rb") as fh:
+            enc = pickle.load(fh)
+    if not isinstance(enc, Mapping) or not enc:
+        raise ValueError(f"{path} does not hold a non-empty gene -> index mapping")
+    return dict(enc)
+
+
+def resolve_gene_encoding(encoding: Union[Mapping, PathLike, None]) -> Encoding:
+    """A dict is returned as is; a path is loaded."""
+    if encoding is None:
+        raise ValueError("a gene encoding (dict or path to one) is required")
+    if isinstance(encoding, Mapping):
+        if not encoding:
+            raise ValueError("the gene encoding is empty")
+        return dict(encoding)
+    return load_gene_encoding(encoding)
+
+
+# ---------------------------------------------------------------------------
+# Inspecting / converting
+# ---------------------------------------------------------------------------
+def _is_int(v) -> bool:
+    return isinstance(v, (int, np.integer)) and not isinstance(v, (bool, np.bool_))
+
+
+def is_index_encoding(encoding: Mapping) -> bool:
+    """True when the values are integers (not one-hot vectors)."""
+    return _is_int(next(iter(encoding.values())))
+
+
+def _as_numpy(v) -> np.ndarray:
+    if hasattr(v, "detach"):                       # torch tensor, without importing torch
+        v = v.detach().cpu().numpy()
+    return np.asarray(v)
+
+
+def to_index_encoding(encoding: Mapping) -> Dict[object, int]:
+    """One-hot vector values -> their position (argmax). Index maps pass through."""
+    if is_index_encoding(encoding):
+        return {k: int(v) for k, v in encoding.items()}
+    out = {}
+    for k, v in encoding.items():
+        a = _as_numpy(v).reshape(-1)
+        if np.count_nonzero(a) != 1:
+            raise ValueError(f"value for {k!r} is not a one-hot vector")
+        out[k] = int(np.argmax(a))
+    return out
+
+
+def num_genes(encoding: Mapping) -> int:
+    """Width of the model input: max index + 1 (index maps) or the vector length."""
+    if is_index_encoding(encoding):
+        return int(max(encoding.values())) + 1
+    return int(_as_numpy(next(iter(encoding.values()))).size)
+
+
+def same_encoding(a: Mapping, b: Mapping) -> bool:
+    """True when two encodings give every gene the same index."""
+    return to_index_encoding(a) == to_index_encoding(b)
+
+
+def encoding_lut(categories: Sequence, encoding: Mapping,
+                 gene_names: Optional[Sequence[str]] = None) -> np.ndarray:
+    """Feature for each distinct value of the feature column, in the order of
+    `categories`: int64 indices (index maps) or a (C, F) array (one-hot maps).
+
+    Integer categories (a cohort store's gene_identity codes) that are not themselves
+    keys are looked up by name: gene_names[code] (the store's root attr gene_names).
+    Raises KeyError naming the genes the encoding lacks -- never silently drops them.
+    """
+    cats = list(categories)
+    keys = cats
+    if cats and _is_int(cats[0]) and not all(c in encoding for c in cats):
+        if gene_names is None:
+            raise KeyError("the feature column holds integer codes but the gene encoding is keyed by "
+                           "gene name -- pass the cohort's gene_names to translate them")
+        bad = [c for c in cats if not 0 <= int(c) < len(gene_names)]
+        if bad:
+            raise KeyError(f"gene codes outside gene_names (0..{len(gene_names) - 1}): {bad[:20]}")
+        keys = [gene_names[int(c)] for c in cats]
+    missing = [k for k in keys if k not in encoding]
+    if missing:
+        raise KeyError(f"{len(missing)} gene(s) are not in the gene encoding, e.g. {missing[:20]} -- "
+                       "rebuild the encoding from the full cohort, or filter these genes at ingest")
+    if is_index_encoding(encoding):
+        return np.fromiter((int(encoding[k]) for k in keys), dtype=np.int64, count=len(keys))
+    rows = [_as_numpy(encoding[k]).reshape(-1) for k in keys]
+    return np.stack(rows) if rows else np.empty((0, num_genes(encoding)))
+
+
+def main(argv: Optional[Iterable[str]] = None) -> None:
+    import argparse
+    p = argparse.ArgumentParser(description="Write the gene encoding of a cohort store.")
+    p.add_argument("cohort_zarr", help="cohort.zarr (root attr gene_names)")
+    p.add_argument("out", help="output .pkl (or .json)")
+    p.add_argument("--drop-controls", action="store_true", help="leave out control / blank codewords")
+    a = p.parse_args(list(argv) if argv is not None else None)
+    enc = encoding_from_cohort(a.cohort_zarr, drop_controls=a.drop_controls)
+    if os.path.exists(a.out):
+        old = load_gene_encoding(a.out)
+        if not same_encoding(old, enc):
+            sys.exit(f"{a.out} exists with a DIFFERENT encoding ({len(old)} genes vs {len(enc)}); "
+                     "graphs and models built with it would no longer match. Remove it to rebuild.")
+        print(f"{a.out}: unchanged ({len(enc)} genes)")
+        return
+    save_gene_encoding(enc, a.out)
+    print(f"{a.out}: {len(enc)} genes")
+
+
+if __name__ == "__main__":
+    main()

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/graph_build.py

@@ -10,14 +10,19 @@ sampler), vectorised so whole samples (1M+ transcripts) are fast:
   fanout (-1 = all neighbours); frontier expansion is vectorised over CSR rows.
 * link_based_subgraph -- num_sampled_edges anchor edges + k-hop expansion (the old
   build_subgraph(subgraph_mode="link_based")).
+* plan_tiles / plan_strips + tile_arrays -- which transcripts make up each tile
+  (core + padding) and the arrays of one tile graph; map_tiles runs tiles in a
+  forked process pool that shares the sample's arrays (SpatialRNA.process(n_procs)).
 
-Conventions (same as before): edge_index rows are (source, target); target is the
-query node and source a neighbour within r -- PyG's flow="source_to_target".
 """
 from __future__ import annotations
 
 import itertools
-from typing import Optional, Sequence, Tuple
+import multiprocessing as mp
+import os
+import sys
+import time
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
 
 import numpy as np
 from scipy.spatial import cKDTree
@@ -321,7 +326,7 @@ class TilePlan:
 def plan_tiles(x: np.ndarray, y: np.ndarray, n_tiles: int, mode: str = "grid") -> TilePlan:
     """Assign every transcript to one of `n_tiles` tiles (mode "grid"; see TilePlan)."""
     if mode != "grid":
-        raise ValueError(f"plan_tiles supports mode='grid' (strips are cut in SpatialRNA.process), not {mode!r}")
+        raise ValueError(f"plan_tiles supports mode='grid' (strips: plan_strips), not {mode!r}")
     x = np.asarray(x, dtype=np.float64)
     y = np.asarray(y, dtype=np.float64)
     n = len(x)
@@ -348,3 +353,135 @@ def plan_tiles(x: np.ndarray, y: np.ndarray, n_tiles: int, mode: str = "grid") -
             tile_of[xo[col_cut[c]:col_cut[c + 1]]] = first + c
         first += cols[r]
     return TilePlan(x, y, tile_of, n_tiles)
+
+
+class StripPlan:
+    """The original SpatialRNA strips: `n_tiles` equal-WIDTH bands along one axis
+    (pandas.cut with bins=n_tiles), padding = transcripts within `pad` below / above
+    the band's own min / max along that axis.
+
+    Same rows, in the same order, as the earlier per-tile pandas code
+    (core rows, then left padding, then right padding, each in original row order),
+    but every tile is cut from arrays sorted once, not by masking the whole sample.
+
+    tile_of[i] = strip of transcript i (-1 if its coordinate is NaN).
+    """
+
+    def __init__(self, v: np.ndarray, tile_of: np.ndarray, n_tiles: int):
+        # keep the column's own dtype: the padding bounds are compared in it, exactly
+        # as the earlier pandas code did (float32 coordinates stay float32)
+        self.v = np.asarray(v)
+        self.tile_of = np.asarray(tile_of, dtype=np.int64)
+        self.n_tiles = n_tiles
+        self._vorder = np.argsort(self.v, kind="stable")
+        self._vsorted = self.v[self._vorder]
+        order = np.argsort(self.tile_of, kind="stable")
+        bounds = np.searchsorted(self.tile_of[order], np.arange(n_tiles + 1))
+        self._cores = [np.sort(order[bounds[t]:bounds[t + 1]]) for t in range(n_tiles)]
+
+    def sizes(self) -> np.ndarray:
+        return np.array([len(c) for c in self._cores])
+
+    def members(self, t: int, pad: float) -> Tuple[np.ndarray, np.ndarray]:
+        core = self._cores[t]
+        if len(core) == 0:
+            return core, core
+        lo, hi = self.v[core].min(), self.v[core].max()
+        vs, vo = self._vsorted, self._vorder
+        left = np.sort(vo[np.searchsorted(vs, lo - pad, "left"):np.searchsorted(vs, lo, "left")])
+        right = np.sort(vo[np.searchsorted(vs, hi, "right"):np.searchsorted(vs, hi + pad, "right")])
+        return core, np.concatenate([left, right])
+
+
+def plan_strips(v: np.ndarray, n_tiles: int) -> StripPlan:
+    """Strips along `v` with exactly pandas.cut(v, bins=n_tiles)'s bin edges."""
+    import pandas as pd
+
+    codes = pd.cut(pd.Series(np.asarray(v)), bins=n_tiles, labels=False)
+    tile_of = codes.fillna(-1).to_numpy().astype(np.int64)
+    return StripPlan(v, tile_of, n_tiles)
+
+
+def tile_arrays(core: np.ndarray, padding: np.ndarray, codes: np.ndarray, lut: np.ndarray,
+                graph_pos: np.ndarray, r: float, max_num_neighbors: Optional[int],
+                store_pos: Optional[np.ndarray] = None, workers: int = -1,
+                method: str = "auto") -> Dict[str, np.ndarray]:
+    """The arrays of one tile graph, from whole-sample arrays (read-only):
+
+    node order: core transcripts, then padding (as members() returns them)
+    x          lut[codes[node]]  (int64 gene index, or a feature row per node)
+    edge_index radius graph over graph_pos[node] (scipy KDTree, radius_graph_kdtree;
+               method "auto" | "ball" | "knn", see there)
+    trans_id   the node's row in the sample (= its index in the cohort store)
+    core_mask  True for the core part
+    pos        store_pos[node] as float32, if store_pos is given
+    """
+    node = np.concatenate([np.asarray(core, dtype=np.int64), np.asarray(padding, dtype=np.int64)])
+    out = {
+        "x": lut[codes[node]],
+        "edge_index": radius_graph_kdtree(graph_pos[node], r=r, max_num_neighbors=max_num_neighbors,
+                                          workers=workers, method=method),
+        "trans_id": node,
+        "core_mask": np.r_[np.ones(len(core), bool), np.zeros(len(padding), bool)],
+    }
+    if store_pos is not None:
+        out["pos"] = np.asarray(store_pos[node], dtype=np.float32)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Forked process pool over tiles
+# ---------------------------------------------------------------------------
+_POOL_TASK: List[Optional[Callable]] = [None]     # set in the parent right before forking
+
+
+def _pool_call(tile):
+    return tile, _POOL_TASK[0](tile)
+
+
+def available_cpus() -> int:
+    """CPUs this process may use (SLURM / taskset aware), not the whole node."""
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:
+        return os.cpu_count() or 1
+
+
+def map_tiles(task: Callable, tiles: Iterable[int], n_procs: int = 1,
+              initializer: Optional[Callable] = None, log: bool = True) -> Dict[int, object]:
+    """task(tile) for every tile; returns {tile: result}.
+
+    n_procs > 1: a 'fork' process pool. `task` (with whatever it closes over -- the
+    whole sample's arrays, the tile plan) is never pickled: it is set as a module
+    global just before the pool forks, so every worker sees the parent's memory
+    copy-on-write. Only tile ids go to the workers and only task's (small) return
+    values come back. Numpy arrays are shared without copying as long as nobody
+    writes to them; keep the shared data in numpy arrays, not in Python objects
+    (touching an object's refcount copies its memory page).
+
+    `initializer` runs once in each worker after the fork (e.g. to reset a library's
+    background thread). Errors in a worker are re-raised here. Without the 'fork'
+    start method (Windows; macOS works) this falls back to running serially.
+    """
+    tiles = list(tiles)
+    n_procs = max(1, min(int(n_procs or 1), len(tiles)))
+    if n_procs > 1 and "fork" not in mp.get_all_start_methods():
+        print("WARNING: no 'fork' start method on this platform -- running tiles serially.",
+              file=sys.stderr, flush=True)
+        n_procs = 1
+    if n_procs == 1:
+        return {t: task(t) for t in tiles}
+
+    out: Dict[int, object] = {}
+    t0 = time.time()
+    _POOL_TASK[0] = task
+    try:
+        with mp.get_context("fork").Pool(n_procs, initializer=initializer) as pool:
+            for i, (tile, res) in enumerate(pool.imap_unordered(_pool_call, tiles, chunksize=1), 1):
+                out[tile] = res
+                if log:
+                    print(f"[{i}/{len(tiles)}] tile {tile} done ({time.time() - t0:.0f}s elapsed, "
+                          f"{n_procs} processes)", flush=True)
+    finally:
+        _POOL_TASK[0] = None
+    return out

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/graph_zarr_io.py

@@ -1,51 +1,13 @@
 """
 SpatialRNA dataset backed by zarr storage.
 
-Why zarr instead of torch.save/.pt for the processed tile graphs:
-  - Chunked + compressed on-disk arrays: a tile's `x` / `edge_index` can be
-    read without ever fully deserializing the tile through pickle, and only
-    the chunks that are actually indexed get decompressed.
-  - Cloud/object-store friendly: a zarr group is a directory of small,
-    independently addressable files, so tiles can live on S3/GCS behind
-    fsspec without a custom loader.
-  - Cross-tool interop: other steps in the pipeline (QC notebooks, R via
-    Rarr, dask) can read the same tiles without a torch_geometric
-    dependency.
-
-This module keeps `SpatialRNA` as a torch_geometric `InMemoryDataset` for
-raw/processed path bookkeeping (`raw_paths`, `processed_dir`, `download`,
-etc.), but replaces the two places that used torch's own (de)serialization
-(`self.save` / `self.load`, and `torch.load` inside `generate_subgraph`)
-with explicit zarr I/O. Because each tile file always holds exactly one
-`Data` object, `len()`/`get()` are overridden directly instead of relying on
-`InMemoryDataset`'s collate/slice machinery, which assumed a torch .pt blob.
-
-Requires: zarr>=2.16 (either the 2.x or 3.x major version -- see the
-`_ZARR_MAJOR` shim below, which is what actually differs between them:
-store construction, `create_dataset` vs `create_array`, and where the
-compressor codec classes live).
-
---- Canonical module (post-consolidation) ---------------------------------
-
-This used to live as a standalone `scripts/spatialrna_zarr.py` (imported via
-sys.path hacks) with a second, incompatible zarr I/O layer duplicated under
-`step1_old_zarr/spatialrna/graph_zarr_io.py` (a `tile_store_path` /
-`read_tile_graph` / `store_exists` / `_write_arrays` / `_read_arrays` API
-used only by the niche-inference step). Both wrote/read the exact same
-on-disk tile layout (`x`, `edge_index`, `trans_id`, `core_mask` arrays in a
-zarr group), so rather than keep two parallel implementations, this module
-is now the single source of truth: the `SpatialRNA` class and
-`save_data_zarr`/`load_data_zarr` below are unchanged, and a thin
-backward-compatible API (`tile_store_path`, `store_exists`,
-`read_tile_graph`, `_write_arrays`, `_read_arrays`, etc. -- see the bottom
-of this file) is provided purely so `niche_zarr_io.py` and the
-inference/GMM scripts can import `spatialrna.graph_zarr_io` without any
-change to their own logic. `step1_old_zarr/` is no longer used.
+Requires: zarr>=2.16 (either the 2.x or 3.x major version).
 """
 from torch_geometric.data import InMemoryDataset, Data
-from .graph_build import plan_tiles, radius_graph_kdtree
+from .graph_build import plan_tiles, plan_strips, tile_arrays, map_tiles, available_cpus
+from .gene_encoding import resolve_gene_encoding, encoding_lut
 from torch_geometric.loader import LinkNeighborLoader
-from typing import Callable, Optional, List, Dict, Any
+from typing import Callable, Optional, List, Dict, Any, Mapping, Union
 import os
 import os.path as osp
 import pandas as pd
@@ -108,6 +70,43 @@ else:
         )
 
 
+def _after_fork_in_child() -> None:
+    """Run in each forked tile worker before it touches torch or zarr.
+
+    A forked child gets a copy of the parent's memory but none of its threads. Two
+    libraries keep state that points at parent threads:
+      * zarr-python 3 runs all I/O on a background asyncio loop thread, created the
+        first time the parent read anything; the child would submit work to a loop
+        nobody runs and hang. Dropping zarr's references makes it start a new loop
+        in the child on first use.
+      * Blosc's internal thread pool -> turned off in the child (zarr writes are
+        already parallel across workers).
+    torch is limited to one thread per worker for the same reason (and to avoid
+    oversubscription; the workers only wrap arrays, they compute nothing in torch).
+    """
+    try:
+        torch.set_num_threads(1)
+    except Exception:
+        pass
+    try:
+        import numcodecs.blosc as _blosc
+        _blosc.use_threads = False
+    except Exception:
+        pass
+    if _ZARR_MAJOR >= 3:
+        try:
+            from zarr.core import sync as _zsync
+            for name in ("loop", "iothread"):
+                ref = getattr(_zsync, name, None)
+                if isinstance(ref, list) and ref:
+                    ref[0] = None
+            for name in ("_lock", "_executor"):
+                if hasattr(_zsync, name):
+                    setattr(_zsync, name, None)
+        except Exception:
+            pass
+
+
 # Rows-per-chunk default. Individual arrays fall back to their own length if
 # the tile is smaller than this (chunks can't exceed an array's own extent).
 _DEFAULT_CHUNK_ROWS = 200_000
@@ -346,9 +345,28 @@ class SpatialRNA(InMemoryDataset):
         load_tile_id (int): Select one tile to load.
         process_tile_ids (List[int]): Select tiles to process (tile graph
             + subgraph generation).
-        one_hot_encoding (dict): Maps gene names to an integer (or a
-            one-hot tensor). Must be identical across every sample analyzed
-            together.
+        one_hot_encoding (dict or path): The gene encoding: gene -> integer index
+            (or -> one-hot tensor), or the path of a saved one (.pkl / .json; see
+            `spatialrna.gene_encoding`). Must be identical across every sample
+            analyzed together. A name-keyed encoding also works on an integer
+            `feature_col` (the cohort store's `gene_identity`): codes are mapped
+            to names through the cohort root attr `gene_names`.
+        n_procs (int): Tiles built in parallel (default 1 = serial). The sample is
+            read and tiled once; N forked worker processes then share it
+            copy-on-write. Peak memory ~ the sample's arrays + N tiles in flight,
+            so size N by the largest tile (radius graph + edge list), not by cores.
+            Output is identical to n_procs=1.
+        radius_method (str): How the radius graph is built -- always on a scipy
+            KDTree (graph_build.radius_graph_kdtree; torch_cluster is not used).
+            "ball": every neighbour within radius_r, then the nearest
+            max_num_neighbors (exact, ties by index). "knn": a k-nearest query
+            bounded by radius_r (same graph up to ties at equal distance; much
+            faster on dense whole-transcriptome data). "auto" (default): knn when
+            most nodes exceed max_num_neighbors, else ball. Stored per tile as attr
+            `radius_method`.
+        graph_workers (int, optional): Threads per radius-graph query (scipy
+            cKDTree `workers`). Default: all CPUs when n_procs=1, else the job's
+            CPUs divided by n_procs (no oversubscription).
         transform, pre_transform: as in `InMemoryDataset`.
         force_reload (bool): Whether to re-process the dataset.
         force_resample (bool): Whether to re-process the tile dataset to
@@ -399,7 +417,10 @@ class SpatialRNA(InMemoryDataset):
         feature_col="gene",
         max_num_neighbors: int = 500,
         process_tile_ids: List[int] = None,
-        one_hot_encoding: dict = None,
+        one_hot_encoding: Union[Mapping, str, os.PathLike] = None,
+        n_procs: int = 1,
+        graph_workers: Optional[int] = None,
+        radius_method: str = "auto",
         subgraph_mode: str = "link_based",
         num_sampled_edges: int = 5000,
         num_seed_nodes: int = 5000,
@@ -434,7 +455,12 @@ class SpatialRNA(InMemoryDataset):
         self.load_tile_id = load_tile_id
         self.pad_hops = pad_hops
         self.max_num_neighbors = max_num_neighbors
-        self.one_hot_encoding = one_hot_encoding
+        assert one_hot_encoding is not None, "one_hot_encoding (a gene encoding dict or its path) is required"
+        self.one_hot_encoding = resolve_gene_encoding(one_hot_encoding)
+        self.n_procs = max(1, int(n_procs or 1))
+        self.graph_workers = graph_workers
+        assert radius_method in ("auto", "ball", "knn"), f"radius_method: auto, ball or knn, not {radius_method!r}"
+        self.radius_method = radius_method
         self.seed = seed
         self.feature_col = feature_col
         self.num_sampled_edges = num_sampled_edges
@@ -445,7 +471,6 @@ class SpatialRNA(InMemoryDataset):
         self.max_subgraph_edges = max_subgraph_edges
 
         assert tile_by_dim in [dim_x, dim_y]
-        assert one_hot_encoding is not None
 
         self.dim_x = dim_x
         self.dim_y = dim_y
@@ -549,6 +574,9 @@ class SpatialRNA(InMemoryDataset):
     def get_padding(self, log: bool, tile_df, full_df, padding: float):
         r"""Padding-area transcripts surrounding a tile's core area.
 
+        Kept for callers of the earlier API; process() now uses
+        graph_build.StripPlan.members, which selects the same rows.
+
         Args:
             log: whether to print logs.
             tile_df: the data.frame for the current tile.
@@ -694,7 +722,23 @@ class SpatialRNA(InMemoryDataset):
             out_path = osp.join(self.subgraph_dir, f"tile{s_id}")
             save_data_zarr(out_path, subgraph_g)
 
-    def process(self) -> None:
+    # --- tile graphs ------------------------------------------------------
+    def _cohort_gene_names(self) -> Optional[List[str]]:
+        try:
+            names = _zopen_read(self.cohort_zarr_path).attrs.get("gene_names")
+        except Exception:
+            names = None
+        return [str(g) for g in names] if names else None
+
+    def process(self, n_procs: Optional[int] = None) -> None:
+        """Build tiles/tile<t> for every tile in process_tile_ids.
+
+        The sample is read once and turned into flat arrays (gene index per
+        transcript, coordinates) and a tile plan; each tile is then cut from those
+        arrays. With n_procs > 1 (default: the n_procs given at construction) the
+        tiles run in forked worker processes that share these arrays.
+        """
+        n_procs = self.n_procs if n_procs is None else max(1, int(n_procs))
         if self.log:
             print("Processing raw file ", self.raw_paths[0])
             print("To create processed files ", self.processed_paths)
@@ -718,6 +762,7 @@ class SpatialRNA(InMemoryDataset):
         if use_z:
             raw_columns = sorted(set(raw_columns) | {self.dim_z})
         coord_cols = [self.dim_x, self.dim_y] + ([self.dim_z] if use_z else [])
+        graph_cols = coord_cols if self.radius_3d else [self.dim_x, self.dim_y]
         raw_data = load_raw_transcripts_zarr(self.raw_paths[0], columns=raw_columns)
         if self.log:
             print(f"Raw data shape {raw_data.shape[0]}")
@@ -731,7 +776,21 @@ class SpatialRNA(InMemoryDataset):
             if self.log:
                 print(f"Building 3-D radius graphs on ({self.dim_x}, {self.dim_y}, {self.dim_z})")
 
-        raw_data["trans_id"] = np.arange(raw_data.shape[0], dtype=np.int64)
+        # --- whole-sample arrays, built once and shared by every tile ----------
+        # gene feature: one encoding lookup per distinct gene, then a per-transcript code
+        cat = pd.Categorical(raw_data[self.feature_col].to_numpy())
+        codes = np.asarray(cat.codes, dtype=np.int64)
+        if (codes < 0).any():
+            raise ValueError(f"{self.sample_name}: {int((codes < 0).sum())} transcripts have no "
+                             f"'{self.feature_col}' value")
+        categories = list(cat.categories)
+        gene_names = (self._cohort_gene_names()
+                      if categories and isinstance(categories[0], (int, np.integer)) else None)
+        lut = encoding_lut(categories, self.one_hot_encoding, gene_names=gene_names)
+        # row i = transcript i = trans_id i (row index into cohort.zarr/<sample>)
+        graph_pos = np.ascontiguousarray(raw_data[graph_cols].to_numpy(np.float64))
+        store_pos = (np.ascontiguousarray(raw_data[coord_cols].to_numpy(np.float32))
+                     if self.store_coords else None)
         padding = self.pad_hops * self.radius_r
 
         if self.tile_mode == "grid":
@@ -742,81 +801,59 @@ class SpatialRNA(InMemoryDataset):
                 print(f"grid tiling: {plan.n_tiles} tiles, {sz.min():,}-{sz.max():,} core transcripts each")
             for tile in self.process_tile_ids:
                 assert 0 <= tile < plan.n_tiles, f"tile {tile} is not in 0:{plan.n_tiles - 1}"
-                core, pad_idx = plan.members(tile, padding)
-                tile_df = pd.concat([raw_data.iloc[core].assign(core_mask=True),
-                                     raw_data.iloc[pad_idx].assign(core_mask=False)])
-                if self.log:
-                    print(f"tile {tile}: core {len(core):,}, core plus padding {tile_df.shape[0]:,}")
-                self._write_tile(tile, tile_df, coord_cols)
-            return
-
-        raw_data["tiles_range"] = pd.cut(
-            raw_data[self.tile_by_dim],
-            bins=self.num_tiles,
-            labels=["tile" + str(x) for x in range(self.num_tiles)],
-        )
+        else:
+            plan = plan_strips(raw_data[self.tile_by_dim].to_numpy(), self.num_tiles)
+            sizes = plan.sizes()
+            for tile in self.process_tile_ids:
+                assert 0 <= tile < plan.n_tiles and sizes[tile] > 0, (
+                    f"tile {tile}: the requested tile contains no transcripts, likely due to the "
+                    "tile id not being in the range 0:(num_tiles-1) -- or an empty strip "
+                    "(a gap in the tissue); graph.tile_mode: grid never makes empty tiles"
+                )
+        del raw_data, cat     # only numpy arrays are shared with the workers
 
-        for tile in self.process_tile_ids:
-            tile_df = raw_data[raw_data.tiles_range == ("tile" + str(tile))]
-            if self.log:
-                print("tile core area shape ", tile_df.shape)
-            assert tile_df.shape[0] > 0, (
-                "The requested tile contains no transcripts, likely due to the "
-                "tile id not being in the range 0:(num_tiles-1) -- or an empty strip "
-                "(a gap in the tissue); graph.tile_mode: grid never makes empty tiles"
-            )
-            tile_df = tile_df.assign(core_mask=True)
-            padding_df = self.get_padding(
-                log=self.log,
-                tile_df=tile_df,
-                full_df=raw_data,
-                padding=padding,
-            )
-            tile_df = pd.concat([tile_df, padding_df])
-            if self.log:
-                print("core area plus paddings, shape ", tile_df.shape)
-            self._write_tile(tile, tile_df, coord_cols)
-
-    def _write_tile(self, tile, tile_df, coord_cols) -> None:
-        """Radius graph over one tile's core + padding transcripts -> tiles/tile<t>."""
-        t0 = time.time()
-        first = next(iter(self.one_hot_encoding.values()))
-        if isinstance(first, (int, np.integer)):
-            # one lookup per distinct gene, not per transcript
-            cat = pd.Categorical(tile_df[self.feature_col].to_numpy())
-            lut = np.fromiter((self.one_hot_encoding[c] for c in cat.categories), dtype=np.int64,
-                              count=len(cat.categories))
-            feature_x = torch.from_numpy(lut[cat.codes])
+        if self.graph_workers:
+            workers = int(self.graph_workers)
         else:
-            feature_x = torch.stack([self.one_hot_encoding[x] for x in tile_df[self.feature_col]])
-        graph_cols = coord_cols if self.radius_3d else [self.dim_x, self.dim_y]
-        # scipy cKDTree radius graph (graph_build.py): same builder as the whole-sample
-        # training graphs, so tile (inference) and sample (training) graphs agree.
-        graph_pos = np.asarray(tile_df[graph_cols], dtype=np.float64)
-        edge_index = torch.as_tensor(
-            radius_graph_kdtree(graph_pos, r=self.radius_r, max_num_neighbors=self.max_num_neighbors),
-            dtype=torch.long)
+            workers = -1 if n_procs == 1 else max(1, available_cpus() // n_procs)
+        if self.log and n_procs > 1:
+            print(f"{len(self.process_tile_ids)} tiles on {n_procs} processes, "
+                  f"{workers} radius-query thread(s) each", flush=True)
+
+        def build_one(tile):
+            t0 = time.time()
+            core, pad_idx = plan.members(tile, padding)
+            arrays = tile_arrays(core, pad_idx, codes, lut, graph_pos, self.radius_r,
+                                 self.max_num_neighbors, store_pos=store_pos, workers=workers,
+                                 method=self.radius_method)
+            t1 = time.time()
+            self._save_tile(tile, arrays, graph_cols, coord_cols)
+            n, n_edges = len(arrays["trans_id"]), arrays["edge_index"].shape[1]
+            if self.log:
+                print(f"tile {tile}: core {len(core):,}, core plus padding {n:,}; {n_edges:,} edges "
+                      f"(mean degree {n_edges / max(n, 1):.1f}), graph {t1 - t0:.0f}s, "
+                      f"write {time.time() - t1:.0f}s", flush=True)
+            return {"core": int(len(core)), "nodes": int(n), "edges": int(n_edges)}
+
+        map_tiles(build_one, self.process_tile_ids, n_procs=n_procs,
+                  initializer=_after_fork_in_child, log=self.log)
 
+    def _save_tile(self, tile, arrays: Dict[str, np.ndarray], graph_cols, coord_cols) -> None:
+        """One tile's arrays (graph_build.tile_arrays) -> tiles/tile<t>."""
         data_1 = Data(
-            x=feature_x,
-            edge_index=edge_index,
-            trans_id=torch.tensor(tile_df.trans_id.to_numpy()),
-            core_mask=torch.tensor(tile_df.core_mask.to_numpy()),
+            x=torch.from_numpy(arrays["x"]),
+            edge_index=torch.from_numpy(arrays["edge_index"]),
+            trans_id=torch.from_numpy(arrays["trans_id"]),
+            core_mask=torch.from_numpy(arrays["core_mask"]),
         )
-        if self.store_coords:
-            data_1.pos = torch.tensor(np.asarray(tile_df[coord_cols], dtype=np.float32))
-
+        if "pos" in arrays:
+            data_1.pos = torch.from_numpy(arrays["pos"])
         out_path = osp.join(self.processed_dir, f"tile{tile}")
-        t1 = time.time()
         save_data_zarr(out_path, data_1)
-        if self.log:
-            n = int(tile_df.shape[0])
-            print(f"tile {tile}: {n:,} nodes, {edge_index.size(1):,} edges "
-                  f"(mean degree {edge_index.size(1) / max(n, 1):.1f}), graph {t1 - t0:.0f}s, "
-                  f"write {time.time() - t1:.0f}s", flush=True)
         meta = _zopen_append(out_path)
         meta.attrs["graph_dims"] = list(graph_cols)
         meta.attrs["tile_mode"] = self.tile_mode
+        meta.attrs["radius_method"] = self.radius_method
         if self.store_coords:
             meta.attrs["pos_columns"] = list(coord_cols)
 

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/train_val.py

@@ -1,21 +1,3 @@
-"""Model-agnostic GNN training and inference on spatial transcript graphs
-(spatialrna 0.2's train_val, for the zarr tile / subgraph stores).
-
-Bring any PyG model with the signature ``model(x, edge_index) -> embeddings``; genes
-are integer codes in ``data.x`` and are one-hot encoded batch by batch.
-
-    from spatialrna import train_val as tv
-    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
-    for epoch in range(5):
-        loss, acc = tv.train_link_prediction(model, graphs, opt, device, num_genes)
-        val_acc = tv.validate_link_prediction(model, val_graphs, device, num_genes)
-    emb, tx_id = tv.embed_tile(model, tile, device, num_genes)   # core transcripts only
-
-`graphs` is any iterable of torch_geometric Data (e.g. spatialrna.datasets.ZarrGraphDataset).
-Link prediction: positive pairs are graph edges (``edge_label_index`` if the graph has
-one, else a random sample of edges), negatives are random partners; the score is the
-dot product of the two embeddings.
-"""
 from __future__ import annotations
 
 import random

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna/viz.py

@@ -1,10 +1,5 @@
 """Spatial maps of transcript labels (clusters / niches / cell types), as pixel images.
 
-Like spatialrna 0.2's plot_pixel: the tissue is cut into px x px um squares and each
-square is coloured by the most frequent label among its transcripts (counts smoothed
-over ~1 neighbouring square first), so every transcript is used and maps stay readable
-at any size. Dots are available too.
-
     from spatialrna import viz
     fig = viz.pixel_map(x, y, labels)                         # one panel
     d = viz.make_data(x, y, label, conf, label_names, niche)  # two label sets

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: spatialrna
-Version: 0.3.0.dev2
+Version: 0.3.0.dev4
 Summary: A package for spatial RNA graph data processing and analysis.
-Author: Ruqian Lyu
-Author-email: Anne Pan <anne.pan@svi.edu.au>
+Author: Yan Pan, Ruqian Lyu
+Maintainer-email: Yan Pan <anne.pan@svi.edu.au>
 License: MIT
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/src/spatialrna.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@ pyproject.toml
 src/spatialrna/__init__.py
 src/spatialrna/cluster.py
 src/spatialrna/datasets.py
+src/spatialrna/gene_encoding.py
 src/spatialrna/graph_build.py
 src/spatialrna/graph_zarr_io.py
 src/spatialrna/readers.py
@@ -14,6 +15,7 @@ src/spatialrna.egg-info/dependency_links.txt
 src/spatialrna.egg-info/requires.txt
 src/spatialrna.egg-info/top_level.txt
 tests/test_cluster_viz.py
+tests/test_parallel_tiles.py
 tests/test_readers.py
 tests/test_tiling.py
 tests/test_zarr_roundtrip.py

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/tests/test_cluster_viz.py

@@ -1,7 +1,3 @@
-"""spatialrna.cluster and spatialrna.viz on synthetic data (no torch needed).
-
-    pytest tests/test_cluster_viz.py
-"""
 import matplotlib
 matplotlib.use("Agg")
 import numpy as np

spatialrna-0.3.0.dev4/tests/test_parallel_tiles.py

@@ -0,0 +1,377 @@
+import os
+import sys
+import tempfile
+import time
+import types
+
+import numpy as np
+import pandas as pd
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(HERE)
+# repo layout (src/spatialrna) or a flat checkout; an installed spatialrna otherwise
+sys.path.insert(0, os.path.join(_ROOT, "src") if os.path.isdir(os.path.join(_ROOT, "src", "spatialrna")) else _ROOT)
+
+
+def _have(*mods):
+    import importlib.util
+    return all(importlib.util.find_spec(m) is not None for m in mods)
+
+
+HAVE_STACK = _have("torch", "torch_geometric", "zarr")
+
+from spatialrna import graph_build as gb          # noqa: E402
+from spatialrna import gene_encoding as ge        # noqa: E402
+
+RNG = np.random.default_rng(7)
+N_GENES = 300
+GENE_NAMES = [f"G{i:03d}" for i in range(N_GENES)]
+
+
+def synthetic_sample(n=60_000):
+    """Two blobs with a gap between them (an empty strip), float32 like the cohort
+    store, some duplicated positions, dense enough that some nodes hit the cap."""
+    a = RNG.normal([0, 0], [60, 40], size=(n // 2, 2))
+    b = RNG.normal([400, 300], [50, 50], size=(n - n // 2, 2))
+    xy = np.vstack([a, b]).astype(np.float32)
+    xy[:200] = xy[200:400]                                 # duplicated positions
+    z = RNG.uniform(0, 10, size=n).astype(np.float32)
+    code = RNG.integers(0, N_GENES, size=n).astype(np.int32)
+    return pd.DataFrame({"x_location": xy[:, 0], "y_location": xy[:, 1], "z_location": z,
+                         "gene_identity": code})
+
+
+# ---------------------------------------------------------------------------
+# The earlier implementation (graph_zarr_io.py of 2026-10-23), verbatim logic
+# ---------------------------------------------------------------------------
+def old_tiles(df, enc, num_tiles, tile_mode, tile_by_dim, radius_r, pad_hops, max_nn,
+              coord_cols, radius_3d, store_coords, tile_ids, feature_col="gene"):
+    raw_data = df.copy()
+    raw_data["trans_id"] = np.arange(raw_data.shape[0], dtype=np.int64)
+    padding = pad_hops * radius_r
+    out = {}
+
+    def write_tile(tile, tile_df):
+        first = next(iter(enc.values()))
+        if isinstance(first, (int, np.integer)):
+            cat = pd.Categorical(tile_df[feature_col].to_numpy())
+            lut = np.fromiter((enc[c] for c in cat.categories), dtype=np.int64, count=len(cat.categories))
+            feature_x = lut[cat.codes]
+        else:
+            feature_x = np.stack([enc[x] for x in tile_df[feature_col]])
+        graph_cols = coord_cols if radius_3d else coord_cols[:2]
+        graph_pos = np.asarray(tile_df[graph_cols], dtype=np.float64)
+        ei = gb.radius_graph_kdtree(graph_pos, r=radius_r, max_num_neighbors=max_nn)
+        d = {"x": feature_x, "edge_index": ei, "trans_id": tile_df.trans_id.to_numpy(),
+             "core_mask": tile_df.core_mask.to_numpy()}
+        if store_coords:
+            d["pos"] = np.asarray(tile_df[coord_cols], dtype=np.float32)
+        out[tile] = d
+
+    if tile_mode == "grid":
+        plan = gb.plan_tiles(raw_data[coord_cols[0]].to_numpy(), raw_data[coord_cols[1]].to_numpy(),
+                             num_tiles, mode="grid")
+        for tile in tile_ids:
+            core, pad_idx = plan.members(tile, padding)
+            tile_df = pd.concat([raw_data.iloc[core].assign(core_mask=True),
+                                 raw_data.iloc[pad_idx].assign(core_mask=False)])
+            write_tile(tile, tile_df)
+        return out
+
+    raw_data["tiles_range"] = pd.cut(raw_data[tile_by_dim], bins=num_tiles,
+                                     labels=["tile" + str(x) for x in range(num_tiles)])
+    for tile in tile_ids:
+        tile_df = raw_data[raw_data.tiles_range == ("tile" + str(tile))]
+        if tile_df.shape[0] == 0:
+            continue
+        tile_df = tile_df.assign(core_mask=True)
+        lm, rm = tile_df[tile_by_dim].min(), tile_df[tile_by_dim].max()
+        fl = (raw_data[tile_by_dim] >= (lm - padding)) & (raw_data[tile_by_dim] < lm)
+        fr = (raw_data[tile_by_dim] > rm) & (raw_data[tile_by_dim] <= (rm + padding))
+        pad_df = pd.concat([raw_data[fl].assign(core_mask=False), raw_data[fr].assign(core_mask=False)])
+        write_tile(tile, pd.concat([tile_df, pad_df]))
+    return out
+
+
+def new_tiles(df, enc, num_tiles, tile_mode, tile_by_dim, radius_r, pad_hops, max_nn,
+              coord_cols, radius_3d, store_coords, tile_ids, feature_col, gene_names=None):
+    cat = pd.Categorical(df[feature_col].to_numpy())
+    codes = np.asarray(cat.codes, dtype=np.int64)
+    lut = ge.encoding_lut(list(cat.categories), enc, gene_names=gene_names)
+    graph_cols = coord_cols if radius_3d else coord_cols[:2]
+    gpos = np.ascontiguousarray(df[graph_cols].to_numpy(np.float64))
+    spos = np.ascontiguousarray(df[coord_cols].to_numpy(np.float32)) if store_coords else None
+    plan = (gb.plan_tiles(df[coord_cols[0]].to_numpy(), df[coord_cols[1]].to_numpy(), num_tiles)
+            if tile_mode == "grid" else gb.plan_strips(df[tile_by_dim].to_numpy(), num_tiles))
+    pad = pad_hops * radius_r
+    out = {}
+    for t in tile_ids:
+        core, p = plan.members(t, pad)
+        if len(core) == 0:
+            continue
+        out[t] = gb.tile_arrays(core, p, codes, lut, gpos, radius_r, max_nn, store_pos=spos)
+    return out
+
+
+def same(a, b, label):
+    assert a.keys() == b.keys(), (label, sorted(a), sorted(b))
+    for t in a:
+        assert a[t].keys() == b[t].keys(), (label, t, a[t].keys(), b[t].keys())
+        for k in a[t]:
+            x, y = np.asarray(a[t][k]), np.asarray(b[t][k])
+            assert x.dtype == y.dtype, (label, t, k, x.dtype, y.dtype)
+            assert x.shape == y.shape and np.array_equal(x, y), (label, t, k, x.shape, y.shape)
+
+
+def test_core_matches_old_code():
+    df = synthetic_sample()
+    order = RNG.permutation(N_GENES)                       # encoding index != gene code
+    enc_named = {g: int(order[i]) for i, g in enumerate(GENE_NAMES)}
+    df_named = df.assign(gene=np.asarray(GENE_NAMES, dtype=object)[df.gene_identity])
+    eye = np.eye(N_GENES, dtype=np.float32)
+    enc_onehot = {g: eye[order[i]] for i, g in enumerate(GENE_NAMES)}
+    cols3 = ["x_location", "y_location", "z_location"]
+    cases = [
+        dict(num_tiles=12, tile_mode="strips", tile_by_dim="y_location", max_nn=20, radius_3d=False, store=True),
+        dict(num_tiles=7, tile_mode="strips", tile_by_dim="x_location", max_nn=None, radius_3d=False, store=False),
+        dict(num_tiles=25, tile_mode="grid", tile_by_dim="y_location", max_nn=20, radius_3d=False, store=True),
+        dict(num_tiles=9, tile_mode="grid", tile_by_dim="y_location", max_nn=15, radius_3d=True, store=True),
+    ]
+    n_cmp = 0
+    for c in cases:
+        common = dict(num_tiles=c["num_tiles"], tile_mode=c["tile_mode"], tile_by_dim=c["tile_by_dim"],
+                      radius_r=6.0, pad_hops=2, max_nn=c["max_nn"], coord_cols=cols3,
+                      radius_3d=c["radius_3d"], store_coords=c["store"], tile_ids=list(range(c["num_tiles"])))
+        old = old_tiles(df_named, enc_named, **common)
+        # new code on the cohort's INTEGER codes + name-keyed encoding (translated via gene_names)
+        new = new_tiles(df, enc_named, feature_col="gene_identity", gene_names=GENE_NAMES, **common)
+        same(old, new, f"index {c}")
+        old1 = old_tiles(df_named, enc_onehot, **common)
+        new1 = new_tiles(df_named, enc_onehot, feature_col="gene", **common)
+        same(old1, new1, f"one-hot {c}")
+        n_cmp += len(old) + len(old1)
+    empty = [t for t in range(12) if t not in old_tiles(df_named, enc_named, 12, "strips", "y_location", 6.0,
+                                                         2, 20, cols3, False, False, list(range(12)))]
+    print(f"  core == old pandas code: {n_cmp} tiles compared (4 configs x 2 encodings; "
+          f"empty strips skipped by both: {empty})")
+
+
+# ---------------------------------------------------------------------------
+# map_tiles
+# ---------------------------------------------------------------------------
+def _private_mb():
+    with open("/proc/self/smaps_rollup") as fh:
+        kb = sum(int(l.split()[1]) for l in fh if l.startswith(("Private_Dirty", "Private_Clean")))
+    return kb / 1024
+
+
+def test_pool_equals_serial_and_shares_memory():
+    df = synthetic_sample(120_000)
+    cat = pd.Categorical(df.gene_identity.to_numpy())
+    codes = np.asarray(cat.codes, dtype=np.int64)
+    lut = np.arange(N_GENES, dtype=np.int64)
+    gpos = np.ascontiguousarray(df[["x_location", "y_location"]].to_numpy(np.float64))
+    plan = gb.plan_tiles(gpos[:, 0], gpos[:, 1], 16)
+    big = np.ones(400 * 2**20 // 8)           # 400 MB the workers read but never write
+    tmp = tempfile.mkdtemp()
+
+    def task(t):
+        core, p = plan.members(t, 12.0)
+        a = gb.tile_arrays(core, p, codes, lut, gpos, 6.0, 20, workers=1)
+        s = float(big[::4096].sum())          # read the shared array
+        np.savez(os.path.join(tmp, f"tile{t}.npz"), **a)
+        return {"edges": a["edge_index"].shape[1], "private_mb": _private_mb(), "s": s}
+
+    t0 = time.time()
+    serial = {t: gb.tile_arrays(*plan.members(t, 12.0), codes, lut, gpos, 6.0, 20, workers=1) for t in range(16)}
+    t_serial = time.time() - t0
+    t0 = time.time()
+    res = gb.map_tiles(task, range(16), n_procs=2, log=False)
+    t_par = time.time() - t0
+    par = {t: dict(np.load(os.path.join(tmp, f"tile{t}.npz"))) for t in range(16)}
+    same(serial, par, "pool vs serial")
+    worst = max(r["private_mb"] for r in res.values())
+    assert worst < 200, f"a worker holds {worst:.0f} MB private memory -- shared arrays were copied"
+    assert gb._POOL_TASK[0] is None
+    print(f"  pool (2 procs) == serial for 16 tiles; worker private memory <= {worst:.0f} MB "
+          f"with a 400 MB + sample shared in the parent; serial {t_serial:.1f}s (graphs only) "
+          f"vs pool {t_par:.1f}s (graphs + writes)")
+
+    def bad(t):
+        if t == 3:
+            raise RuntimeError("tile 3 failed")
+        return t
+    try:
+        gb.map_tiles(bad, range(6), n_procs=2, log=False)
+        raise AssertionError("worker error was swallowed")
+    except RuntimeError as e:
+        assert "tile 3" in str(e)
+    assert gb.map_tiles(lambda t: t * 2, [4], n_procs=8) == {4: 8}     # n_procs capped by #tiles
+    print("  worker errors propagate; n_procs capped at the number of tiles")
+
+
+# ---------------------------------------------------------------------------
+# The real SpatialRNA.process() with stand-in torch / PyG / zarr
+# ---------------------------------------------------------------------------
+def _install_stubs():
+    torch = types.ModuleType("torch")
+    torch.from_numpy = lambda a: a
+    torch.set_num_threads = lambda n: None
+    torch.is_tensor = lambda v: False
+    torch.long, torch.bool = "long", "bool"
+    torch.dtype = object
+    torch.Tensor = np.ndarray
+    sys.modules["torch"] = torch
+
+    pyg = types.ModuleType("torch_geometric")
+    data = types.ModuleType("torch_geometric.data")
+
+    class Data(dict):
+        def __init__(self, **kw):
+            super().__init__(**kw)
+
+        def __setattr__(self, k, v):
+            self[k] = v
+
+    class InMemoryDataset:
+        @property
+        def raw_paths(self):
+            return [os.path.join(self.raw_dir, f) for f in self.raw_file_names]
+
+        @property
+        def processed_paths(self):
+            return [os.path.join(self.processed_dir, f) for f in self.processed_file_names]
+
+        def __init__(self, root, *a, **k):
+            self.root = root
+            self.process()
+
+    data.Data, data.InMemoryDataset = Data, InMemoryDataset
+    loader = types.ModuleType("torch_geometric.loader")
+    loader.LinkNeighborLoader = object
+    sys.modules.update({"torch_geometric": pyg, "torch_geometric.data": data,
+                        "torch_geometric.loader": loader})
+    zarr = types.ModuleType("zarr")
+    zarr.__version__ = "3.1.0"
+    codecs = types.ModuleType("zarr.codecs")
+    codecs.BloscCodec = lambda **k: None
+    zarr.codecs = codecs
+    sys.modules.update({"zarr": zarr, "zarr.codecs": codecs})
+
+
+def test_spatialrna_process():
+    """Without torch / PyG / zarr only: stand-ins replace them in sys.modules, which must
+    never happen in an environment that has the real ones (see the _real test)."""
+    if HAVE_STACK:
+        print("  skipped (real torch / PyG / zarr installed: test_spatialrna_process_real runs instead)")
+        return
+    _install_stubs()
+    from spatialrna import graph_zarr_io as gz
+    df = synthetic_sample(80_000)
+    order = RNG.permutation(N_GENES)
+    enc = {g: int(order[i]) for i, g in enumerate(GENE_NAMES)}
+    tmp = tempfile.mkdtemp()
+    enc_path = ge.save_gene_encoding(enc, os.path.join(tmp, "gene_encoding.pkl"))
+
+    gz.load_raw_transcripts_zarr = lambda path, columns=None: df[columns].copy()
+
+    class Attrs:
+        def __init__(self):
+            self.attrs = {"gene_names": GENE_NAMES}
+
+        def array_keys(self):
+            return list(df.columns)
+    gz._zopen_read = lambda path: Attrs()
+
+    results = {}
+    for n_procs in (1, 2):
+        out = os.path.join(tmp, f"S1_{n_procs}.zarr")
+
+        def save_tile(self, tile, arrays, graph_cols, coord_cols, out=out):
+            os.makedirs(out, exist_ok=True)
+            np.savez(os.path.join(out, f"tile{tile}.npz"), **arrays)
+        gz.SpatialRNA._save_tile = save_tile
+        gz.SpatialRNA(root=out, sample_name="S1", cohort_zarr_path=tmp, radius_r=6.0,
+                      dim_x="x_location", dim_y="y_location", dim_z="z_location", tile_mode="grid", tile_by_dim="y_location",
+                      store_coords=True, num_tiles=20, process_tile_ids=list(range(20)),
+                      feature_col="gene_identity", max_num_neighbors=20,
+                      one_hot_encoding=enc_path, n_procs=n_procs, load_type="blank", log=False)
+        results[n_procs] = {t: dict(np.load(os.path.join(out, f"tile{t}.npz"))) for t in range(20)}
+    same(results[1], results[2], "SpatialRNA.process n_procs 1 vs 2")
+    # gene index = encoding of gene_names[code], not the raw code
+    t0 = results[1][0]
+    exp = np.array([enc[GENE_NAMES[c]] for c in df.gene_identity.to_numpy()[t0["trans_id"]]])
+    assert np.array_equal(t0["x"], exp)
+    # every transcript is core in exactly one tile
+    core_ids = np.concatenate([r["trans_id"][r["core_mask"]] for r in results[1].values()])
+    assert np.array_equal(np.sort(core_ids), np.arange(len(df)))
+    print("  SpatialRNA.process(): n_procs=1 == n_procs=2 (20 grid tiles, 3-D pos); encoding loaded "
+          "from .pkl path; gene_identity codes mapped through gene_names; every transcript core once")
+
+
+def test_spatialrna_process_real():
+    """End to end with the real stack: cohort.zarr -> SpatialRNA tiles with n_procs=1 and
+    n_procs=2 (zarr written from forked workers) must be identical."""
+    if not HAVE_STACK:
+        print("  skipped (needs torch, torch_geometric and zarr)")
+        return
+    import zarr
+    from spatialrna import graph_zarr_io as gz
+    df = synthetic_sample(80_000)
+    tmp = tempfile.mkdtemp()
+    cohort = os.path.join(tmp, "cohort.zarr")
+    zarr.open_group(cohort, mode="w").attrs["gene_names"] = GENE_NAMES
+    gz.save_raw_transcripts_zarr(os.path.join(cohort, "S1"), df)
+    order = RNG.permutation(N_GENES)
+    enc_path = ge.save_gene_encoding({g: int(order[i]) for i, g in enumerate(GENE_NAMES)},
+                                      os.path.join(tmp, "gene_encoding.pkl"))
+    out = {}
+    for n_procs in (1, 2):
+        store = os.path.join(tmp, f"run{n_procs}", "S1.zarr")
+        gz.ensure_sample_store(store, "S1")
+        t0 = time.time()
+        gz.SpatialRNA(root=store, sample_name="S1", cohort_zarr_path=cohort, radius_r=6.0,
+                      dim_x="x_location", dim_y="y_location", tile_by_dim="y_location", dim_z="z_location",
+                      tile_mode="grid", store_coords=True, num_tiles=12, process_tile_ids=list(range(12)),
+                      feature_col="gene_identity", max_num_neighbors=20, one_hot_encoding=enc_path,
+                      n_procs=n_procs, load_type="blank", log=False)
+        out[n_procs] = {t: {k: np.asarray(v[:]) for k, v in zarr.open_group(
+            os.path.join(store, "tiles", f"tile{t}"), mode="r").arrays()} for t in range(12)}
+        print(f"  n_procs={n_procs}: 12 tiles in {time.time() - t0:.1f}s")
+    same(out[1], out[2], "real SpatialRNA n_procs 1 vs 2")
+    print("  real zarr + torch: n_procs=1 == n_procs=2 (forked workers wrote every tile)")
+
+
+def test_gene_encoding():
+    tmp = tempfile.mkdtemp()
+    enc = ge.build_gene_encoding(["CD3E", "KRT8", "NegControlProbe_1", "COL1A1"], drop_controls=False)
+    assert enc == {"CD3E": 0, "KRT8": 1, "NegControlProbe_1": 2, "COL1A1": 3}
+    for ext in ("pkl", "json"):
+        p = ge.save_gene_encoding(enc, os.path.join(tmp, f"e.{ext}"))
+        assert ge.load_gene_encoding(p) == enc
+    assert ge.resolve_gene_encoding(enc) == enc and ge.num_genes(enc) == 4
+    eye = np.eye(4, dtype=np.float32)
+    onehot = {g: eye[i] for g, i in enc.items()}
+    assert ge.to_index_encoding(onehot) == enc and ge.same_encoding(onehot, enc) and ge.num_genes(onehot) == 4
+    ge.save_gene_encoding(onehot, os.path.join(tmp, "oh.pkl"))                 # old-style one-hot pkl
+    assert ge.same_encoding(ge.load_gene_encoding(os.path.join(tmp, "oh.pkl")), enc)
+    for bad, msg in ((lambda: ge.build_gene_encoding(["A", "B", "A"]), "duplicate"),
+                     (lambda: ge.encoding_lut(["CD3E", "MS4A1"], enc), "MS4A1"),
+                     (lambda: ge.encoding_lut([0, 7], enc, gene_names=list(enc)), "outside"),
+                     (lambda: ge.encoding_lut([0, 1], enc), "gene_names"),
+                     (lambda: ge.save_gene_encoding(onehot, os.path.join(tmp, "x.json")), "integers")):
+        try:
+            bad()
+            raise AssertionError(f"no error for {msg}")
+        except (ValueError, KeyError) as e:
+            assert msg in str(e), (msg, e)
+    assert list(ge.encoding_lut([3, 0], enc, gene_names=list(enc))) == [3, 0]
+    print("  gene_encoding: build / .pkl + .json round trip / one-hot -> index / clear errors")
+
+
+if __name__ == "__main__":
+    for f in (test_gene_encoding, test_core_matches_old_code, test_pool_equals_serial_and_shares_memory,
+              test_spatialrna_process, test_spatialrna_process_real):
+        print(f.__name__)
+        f()
+    print("ALL PASSED")

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/tests/test_readers.py

@@ -1,7 +1,3 @@
-"""spatialrna.readers on tiny Xenium- and Atera-shaped transcripts.zarr.zip files.
-
-    pytest tests/test_readers.py
-"""
 import json
 import os
 

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/tests/test_tiling.py

@@ -1,7 +1,3 @@
-"""spatialrna.graph_build.plan_tiles: balanced grid tiles (numpy + scipy only).
-
-    pytest tests/test_tiling.py
-"""
 import numpy as np
 from scipy.spatial import cKDTree
 

{spatialrna-0.3.0.dev2 → spatialrna-0.3.0.dev4}/tests/test_zarr_roundtrip.py

@@ -1,19 +1,3 @@
-"""
-Smoke test for spatialrna.graph_zarr_io's save/load round-trip on zarr 3.x.
-
-Run (with the `spatialrna` package installed -- `pip install -e .` from the
-repo root) BEFORE pointing the dataset/training scripts at real data:
-
-    pytest tests/test_zarr_roundtrip.py
-    # or standalone:
-    python tests/test_zarr_roundtrip.py
-
-It builds a tiny fake tile (10 nodes, a handful of edges, a bool mask),
-writes it with `save_data_zarr`, reads it back with `load_data_zarr`, and
-asserts every tensor matches exactly (dtype and values). It also checks
-`zarr_group_exists` and the `keys=` partial-read path used by the training
-script (reading only `x`/`edge_index` without touching `core_mask`).
-"""
 import shutil
 import tempfile
 import os.path as osp