Merge pull request #85 from Genentech/handle-inf

Handle infinite values in addition to nans

Author: joshua-gould

scallops/features/preprocessing.py

@@ -54,14 +54,16 @@ def _transform_feature_group(x):
 
 def filter_data(
     data: anndata.AnnData,
-    max_fraction_nans: float | None = 0.25,
+    max_fraction_not_finite: float | None = 0.25,
     min_variance: float | None = 0.1,
     by: str | Sequence | None = None,
 ) -> anndata.AnnData:
-    """Filter cells using `max_fraction_nans` then filter features using `min_variance`
+    """Filter cells using `max_fraction_not_finite` then filter features using
+    `min_variance`
 
     :param data: AnnData object
-    :param max_fraction_nans: Keep cells with <= `max_fraction_nans` missing values
+    :param max_fraction_not_finite: Keep cells with <= `max_fraction_not_finite`
+    missing or infinite values
     :param min_variance: Keep features with variance >= `min_variance`
     :param by: Column(s) in `data.obs` to stratify by when computing variance. If
     provided, the median variance is used for filtering.
@@ -70,10 +72,10 @@ def filter_data(
     xp = get_namespace(data.X)
     keep_cells = None
     keep_features = None
-    if max_fraction_nans is not None:
-        nan_counts_per_cell = xp.isnan(data.X).sum(axis=1)
-        max_nans = int(data.shape[1] * max_fraction_nans)
-        keep_cells = nan_counts_per_cell <= max_nans
+    if max_fraction_not_finite is not None:
+        invalid_counts_per_cell = (~xp.isfinite(data.X)).sum(axis=1)
+        max_counts = int(data.shape[1] * max_fraction_not_finite)
+        keep_cells = invalid_counts_per_cell <= max_counts
     if min_variance is not None:
         if by is not None:
             if isinstance(keep_cells, da.Array):
@@ -102,7 +104,7 @@ def filter_data(
                 if keep_cells is not None
                 else xp.var(data.X, axis=0)
             )
-        keep_features = variance >= min_variance
+        keep_features = (variance >= min_variance) & (xp.isfinite(variance))
 
     if isinstance(data.X, da.Array):
         keep_features, keep_cells = dask.compute(keep_features, keep_cells)

scallops/tests/test_features_preprocessing.py

@@ -28,14 +28,14 @@ def test_filter_data(use_dask, by):
     adata.X[1, 0] = 100
     adata.X[0, 0] = np.nan
     # np.var(adata.X, axis=0) array([nan,  5.], dtype=float32)
-    test_nan_filter = filter_data(adata, max_fraction_nans=0, min_variance=None)
+    test_nan_filter = filter_data(adata, max_fraction_not_finite=0, min_variance=None)
     assert test_nan_filter.shape == (3, 2)
     # np.var(adata.X, axis=0) # array([nan,  5.]
     # np.var(adata[adata.obs['well'] == 'well1'].X, axis=0)  # array([nan,  4.])
     # np.var(adata[adata.obs['well'] == 'well2'].X, axis=0)  # array([2209.,    4.]
-    d1 = filter_data(adata, max_fraction_nans=None, min_variance=0, by=by)
+    d1 = filter_data(adata, max_fraction_not_finite=None, min_variance=0, by=by)
     # np.var(adata[1:].X, axis=0)  array([2006.2222, 2.6666667]
-    d2 = filter_data(adata, max_fraction_nans=0, min_variance=5, by=by)
+    d2 = filter_data(adata, max_fraction_not_finite=0, min_variance=5, by=by)
 
     assert d1.shape == (4, 1)
     assert d2.shape == (3, 1)