write backbone

Author: quentinblampey

spatialdata_xenium_explorer/__init__.py

@@ -2,4 +2,8 @@
 
 __version__ = importlib.metadata.version("spatialdata_xenium_explorer")
 
+from .images import write_multiscale
+from .points import write_transcripts
+from .table import write_cell_categories, write_gene_counts
+from .shapes import write_polygons
 from .converter import write

spatialdata_xenium_explorer/constants.py

@@ -1,15 +1,24 @@
-class Versions:
-    EXPERIMENT = [2, 0]
-    GROUPS = [5, 0]
-    CELL_CATEGORIES = [1, 0]
+class FileNames:
+    IMAGE = "morphology.ome.tif"
+    POINTS = "transcripts.zarr.zip"
+    SHAPES = "cells.zarr.zip"
+    TABLE = "cell_feature_matrix.zarr.zip"
+    CELL_CATEGORIES = "analysis.zarr.zip"
+    METADATA = "experiment.xenium"
 
 
 class ExplorerConstants:
     GRID_SIZE = 250
     QUALITY_SCORE = 40
 
 
-def cell_categories_attrs():
+class Versions:
+    EXPERIMENT = [2, 0]
+    GROUPS = [5, 0]
+    CELL_CATEGORIES = [1, 0]
+
+
+def cell_categories_attrs() -> dict:
     return {
         "major_version": Versions.CELL_CATEGORIES[0],
         "minor_version": Versions.CELL_CATEGORIES[1],
@@ -19,7 +28,7 @@ def cell_categories_attrs():
     }
 
 
