[feat] add option to generate synthetic coords

docs/release_notes/index.md

@@ -21,6 +21,8 @@ is available in the [commit logs](https://github.com/scverse/scvi-tools/commits/
     in `batch_representation="embedding"` to {class}`scvi.model.SCVI` {pr}`2576`.
 -   Add experimental mixin classes {class}`scvi.model.base.EmbeddingMixin` and
     {class}`scvi.module.base.EmbeddingModuleMixin` {pr}`2576`.
+-   Add option to generate synthetic spatial coordinates in {func}`scvi.data.synthetic_iid` with
+    argument `generate_coordinates` {pr}`2603`.
 
 #### Changed
 

scvi/data/_built_in_data/_synthetic.py

@@ -22,13 +22,15 @@ def _generate_synthetic(
     n_labels: int,
     dropout_ratio: float,
     sparse_format: Optional[str],
+    generate_coordinates: bool,
     return_mudata: bool,
     batch_key: str = "batch",
     labels_key: str = "labels",
     rna_key: str = "rna",
     protein_expression_key: str = "protein_expression",
     protein_names_key: str = "protein_names",
     accessibility_key: str = "accessibility",
+    coordinates_key: str = "coordinates",
 ) -> AnnOrMuData:
     n_obs = batch_size * n_batches
 
@@ -61,6 +63,9 @@ def sparsify_data(data: np.ndarray):
         labels = np.random.randint(0, n_labels, size=(n_obs,))
         labels = np.array([f"label_{i}" for i in labels])
 
+    if generate_coordinates:
+        coords = np.random.normal(size=(n_obs, 2))
+
     adata = AnnData(rna)
     if return_mudata:
         mod_dict = {rna_key: adata}
@@ -83,5 +88,7 @@ def sparsify_data(data: np.ndarray):
     adata.obs[batch_key] = pd.Categorical(batch)
     if n_labels > 0:
         adata.obs[labels_key] = pd.Categorical(labels)
+    if generate_coordinates:
+        adata.obsm[coordinates_key] = coords
 
     return adata

scvi/data/_datasets.py

@@ -515,6 +515,7 @@ def synthetic_iid(
     n_labels: int = 3,
     dropout_ratio: float = 0.7,
     sparse_format: str | None = None,
+    generate_coordinates: bool = False,
     return_mudata: bool = False,
 ) -> AnnOrMuData:
     """Synthetic multimodal dataset.
@@ -551,6 +552,8 @@ def synthetic_iid(
         * `None`: Store as a dense :class:`numpy.ndarray`.
         * `"csr_matrix"`: Store as a :class:`scipy.sparse.csr_matrix`.
         * `"csc_matrix"`: Store as a :class:`scipy.sparse.csc_matrix`.
+    generate_coordinates
+        Whether to generate spatial coordinates for the cells.
     return_mudata
         Returns a :class:`~mudata.MuData` if `True`, else :class:`~anndata.AnnData`.
 
@@ -563,6 +566,8 @@ def synthetic_iid(
     * `.obsm["protein_expression"]`: Protein expression matrix.
     * `.uns["protein_names"]`: Array of protein names.
     * `.obsm["accessibility"]`: Accessibility expression matrix.
+    * `.obsm["coordinates"]`: Spatial coordinates for the cells if ``generate_coordinates`` is
+      ``True``.
 
     :class:`~mudata.MuData` (if `return_mudata=True`) with the following fields:
 
@@ -571,6 +576,8 @@ def synthetic_iid(
     * `.mod["rna"]`: RNA expression data.
     * `.mod["protein_expression"]`: Protein expression data.
     * `.mod["accessibility"]`: Accessibility expression data.
+    * `.obsm["coordinates"]`: Spatial coordinates for the cells if ``generate_coordinates`` is
+      ``True``.
 
     Examples
     --------
@@ -591,6 +598,7 @@ def synthetic_iid(
         n_labels=n_labels,
         dropout_ratio=dropout_ratio,
         sparse_format=sparse_format,
+        generate_coordinates=generate_coordinates,
         return_mudata=return_mudata,
     )
 

tests/data/test_synthetic_iid.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import numpy as np
 import pytest
 
 from scvi.data import synthetic_iid
@@ -8,3 +9,11 @@
 @pytest.mark.parametrize("sparse_format", ["csr_matrix", "csc_matrix", None])
 def test_synthetic_iid_sparse_format(sparse_format: Optional[str]):
     _ = synthetic_iid(sparse_format=sparse_format)
+
+
+@pytest.mark.parametrize("return_mudata", [False, True])
+def test_synthetic_iid_coords(return_mudata: bool):
+    adata = synthetic_iid(return_mudata=return_mudata, generate_coordinates=True)
+    assert "coordinates" in adata.obsm
+    assert isinstance(adata.obsm["coordinates"], np.ndarray)
+    assert adata.obsm["coordinates"].shape == (adata.n_obs, 2)