min_cells argument to load_and_normalize()

src/vsparse/_rapid_load.py

@@ -188,6 +188,24 @@ def _weighted_bincount(
     return partial.sum(axis=0)
 
 
+@numba.njit(cache=True, parallel=True)
+def _gene_detection_counts(
+    indices: np.ndarray, data: np.ndarray, n_genes: int, nthreads: int
+) -> np.ndarray:
+    """Number of cells with positive expression for each gene."""
+    n = indices.shape[0]
+    chunk = (n + nthreads - 1) // nthreads
+    partial = np.zeros((nthreads, n_genes), dtype=np.int64)
+    for t in numba.prange(nthreads):  # ty: ignore[not-iterable]
+        start = t * chunk
+        end = min(n, start + chunk)
+        local = partial[t]
+        for k in range(start, end):
+            if data[k] > 0:
+                local[indices[k]] += 1
+    return partial.sum(axis=0)
+
+
 # -- fused row+column filter/compaction --------------------------------------
 
 
@@ -348,16 +366,18 @@ def load_and_normalize(
     *,
     min_cell_counts: float = 10.0,
     gene_threshold: float = 0.0,
+    min_cells: int | None = None,
     obs_filter: Callable[[pd.DataFrame], object] | None = None,
     x_key: str = "X",
 ) -> ad.AnnData:
     """Load, filter, and depth-normalize a VCSR/IVCSR-backed ``.h5ad`` file.
 
     Reproduces ``parafac2.normalize.prepare_dataset``: cells with total
     counts <= ``min_cell_counts`` and genes with total counts <=
-    ``gene_threshold * n_cells`` (both measured on the raw, unfiltered
-    counts after any ``obs_filter``, matching the reference implementation)
-    are dropped; the remaining matrix is row-normalized to the median per-cell depth, then
+    ``gene_threshold * n_cells`` are dropped. When ``min_cells`` is given,
+    genes expressed in fewer than ``min_cells`` cells are also dropped. Gene
+    filters are measured on the raw counts after any ``obs_filter``.
+    The remaining matrix is row-normalized to the median per-cell depth, then
     column-normalized by gene sum, then transformed as ``log10(1000x + 1)``.
     Surrounding metadata (``obs``, ``var``, ``obsm``, etc.) is sliced to
     match the retained cells and genes.
@@ -374,6 +394,9 @@ def load_and_normalize(
         Minimum threshold fraction for gene inclusion, as in
         ``parafac2.normalize.prepare_dataset``: genes with total raw counts
         <= ``gene_threshold * n_cells`` are dropped.
+    min_cells
+        Optional gene filter. Genes expressed in fewer than this
+        many cells are dropped. Expression is defined as a raw count > 0.
     obs_filter
         Optional callable receiving ``obs`` and returning a one-dimensional
         boolean mask. When provided, rows are subset before cell filtering,
@@ -405,6 +428,12 @@ def load_and_normalize(
     import h5py
     import hdf5plugin  # noqa: F401  -- registers the Blosc2 HDF5 filter
 
+    if min_cells is not None:
+        if isinstance(min_cells, bool) or not isinstance(min_cells, int):
+            raise TypeError("min_cells must be an integer or None")
+        if min_cells < 0:
+            raise ValueError("min_cells must be non-negative")
+
     with h5py.File(Path(path), "r") as f:
         g = f[x_key]
         shape = (int(g.attrs["shape"][0]), int(g.attrs["shape"][1]))
@@ -437,6 +466,11 @@ def load_and_normalize(
 
         gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads())
         gene_mask = gene_totals_raw > (gene_threshold * n_cells)
+        if min_cells is not None:
+            gene_detection_counts = _gene_detection_counts(
+                indices, data, n_genes, numba.get_num_threads()
+            )
+            gene_mask &= gene_detection_counts >= min_cells
         metadata_cell_mask = cell_mask
     else:
         obs = kwargs.get("obs")
@@ -463,6 +497,11 @@ def load_and_normalize(
 
         gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads())
         gene_mask = gene_totals_raw > (gene_threshold * selected_rows.shape[0])
+        if min_cells is not None:
+            gene_detection_counts = _gene_detection_counts(
+                indices, data, n_genes, numba.get_num_threads()
+            )
+            gene_mask &= gene_detection_counts >= min_cells
 
         metadata_cell_mask = np.zeros(n_cells, dtype=np.bool_)
         metadata_cell_mask[selected_rows[cell_mask]] = True

tests/test_rapid_load.py

@@ -9,11 +9,18 @@
 from vsparse._rapid_load import load_and_normalize, load_packed
 
 
-def _reference_prepare(dense: np.ndarray, min_cell_counts: float, gene_threshold: float):
+def _reference_prepare(
+    dense: np.ndarray,
+    min_cell_counts: float,
+    gene_threshold: float,
+    min_cells: int | None = None,
+):
     """Direct translation of parafac2.normalize.prepare_dataset's math, on a plain array."""
     X = sp.csr_array(dense)
     cell_mask = np.ravel(X.sum(axis=1)) > min_cell_counts
     gene_mask = np.ravel(X.sum(axis=0)) > (gene_threshold * X.shape[0])
+    if min_cells is not None:
+        gene_mask &= np.ravel((X > 0).sum(axis=0)) >= min_cells
     Xf = sp.csr_array(X[cell_mask][:, gene_mask])
     Xf.data = Xf.data.astype(np.float32)
 
@@ -71,6 +78,90 @@ def test_no_filtering_keeps_everything(tmp_path, rng):
     assert result.n_vars == dense.shape[1]
 
 
+def test_min_cells_filters_genes_by_detection_count(tmp_path):
+    dense = np.array(
+        [
+            [5, 1, 1, 1],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 1],
+        ],
+        dtype=np.float64,
+    )
+    path = _write_ivcsr(tmp_path, dense)
+
+    result = load_and_normalize(
+        path,
+        min_cell_counts=-1.0,
+        gene_threshold=0.0,
+        min_cells=3,
+    )
+    ref_X, _, ref_genes = _reference_prepare(dense, -1.0, 0.0, min_cells=3)
+
+    assert list(result.var_names) == [str(i) for i in ref_genes]
+    assert list(result.var_names) == ["2", "3"]
+    assert isinstance(result.X, sp.csr_array)
+    np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5)
+
+
+def test_min_cells_uses_obs_filtered_cohort(tmp_path):
+    import pandas as pd
+
+    dense = np.array(
+        [
+            [1, 1, 1],
+            [1, 0, 1],
+            [0, 1, 1],
+            [1, 0, 1],
+            [0, 0, 1],
+            [1, 0, 1],
+        ],
+        dtype=np.float64,
+    )
+    obs = pd.DataFrame(
+        {"condition": ["control", "treated"] * 3},
+        index=[f"cell_{i}" for i in range(dense.shape[0])],
+    )
+    adata = ad.AnnData(X=sp.csr_array(dense), obs=obs)
+    path = tmp_path / "min_cells_obs_filter.ivcsr.h5ad"
+    VCSCAnnData.from_anndata(adata, format="csr").write_h5ad(path, format="ivcsc")
+
+    obs_mask = np.asarray(obs["condition"] == "control")
+    result = load_and_normalize(
+        path,
+        min_cell_counts=-1.0,
+        gene_threshold=0.0,
+        min_cells=2,
+        obs_filter=lambda obs: obs["condition"] == "control",
+    )
+    ref_X, _, ref_genes = _reference_prepare(
+        dense[obs_mask], -1.0, 0.0, min_cells=2
+    )
+
+    assert list(result.var_names) == [str(i) for i in ref_genes]
+    assert list(result.var_names) == ["1", "2"]
+    assert isinstance(result.X, sp.csr_array)
+    np.testing.assert_allclose(result.X.toarray(), ref_X.toarray(), rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("min_cells", "error", "match"),
+    [
+        (-1, ValueError, "non-negative"),
+        (1.5, TypeError, "integer or None"),
+        (True, TypeError, "integer or None"),
+    ],
+)
+def test_min_cells_validation(tmp_path, rng, min_cells, error, match):
+    dense = rng.integers(0, 5, size=(10, 4)).astype(np.float64)
+    path = _write_ivcsr(tmp_path, dense)
+
+    with pytest.raises(error, match=match):
+        load_and_normalize(path, min_cells=min_cells)
+
+
 def test_obs_filter_matches_subset_reference(tmp_path):
     import pandas as pd