Spacelab 0.28.0__py3-none-any.whl

spacelab/hlod/hashing.py

@@ -0,0 +1,38 @@
+"""splitmix64, vectorized over uint64 arrays.
+
+hash(k) is the first output of a splitmix64 generator whose state is k, so
+hash(seed_state) reproduces the published sequence for that seed.
+"""
+
+import numpy as np
+
+GOLDEN = 0x9E3779B97F4A7C15
+_M1 = 0xBF58476D1CE4E5B9
+_M2 = 0x94D049BB133111EB
+_MASK = (1 << 64) - 1
+
+
+def splitmix64_scalar(k: int) -> int:
+    z = (k + GOLDEN) & _MASK
+    z = ((z ^ (z >> 30)) * _M1) & _MASK
+    z = ((z ^ (z >> 27)) * _M2) & _MASK
+    return z ^ (z >> 31)
+
+
+def splitmix64(k: np.ndarray) -> np.ndarray:
+    z = np.asarray(k, dtype=np.uint64)
+    with np.errstate(over="ignore"):  # wrapping uint64 arithmetic is the algorithm
+        z = z + np.uint64(GOLDEN)
+        z = (z ^ (z >> np.uint64(30))) * np.uint64(_M1)
+        z = (z ^ (z >> np.uint64(27))) * np.uint64(_M2)
+    return z ^ (z >> np.uint64(31))
+
+
+def gene_xy_key(gene_code: int, x: np.ndarray, y: np.ndarray) -> np.ndarray:
+    """Fold the bytes of (gene_code uint32, x float32, y float32) into a uint64
+    key, for sources without a transcript id. Duplicate (x, y) within a gene
+    share a key, so they land in the same level."""
+    xb = np.ascontiguousarray(x, dtype="<f4").view("<u4").astype(np.uint64)
+    yb = np.ascontiguousarray(y, dtype="<f4").view("<u4").astype(np.uint64)
+    hi = (np.uint64(gene_code) << np.uint64(32)) | xb
+    return hi ^ splitmix64(yb)

spacelab/hlod/hilbert.py

