spatialscape 0.3.0__py3-none-any.whl

spatialscape/manifest.py

@@ -0,0 +1,162 @@
+"""Output schema: manifest.json written into every bundle.
+
+Field names are camelCase on purpose so the JSON matches the TypeScript types in
+app/src/data/manifest.ts one-to-one.
+"""
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+FORMAT_VERSION = 1
+
+Platform = Literal["visium", "visium_hd", "xenium", "merfish", "snrnaseq", "other"]
+SampleKind = Literal["spatial", "embedding"]
+LayoutMode = Literal["grid", "strip"]
+BUILTIN_COLORMAPS = ["viridis", "magma", "inferno", "plasma", "cividis", "turbo", "greys", "rdbu"]
+
+
+class _Model(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class Vocabulary(_Model):
+    categories: list[str]
+    colors: list[str]
+    aliases: dict[str, str] = Field(default_factory=dict)
+
+
+class CategoricalField(_Model):
+    id: str
+    name: str
+    type: Literal["categorical"] = "categorical"
+    vocabulary: str
+    description: str | None = None
+
+
+class ContinuousField(_Model):
+    id: str
+    name: str
+    type: Literal["continuous"] = "continuous"
+    range: tuple[float, float] | None = None
+    colormap: str | None = None
+    description: str | None = None
+
+
+ObsField = Annotated[CategoricalField | ContinuousField, Field(discriminator="type")]
+
+
+class ColorSpec(_Model):
+    kind: Literal["gene", "field"]
+    gene: str | None = None
+    field: str | None = None
+
+
+class Layout(_Model):
+    mode: LayoutMode = "grid"
+    gutterFraction: float = 0.1
+    order: list[str] = Field(default_factory=list)
+    columns: int | None = None  # grid: a fixed number of columns
+
+
+class ImageChannel(_Model):
+    name: str
+    color: str
+    window: tuple[float, float]
+
+
+class SampleImage(_Model):
+    id: str
+    name: str
+    path: str
+    kind: Literal["rgb", "multichannel"] = "rgb"
+    pixelSize: float
+    translate: tuple[float, float] = (0.0, 0.0)
+    size: tuple[int, int]  # level-0 (height, width) in pixels
+    channels: list[ImageChannel] | None = None
+    defaultOpacity: float = 1.0
+
+
+class ExprSpec(_Model):
+    kind: Literal["u8", "csc"] = "u8"
+    sharded: bool = True
+    shardGenes: int = 512
+    normalized: Literal["lognorm", "log1p_cp10k"] | None = None  # as stored by the source, or log1p(CP10k) applied by the build
+
+
+class PolygonInfo(_Model):
+    path: str = "polygons"  # polygons.i16 + polygons.offsets.u32
+    cells: int = 0  # cells that have a polygon
+    vertices: int = 0
+    scale: float = 0.1  # µm per int16 unit (deltas from the cell centroid)
+
+
+class TrackSet(_Model):
+    field: str
+    categories: dict[str, str]  # category label -> bigWig path relative to the sample folder
+
+
+class TracksInfo(_Model):
+    """Genome tracks of an ATAC sample (see tracks.py): bigWigs of insertions per million per category."""
+
+    genome: str
+    bin: int
+    all: str | None = None  # bigWig of every cell: built from fragments, or given with bigwigs
+    units: str = "insertions per million"  # what the tracks measure, for the panel's scale
+    by: list[TrackSet] = Field(default_factory=list)
+    peaks: str | None = None  # BED of the sample's peak features
+    genes: str | None = None  # folder of gene models: index.json (names per chromosome) + <chrom>.json (see genes.py)
+    links: str | None = None  # BEDPE of peak-gene links
+    chromSizes: str = "tracks/chrom.sizes"
+
+
+class Sample(_Model):
+    id: str
+    name: str
+    platform: Platform
+    kind: SampleKind = "spatial"
+    group: str | None = None
+    nObs: int
+    nGenes: int
+    bbox: tuple[float, float, float, float]
+    toDataset: list[float] = Field(default_factory=lambda: [1, 0, 0, 0, 1, 0, 0, 0, 1])
+    pointRadius: float
+    expr: ExprSpec = Field(default_factory=ExprSpec)
+    hasF16: bool = False
+    fields: list[str] = Field(default_factory=list)
+    images: list[SampleImage] = Field(default_factory=list)
+    outlines: list[str] = Field(default_factory=list)  # categorical field ids with outlines/<field>.json
+    polygons: PolygonInfo | None = None
+    idBlock: int = 65536
+    suggestedGene: str | None = None  # most variable gene in this sample; the dataset's default when none is configured
+    tracks: TracksInfo | None = None  # genome tracks (ATAC samples)
+
+
+class FeatureGroup(_Model):
+    id: str
+    name: str
+    units: str | None = None
+    count: int = 0
+
+
+class Manifest(_Model):
+    formatVersion: Literal[1] = FORMAT_VERSION
+    id: str
+    name: str
+    description: str | None = None
+    defaultGene: str | None = None
+    defaultColor: ColorSpec
+    layout: Layout = Field(default_factory=Layout)
+    colormaps: list[str] = Field(default_factory=lambda: list(BUILTIN_COLORMAPS))
+    vocabularies: dict[str, Vocabulary] = Field(default_factory=dict)
+    fields: list[ObsField] = Field(default_factory=list)
+    featureGroups: list[FeatureGroup] = Field(default_factory=list)
+    samples: list[Sample] = Field(default_factory=list)
+    thumbnailSample: str | None = None  # sample drawn on the dataset thumbnail (default: first spatial sample in layout order)
+
+    def field_by_id(self, fid: str) -> CategoricalField | ContinuousField | None:
+        for f in self.fields:
+            if f.id == fid:
+                return f
+        return None

spatialscape/obs.py

@@ -0,0 +1,48 @@
+"""Per-observation arrays: positions, shuffle order, ids, categorical codes, continuous values."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import zarr
+from zarr.codecs import ZstdCodec
+
+MAX_CHUNK = 1 << 20
+
+
+def _chunk(n: int) -> int:
+    return max(1, min(n, MAX_CHUNK))
+
+
+def write_obs(
+    out_dir: Path,
+    *,
+    xy: np.ndarray,
+    order: np.ndarray,
+    ids: list[str],
+    cats: dict[str, np.ndarray],
+    nums: dict[str, np.ndarray],
+    id_block: int = 65536,
+) -> None:
+    n = xy.shape[0]
+    store = zarr.storage.LocalStore(str(out_dir / "obs.zarr"))
+    root = zarr.open_group(store, mode="w", zarr_format=3)
+    z = ZstdCodec(level=5)
+    a = root.create_array("xy", shape=(n, 2), chunks=(_chunk(n), 2), dtype="float32", compressors=[z])
+    a[:] = xy.astype(np.float32)
+    o = root.create_array("order", shape=(n,), chunks=(_chunk(n),), dtype="uint32", compressors=[z])
+    o[:] = order.astype(np.uint32)
+    cg = root.create_group("cat")
+    for fid, codes in cats.items():
+        arr = cg.create_array(fid, shape=(n,), chunks=(_chunk(n),), dtype="uint16", compressors=[z])
+        arr[:] = codes.astype(np.uint16)
+    ng = root.create_group("num")
+    for fid, vals in nums.items():
+        arr = ng.create_array(fid, shape=(n,), chunks=(_chunk(n),), dtype="float32", compressors=[z])
+        arr[:] = vals.astype(np.float32)
+    ids_dir = out_dir / "ids"
+    ids_dir.mkdir(parents=True, exist_ok=True)
+    for b, start in enumerate(range(0, n, id_block)):
+        with open(ids_dir / f"{b}.json", "w") as fh:
+            json.dump(ids[start : start + id_block], fh, separators=(",", ":"))

spatialscape/outlines.py

@@ -0,0 +1,145 @@
+"""Boundary polylines for categorical fields (e.g. spatial domains), computed from a bundle's obs arrays.
+
+Approach: rasterize labels onto a grid (nearest point within ~0.75 x spacing), majority-smooth,
+trace each category with marching squares, simplify (RDP), and store as µm polylines.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+from scipy import ndimage
+from scipy.spatial import cKDTree
+
+MAX_GRID_CELLS = 2_000_000
+BACKGROUND = -1
+
+
+def _median_nn(xy: np.ndarray, sample: int = 3000, seed: int = 0) -> float:
+    rng = np.random.default_rng(seed)
+    idx = rng.choice(len(xy), size=min(sample, len(xy)), replace=False)
+    d, _ = cKDTree(xy).query(xy[idx], k=2)
+    return float(np.median(d[:, 1]))
+
+
+def _rdp(points: np.ndarray, tol: float) -> np.ndarray:
+    """Iterative Ramer-Douglas-Peucker for an (n, 2) polyline."""
+    n = len(points)
+    if n < 3:
+        return points
+    keep = np.zeros(n, dtype=bool)
+    keep[0] = keep[-1] = True
+    stack = [(0, n - 1)]
+    while stack:
+        a, b = stack.pop()
+        if b <= a + 1:
+            continue
+        seg = points[b] - points[a]
+        L = np.hypot(*seg)
+        pts = points[a + 1 : b]
+        if L == 0:
+            d = np.hypot(*(pts - points[a]).T)
+        else:
+            d = np.abs(seg[0] * (points[a][1] - pts[:, 1]) - (points[a][0] - pts[:, 0]) * seg[1]) / L
+        i = int(np.argmax(d))
+        if d[i] > tol:
+            k = a + 1 + i
+            keep[k] = True
+            stack.append((a, k))
+            stack.append((k, b))
+    return points[keep]
+
+
+def compute_outlines(
+    xy: np.ndarray,
+    codes: np.ndarray,
+    n_categories: int,
+    *,
+    skip_codes: set[int] | None = None,
+    bin_um: float | None = None,
+    min_cells: int = 25,
+    smooth_um: float = 150.0,
+    min_feature_um: float = 300.0,
+) -> dict:
+    """Return {"bin": bin_um, "paths": [{"c": code, "p": [[x, y], ...]}, ...]} in the xy units (µm)."""
+    from skimage.measure import find_contours
+
+    skip_codes = skip_codes or set()
+    xy = np.asarray(xy, dtype=np.float64)
+    codes = np.asarray(codes).astype(np.int64)
+    nn = _median_nn(xy)
+    # unlabeled cells (e.g. "NA") are treated as unknown: dropped so neighbours fill their space
+    if skip_codes:
+        keep = ~np.isin(codes, list(skip_codes))
+        if keep.sum() >= min_cells:
+            xy, codes = xy[keep], codes[keep]
+    b = float(bin_um) if bin_um else max(0.5, 0.5 * nn)
+    x0, y0 = xy.min(axis=0)
+    x1, y1 = xy.max(axis=0)
+    while ((x1 - x0) / b + 3) * ((y1 - y0) / b + 3) > MAX_GRID_CELLS:
+        b *= 1.5
+    w = int(np.ceil((x1 - x0) / b)) + 3
+    h = int(np.ceil((y1 - y0) / b)) + 3
+    gx = x0 - b + (np.arange(w) + 0.5) * b
+    gy = y0 - b + (np.arange(h) + 0.5) * b
+    GX, GY = np.meshgrid(gx, gy)
+    tree = cKDTree(xy)
+    # fill radius: 2x the median spacing closes gaps inside tissue but keeps real holes/edges as background
+    r = max(2.0 * nn, 0.75 * b * np.sqrt(2))
+    d, i = tree.query(np.column_stack([GX.ravel(), GY.ravel()]), k=1, distance_upper_bound=r)
+    lab = np.full(GX.size, BACKGROUND, dtype=np.int64)
+    ok = np.isfinite(d)
+    lab[ok] = codes[i[ok]]
+    lab = lab.reshape(h, w)
+
+    # majority smoothing over a (smooth_um)^2 window (one-hot sums, argmax); background competes too
+    counts = np.bincount(codes, minlength=n_categories)
+    cats = [c for c in range(n_categories) if counts[c] >= min_cells and c not in skip_codes]
+    smooth = int(max(3, round(smooth_um / b)) | 1)  # odd window size in bins
+    if smooth > 1:
+        best = np.zeros((h, w), dtype=np.float32)
+        best_lab = np.full((h, w), BACKGROUND, dtype=np.int64)
+        for c in [BACKGROUND, *cats]:
+            s = ndimage.uniform_filter((lab == c).astype(np.float32), size=smooth, mode="constant")
+            upd = s > best
+            best[upd] = s[upd]
+            best_lab[upd] = c
+        lab = best_lab
+
+    import warnings
+
+    from skimage.morphology import remove_small_holes, remove_small_objects
+
+    # drop islands/holes smaller than (min_feature_um)^2 so outlines follow regions, not single cells
+    min_px = int(max(16, (min_feature_um / b) ** 2))
+    paths = []
+    tol = 0.35 * b
+    for c in cats:
+        m = lab == c
+        with warnings.catch_warnings():  # skimage 0.26 renamed these parameters; both spellings still work
+            warnings.simplefilter("ignore", FutureWarning)
+            m = remove_small_objects(m, min_size=min_px)
+            m = remove_small_holes(m, area_threshold=min_px)
+        mask = np.pad(m.astype(np.float32), 1)
+        for cont in find_contours(mask, 0.5):
+            if len(cont) < 6:
+                continue
+            pts = np.column_stack([x0 - b + (cont[:, 1] - 1 + 0.5) * b, y0 - b + (cont[:, 0] - 1 + 0.5) * b])
+            pts = _rdp(pts, tol)
+            if len(pts) < 4:
+                continue
+            per = np.hypot(*np.diff(pts, axis=0).T).sum()
+            if per < 10 * b:
+                continue
+            paths.append({"c": int(c), "p": np.round(pts, 1).tolist()})
+    return {"bin": round(b, 3), "paths": paths}
+
+
+def write_outlines(out_dir: Path, field: str, data: dict) -> Path:
+    d = Path(out_dir) / "outlines"
+    d.mkdir(parents=True, exist_ok=True)
+    p = d / f"{field}.json"
+    with open(p, "w") as fh:
+        json.dump(data, fh, separators=(",", ":"))
+    return p

spatialscape/polygons.py

@@ -0,0 +1,85 @@
+"""Per-cell boundary polygons packed as int16 deltas from the cell centroid.
+
+Files (sample-local µm frame, bundle row order):
+  polygons.offsets.u32  uint32[n_obs + 1]  vertex start index per cell (zero-length = no polygon)
+  polygons.i16          int16[n_vertices * 2]  (dx, dy) * 10 relative to the cell's xy
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+
+SCALE = 0.1  # µm per unit
+
+
+def read_polygon_table(path: Path, id_col: str, x_col: str, y_col: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    t = pq.read_table(path, columns=[id_col, x_col, y_col])
+    ids = np.asarray(t.column(id_col).to_pylist(), dtype=str)
+    x = t.column(x_col).to_numpy(zero_copy_only=False).astype(np.float64)
+    y = t.column(y_col).to_numpy(zero_copy_only=False).astype(np.float64)
+    return ids, x, y
+
+
+def read_polygon_obsm(adata, key: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Flatten an obsm array of shape (n_obs, n_vertices, 2) into (ids, x, y); NaN-padded vertices are dropped."""
+    arr = np.asarray(adata.obsm[key], dtype=np.float64)
+    if arr.ndim != 3 or arr.shape[2] != 2:
+        raise ValueError(f"obsm/{key}: expected shape (n_obs, n_vertices, 2), got {arr.shape}")
+    _n, v, _ = arr.shape
+    ids = np.repeat(np.asarray(adata.obs_names, dtype=str), v)
+    x, y = arr[:, :, 0].ravel(), arr[:, :, 1].ravel()
+    ok = np.isfinite(x) & np.isfinite(y)
+    return ids[ok], x[ok], y[ok]
+
+
+def apply_affine(x: np.ndarray, y: np.ndarray, affine) -> tuple[np.ndarray, np.ndarray]:
+    a, b, c, d, e, f = (float(v) for v in affine)
+    return a * x + b * y + c, d * x + e * y + f
+
+
+def pack_polygons(
+    obs_ids: list[str],
+    xy_um: np.ndarray,
+    poly_ids: np.ndarray,
+    poly_xy_um: np.ndarray,
+    *,
+    max_vertices: int = 24,
+) -> tuple[np.ndarray, np.ndarray, int, float]:
+    """Return (offsets uint32[n+1], deltas int16[nv*2], n_cells_with_polygon, median_centroid_offset_um) in bundle
+    row order. The last value is the median distance between a polygon's mean vertex and its cell's xy; a large
+    value means the polygons are in a different frame than the coordinates."""
+    order = np.argsort(poly_ids, kind="stable")
+    poly_ids = poly_ids[order]
+    poly_xy_um = poly_xy_um[order]
+    uniq, start, counts = np.unique(poly_ids, return_index=True, return_counts=True)
+    lookup = {u: (int(s), int(c)) for u, s, c in zip(uniq, start, counts, strict=True)}
+    n = len(obs_ids)
+    offsets = np.zeros(n + 1, dtype=np.uint32)
+    chunks: list[np.ndarray] = []
+    total = 0
+    with_poly = 0
+    dev: list[float] = []
+    for i, cid in enumerate(obs_ids):
+        hit = lookup.get(str(cid))
+        if hit:
+            s, c = hit
+            pts = poly_xy_um[s : s + c]
+            dev.append(float(np.hypot(*(pts.mean(axis=0) - xy_um[i]))))
+            if c > max_vertices:
+                pts = pts[np.linspace(0, c - 1, max_vertices).round().astype(int)]
+            d = np.rint((pts - xy_um[i]) / SCALE)
+            d = np.clip(d, -32767, 32767).astype(np.int16)
+            chunks.append(d.ravel())
+            total += len(pts)
+            with_poly += 1
+        offsets[i + 1] = total
+    deltas = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.int16)
+    return offsets, deltas, with_poly, (float(np.median(dev)) if dev else 0.0)
+
+
+def write_polygons(out_dir: Path, offsets: np.ndarray, deltas: np.ndarray) -> None:
+    (Path(out_dir) / "polygons.offsets.u32").write_bytes(offsets.astype("<u4").tobytes())
+    (Path(out_dir) / "polygons.i16").write_bytes(deltas.astype("<i2").tobytes())