Build, freeze and split cohorts

Build, freeze and split cohorts#

A SpatialDataCohort is the unit you work with: an ordered set of dataset records you can combine, pin to an exact membership, download reproducibly and split into train, validation and test folds without leaking a study across folds. Everything below runs against the live catalog.

pip install git+https://github.com/theislab/spatialdata-db-client

import tempfile
from pathlib import Path

import sddb
from sddb import CohortDriftError, CohortManifest, Manifest, SpatialDataCohort, SplitManifest

cat = sddb.Catalog()
work = Path(tempfile.mkdtemp())  # scratch directory for the files this guide writes

Combine and refine cohorts#

Cohorts support | (union, deduplicated by uid) and - (difference). filter() keeps rows matching column=value (a list means any of), and where() keeps rows where a boolean mask over .members is true.

visium = cat.query(organism="human", tissue="lung", technology="Visium")
visium_hd = cat.query(organism="human", tissue="lung", technology="VisiumHD")

both = visium | visium_hd
print(len(visium), "+", len(visium_hd), "->", len(both))
print("difference:", len(both - visium_hd), "Visium-only")

37 + 7 -> 44
difference: 37 Visium-only

cancer = both.filter(disease="lung cancer")
deep = both.where(both.members["n_obs"] > 20_000)
print(len(cancer), "lung cancer;", len(deep), "with more than 20,000 observations")

6 lung cancer; 6 with more than 20,000 observations

Freeze a cohort#

freeze() pins the selection to its exact members (full uids, version suffix included) and gives it a deterministic hash. It is offline and cheap. We use four small stores (about 10 MB each):

small = cat.cohort(uids=["xVK6T0rajFTgOKci0000", "cgmLrEw3hgv7yhZf0000", "gtLUT2zUAH9l9Uzi0002", "cdJ4ig3BNiTZMC3T0002"])
manifest = small.freeze()
print(manifest.hash)
print(manifest.filter)

sha256:fa97d3eadb3b3a9fed4708dc372ece91e83164644721274f5ec36d8b0128e4e9
{'uid': ['cdJ4ig3BNiTZMC3T0002', 'cgmLrEw3hgv7yhZf0000', 'gtLUT2zUAH9l9Uzi0002', 'xVK6T0rajFTgOKci0000']}

CohortManifest.write() saves JSON (large member lists spill into a parquet sidecar), CohortManifest.load() reads it back and SpatialDataCohort.from_manifest() rebuilds exactly the pinned members, offline, with no catalog query.

path = manifest.write(work / "cohort.json")
restored = SpatialDataCohort.from_manifest(CohortManifest.load(path))
print(restored)
print("same members:", list(restored.members["uid"]) == list(small.members["uid"]))
print("same hash:", restored.freeze().hash == manifest.hash)

<SpatialDataCohort: 4 datasets>
same members: True
same hash: True

Pass catalog= to from_manifest to also be warned about members that a newer catalog supersedes.

restored = SpatialDataCohort.from_manifest(CohortManifest.load(path), catalog=cat)

Check a cohort is still valid#

  • verify() is a liveness check: it raises CohortDriftError if a pinned store no longer resolves.

  • check_updates(catalog) returns {pinned_uid: current_uid} for members a newer version supersedes.

small.verify()  # returns None when every store resolves
small.check_updates(cat)  # {} means nothing has been superseded

{}

To see the failure mode, point one member at a store that does not exist.

broken = small.members.copy()
broken.loc[0, "zarr_url"] = broken.loc[0, "zarr_url"].replace(".zarr", "-missing.zarr")
try:
    SpatialDataCohort(broken).verify()
except CohortDriftError as err:
    print("CohortDriftError:", err.drifted)

CohortDriftError: {'gtLUT2zUAH9l9Uzi0002': 'store does not resolve'}

Download reproducibly#

download(dest) copies every store to dest/<uid>.zarr (resumable, parallel) and writes dest/manifest.json. pin_versions=True records the catalog version in that manifest, so downloading into a directory pinned to another catalog version raises ManifestVersionMismatch unless you pass allow_version_change=True. Pinning only records a version when the cohort came from a versioned catalog (Catalog(version=...)); for the current catalog it is recorded as None.

Two of the small mouse stores keep this fast.

mouse = small.filter(organism="mouse")
dl = mouse.download(work / "data", pin_versions=True)
print("catalog_version:", dl.catalog_version)
for e in dl.entries:
    print(e.uid, e.status, f"{e.size_bytes / 1e6:.1f} MB")

catalog_version: None
gtLUT2zUAH9l9Uzi0002 complete 10.5 MB
cdJ4ig3BNiTZMC3T0002 complete 12.3 MB

A Manifest records the exact URLs that were fetched. refetch() replays them into a new directory, using the recorded URLs rather than the current catalog.

recorded = Manifest.read(work / "data" / "manifest.json")
copy = recorded.refetch(work / "data_copy")
print([(e.uid, e.status) for e in copy.entries])

[('gtLUT2zUAH9l9Uzi0002', 'complete'), ('cdJ4ig3BNiTZMC3T0002', 'complete')]

Only URLs are replayed: a store republished at the same URL is not detected.

Split without leaking studies#

split(by="study_id", train=, val=, test=, seed=) assigns whole studies to one fold, so no study appears in two folds. It is deterministic for a given seed, independent of row order, and refuses to create empty folds. Datasets without a study_id become their own group. The result is {fold: SpatialDataCohort}.

folds = both.split(by="study_id", train=0.8, val=0.1, test=0.1, seed=42)
print({name: len(fold) for name, fold in folds.items()})

studies = {name: set(fold.members["study_id"].dropna()) for name, fold in folds.items()}
print("study overlap train/test:", studies["train"] & studies["test"])

{'train': 36, 'val': 4, 'test': 4}
study overlap train/test: set()

Every fold carries the same SplitManifest, pinned to the parent cohort’s hash. Write it next to your experiment and load it later to recover the identical assignment.

sm = folds["train"].split_manifest
sm.write(work / "split.json")
again = SplitManifest.load(work / "split.json")
print(again.by, again.seed, again.ratios)
print("parent hash matches cohort:", again.parent_hash == both.freeze().hash)
print("same assignments:", again.assignments == sm.assignments)

study_id 42 {'test': 0.1, 'train': 0.8, 'val': 0.1}
parent hash matches cohort: True
same assignments: True

print("Catalog generated at:", cat.generated_at)

Catalog generated at: 2026-10-18