@@ -0,0 +1,91 @@
+"""Hilbert curve index on a 2^order x 2^order grid.
+
+The standard iterative xy2d (Wikipedia's), fixed so the TypeScript client
+(src/hlod/hilbert.ts) produces identical indices. Its prefix property is what
+the store relies on: the depth-m cell of a point is the top 2m bits of its
+order-K index, so every quadtree cell is one contiguous run of h-sorted rows.
+"""
+
+import numpy as np
+
+# Rows per vectorized block; bounds the temporaries to a few hundred MB.
+BLOCK = 10_000_000
+
+
+def xy2d_scalar(order: int, x: int, y: int) -> int:
+    n = 1 << order
+    d = 0
+    s = n >> 1
+    while s > 0:
+        rx = 1 if x & s else 0
+        ry = 1 if y & s else 0
+        d += s * s * ((3 * rx) ^ ry)
+        if ry == 0:
+            if rx == 1:
+                x = n - 1 - x  # n, not s
+                y = n - 1 - y
+            x, y = y, x
+        s >>= 1
+    return d
+
+
+def d2xy_scalar(order: int, d: int) -> tuple[int, int]:
+    """Inverse of xy2d_scalar; used by the tests only."""
+    x = y = 0
+    t = d
+    s = 1
+    n = 1 << order
+    while s < n:
+        rx = 1 & (t // 2)
+        ry = 1 & (t ^ rx)
+        if ry == 0:
+            if rx == 1:
+                x = s - 1 - x
+                y = s - 1 - y
+            x, y = y, x
+        x += s * rx
+        y += s * ry
+        t //= 4
+        s <<= 1
+    return x, y
+
+
+def _xy2d_block(order: int, x: np.ndarray, y: np.ndarray) -> np.ndarray:
+    n = np.uint32((1 << order) - 1)
+    x = x.astype(np.uint32, copy=True)
+    y = y.astype(np.uint32, copy=True)
+    # uint64: at order 16 the last step can add 3 * 4^15 to a sum near 4^16.
+    d = np.zeros(x.shape, dtype=np.uint64)
+    s = 1 << (order - 1)
+    while s > 0:
+        rx = (x & np.uint32(s)) != 0
+        ry = (y & np.uint32(s)) != 0
+        d += np.uint64(s * s) * ((3 * rx.astype(np.uint64)) ^ ry.astype(np.uint64))
+        flip = ~ry & rx
+        x[flip] = n - x[flip]
+        y[flip] = n - y[flip]
+        swap = ~ry
+        x[swap], y[swap] = y[swap], x[swap].copy()
+        s >>= 1
+    return d
+
+
+def xy2d(order: int, x: np.ndarray, y: np.ndarray) -> np.ndarray:
+    """Vectorized xy2d. uint32 for order <= 16, else uint64."""
+    if order < 1:
+        return np.zeros(np.shape(x), dtype=np.uint32)
+    x = np.asarray(x)
+    y = np.asarray(y)
+    out_dtype = np.uint32 if order <= 16 else np.uint64
+    out = np.empty(x.shape, dtype=out_dtype)
+    flat_x, flat_y, flat_out = x.reshape(-1), y.reshape(-1), out.reshape(-1)
+    for a in range(0, flat_x.size, BLOCK):
+        flat_out[a:a + BLOCK] = _xy2d_block(order, flat_x[a:a + BLOCK], flat_y[a:a + BLOCK])
+    return out
+
+
+def quantize(v: np.ndarray, v0: float, size: float, bits: int) -> np.ndarray:
+    """floor((v - v0) / size * 2^bits), clamped to [0, 2^bits - 1], as uint32."""
+    n = 1 << bits
+    q = np.floor((np.asarray(v, dtype=np.float64) - v0) / size * n)
+    return np.clip(q, 0, n - 1).astype(np.uint32)

spacelab/hlod/read.py

@@ -0,0 +1,150 @@
+"""Analysis-side reader. The store is plain Zarr (zarr.open(path)["coords"][:]
+works); this adds the gene, level and region lookups the sort order makes cheap.
+
+    store = hlod.open("points.zarr")      # local path, s3://, https://
+    g = store.gene("CD3E")
+    g.coords()                            # (N_g, 2) float32, one contiguous slice
+    g.coords(levels=2)                    # levels 0..2: a uniform subsample of fraction F[2]
+    g.table(columns=["cell_id", "qv"])    # pyarrow Table, co-sorted
+    g.roi(xmin, ymin, xmax, ymax)         # every transcript of the gene in the box
+    store.roi(xmin, ymin, xmax, ymax)     # across genes, with a gene column
+"""
+
+from collections.abc import Sequence
+from functools import cached_property
+
+import numpy as np
+import pyarrow as pa
+import zarr
+
+from .build import FORMAT
+from .cells import DEFAULT_MAX_GAP_ROWS, Extent, block_offsets, cells_in_box, coalesce, ranges_for
+
+
+def open(path: str, storage_options: dict | None = None) -> "Store":  # noqa: A001
+    return Store(path, storage_options)
+
+
+class Store:
+    def __init__(self, path: str, storage_options: dict | None = None):
+        self.root = zarr.open_group(path, mode="r", storage_options=storage_options)
+        self.attrs = dict(self.root.attrs)
+        if self.attrs.get("format") != FORMAT:
+            raise ValueError(f"{path} is not a {FORMAT} store")
+        self.genes: list[str] = list(self.attrs["genes"])
+        self.columns: list[str] = list(self.attrs.get("columns", []))
+        self.categories: dict[str, str] = dict(self.attrs.get("categories", {}))
+        self.extent = Extent.from_attrs(self.attrs)
+        self.K, self.g, self.P, self.Dmax = (int(self.attrs[k]) for k in ("K", "g", "P", "Dmax"))
+        self.B = block_offsets(self.g, self.Dmax)
+        self._gene_of = {name: i for i, name in enumerate(self.genes)}
+
+    @cached_property
+    def gene_offsets(self) -> np.ndarray:
+        return self.root["index/gene_offsets"][:].astype(np.int64)
+
+    @cached_property
+    def level_offsets(self) -> np.ndarray:
+        return self.root["index/level_offsets"][:].astype(np.int64)
+
+    @property
+    def n_rows(self) -> int:
+        return int(self.gene_offsets[-1])
+
+    def gene(self, gene: str | int) -> "Gene":
+        gi = gene if isinstance(gene, (int, np.integer)) else self._gene_of.get(gene)
+        if gi is None or not 0 <= gi < len(self.genes):
+            raise KeyError(gene)
+        return Gene(self, int(gi))
+
+    def fraction(self, level: int, n_gene: int) -> float:
+        """F[level]: the share of a gene's rows in levels 0..level."""
+        return 1.0 if level >= self.Dmax or not n_gene else min(1.0, self.P * 4 ** level / n_gene)
+
+    def read_rows(self, name: str, start: int, end: int) -> np.ndarray:
+        return self.root[name][start:end]
+
+    @cached_property
+    def _category_names(self) -> dict[str, np.ndarray]:
+        return {c: np.asarray(self.root[a][:], dtype=object) for c, a in self.categories.items()}
+
+    def column_array(self, name: str, values: np.ndarray) -> pa.Array:
+        if name in self.categories:
+            names = pa.array(self._category_names[name].tolist(), type=pa.string())
+            return pa.DictionaryArray.from_arrays(pa.array(values.astype(np.int32)), names)
+        return pa.array(values)
+
+    def roi(self, xmin, ymin, xmax, ymax, genes: Sequence[str] | None = None, exact: bool = True,
+            columns: Sequence[str] = ()) -> pa.Table:
+        names = self.genes if genes is None else list(genes)
+        tables = []
+        for name in names:
+            t = self.gene(name).roi(xmin, ymin, xmax, ymax, exact=exact, columns=columns)
+            tables.append(t.append_column("gene", pa.array([name] * t.num_rows, pa.string())))
+        if not tables:
+            return pa.table({"x": pa.array([], pa.float32()), "y": pa.array([], pa.float32()),
+                             "gene": pa.array([], pa.string())})
+        return pa.concat_tables(tables)
+
+
+class Gene:
+    def __init__(self, store: Store, gi: int):
+        self.store = store
+        self.index = gi
+        self.name = store.genes[gi]
+        self.start = int(store.gene_offsets[gi])
+        self.end = int(store.gene_offsets[gi + 1])
+        self.level_offsets = store.level_offsets[gi]
+
+    def __len__(self) -> int:
+        return self.end - self.start
+
+    @cached_property
+    def cell_index(self) -> np.ndarray:
+        return self.store.root["index/cell_index"][self.index].astype(np.int64)
+
+    def _stop(self, levels: int | None) -> int:
+        if levels is None or levels >= self.store.Dmax:
+            return self.end
+        return int(self.level_offsets[levels + 1])
+
+    def coords(self, levels: int | None = None) -> np.ndarray:
+        return self.store.read_rows("coords", self.start, self._stop(levels))
+
+    def table(self, columns: Sequence[str] | None = None, levels: int | None = None) -> pa.Table:
+        stop = self._stop(levels)
+        return self._table([(self.start, stop)], columns)
+
+    def _table(self, spans, columns) -> pa.Table:
+        columns = self.store.columns if columns is None else list(columns)
+        xy = (np.concatenate([self.store.read_rows("coords", s, e) for s, e in spans])
+              if spans else np.empty((0, 2), np.float32))
+        out = {"x": pa.array(xy[:, 0]), "y": pa.array(xy[:, 1])}
+        for c in columns:
+            vals = np.concatenate([self.store.read_rows(c, s, e) for s, e in spans]) if spans \
+                else np.empty(0, self.store.root[c].dtype)
+            out[c] = self.store.column_array(c, vals)
+        return pa.table(out)
+
+    def roi_ranges(self, xmin, ymin, xmax, ymax,
+                   max_gap_rows: int = DEFAULT_MAX_GAP_ROWS) -> list[tuple[int, int]]:
+        """Coalesced absolute row ranges covering the box at every level; the
+        same cellsInBox -> rangesFor -> coalesce the viewer runs."""
+        s = self.store
+        ranges = []
+        for d in range(s.Dmax + 1):
+            cells = cells_in_box((xmin, ymin, xmax, ymax), d + s.g, s.extent)
+            ranges.extend(coalesce(ranges_for(self.cell_index, self.start, d, cells, s.g, s.Dmax),
+                                   max_gap_rows))
+        return [(r.start, r.end) for r in coalesce(ranges, max_gap_rows)]
+
+    def roi(self, xmin, ymin, xmax, ymax, exact: bool = True,
+            columns: Sequence[str] = ()) -> pa.Table:
+        """Transcripts of this gene in the box; never reads the gene's full array.
+        exact=False keeps the overhang of the cells the box touches."""
+        t = self._table(self.roi_ranges(xmin, ymin, xmax, ymax), columns)
+        if exact and t.num_rows:
+            x = t.column("x").to_numpy()
+            y = t.column("y").to_numpy()
+            t = t.filter(pa.array((x >= xmin) & (x <= xmax) & (y >= ymin) & (y <= ymax)))
+        return t

spacelab/hlod/verify.py

@@ -0,0 +1,142 @@
+"""Structural and statistical checks a store must pass before it is published."""
+
+import math
+
+import numpy as np
+
+from .build import FORMAT_VERSION
+from .hilbert import quantize, xy2d
+from .read import Store
+
+
+def ks_2samp(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
+    """Two-sample Kolmogorov-Smirnov statistic and asymptotic p-value
+    (Numerical Recipes' kstwo), so verify needs no scipy."""
+    a = np.sort(a)
+    b = np.sort(b)
+    both = np.concatenate([a, b])
+    cdf_a = np.searchsorted(a, both, side="right") / a.size
+    cdf_b = np.searchsorted(b, both, side="right") / b.size
+    d = float(np.abs(cdf_a - cdf_b).max())
+    en = math.sqrt(a.size * b.size / (a.size + b.size))
+    lam = (en + 0.12 + 0.11 / en) * d
+    if lam < 1e-3:
+        return d, 1.0
+    p = 2 * sum((-1) ** (j - 1) * math.exp(-2 * j * j * lam * lam) for j in range(1, 101))
+    return d, float(min(max(p, 0.0), 1.0))
+
+
+class Report:
+    def __init__(self):
+        self.failures: list[str] = []
+        self.notes: list[str] = []
+
+    def check(self, ok: bool, msg: str):
+        if not ok:
+            self.failures.append(msg)
+
+    @property
+    def ok(self) -> bool:
+        return not self.failures
+
+
+def verify(path: str, roi_samples: int = 50, seed: int = 7, ks_p: float = 0.001,
+           binomial_sd: float = 4.0) -> Report:
+    store = Store(path)
+    rep = Report()
+    a = store.attrs
+    rep.check(a.get("format_version") == FORMAT_VERSION, f"format_version {a.get('format_version')}")
+    K, g, dmax, P = store.K, store.g, store.Dmax, store.P
+    n = store.n_rows
+    G = len(store.genes)
+    go, lo = store.gene_offsets, store.level_offsets
+    rep.check(go.shape == (G + 1,) and go[0] == 0, "gene_offsets shape/start")
+    rep.check(bool((np.diff(go) >= 0).all()), "gene_offsets not monotone")
+    rep.check(store.root["coords"].shape == (n, 2), "coords rows != gene_offsets[G]")
+    rep.check(lo.shape == (G, dmax + 2), "level_offsets shape")
+    for c in store.columns:
+        rep.check(store.root[c].shape[0] == n, f"column {c} has the wrong length")
+
+    if "original_index" in store.root:
+        orig = store.root["original_index"][:]
+        rep.check(bool((np.sort(orig) == np.arange(n, dtype=orig.dtype)).all()),
+                  "original_index is not a permutation: some input row is missing or repeated")
+        del orig
+
+    ext = store.extent
+    B = store.B
+    cell_index = store.root["index/cell_index"][:].astype(np.int64)
+    use_binomial = a.get("hash", {}).get("key") == "transcript_id"
+    rows_per_level = np.zeros(dmax + 1, dtype=np.int64)
+    for gi, name in enumerate(store.genes):
+        s, e = int(go[gi]), int(go[gi + 1])
+        ng = e - s
+        rel = lo[gi] - s
+        rep.check(lo[gi, 0] == s and lo[gi, -1] == e, f"{name}: level_offsets don't span the gene")
+        rep.check(bool((np.diff(rel) >= 0).all()), f"{name}: level_offsets not monotone")
+        rows_per_level += np.diff(rel)
+        xy = store.read_rows("coords", s, e)
+        h = xy2d(K, quantize(xy[:, 0], ext.x0, ext.size, K), quantize(xy[:, 1], ext.y0, ext.size, K))
+        prev_f = 0.0
+        for d in range(dmax + 1):
+            ls, le = int(rel[d]), int(rel[d + 1])
+            hd = h[ls:le]
+            rep.check(bool((np.diff(hd.astype(np.int64)) >= 0).all()),
+                      f"{name} level {d}: h not sorted")
+            block = cell_index[gi, B[d]:B[d + 1]]
+            rep.check(block[0] == ls and block[-1] == le,
+                      f"{name} level {d}: cell_index block disagrees with level_offsets")
+            rep.check(bool((np.diff(block) >= 0).all()), f"{name} level {d}: cell_index not monotone")
+            m = d + g
+            expect = np.repeat(np.arange(4 ** m), np.diff(block))
+            got = (hd >> np.uint32(2 * (K - m))).astype(np.int64)
+            rep.check(expect.size == got.size and bool((expect == got).all()),
+                      f"{name} level {d}: rows outside the cell cell_index assigns them")
+            f = store.fraction(d, ng)
+            if use_binomial and ng:
+                p = f - prev_f
+                sd = math.sqrt(ng * p * (1 - p))
+                rep.check(abs((le - ls) - ng * p) <= binomial_sd * sd + 1,
+                          f"{name} level {d}: {le - ls} rows, expected {ng * p:.0f} ± {sd:.0f}")
+            prev_f = f
+        # A level prefix that is a real subsample should look like the whole gene.
+        for d in range(dmax):
+            k = int(rel[d + 1])
+            if 1000 <= k < ng:
+                for axis, label in ((0, "x"), (1, "y")):
+                    _, p = ks_2samp(xy[:k, axis], xy[:, axis])
+                    rep.check(p > ks_p, f"{name} levels 0..{d}: KS on {label} p={p:.2g}")
+
+    rep.notes.append(f"rows per level: {rows_per_level.tolist()}")
+    rep.notes.append(f"index bytes per gene: {cell_index.shape[1] * 4 + (dmax + 2) * 8 + 8:,}")
+    if cell_index.size:
+        rep.notes.append(f"largest cell: {int(np.diff(cell_index, axis=1).max()):,} rows")
+
+    rng = np.random.default_rng(seed)
+    if G and roi_samples:
+        for i in range(roi_samples):
+            gi = int(rng.integers(G))
+            gene = store.gene(gi)
+            box = _random_box(rng, ext, i)
+            got = gene.roi(*box, exact=True)
+            xy = gene.coords()
+            want = xy[(xy[:, 0] >= box[0]) & (xy[:, 0] <= box[2])
+                      & (xy[:, 1] >= box[1]) & (xy[:, 1] <= box[3])]
+            gx = np.stack([got.column("x").to_numpy(), got.column("y").to_numpy()], axis=1)
+            same = gx.shape == want.shape and bool(
+                (gx[np.lexsort(gx.T)] == want[np.lexsort(want.T)]).all())
+            rep.check(same, f"roi {box} on {gene.name}: {len(gx)} rows, brute force {len(want)}")
+    return rep
+
+
+def _random_box(rng, ext, i: int):
+    """Mostly in-extent boxes of varied size, with every 10th smaller than a
+    deep cell and every 7th wholly outside the extent."""
+    if i % 7 == 6:
+        x = ext.x0 + ext.size * (1.5 + rng.random())
+        return (x, ext.y0, x + ext.size * 0.1, ext.y0 + ext.size * 0.1)
+    scale = 1e-4 if i % 10 == 9 else rng.uniform(0.01, 0.6)
+    w = ext.size * scale
+    x = ext.x0 + rng.uniform(-0.05, 1.0) * ext.size
+    y = ext.y0 + rng.uniform(-0.05, 1.0) * ext.size
+    return (x, y, x + w, y + w * rng.uniform(0.5, 2.0))