-def cell_summary_attrs():
+def cell_summary_attrs() -> dict:
     return {
         "column_descriptions": [
             "Cell centroid in X",
@@ -42,7 +51,7 @@ def cell_summary_attrs():
     }
 
 
-def group_attrs():
+def group_attrs() -> dict:
     return {
         "major_version": Versions.GROUPS[0],
         "minor_version": Versions.GROUPS[1],
@@ -57,13 +66,17 @@ def group_attrs():
     }
 
 
-def experiment_dict():
+def experiment_dict(run_name: str, region_name: str, num_cells: int) -> dict:
     return {
         "major_version": Versions.EXPERIMENT[0],
         "minor_version": Versions.EXPERIMENT[1],
+        "run_name": run_name,
+        "region_name": region_name,
+        "experiment_uuid": "N/A",
+        "panel_tissue_type": "N/A",
         "run_start_time": "N/A",
         "preservation_method": "N/A",
-        "num_cells": 0,
+        "num_cells": num_cells,
         "transcripts_per_cell": 0,
         "transcripts_per_100um": 0,
         "cassette_name": "N/A",
@@ -98,7 +111,7 @@ def experiment_dict():
     }
 
 
-def image_options():
+def image_options() -> dict:
     return {
         "photometric": "minisblack",
         "tile": (1024, 1024),

spatialdata_xenium_explorer/converter.py

@@ -1,17 +1,41 @@
 import json
+from pathlib import Path
 
-from .constants import EXPERIMENT
+from spatialdata import SpatialData
 
+from . import (
+    write_cell_categories,
+    write_gene_counts,
+    write_multiscale,
+    write_polygons,
+    write_transcripts,
+)
+from .constants import FileNames, experiment_dict
 
-def write_experiment(path, run_name, tissue, region_name, uuid):
-    EXPERIMENT["run_name"] = run_name
-    EXPERIMENT["panel_tissue_type"] = tissue
-    EXPERIMENT["region_name"] = region_name
-    EXPERIMENT["experiment_uuid"] = uuid
 
-    with open(path, "w") as f:
+def write(path: str, sdata: SpatialData, image_key: str, gene_column: str) -> None:
+    path: Path = Path(path)
+    assert (
+        not path.exists() or path.is_dir()
+    ), f"A path to an existing file was provided. It should be a path to a directory."
+
+    path.mkdir(parents=True, exist_ok=True)
+
+    adata = sdata.table
+
+    EXPERIMENT = experiment_dict(..., ..., adata.n_obs)
+    with open(path / FileNames.METADATA, "w") as f:
         json.dump(EXPERIMENT, f, indent=4)
 
+    write_gene_counts(path / FileNames.TABLE, adata)
+    write_cell_categories(path / FileNames.CELL_CATEGORIES, adata)
+
+    polygons = sdata.shapes["..."]
+    # TODO: transform polygon coords to pixel
+    write_polygons(path / FileNames.SHAPES, polygons)
+
+    # TODO : make it memory efficient
+    write_multiscale(path / FileNames.IMAGE, sdata.images[image_key])
 
-def write(path, sdata):
-    ...
+    df = sdata.points["..."]
+    write_transcripts(path / FileNames.POINTS, df, gene_column)

spatialdata_xenium_explorer/images.py

@@ -7,7 +7,7 @@
 from .constants import image_metadata, image_options
 
 
-def to_uint8(arr):
+def to_uint8(arr: np.ndarray) -> np.ndarray:
     print(f"Writing image of shape {arr.shape}")
     return (arr // 256).astype(np.uint8)
 

spatialdata_xenium_explorer/points.py

@@ -16,15 +16,13 @@ def subsample_indices(n_samples, factor: int = 4):
 def write_transcripts(
     path: Path,
     df: pd.DataFrame,
-    x: str = "x",
-    y: str = "y",
     gene: str = "gene",
     max_levels: int = 15,
 ):
     num_transcripts = len(df)
     df[gene] = df[gene].astype("category")
 
-    location = df[[x, y]]
+    location = df[["x", "y"]]
     location = np.concatenate([location, np.zeros((num_transcripts, 1))], axis=1)
 
     xmax, ymax = location[:, :2].max(axis=0)

spatialdata_xenium_explorer/shapes.py

@@ -1,25 +1,42 @@
 from math import ceil
 from pathlib import Path
+from typing import Iterable
 
 import numpy as np
 import zarr
 from shapely.geometry import Polygon
 
 from .constants import cell_summary_attrs, group_attrs
-from .utils import pad
 
 
-def write_polygons(path: Path, polygons: list[Polygon], area: np.ndarray) -> None:
-    coordinates = np.stack([pad(p, 3, 13) for p in polygons])
+def pad_polygon(polygon: Polygon, max_vertices: int, tolerance: float = 1) -> np.ndarray:
+    n_vertices = len(polygon.exterior.coords)
+    assert n_vertices >= 3
+
+    coords = polygon.exterior.coords._coords
+
+    if n_vertices == max_vertices:
+        return coords.flatten()
+
+    if n_vertices < max_vertices:
+        return np.pad(
+            coords, ((0, max_vertices - n_vertices), (0, 0)), mode="edge"
+        ).flatten()
+
+    # TODO: improve it: how to choose the right tolerance?
+    polygon = polygon.simplify(tolerance=tolerance)
+    return pad_polygon(polygon, max_vertices, tolerance + 1)
+
+
+def write_polygons(path: Path, polygons: Iterable[Polygon], max_vertices: int) -> None:
+    coordinates = np.stack([pad_polygon(p, max_vertices) for p in polygons])
     coordinates /= 4.705882
 
     num_cells = len(coordinates)
     cells_fourth = ceil(num_cells / 4)
     cells_half = ceil(num_cells / 2)
 
-    CELLS_SUMMARY_ATTRS = cell_summary_attrs()
     GROUP_ATTRS = group_attrs()
-
     GROUP_ATTRS["number_cells"] = num_cells
 
     polygon_vertices = np.stack([coordinates, coordinates])
@@ -42,14 +59,14 @@ def write_polygons(path: Path, polygons: list[Polygon], area: np.ndarray) -> Non
         g.array("cell_id", cell_id, dtype="uint32", chunks=(cells_half, 1))
 
         cell_summary = np.zeros((num_cells, 7))
-        cell_summary[:, 2] = area
+        cell_summary[:, 2] = [p.area for p in polygons]
         g.array(
             "cell_summary",
             cell_summary,
             dtype="float64",
             chunks=(num_cells, 1),
         )
-        g["cell_summary"].attrs.put(CELLS_SUMMARY_ATTRS)
+        g["cell_summary"].attrs.put(cell_summary_attrs())
 
         g.array(
             "polygon_num_vertices",

spatialdata_xenium_explorer/table.py

@@ -1,11 +1,11 @@
-import anndata
 import numpy as np
 import zarr
+from anndata import AnnData
 
 from .constants import cell_categories_attrs
 
 
-def write_gene_counts(path, adata: anndata.AnnData):
+def write_gene_counts(path: str, adata: AnnData) -> None:
     counts = adata.layers["counts"]
 
     feature_keys = list(adata.var_names) + ["Total transcripts"]
@@ -54,7 +54,9 @@ def write_gene_counts(path, adata: anndata.AnnData):
         cells_group.array("indptr", indptr, dtype="uint32", chunks=indptr.shape)
 
 
-def add_group(root: zarr.Group, index: int, values: np.ndarray, categories: list[str]):
+def _write_categorical_column(
+    root: zarr.Group, index: int, values: np.ndarray, categories: list[str]
+) -> None:
     group = root.create_group(index)
     values_indices = [np.where(values == cat)[0] for cat in categories]
     values_cum_len = np.cumsum([len(indices) for indices in values_indices])
@@ -66,23 +68,23 @@ def add_group(root: zarr.Group, index: int, values: np.ndarray, categories: list
     group.array("indptr", indptr, dtype="uint32", chunks=(len(indptr),))
 
 
-def write_cell_categories(path: str, adata: anndata.AnnData):
-    categorical_columns = [
-        name for name, cat in adata.obs.dtypes.items() if cat == "category"
-    ]
+def write_cell_categories(path: str, adata: AnnData) -> None:
+    cat_columns = [name for name, cat in adata.obs.dtypes.items() if cat == "category"]
+
+    print(f"Saving {len(cat_columns)} cell categories: {', '.join(cat_columns)}")
 
     ATTRS = cell_categories_attrs()
-    ATTRS["number_groupings"] = len(categorical_columns)
+    ATTRS["number_groupings"] = len(cat_columns)
 
     with zarr.ZipStore(path, mode="w") as store:
         g = zarr.group(store=store)
         cell_groups = g.create_group("cell_groups")
 
-        for i, name in enumerate(categorical_columns):
+        for i, name in enumerate(cat_columns):
             categories = list(adata.obs[name].cat.categories)
             ATTRS["grouping_names"].append(name)
             ATTRS["group_names"].append(categories)
 
-            add_group(cell_groups, i, adata.obs[name], categories)
+            _write_categorical_column(cell_groups, i, adata.obs[name], categories)
 
         cell_groups.attrs.put(ATTRS)