Use cohorts for machine learning

Use cohorts for machine learning#

This guide goes from a cohort to training-ready data: implement a small TaskAdapter, pull one concatenated AnnData table out of several stores without downloading images, and loop over the train fold of a study-level split. It uses four small stores (about 10 MB each):

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

import sddb

cat = sddb.Catalog()
cohort = cat.cohort(uids=["xVK6T0rajFTgOKci0000", "cgmLrEw3hgv7yhZf0000", "gtLUT2zUAH9l9Uzi0002", "cdJ4ig3BNiTZMC3T0002"])
cohort

<SpatialDataCohort: 4 datasets>

One table across many stores#

to_anndata(table=...) reads only each member’s tables/<name> group (no images, no shapes) and concatenates them with an outer join. obs["uid"] says which dataset a row came from and obs["study_id"] which study.

visium = cohort.filter(technology="Visium")
adata = visium.to_anndata(table="table")
print(adata)
adata.obs["uid"].value_counts().reindex([d.uid for d in visium])

AnnData object with n_obs × n_vars = 5662 × 37532
    obs: 'in_tissue', 'array_row', 'array_col', 'spot_id', 'region', 'uid', 'study_id'
    obsm: 'spatial'
    layers: 'csc_counts'

uid
gtLUT2zUAH9l9Uzi0002     459
cdJ4ig3BNiTZMC3T0002     211
cgmLrEw3hgv7yhZf0000    4992
Name: count, dtype: int64

Open stores as SpatialData#

load(elements=[...]) opens every member as SpatialData but copies only the listed elements. sdatas opens every member completely (download first, so mind the size).

sdatas = visium.load(elements=["tables/table"])
{d.uid: s.tables["table"].shape for d, s in zip(visium, sdatas, strict=True)}

{'gtLUT2zUAH9l9Uzi0002': (459, 19465),
 'cdJ4ig3BNiTZMC3T0002': (211, 19465),
 'cgmLrEw3hgv7yhZf0000': (4992, 18085)}

full = visium[:1].sdatas[0]
print("images:", len(full.images), "| tables:", list(full.tables))

images: 3 | tables: ['table']

A minimal TaskAdapter#

TaskAdapter is a protocol: the client defines the contract and your task repository implements it. An adapter declares its name, version and config, states its requirements(), and provides three hooks:

  • validate_meta(meta): cheap, metadata-only check (advisory), called by the task framework with an ObjectMeta (the element inventory, read without any array chunk).

  • validate(sdata): authoritative check on the opened object.

  • build(sdata): produce the task item for one object.

Each hook returns a list of issue strings, empty when all is well (build returns the item).

from typing import Any

from sddb import TaskAdapter


class SpotSummaryAdapter:
    """Per-slide summary: number of spots, genes and total counts."""

    name = "spot-summary"
    version = "0.1.0"

    def __init__(self, min_spots: int = 1_000) -> None:
        self.config: dict[str, Any] = {"min_spots": min_spots}

    def requirements(self) -> dict[str, Any]:
        return {"tables": 1, "elements": ["tables/*"]}

    def validate_meta(self, meta) -> list[str]:
        return [] if meta.tables else ["no table in store"]

    def validate(self, sdata) -> list[str]:
        table = next(iter(sdata.tables.values()))
        n = self.config["min_spots"]
        return [] if table.n_obs >= n else [f"only {table.n_obs} spots (< {n})"]

    def build(self, sdata) -> dict[str, float]:
        table = next(iter(sdata.tables.values()))
        return {"n_spots": table.n_obs, "n_genes": table.n_vars, "total_counts": float(table.X.sum())}


adapter = SpotSummaryAdapter()
isinstance(adapter, TaskAdapter)

True

Loop over a split’s train fold#

split() is group-safe: every study lands in exactly one fold. Datasets without a study_id form their own group. Each fold is itself a cohort, so you can iterate it, check each object with the adapter and build items. Members keep catalog order throughout this guide, and outputs are keyed by uid.

folds = cohort.split(by="study_id", train=0.5, val=0.25, test=0.25, seed=0)
{name: [d.uid for d in fold] for name, fold in folds.items()}

{'train': ['cdJ4ig3BNiTZMC3T0002', 'cgmLrEw3hgv7yhZf0000'],
 'val': ['xVK6T0rajFTgOKci0000'],
 'test': ['gtLUT2zUAH9l9Uzi0002']}

items = {}
for ds in folds["train"]:
    table = ds.default_table or "table"
    sdata = ds.load(lazy=False, elements=[f"tables/{table}"])
    issues = adapter.validate(sdata)
    if issues:
        print(ds.uid, "skipped:", issues)
        continue
    items[ds.uid] = adapter.build(sdata)
items

cdJ4ig3BNiTZMC3T0002 skipped: ['only 211 spots (< 1000)']

{'cgmLrEw3hgv7yhZf0000': {'n_spots': 4992,
  'n_genes': 18085,
  'total_counts': 1069313.0}}

Because the split manifest travels with the fold, you can store it next to your results and recover the exact assignment later (see the cohorts guide).

folds["train"].split_manifest.assignments

{'gtLUT2zUAH9l9Uzi0002': 'test',
 'cdJ4ig3BNiTZMC3T0002': 'train',
 'cgmLrEw3hgv7yhZf0000': 'train',
 'xVK6T0rajFTgOKci0000': 'val'}

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

Catalog generated at: 2026-10-18