Draft SpatialData.filter()

CHANGELOG.md

@@ -10,11 +10,15 @@ and this project adheres to [Semantic Versioning][].
 
 ## [0.x.x] - 2024-xx-xx
 
+### Minor
+
+-   Added `SpatialData.filter()` method for subsetting by `obs` and `var` @aeisenbarth
+
 ## [0.2.1] - 2024-07-04
 
 ### Minor
 
--  Relaxing `spatial-image` package requirement #616
+-   Relaxing `spatial-image` package requirement #616
 
 ## [0.2.0] - 2024-07-03
 

src/spatialdata/_core/spatialdata.py

@@ -3,11 +3,12 @@
 import hashlib
 import os
 import warnings
-from collections.abc import Generator
+from collections.abc import Generator, Iterable
 from itertools import chain
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
+import numpy as np
 import pandas as pd
 import zarr
 from anndata import AnnData
@@ -2061,6 +2062,74 @@ def subset(
         )
         return SpatialData(**elements_dict, tables=tables)
 
+    def filter(
+        self,
+        elements: Iterable[str] | None = None,
+        regions: Iterable[str] | None = None,
+        obs_keys: Iterable[str] | None = None,
+        var_keys: Iterable[str] | None = None,
+        var_names: Iterable[str] | None = None,
+        layers: Iterable[str] | None = None,
+        region_key: str = "region",
+        instance_key: str = "instance_id",
+    ) -> SpatialData:
+        """
+        Filter a SpatialData object to contain only specified elements or table entries.
+
+        Args:
+            sdata: A SpatialData object
+            elements: Names of elements to include. Defaults to [].
+            regions: Regions to include in the table. Defaults to regions of all selected elements.
+            obs_keys: Names of obs columns to include. Defaults to [].
+            var_keys: Names of var columns to include. Defaults to [].
+            var_names: Names of variables (X columns) to include. Defaults to [].
+            layers: Names of X layers to include. Defaults to [].
+
+        Returns
+        -------
+            A new SpatialData instance
+        """
+        elements = [] if elements is None else list(elements)
+
+        sdata_subset = self.subset(element_names=elements, filter_tables=True) if elements else SpatialData()
+        # Ensure the returned SpatialData is not backed to the original reference dataset,
+        # so that it can be safely modified.
+        assert not sdata_subset.is_backed()
+        # Further filtering on the table
+        if (table := sdata_subset.tables.get("table")) is not None:
+            regions = elements if regions is None else regions
+            obs_keys = [] if obs_keys is None else list(obs_keys)
+            if instance_key not in obs_keys:
+                obs_keys.insert(0, instance_key)
+            if region_key not in obs_keys:
+                obs_keys.insert(0, region_key)
+            var_keys = [] if var_keys is None else var_keys
+            var_names = [] if var_names is None else var_names
+            # Preserve order by checking "isin" instead of slicing. Also guarantees no duplicates.
+            table_subset = table[
+                table.obs[region_key].isin(regions),
+                table.var_names.isin(var_names),
+            ]
+            layers_subset = (
+                {key: layer for key, layer in table_subset.layers.items() if key in layers}
+                if table_subset.layers is not None and len(var_names) > 0
+                else None
+            )
+            table_subset = TableModel.parse(
+                AnnData(
+                    X=table_subset.X if len(var_names) > 0 else None,
+                    obs=table_subset.obs.loc[:, table_subset.obs.columns.isin(obs_keys)],
+                    var=table_subset.var.loc[:, table_subset.var.columns.isin(var_keys)],
+                    layers=layers_subset,
+                ),
+                region_key=region_key,
+                instance_key=instance_key,
+                region=np.unique(table_subset.obs[region_key]).tolist(),
+            )
+            del sdata_subset.tables["table"]
+            sdata_subset.tables["table"] = table_subset
+        return sdata_subset
+
     def __getitem__(self, item: str) -> SpatialElement:
         """
         Return the element with the given name.