version 0, WIP

Author: quentinblampey

README.md

@@ -17,4 +17,6 @@ sdata = spatialdata.read_zarr("...")
 image_key = "..." # The name of the MultiscaleSpatialImage to be exported
 
 spatialdata_xenium_explorer.write("/path/to/directory", sdata, image_key)
-```
+```
+
+This will create up to 6 files, among which a file called `experiment.xenium`. Double-click on this file to open it on the [Xenium Explorer](https://www.10xgenomics.com/support/software/xenium-explorer/downloads) (make sure you have the latest version of the software).

pyproject.toml

@@ -26,7 +26,7 @@ testpaths = ["tests"]
 python_files = "test_*.py"
 
 [tool.black]
-line-length = 90
+line-length = 100
 include = '\.pyi?$'
 exclude = '''
 /(

spatialdata_xenium_explorer/converter.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 from spatialdata import SpatialData
+from spatialdata.transformations import Identity, set_transformation
 
 from . import (
     write_cell_categories,
@@ -13,29 +14,63 @@
 from .constants import FileNames, experiment_dict
 
 
-def write(path: str, sdata: SpatialData, image_key: str, gene_column: str) -> None:
+def _order_instances(sdata: SpatialData, shapes_key: str):
+    adata = sdata.table
+
+    instance_key = adata.uns["spatialdata_attrs"]["instance_key"]
+    region_key = adata.uns["spatialdata_attrs"]["region_key"]
+
+    adata = adata[adata.obs[region_key] == shapes_key].copy()
+    adata.obs.set_index(instance_key, inplace=True)
+    adata = adata[sdata.shapes[shapes_key].index].copy()
+    return adata
+
+
+def write(
+    path: str,
+    sdata: SpatialData,
+    image_key: str,
+    shapes_key: str,
+    points_key: str,
+    gene_column: str,
+    layer: str | None = None,
+) -> None:
+    """
+    Transform a SpatialData object into inputs for the Xenium Explorer.
+    Currently only images of type MultiscaleSpatialImage are supported.
+
+    Args:
+        path: Path to the directory where files will be saved.
+        sdata: SpatialData object.
+        image_key: Name of the image or interest.
+        shapes_key: Name of the cell shapes.
+        points_key: Name of the transcripts (key of `sdata.points`).
+        gene_column: Column name of the points dataframe containing the gene names.
+        layer: Layer of `sdata.table` where the gene counts are saved. If `None`, uses `sdata.table.X`.
+    """
     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
+    adata = _order_instances(sdata, shapes_key)
 
-    EXPERIMENT = experiment_dict(..., ..., adata.n_obs)
+    EXPERIMENT = experiment_dict(image_key, shapes_key, 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_gene_counts(path / FileNames.TABLE, adata, layer)
     write_cell_categories(path / FileNames.CELL_CATEGORIES, adata)
 
-    polygons = sdata.shapes["..."]
-    # TODO: transform polygon coords to pixel
-    write_polygons(path / FileNames.SHAPES, polygons)
+    pixels_cs = "__pixels"
+    set_transformation(sdata.images[image_key], Identity(), pixels_cs)
+
+    gdf = sdata.transform_element_to_coordinate_system(sdata.shapes[shapes_key], pixels_cs)
+    write_polygons(path / FileNames.SHAPES, gdf.geometry)
 
-    # TODO : make it memory efficient
     write_multiscale(path / FileNames.IMAGE, sdata.images[image_key])
 
-    df = sdata.points["..."]
+    df = sdata.transform_element_to_coordinate_system(sdata.points[points_key], pixels_cs)
     write_transcripts(path / FileNames.POINTS, df, gene_column)

spatialdata_xenium_explorer/images.py

@@ -1,19 +1,24 @@
-from pathlib import Path
-
 import numpy as np
 import tifffile as tf
 from multiscale_spatial_image import MultiscaleSpatialImage
 
 from .constants import image_metadata, image_options
 
 
-def to_uint8(arr: np.ndarray) -> np.ndarray:
-    print(f"Writing image of shape {arr.shape}")
-    return (arr // 256).astype(np.uint8)
+def _astype_uint8(arr: np.ndarray) -> np.ndarray:
+    assert np.issubdtype(
+        arr.dtype, np.integer
+    ), f"The image dtype has to be an integer dtype. Found {arr.dtype}"
+
+    if arr.dtype == np.uint8:
+        return arr
+
+    factor = np.iinfo(np.uint8).max / np.iinfo(arr.dtype).max
+    return (arr * factor).astype(np.uint8)
 
 
 def write_multiscale(
-    output_path: Path,
+    path: str,
     multiscale: MultiscaleSpatialImage,
     pixelsize: float = 0.2125,
 ):
@@ -22,9 +27,10 @@ def write_multiscale(
 
     metadata = image_metadata(channel_names, pixelsize)
 
-    with tf.TiffWriter(output_path, bigtiff=True) as tif:
+    # TODO : make it memory efficient
+    with tf.TiffWriter(path, bigtiff=True) as tif:
         tif.write(
-            to_uint8(multiscale[scale_names[0]]["image"].values),
+            _astype_uint8(multiscale[scale_names[0]]["image"].values),
             subifds=len(scale_names) - 1,
             resolution=(1e4 / pixelsize, 1e4 / pixelsize),
             metadata=metadata,
@@ -33,7 +39,7 @@ def write_multiscale(
 
         for i, scale in enumerate(scale_names[1:]):
             tif.write(
-                to_uint8(multiscale[scale]["image"].values),
+                _astype_uint8(multiscale[scale]["image"].values),
                 subfiletype=1,
                 resolution=(
                     1e4 * 2 ** (i + 1) / pixelsize,

spatialdata_xenium_explorer/points.py

@@ -1,8 +1,8 @@
 from math import ceil
 from pathlib import Path
 
+import dask.dataframe as dd
 import numpy as np
-import pandas as pd
 import zarr
 
 from .constants import ExplorerConstants
@@ -15,10 +15,13 @@ def subsample_indices(n_samples, factor: int = 4):
 
 def write_transcripts(
     path: Path,
-    df: pd.DataFrame,
+    df: dd.DataFrame,
     gene: str = "gene",
     max_levels: int = 15,
 ):
+    # TODO: make everything using dask instead of pandas
+    df = df.compute()
+
     num_transcripts = len(df)
     df[gene] = df[gene].astype("category")
 
@@ -178,5 +181,3 @@ def write_transcripts(
             transcript_id = transcript_id[sub_indices]
 
         grids.attrs.put(GRIDS_ATTRS)
-
-        print(g.info)

spatialdata_xenium_explorer/table.py

@@ -1,12 +1,14 @@
 import numpy as np
 import zarr
 from anndata import AnnData
+from scipy.sparse import csr_matrix
 
 from .constants import cell_categories_attrs
 
 
-def write_gene_counts(path: str, adata: AnnData) -> None:
-    counts = adata.layers["counts"]
+def write_gene_counts(path: str, adata: AnnData, layer: str | None) -> None:
+    counts = adata.X if layer is None else adata.layers[layer]
+    counts = csr_matrix(counts)
 
     feature_keys = list(adata.var_names) + ["Total transcripts"]
     feature_ids = feature_keys
@@ -69,6 +71,7 @@ def _write_categorical_column(
 
 
 def write_cell_categories(path: str, adata: AnnData) -> None:
+    # TODO: consider also columns that can be transformed to a categorical column?
     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)}")