Spacelab 0.28.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. spacelab/__init__.py +6 -0
  2. spacelab/__main__.py +3 -0
  3. spacelab/cli.py +97 -0
  4. spacelab/hlod/__init__.py +8 -0
  5. spacelab/hlod/__main__.py +5 -0
  6. spacelab/hlod/build.py +349 -0
  7. spacelab/hlod/cells.py +82 -0
  8. spacelab/hlod/cli.py +90 -0
  9. spacelab/hlod/hashing.py +38 -0
  10. spacelab/hlod/hilbert.py +91 -0
  11. spacelab/hlod/read.py +150 -0
  12. spacelab/hlod/verify.py +142 -0
  13. spacelab/prep_spatialdata.py +700 -0
  14. spacelab/server.py +130 -0
  15. spacelab/static/AI_lens_black.png +0 -0
  16. spacelab/static/ai-new-logo.png +0 -0
  17. spacelab/static/aifi_logo.png +0 -0
  18. spacelab/static/assets/AllenInstitutePlusText-Bold-DxLtQ4oE.woff2 +0 -0
  19. spacelab/static/assets/AllenInstitutePlusText-BoldItalic-mxzHQhQO.woff2 +0 -0
  20. spacelab/static/assets/AllenInstitutePlusText-Italic-CaKmxpsW.woff2 +0 -0
  21. spacelab/static/assets/AllenInstitutePlusText-Light-DUThXl1m.woff2 +0 -0
  22. spacelab/static/assets/AllenInstitutePlusText-LightItalic-CxfKknzG.woff2 +0 -0
  23. spacelab/static/assets/AllenInstitutePlusText-Regular-B0ryviHQ.woff2 +0 -0
  24. spacelab/static/assets/AllenInstitutePlusText-Semibold-3PoiDF1d.woff2 +0 -0
  25. spacelab/static/assets/AllenInstitutePlusText-SemiboldItalic-uXaVqnFo.woff2 +0 -0
  26. spacelab/static/assets/blosc-BKQbfHyB.js +18 -0
  27. spacelab/static/assets/blosc-j6ygecpn.js +18 -0
  28. spacelab/static/assets/chunk-INHXZS53-D3tQiqtZ.js +2 -0
  29. spacelab/static/assets/chunk-INHXZS53-QjxmivJp.js +2 -0
  30. spacelab/static/assets/deflate-BRVR_fNk.js +2 -0
  31. spacelab/static/assets/hlodTranscripts.worker-CjhNLk56.js +2 -0
  32. spacelab/static/assets/index-BT-mDgg2.js +5025 -0
  33. spacelab/static/assets/index-CNIqsj6m.css +1 -0
  34. spacelab/static/assets/jpeg-JUDfHHDo.js +2 -0
  35. spacelab/static/assets/lerc-CXQdKTps.js +2 -0
  36. spacelab/static/assets/lz4-C2MQfzIA.js +16 -0
  37. spacelab/static/assets/lz4-rzy_Siro.js +16 -0
  38. spacelab/static/assets/lzw-DBCCrYEo.js +2 -0
  39. spacelab/static/assets/packbits-LrtZ2mhj.js +2 -0
  40. spacelab/static/assets/pako.esm-C69tvH9t.js +2 -0
  41. spacelab/static/assets/polygons.worker-ClOPdWlW.js +2 -0
  42. spacelab/static/assets/raw-lPmjnySZ.js +2 -0
  43. spacelab/static/assets/spatialdataExpression.worker-D4W_OHS3.js +4 -0
  44. spacelab/static/assets/spatialdataShapes.worker-DQV1zeTt.js +974 -0
  45. spacelab/static/assets/spatialdataTranscripts.worker-DH9L-MwZ.js +974 -0
  46. spacelab/static/assets/transcripts.worker-CtxyO4n3.js +2 -0
  47. spacelab/static/assets/umap.worker-BPHuoLPV.js +2 -0
  48. spacelab/static/assets/webimage-xg0Pp5Jv.js +2 -0
  49. spacelab/static/assets/zstd-Cp7BEDDH.js +16 -0
  50. spacelab/static/assets/zstd-WlDjm0Vu.js +16 -0
  51. spacelab/static/blurred-bg.jpg +0 -0
  52. spacelab/static/favicon.ico +0 -0
  53. spacelab/static/index.html +22 -0
  54. spacelab/static/logo192.png +0 -0
  55. spacelab/static/logo512.png +0 -0
  56. spacelab/static/manifest.json +15 -0
  57. spacelab/static/robots.txt +3 -0
  58. spacelab-0.28.0.dist-info/METADATA +30 -0
  59. spacelab-0.28.0.dist-info/RECORD +62 -0
  60. spacelab-0.28.0.dist-info/WHEEL +4 -0
  61. spacelab-0.28.0.dist-info/entry_points.txt +2 -0
  62. spacelab-0.28.0.dist-info/licenses/LICENSE +21 -0
spacelab/__init__.py ADDED
@@ -0,0 +1,6 @@
1
+ from importlib.metadata import PackageNotFoundError, version
2
+
3
+ try:
4
+ __version__ = version("spacelab")
5
+ except PackageNotFoundError: # imported from a source tree, not an install
6
+ __version__ = "0+unknown"
spacelab/__main__.py ADDED
@@ -0,0 +1,3 @@
1
+ from spacelab.cli import main
2
+
3
+ main()
spacelab/cli.py ADDED
@@ -0,0 +1,97 @@
1
+ import argparse
2
+ import os
3
+ import shutil
4
+ import subprocess
5
+ import sys
6
+ import threading
7
+ import webbrowser
8
+ from pathlib import Path
9
+
10
+ from spacelab import __version__
11
+ from spacelab.server import STATIC_DIR, make_server
12
+
13
+ PREP_SCRIPT = Path(__file__).parent / "prep_spatialdata.py"
14
+
15
+ MANIFEST = "tissue-set-manifest.json"
16
+
17
+
18
+ def prep(args: list[str]) -> None:
19
+ if not PREP_SCRIPT.is_file():
20
+ sys.exit(f"error: prep script not bundled at {PREP_SCRIPT} (broken install?)")
21
+ # uv resolves the script's inline dependencies (zarr 3, shapely, ...) on the fly.
22
+ uv = shutil.which("uv")
23
+ cmd = [uv, "run", "--script", str(PREP_SCRIPT), *args] if uv else [sys.executable, str(PREP_SCRIPT), *args]
24
+ try:
25
+ sys.exit(subprocess.call(cmd))
26
+ except KeyboardInterrupt:
27
+ sys.exit(130)
28
+
29
+
30
+ def hlod(args: list[str]) -> None:
31
+ try:
32
+ import numpy, pyarrow, zarr # noqa: F401, E401
33
+ except ImportError:
34
+ # Same trick as prep: let uv supply the hlod extra's dependencies.
35
+ uv = shutil.which("uv")
36
+ if not uv:
37
+ sys.exit("error: `spacelab hlod` needs the hlod extra: pip install 'spacelab[hlod]'")
38
+ env = {**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)}
39
+ cmd = [uv, "run", "--no-project", "--with", "numpy>=1.26", "--with", "pyarrow>=14",
40
+ "--with", "zarr>=3.1", "python", "-m", "spacelab.hlod", *args]
41
+ try:
42
+ sys.exit(subprocess.call(cmd, env=env))
43
+ except KeyboardInterrupt:
44
+ sys.exit(130)
45
+ from spacelab.hlod.cli import main as hlod_main
46
+
47
+ sys.exit(hlod_main(args))
48
+
49
+
50
+ def main(argv: list[str] | None = None) -> None:
51
+ argv = sys.argv[1:] if argv is None else argv
52
+ # Handled before argparse: the serve command's optional positional would swallow "prep".
53
+ if argv and argv[0] == "prep":
54
+ prep(argv[1:])
55
+ return
56
+ if argv and argv[0] == "hlod":
57
+ hlod(argv[1:])
58
+ return
59
+ parser = argparse.ArgumentParser(
60
+ prog="spacelab",
61
+ description="Serve a converted spatial dataset in the browser. "
62
+ "Run `spacelab prep --help` to add fast-loading sidecars to a SpatialData store, "
63
+ "or `spacelab hlod --help` for the transcript LOD store.",
64
+ )
65
+ parser.add_argument(
66
+ "data_dir", nargs="?", default=".", help="*_app_ready directory (default: cwd)"
67
+ )
68
+ parser.add_argument("--port", type=int, default=3000)
69
+ parser.add_argument("--host", default="127.0.0.1")
70
+ parser.add_argument("--no-browser", action="store_true", help="don't open a browser tab")
71
+ parser.add_argument("--version", action="version", version=f"spacelab {__version__}")
72
+ args = parser.parse_args(argv)
73
+
74
+ data_dir = Path(args.data_dir).expanduser().resolve()
75
+ # Without a manifest the app 404s on it and falls back to its connect-a-bucket page.
76
+ has_manifest = (data_dir / MANIFEST).is_file()
77
+ if not has_manifest:
78
+ print(f"No {MANIFEST} found in {data_dir}, so defaulting to cloud connect for data.")
79
+ if not (STATIC_DIR / "index.html").is_file():
80
+ sys.exit(f"error: frontend not bundled at {STATIC_DIR} (broken install?)")
81
+
82
+ try:
83
+ server = make_server(str(data_dir), args.host, args.port)
84
+ except OSError as e:
85
+ sys.exit(f"error: cannot bind {args.host}:{args.port} ({e}); try --port")
86
+
87
+ url = f"http://{'localhost' if args.host in ('127.0.0.1', '0.0.0.0') else args.host}:{args.port}"
88
+ target = f"{data_dir} " if has_manifest else ""
89
+ print(f"🚀 spacelab v{__version__} launching {target}on {url} (Ctrl+C to stop)", flush=True)
90
+ if not args.no_browser:
91
+ threading.Timer(0.5, webbrowser.open, [url]).start()
92
+ try:
93
+ server.serve_forever()
94
+ except KeyboardInterrupt:
95
+ print("\nstopped")
96
+ finally:
97
+ server.server_close()
@@ -0,0 +1,8 @@
1
+ """Hilbert additive-LOD transcript store: build, read and verify.
2
+
3
+ Needs the `hlod` extra (numpy, pyarrow, zarr>=3): pip install 'spacelab[hlod]'.
4
+ """
5
+
6
+ from .read import Gene, Store, open
7
+
8
+ __all__ = ["Gene", "Store", "open"]
@@ -0,0 +1,5 @@
1
+ import sys
2
+
3
+ from .cli import main
4
+
5
+ sys.exit(main())
spacelab/hlod/build.py ADDED
@@ -0,0 +1,349 @@
1
+ """Build a Hilbert additive-LOD point store from a transcript table.
2
+
3
+ Rows are sorted by (gene, level, h): a gene is one contiguous slice, levels
4
+ 0..d of a gene are one contiguous prefix of it, and within a level every
5
+ quadtree cell at the level's index depth is one contiguous run. Levels are
6
+ disjoint random subsets chosen by hashing a stable per-row key, so levels
7
+ 0..d together are a uniform random sample of fraction F[d].
8
+ """
9
+
10
+ import math
11
+ import os
12
+ import shutil
13
+ import tempfile
14
+ import time
15
+ from collections.abc import Callable, Sequence
16
+
17
+ import numpy as np
18
+ import pyarrow as pa
19
+ import pyarrow.parquet as pq
20
+ import zarr
21
+ from zarr.codecs import ZstdCodec
22
+
23
+ from .cells import block_offsets
24
+ from .hashing import gene_xy_key, splitmix64
25
+ from .hilbert import quantize, xy2d
26
+
27
+ # Written to the root attrs. The browser reader (src/hlod/store.ts) refuses a
28
+ # store with any other format or version, read.py checks the format, and
29
+ # `hlod verify` the version; bump the version on any layout change.
30
+ FORMAT = "hilbert-lod-points"
31
+ FORMAT_VERSION = 1
32
+ # Cap on the deepest level, Dmax. Each level is 4x the previous, so 6 levels past
33
+ # level 0 reach P * 4^6 (~41M rows at the default P); a larger gene puts the rest
34
+ # in its last level. Bounds the per-gene cell table, 4^(d+g) entries per level d.
35
+ MAX_DMAX = 6
36
+ # Array names the builder writes itself; an extra column with one of these names
37
+ # would overwrite it.
38
+ RESERVED = {"coords", "index", "original_index"}
39
+
40
+ DEFAULTS = dict(
41
+ # Level-0 rows per gene: the subsample drawn when the whole tissue is in view.
42
+ # Smaller means sparser zoomed-out views and more levels.
43
+ P=10_000,
44
+ # Hilbert grid resolution: 2^K x 2^K over the square extent. 16 keeps h in uint32.
45
+ K=16,
46
+ # Extra quadtree depth for the cell index: level d is indexed at depth d + g,
47
+ # so level 0 already splits into 4^g cells for viewport queries.
48
+ g=1,
49
+ # Seed for the per-row hash that assigns rows to levels; fixed so rebuilds of
50
+ # the same input give the same store.
51
+ seed=42,
52
+ # Rows per zarr chunk, the unit of compression and of a byte-range read.
53
+ chunk=16_384,
54
+ # Rows per shard file (a multiple of chunk); fewer, larger files to serve.
55
+ shard=4_194_304,
56
+ # Where each shard file keeps its chunk index ("start" or "end", zarr's
57
+ # default); recorded in the array metadata, and the reader handles both.
58
+ index_location="end",
59
+ )
60
+
61
+
62
+ class BuildError(Exception):
63
+ """Bad input or options; the message is meant for the user."""
64
+
65
+
66
+ def choose_dmax(gene_counts: np.ndarray, P: int) -> int:
67
+ """Deepest level any gene needs for its last level to hold all its rows,
68
+ in 1..MAX_DMAX."""
69
+ deepest = 0
70
+ for n in gene_counts:
71
+ if n > P:
72
+ deepest = max(deepest, math.ceil(math.log(n / P, 4)))
73
+ return max(1, min(MAX_DMAX, deepest))
74
+
75
+
76
+ def level_thresholds(n_gene: int, P: int, dmax: int) -> list[int]:
77
+ """T[d] for d < dmax as uint64 cut points: a row is in level d when
78
+ T[d-1] <= r < T[d]. Integer comparison avoids float rounding at the edges."""
79
+ out = []
80
+ for d in range(dmax):
81
+ frac = min(1.0, P * 4 ** d / n_gene) if n_gene else 1.0
82
+ out.append(min(math.floor(frac * 2 ** 64), 2 ** 64 - 1))
83
+ return out
84
+
85
+
86
+ def assign_levels(r: np.ndarray, n_gene: int, P: int, dmax: int) -> np.ndarray:
87
+ """Level (uint8) of each row of one gene, from its uint64 hash `r`."""
88
+ t = level_thresholds(n_gene, P, dmax)
89
+ level = np.searchsorted(np.array(t, dtype=np.uint64), r, side="right").astype(np.uint8)
90
+ # A level whose cumulative fraction is 1 takes every remaining row, including
91
+ # the one hash value (2^64 - 1) the clamped threshold would push past it.
92
+ full = next((d for d in range(dmax) if n_gene and P * 4 ** d >= n_gene), None)
93
+ if full is not None:
94
+ np.minimum(level, full, out=level)
95
+ return level
96
+
97
+
98
+ def _read_table(source, cols: list[str]) -> pa.Table:
99
+ """`cols` of a pyarrow Table or parquet file/directory; BuildError if any is missing."""
100
+ if isinstance(source, pa.Table):
101
+ missing = [c for c in cols if c not in source.column_names]
102
+ if missing:
103
+ raise BuildError(f"input has no column(s) {missing}")
104
+ return source.select(cols)
105
+ schema_names = set(pq.read_schema(_first_parquet(source)).names)
106
+ missing = [c for c in cols if c not in schema_names]
107
+ if missing:
108
+ raise BuildError(f"{source} has no column(s) {missing} (has: {sorted(schema_names)})")
109
+ return pq.read_table(source, columns=cols)
110
+
111
+
112
+ def _first_parquet(path: str) -> str:
113
+ """`path` itself if a file, else its first .parquet part (for reading the schema)."""
114
+ if os.path.isfile(path):
115
+ return path
116
+ parts = sorted(f for f in os.listdir(path) if f.endswith(".parquet"))
117
+ if not parts:
118
+ raise BuildError(f"{path} has no parquet files")
119
+ return os.path.join(path, parts[0])
120
+
121
+
122
+ def _dictionary(column: pa.ChunkedArray) -> tuple[np.ndarray, list[str]]:
123
+ """(codes, names) with names sorted, so a code is the position in `names`.
124
+ Each parquet part carries its own dictionary, so they are unified first.
125
+ A null becomes the empty string."""
126
+ column = column.cast(pa.string()) if not pa.types.is_dictionary(column.type) else column
127
+ if not pa.types.is_dictionary(column.type):
128
+ column = column.dictionary_encode(null_encoding="encode")
129
+ arr = column.unify_dictionaries().combine_chunks()
130
+ if arr.null_count:
131
+ arr = pa.DictionaryArray.from_arrays(
132
+ arr.indices.fill_null(len(arr.dictionary)),
133
+ pa.concat_arrays([arr.dictionary.cast(pa.string()), pa.array([""])]))
134
+ raw = ["" if v is None else str(v) for v in arr.dictionary.to_pylist()]
135
+ order = sorted(range(len(raw)), key=lambda i: raw[i])
136
+ names, remap = [], np.empty(len(raw), dtype=np.int64)
137
+ for i in order: # merge duplicate names (a unified dictionary can repeat after casting)
138
+ if not names or names[-1] != raw[i]:
139
+ names.append(raw[i])
140
+ remap[i] = len(names) - 1
141
+ codes = remap[arr.indices.to_numpy(zero_copy_only=False).astype(np.int64)]
142
+ return codes, names
143
+
144
+
145
+ def _numeric(column: pa.ChunkedArray, dtype=None) -> np.ndarray:
146
+ """Column as a numpy array, float nulls as NaN, cast to `dtype` if given."""
147
+ arr = column.combine_chunks()
148
+ if arr.null_count and pa.types.is_floating(arr.type):
149
+ arr = arr.fill_null(float("nan"))
150
+ out = arr.to_numpy(zero_copy_only=False)
151
+ return out.astype(dtype) if dtype is not None else out
152
+
153
+
154
+ def _is_categorical(t: pa.DataType) -> bool:
155
+ """String-like columns, stored as codes plus a `<column>_categories` array."""
156
+ return pa.types.is_string(t) or pa.types.is_large_string(t) or pa.types.is_dictionary(t)
157
+
158
+
159
+ def build(source, output: str, *, gene_col: str, x_col: str, y_col: str,
160
+ id_col: str | None = None, columns: Sequence[str] = (),
161
+ P: int = DEFAULTS["P"], K: int = DEFAULTS["K"], g: int = DEFAULTS["g"],
162
+ seed: int = DEFAULTS["seed"], chunk: int = DEFAULTS["chunk"],
163
+ shard: int = DEFAULTS["shard"], index_location: str = DEFAULTS["index_location"],
164
+ original_index: bool = True, tmpdir: str | None = None,
165
+ extra_attrs: dict | None = None, log: Callable[[str], None] = print) -> dict:
166
+ """Write the store to `output` (replaced atomically via `output.tmp`).
167
+ `source` is a parquet file, a directory of parquet parts, or a pyarrow Table.
168
+ Returns a summary dict (also logged)."""
169
+ if shard % chunk:
170
+ raise BuildError(f"shard ({shard}) must be a multiple of chunk ({chunk})")
171
+ if index_location not in ("start", "end"):
172
+ raise BuildError("index_location must be 'start' or 'end'")
173
+ if not 1 <= K <= 16:
174
+ raise BuildError("K must be in 1..16 (h is stored in uint32)")
175
+ columns = list(columns)
176
+ bad = RESERVED.intersection(columns)
177
+ if bad:
178
+ raise BuildError(f"column name(s) {sorted(bad)} are reserved")
179
+ t0 = time.time()
180
+
181
+ read_cols = list(dict.fromkeys([gene_col, x_col, y_col, *([id_col] if id_col else []),
182
+ *columns]))
183
+ tbl = _read_table(source, read_cols)
184
+ n = tbl.num_rows
185
+ gene_codes, genes = _dictionary(tbl.column(gene_col))
186
+ # Quantize the float32 values actually stored, so verify (which only sees the
187
+ # stored coords) recomputes exactly the h the build sorted by.
188
+ x = _numeric(tbl.column(x_col), np.float32)
189
+ y = _numeric(tbl.column(y_col), np.float32)
190
+ if n and not (np.isfinite(x).all() and np.isfinite(y).all()):
191
+ raise BuildError("x/y contain NaN or infinite values")
192
+ ids = _numeric(tbl.column(id_col)).astype(np.uint64) if id_col else None
193
+
194
+ extras: dict[str, np.ndarray] = {}
195
+ categories: dict[str, list[str]] = {}
196
+ for c in columns:
197
+ col = tbl.column(c)
198
+ if _is_categorical(col.type):
199
+ codes, names = _dictionary(col)
200
+ extras[c] = codes.astype(np.uint32)
201
+ categories[c] = names
202
+ else:
203
+ extras[c] = _numeric(col)
204
+ del tbl
205
+
206
+ x0 = float(x.min()) if n else 0.0
207
+ y0 = float(y.min()) if n else 0.0
208
+ span = max(float(x.max()) - x0, float(y.max()) - y0) if n else 0.0
209
+ size = span * (1 + 1e-9) if span > 0 else 1.0
210
+
211
+ G = len(genes)
212
+ counts = np.bincount(gene_codes, minlength=G).astype(np.int64)
213
+ dmax = choose_dmax(counts, P)
214
+ if dmax + g > K:
215
+ raise BuildError(f"index depth {dmax + g} exceeds K={K}")
216
+ B = block_offsets(g, dmax)
217
+ S = int(B[-1])
218
+ gene_offsets = np.concatenate([[0], np.cumsum(counts)]).astype(np.uint64)
219
+ order_g = np.argsort(gene_codes, kind="stable")
220
+ del gene_codes
221
+
222
+ tmp_root = tempfile.mkdtemp(prefix="hlod-", dir=tmpdir)
223
+ try:
224
+ def mm(name, dtype, shape):
225
+ return np.lib.format.open_memmap(os.path.join(tmp_root, f"{name}.npy"), mode="w+",
226
+ dtype=dtype, shape=shape)
227
+
228
+ out_coords = mm("coords", np.float32, (n, 2))
229
+ out_orig = mm("original_index", np.uint64, (n,)) if original_index else None
230
+ out_extra = {c: mm(f"col_{i}", a.dtype, (n,)) for i, (c, a) in enumerate(extras.items())}
231
+ level_offsets = np.zeros((G, dmax + 2), dtype=np.uint64)
232
+ cell_index = np.zeros((G, S), dtype=np.uint32)
233
+ level_rows = np.zeros((G, dmax + 1), dtype=np.int64)
234
+
235
+ for gi in range(G):
236
+ a, b = int(gene_offsets[gi]), int(gene_offsets[gi + 1])
237
+ idx = order_g[a:b]
238
+ gx, gy = x[idx], y[idx]
239
+ h = xy2d(K, quantize(gx, x0, size, K), quantize(gy, y0, size, K))
240
+ key = ids[idx] if ids is not None else gene_xy_key(gi, gx, gy)
241
+ r = splitmix64(key ^ np.uint64(seed))
242
+ level = assign_levels(r, b - a, P, dmax)
243
+ order = np.lexsort((h, level))
244
+ h, level, idx = h[order], level[order], idx[order]
245
+
246
+ out_coords[a:b, 0] = gx[order]
247
+ out_coords[a:b, 1] = gy[order]
248
+ if out_orig is not None:
249
+ out_orig[a:b] = idx
250
+ for c, arr in extras.items():
251
+ out_extra[c][a:b] = arr[idx]
252
+
253
+ per_level = np.bincount(level, minlength=dmax + 1)
254
+ level_rows[gi] = per_level
255
+ rel = np.concatenate([[0], np.cumsum(per_level)])
256
+ level_offsets[gi] = a + rel
257
+ for d in range(dmax + 1):
258
+ m = d + g
259
+ cells = (h[rel[d]:rel[d + 1]] >> np.uint32(2 * (K - m))).astype(np.int64)
260
+ hist = np.bincount(cells, minlength=4 ** m)
261
+ cell_index[gi, B[d]:B[d + 1]] = rel[d] + np.concatenate([[0], np.cumsum(hist)])
262
+
263
+ attrs = {
264
+ "format": FORMAT,
265
+ "format_version": FORMAT_VERSION,
266
+ "K": K, "g": g, "P": P, "Dmax": dmax,
267
+ "extent": {"x0": x0, "y0": y0, "size": size},
268
+ "hash": {"algo": "splitmix64", "seed": seed,
269
+ "key": "transcript_id" if id_col else "gene_xy"},
270
+ "genes": genes,
271
+ "columns": columns,
272
+ "categories": {c: f"{c}_categories" for c in categories},
273
+ "coordinate_units": "micron",
274
+ "chunk": chunk, "shard": shard, "index_location": index_location,
275
+ "source": {"path": source if isinstance(source, str) else None, "rows": n},
276
+ **(extra_attrs or {}),
277
+ }
278
+ _write_store(output, attrs, tmp_root, out_coords, out_orig, out_extra, categories,
279
+ gene_offsets, level_offsets, cell_index, chunk, shard, index_location)
280
+ finally:
281
+ shutil.rmtree(tmp_root, ignore_errors=True)
282
+
283
+ index_bytes = S * 4 + (dmax + 2) * 8 + 8
284
+ largest = int(np.diff(cell_index.astype(np.int64), axis=1).max()) if G and S > 1 else 0
285
+ summary = {"path": output, "rows": n, "genes": G, "Dmax": dmax, "S": S,
286
+ "index_bytes_per_gene": index_bytes, "largest_cell_rows": largest,
287
+ "rows_per_level": level_rows.sum(axis=0).tolist(),
288
+ "seconds": round(time.time() - t0, 1)}
289
+ log(f"hlod: {n:,} rows, {G} genes, Dmax={dmax}, rows/level={summary['rows_per_level']}, "
290
+ f"index {index_bytes:,} B/gene, largest cell {largest:,} rows, {summary['seconds']}s")
291
+ return summary
292
+
293
+
294
+ def _sharded(group, name, data: np.ndarray, chunk_rows: int, shard_rows: int,
295
+ index_location: str, compress: bool):
296
+ """Write `data` as a sharded zarr array, chunked along rows. Chunk and shard
297
+ sizes are clamped to the array."""
298
+ rows = data.shape[0]
299
+ chunk_rows = max(1, min(chunk_rows, rows))
300
+ # Whole multiple of the chunk; a shard larger than the array just holds every chunk.
301
+ shard_rows = max(chunk_rows, min(shard_rows, -(-rows // chunk_rows) * chunk_rows))
302
+ shard_rows -= shard_rows % chunk_rows
303
+ tail = data.shape[1:]
304
+ arr = group.create_array(
305
+ name, shape=data.shape, dtype=data.dtype,
306
+ chunks=(chunk_rows, *tail), shards={"shape": (shard_rows, *tail),
307
+ "index_location": index_location},
308
+ compressors=ZstdCodec(level=3) if compress else None,
309
+ # Keep every chunk present so uncompressed byte offsets stay contiguous.
310
+ config={"write_empty_chunks": True},
311
+ )
312
+ # Whole shards at a time: a partial shard write is a read-modify-write.
313
+ for s in range(0, rows, shard_rows):
314
+ arr[s:s + shard_rows] = np.asarray(data[s:s + shard_rows])
315
+ return arr
316
+
317
+
318
+ def _write_store(output, attrs, tmp_root, coords, orig, extra, categories, gene_offsets,
319
+ level_offsets, cell_index, chunk, shard, index_location):
320
+ """Write every array and the attrs into `output`.tmp, consolidate the metadata,
321
+ then move it to `output`."""
322
+ staging = output.rstrip("/") + ".tmp"
323
+ shutil.rmtree(staging, ignore_errors=True)
324
+ root = zarr.open_group(staging, mode="w", zarr_format=3)
325
+ _sharded(root, "coords", coords, chunk, shard, index_location, compress=False)
326
+ if orig is not None:
327
+ _sharded(root, "original_index", orig, chunk, shard, index_location, compress=True)
328
+ for c, arr in extra.items():
329
+ _sharded(root, c, arr, chunk, shard, index_location, compress=True)
330
+ if c in categories:
331
+ names = np.array(categories[c], dtype=object)
332
+ cat = root.create_array(f"{c}_categories", shape=names.shape, dtype=str,
333
+ chunks=(max(1, min(names.size, 1 << 20)),),
334
+ compressors=ZstdCodec(level=3))
335
+ if names.size:
336
+ cat[:] = names
337
+ index = root.create_group("index")
338
+ G, S = cell_index.shape
339
+ # Uncompressed so the client reads them with the same byte-range reader as coords.
340
+ _sharded(index, "gene_offsets", gene_offsets, gene_offsets.size, gene_offsets.size,
341
+ index_location, compress=False)
342
+ _sharded(index, "level_offsets", level_offsets, max(1, G), max(1, G), index_location,
343
+ compress=False)
344
+ _sharded(index, "cell_index", cell_index, 1, max(1, G), index_location, compress=False)
345
+ root.attrs.update(attrs)
346
+ zarr.consolidate_metadata(staging)
347
+ if os.path.exists(output):
348
+ shutil.rmtree(output)
349
+ os.replace(staging, output)
spacelab/hlod/cells.py ADDED
@@ -0,0 +1,82 @@
1
+ """Quadtree-cell lookups shared by build, read and verify; src/hlod/cells.ts is
2
+ a port of this module and both are tested against the same fixtures."""
3
+
4
+ from dataclasses import dataclass, field
5
+
6
+ import numpy as np
7
+
8
+ from .hilbert import xy2d
9
+
10
+ DEFAULT_MAX_GAP_ROWS = 32_768 # two chunks
11
+
12
+
13
+ @dataclass(frozen=True)
14
+ class Extent:
15
+ x0: float
16
+ y0: float
17
+ size: float
18
+
19
+ @classmethod
20
+ def from_attrs(cls, attrs: dict) -> "Extent":
21
+ e = attrs["extent"]
22
+ return cls(float(e["x0"]), float(e["y0"]), float(e["size"]))
23
+
24
+
25
+ def block_offsets(g: int, dmax: int) -> np.ndarray:
26
+ """B[d]: where level d's block starts in a gene's cell_index row; B[dmax + 1]
27
+ is the row length S."""
28
+ sizes = [4 ** (d + g) + 1 for d in range(dmax + 1)]
29
+ return np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64)
30
+
31
+
32
+ def cell_coord(v: float, v0: float, size: float, depth: int) -> int:
33
+ n = 1 << depth
34
+ return int(min(max(np.floor((v - v0) / size * n), 0), n - 1))
35
+
36
+
37
+ def cells_in_box(box, depth: int, extent: Extent) -> np.ndarray:
38
+ """Hilbert indices of every depth-`depth` cell the box touches, ascending.
39
+ A box entirely outside the extent touches none; one partly outside is
40
+ clamped to it."""
41
+ xmin, ymin, xmax, ymax = box
42
+ x1 = extent.x0 + extent.size
43
+ y1 = extent.y0 + extent.size
44
+ if xmax < extent.x0 or ymax < extent.y0 or xmin > x1 or ymin > y1 or xmin > xmax or ymin > ymax:
45
+ return np.empty(0, dtype=np.int64)
46
+ i0, i1 = (cell_coord(v, extent.x0, extent.size, depth) for v in (xmin, xmax))
47
+ j0, j1 = (cell_coord(v, extent.y0, extent.size, depth) for v in (ymin, ymax))
48
+ ii, jj = np.meshgrid(np.arange(i0, i1 + 1), np.arange(j0, j1 + 1), indexing="ij")
49
+ return np.sort(xy2d(depth, ii.ravel(), jj.ravel()).astype(np.int64))
50
+
51
+
52
+ @dataclass
53
+ class Range:
54
+ start: int # absolute rows, [start, end)
55
+ end: int
56
+ cells: list[int] = field(default_factory=list)
57
+
58
+
59
+ def ranges_for(cell_index_row: np.ndarray, gene_start: int, level: int, cells, g: int,
60
+ dmax: int) -> list[Range]:
61
+ """Absolute row range of each listed cell of `level`, empty cells dropped."""
62
+ b = int(block_offsets(g, dmax)[level])
63
+ out = []
64
+ for c in cells:
65
+ s = int(cell_index_row[b + c])
66
+ e = int(cell_index_row[b + c + 1])
67
+ if e > s:
68
+ out.append(Range(gene_start + s, gene_start + e, [int(c)]))
69
+ return out
70
+
71
+
72
+ def coalesce(ranges: list[Range], max_gap_rows: int = DEFAULT_MAX_GAP_ROWS) -> list[Range]:
73
+ """Merge ranges (sorted by start) whose gap is <= max_gap_rows; each merged
74
+ range keeps the cells it covers so a read can be split back per cell."""
75
+ out: list[Range] = []
76
+ for r in sorted(ranges, key=lambda r: r.start):
77
+ if out and r.start - out[-1].end <= max_gap_rows:
78
+ out[-1].end = max(out[-1].end, r.end)
79
+ out[-1].cells.extend(r.cells)
80
+ else:
81
+ out.append(Range(r.start, r.end, list(r.cells)))
82
+ return out
spacelab/hlod/cli.py ADDED
@@ -0,0 +1,90 @@
1
+ """spacelab hlod build | verify | fixtures"""
2
+
3
+ import argparse
4
+ import json
5
+ import random
6
+ import sys
7
+
8
+ from .build import DEFAULTS, BuildError, build
9
+ from .hilbert import xy2d_scalar
10
+
11
+
12
+ def _build(args) -> int:
13
+ columns = [c for c in (args.columns or "").split(",") if c]
14
+ try:
15
+ build(args.input, args.output, gene_col=args.gene_col, x_col=args.x_col,
16
+ y_col=args.y_col, id_col=args.id_col or None, columns=columns, P=args.P, K=args.K,
17
+ g=args.g, seed=args.seed, chunk=args.chunk, shard=args.shard,
18
+ index_location=args.index_location, tmpdir=args.tmpdir)
19
+ except BuildError as e:
20
+ print(f"error: {e}", file=sys.stderr)
21
+ return 1
22
+ return 0
23
+
24
+
25
+ def _verify(args) -> int:
26
+ from .verify import verify
27
+
28
+ rep = verify(args.store, roi_samples=args.roi_samples, seed=args.seed)
29
+ for note in rep.notes:
30
+ print(note)
31
+ for f in rep.failures:
32
+ print(f"FAIL {f}", file=sys.stderr)
33
+ print("verify: ok" if rep.ok else f"verify: {len(rep.failures)} failure(s)")
34
+ return 0 if rep.ok else 1
35
+
36
+
37
+ def fixtures(n: int = 1000, seed: int = 0) -> list[dict]:
38
+ rng = random.Random(seed)
39
+ out = []
40
+ for i in range(n):
41
+ order = 16 if i < 100 else rng.randint(1, 16)
42
+ side = 1 << order
43
+ x, y = rng.randrange(side), rng.randrange(side)
44
+ out.append({"order": order, "x": x, "y": y, "d": xy2d_scalar(order, x, y)})
45
+ return out
46
+
47
+
48
+ def _fixtures(args) -> int:
49
+ with open(args.out, "w") as f:
50
+ json.dump(fixtures(), f, separators=(",", ":"))
51
+ f.write("\n")
52
+ print(f"wrote {args.out}")
53
+ return 0
54
+
55
+
56
+ def main(argv=None) -> int:
57
+ ap = argparse.ArgumentParser(prog="spacelab hlod",
58
+ description="Hilbert additive-LOD transcript store")
59
+ sub = ap.add_subparsers(dest="cmd", required=True)
60
+
61
+ b = sub.add_parser("build", help="build a store from a parquet table")
62
+ b.add_argument("--input", required=True, help="parquet file or directory of parts")
63
+ b.add_argument("--output", required=True)
64
+ b.add_argument("--gene-col", default="feature_name")
65
+ b.add_argument("--x-col", default="x")
66
+ b.add_argument("--y-col", default="y")
67
+ b.add_argument("--id-col", default="transcript_id", help="'' to hash (gene, x, y) instead")
68
+ b.add_argument("--columns", default="", help="comma-separated co-sorted extra columns")
69
+ for flag in ("P", "K", "g", "seed", "chunk", "shard"):
70
+ b.add_argument(f"--{flag}", type=int, default=DEFAULTS[flag])
71
+ b.add_argument("--index-location", choices=["start", "end"], default=DEFAULTS["index_location"])
72
+ b.add_argument("--tmpdir", default=None)
73
+ b.set_defaults(fn=_build)
74
+
75
+ v = sub.add_parser("verify", help="check a built store; exits non-zero on failure")
76
+ v.add_argument("store")
77
+ v.add_argument("--roi-samples", type=int, default=50)
78
+ v.add_argument("--seed", type=int, default=7)
79
+ v.set_defaults(fn=_verify)
80
+
81
+ f = sub.add_parser("fixtures", help="write Hilbert test vectors for the TypeScript tests")
82
+ f.add_argument("--out", required=True)
83
+ f.set_defaults(fn=_fixtures)
84
+
85
+ args = ap.parse_args(argv)
86
+ return args.fn(args)
87
+
88
+
89
+ if __name__ == "__main__":
90
+ sys.exit(main())