Add sparse.COO support to harmonize_missing_values and infer_feature_types

docs/conf.py

@@ -159,6 +159,7 @@
     ("py:class", "anndata.compat.CupySparseMatrix"),
     ("py:class", "sparse.numba_backend._coo.core.COO"),
     ("py:class", "sparse._coo.core.COO"),
+    ("py:class", "sparse.COO"),
     ("py:class", "pandas.core.arrays.base.ExtensionArray"),
 ]
 

src/ehrdata/_feature_types.py

@@ -1,25 +1,31 @@
 from __future__ import annotations
 
 import warnings
-from functools import wraps
+from functools import singledispatch, wraps
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import pandas as pd
+import scipy.sparse as sp
+import sparse
 from dateutil.parser import isoparse  # type: ignore
 from fast_array_utils.conv import to_dense
 from rich import print
 from rich.tree import Tree
 from scipy.sparse import issparse
 
 from ehrdata._logger import logger
+from ehrdata._types import CSArray, CSMatrix
 from ehrdata.core.constants import CATEGORICAL_TAG, DATE_TAG, FEATURE_TYPE_KEY, MISSING_VALUES, NUMERIC_TAG
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
     from ehrdata import EHRData
 
+# sparse array types
+_FAST_SPARSE_TYPES = (sparse.COO, CSMatrix, CSArray)
+
 
 def _detect_feature_type(
     col: pd.Series,
@@ -70,6 +76,126 @@ def _detect_feature_type(
     return NUMERIC_TAG, False  # type: ignore
 
 
+def _classify_sparse_binary_columns(
+    var_coord: np.ndarray,
+    data: np.ndarray,
+    n_vars: int,
+    entries_per_var: int,
+    *,
+    fill_value,
+    fill_is_nan: bool,
+) -> np.ndarray:
+    """Per variable, decides if its values (explicit + implicit) are exactly `{0, 1}`."""
+    is_float = np.issubdtype(data.dtype, np.floating)
+    nan_mask = np.isnan(data) if is_float else np.zeros(data.shape, dtype=bool)
+
+    has_zero = np.zeros(n_vars, dtype=bool)
+    has_one = np.zeros(n_vars, dtype=bool)
+    has_other = np.zeros(n_vars, dtype=bool)
+    stored_count = np.zeros(n_vars, dtype=np.int64)
+    non_nan_count = np.zeros(n_vars, dtype=np.int64)
+
+    np.logical_or.at(has_zero, var_coord, (data == 0) & ~nan_mask)
+    np.logical_or.at(has_one, var_coord, (data == 1) & ~nan_mask)
+    np.logical_or.at(has_other, var_coord, ~nan_mask & (data != 0) & (data != 1))
+    np.add.at(stored_count, var_coord, 1)
+    np.add.at(non_nan_count, var_coord, ~nan_mask)
+
+    has_implicit = stored_count < entries_per_var
+    if not fill_is_nan:
+        non_nan_count += has_implicit * (entries_per_var - stored_count)
+        if fill_value == 0:
+            has_zero |= has_implicit
+        elif fill_value == 1:
+            has_one |= has_implicit
+        else:
+            has_other |= has_implicit
+
+    all_nan = np.nonzero(non_nan_count == 0)[0]
+    if all_nan.size:
+        err_msg = f"Feature at position {all_nan[0]} contains only NaN values. Please drop this feature to infer the feature type."
+        raise ValueError(err_msg)
+
+    return has_zero & has_one & ~has_other
+
+
+def _binary_mask_to_results(
+    is_binary: np.ndarray, binary_as: Literal["categorical", "numeric"]
+) -> list[tuple[Literal["categorical", "numeric"], bool]]:
+    binary_result = (NUMERIC_TAG, False) if binary_as == "numeric" else (CATEGORICAL_TAG, True)
+    numeric_result = (NUMERIC_TAG, False)
+    return [binary_result if b else numeric_result for b in is_binary]
+
+
+def _detect_feature_types_sparse_coo_non_numeric(
+    X: sparse.COO,
+    *,
+    binary_as: Literal["categorical", "numeric"] = "categorical",
+) -> list[tuple[Literal["date", "categorical", "numeric"], bool]]:
+    """Detect the feature type of every variable of a non-numeric `sparse.COO` array."""
+    n_vars = X.shape[1]
+    entries_per_var = X.size // n_vars
+
+    order = np.argsort(X.coords[1], kind="stable")
+    sorted_vars = X.coords[1][order]
+    sorted_data = X.data[order]
+    starts = np.searchsorted(sorted_vars, np.arange(n_vars))
+    ends = np.searchsorted(sorted_vars, np.arange(n_vars), side="right")
+
+    results = []
+    for j in range(n_vars):
+        if ends[j] - starts[j] < entries_per_var:
+            # rare: a genuine implicit fill value among non-numeric data, densify just this column
+            col = pd.Series(X[:, j].todense().reshape(-1))
+        else:
+            col = pd.Series(sorted_data[starts[j] : ends[j]])
+        results.append(_detect_feature_type(col, binary_as=binary_as))
+
+    return results
+
+
+@singledispatch
+def _detect_feature_types_sparse(X, *, binary_as: Literal["categorical", "numeric"] = "categorical"):
+    """Detect the feature type of every variable (axis 1) of a sparse array, without densifying."""
+    msg = f"_detect_feature_types_sparse does not support array type {type(X)}."
+    raise NotImplementedError(msg)
+
+
+@_detect_feature_types_sparse.register(sparse.COO)
+def _(X: sparse.COO, *, binary_as: Literal["categorical", "numeric"] = "categorical"):
+    # sparse.COO isn't restricted to numeric dtypes in memory (only ehrdata's on-disk binsparse
+    # writer enforces that); route non-numeric arrays through the dedicated per-column path.
+    if not (np.issubdtype(X.dtype, np.number) or np.issubdtype(X.dtype, np.bool_)):
+        return _detect_feature_types_sparse_coo_non_numeric(X, binary_as=binary_as)
+
+    n_vars = X.shape[1]
+    entries_per_var = X.size // n_vars  # rows * time steps (for 3D) per variable
+    fill_is_nan = np.issubdtype(X.dtype, np.floating) and np.isnan(X.fill_value)
+    is_binary = _classify_sparse_binary_columns(
+        X.coords[1], X.data, n_vars, entries_per_var, fill_value=X.fill_value, fill_is_nan=fill_is_nan
+    )
+    return _binary_mask_to_results(is_binary, binary_as)
+
+
+@_detect_feature_types_sparse.register(sp.csr_array)
+@_detect_feature_types_sparse.register(sp.csr_matrix)
+def _(X, *, binary_as: Literal["categorical", "numeric"] = "categorical"):
+    # CSR stores each entry's column index in .indices
+    n_vars = X.shape[1]
+    is_binary = _classify_sparse_binary_columns(X.indices, X.data, n_vars, X.shape[0], fill_value=0, fill_is_nan=False)
+    return _binary_mask_to_results(is_binary, binary_as)
+
+
+@_detect_feature_types_sparse.register(sp.csc_array)
+@_detect_feature_types_sparse.register(sp.csc_matrix)
+def _(X, *, binary_as: Literal["categorical", "numeric"] = "categorical"):
+    # CSC stores entries grouped by column via .indptr
+    n_vars = X.shape[1]
+    var_coord = np.repeat(np.arange(n_vars), np.diff(X.indptr))
+    is_binary = _classify_sparse_binary_columns(var_coord, X.data, n_vars, X.shape[0], fill_value=0, fill_is_nan=False)
+    return _binary_mask_to_results(is_binary, binary_as)
+
+
 def infer_feature_types(
     edata: EHRData,
     *,
@@ -120,24 +246,31 @@ def infer_feature_types(
             stacklevel=2,
         )
 
-    if issparse(X):
-        X = to_dense(X)
+    # The non-densifying fast path applies to the array types registered on _detect_feature_types_sparse
+    # (sparse.COO handles numeric and non-numeric dtypes internally; scipy.sparse is always numeric).
+    if isinstance(X, _FAST_SPARSE_TYPES):
+        sparse_results = _detect_feature_types_sparse(X, binary_as=binary_as)
+    else:
+        sparse_results = None
 
-    if X.ndim == 3:
-        n_obs, n_vars, n_t = X.shape
-        X = X.transpose(0, 2, 1).reshape(n_obs * n_t, n_vars)
+        if X.ndim == 3:
+            n_obs, n_vars, n_t = X.shape
+            X = X.transpose(0, 2, 1).reshape(n_obs * n_t, n_vars)
 
-    df = pd.DataFrame(X.reshape(-1, edata.shape[1]), columns=edata.var_names)
+        df = pd.DataFrame(X.reshape(-1, edata.shape[1]), columns=edata.var_names)
 
-    for feature in edata.var_names:
+    for i, feature in enumerate(edata.var_names):
         if (
             FEATURE_TYPE_KEY in edata.var
             and edata.var[FEATURE_TYPE_KEY][feature] is not None
             and not pd.isna(edata.var[FEATURE_TYPE_KEY][feature])
         ):
             feature_types[feature] = edata.var[FEATURE_TYPE_KEY][feature]
         else:
-            feature_types[feature], raise_warning = _detect_feature_type(df[feature], binary_as=binary_as)
+            if sparse_results is not None:
+                feature_types[feature], raise_warning = sparse_results[i]
+            else:
+                feature_types[feature], raise_warning = _detect_feature_type(df[feature], binary_as=binary_as)
             if raise_warning:
                 uncertain_features.append(feature)
 
@@ -291,21 +424,98 @@ def replace_feature_types(
     edata.var.loc[features, FEATURE_TYPE_KEY] = corrected_type
 
 
+@singledispatch
+def _harmonize_missing_values_numeric(X, *, var_names: pd.Index, vars: Iterable[str] | None) -> tuple[object, bool]:
+    """Treat the implicit zero fill value of a numeric array as missing, if that array type is ambiguous about it."""
+    return X, False
+
+
+@_harmonize_missing_values_numeric.register(sparse.COO)
+def _(X: sparse.COO, *, var_names: pd.Index, vars: Iterable[str] | None) -> tuple[sparse.COO, bool]:
+    """Swap a numeric :class:`sparse.COO` array's zero fill value to `np.nan`, keeping `vars` columns real."""
+    if np.issubdtype(X.dtype, np.bool_) or X.fill_value != 0:
+        return X, False
+
+    excluded = (
+        np.asarray([var_names.get_loc(v) for v in vars], dtype=np.intp)
+        if vars is not None
+        else np.array([], dtype=np.intp)
+    )
+
+    var_coord = X.coords[1]
+    is_excluded = np.isin(var_coord, excluded) if excluded.size else np.zeros(var_coord.shape, dtype=bool)
+
+    data = X.data.astype(np.float64, copy=True)
+    data[(data == 0) & ~is_excluded] = np.nan
+    coords = X.coords
+
+    if excluded.size:
+        # Materialize every position of the excluded columns that isn't already stored explicitly,
+        # since those positions would otherwise silently read back as the new np.nan fill value.
+        excluded_shape = list(X.shape)
+        excluded_shape[1] = excluded.size
+        stored = np.zeros(excluded_shape, dtype=bool)
+
+        pos_in_excluded = np.full(X.shape[1], -1, dtype=np.intp)
+        pos_in_excluded[excluded] = np.arange(excluded.size)
+
+        if is_excluded.any():
+            index = (
+                coords[0][is_excluded],
+                pos_in_excluded[var_coord[is_excluded]],
+                *(coords[ax][is_excluded] for ax in range(2, X.ndim)),
+            )
+            stored[index] = True
+
+        missing = np.argwhere(~stored)
+        if missing.size:
+            missing = missing.T.copy()
+            missing[1] = excluded[missing[1]]
+            coords = np.concatenate([coords, missing], axis=1)
+            data = np.concatenate([data, np.zeros(missing.shape[1], dtype=np.float64)])
+
+    return sparse.COO(coords, data, shape=X.shape, fill_value=np.nan), True
+
+
+@_harmonize_missing_values_numeric.register(sp.csr_array)
+@_harmonize_missing_values_numeric.register(sp.csr_matrix)
+@_harmonize_missing_values_numeric.register(sp.csc_array)
+@_harmonize_missing_values_numeric.register(sp.csc_matrix)
+def _(X, *, var_names: pd.Index, vars: Iterable[str] | None) -> tuple[object, bool]:
+    err_msg = (
+        f"ed.harmonize_missing_values cannot treat the implicit zero fill value as missing for "
+        f"array type {type(X)}. Please convert to a dense array or sparse.COO first."
+    )
+    raise NotImplementedError(err_msg)
+
+
 def harmonize_missing_values(
     edata: EHRData,
     *,
     layer: str | None = None,
     missing_values: Iterable[str] | None = ["nan", "np.nan", "<NA>", "pd.NA"],
+    vars: Iterable[str] | None = None,
     copy: bool = False,
 ) -> EHRData | None:
     """Harmonize missing values in the :class:`~ehrdata.EHRData` object.
 
     This function will replace strings that are considered to represent missing values with `np.nan`.
 
+    For a numeric :class:`sparse.COO` layer, the implicit fill value `0` is ambiguous between "not
+    measured" and "measured as zero". This function treats it as missing for every variable except
+    those listed in `vars`.
+    Columns not in `vars` stay sparse; columns in `vars` keep `0` as a real value,
+    which requires materializing their previously-implicit zeros as explicit stored entries so they
+    don't pick up the new `np.nan` fill value. Boolean layers (sparse or dense) are also
+    left untouched, as `0`/`False` is not ambiguous for them.
+
     Args:
         edata: Data object.
         layer: The layer to use from the :class:`~ehrdata.EHRData` object. If `None`, the `X` layer is used.
         missing_values: The strings that are considered to represent missing values and should be replaced with np.nan
+        vars: For a sparse :class:`sparse.COO` layer only: variable names whose `0` values represent a
+            real, measured zero rather than a missing value, and so are excluded from the
+            zero-as-missing treatment described above.
         copy: Whether to return a copy of the :class:`~ehrdata.EHRData` object with the missing values replaced.
 
     Examples:
@@ -317,10 +527,31 @@ def harmonize_missing_values(
         edata = edata.copy()
     X = edata.X if layer is None else edata.layers[layer]
 
-    # note that every scipy sparse array is of a numeric dtype and will enter this if block
-    # further sparse.COO, while being allowed in theory to be str dtype, is only allowed numeric dtypes under binsparse specification which we follow closely
+    # reject non numeric sparse.COO explicitly rather than silently densifying
+    if isinstance(X, sparse.COO) and not (np.issubdtype(X.dtype, np.number) or np.issubdtype(X.dtype, np.bool_)):
+        err_msg = (
+            f"ed.harmonize_missing_values can only be used on a numeric sparse.COO layer "
+            f"({'X' if layer is None else layer!r} has dtype {X.dtype})."
+        )
+        raise NotImplementedError(err_msg)
+
+    # every scipy sparse array is of a numeric dtype and will enter this if block
     if np.issubdtype(X.dtype, np.number) or np.issubdtype(X.dtype, np.bool_):
-        logger.debug(f"ed.harmonize_missing_values does not affect numeric layer {'X' if layer is None else layer}.")
+        harmonized, changed = _harmonize_missing_values_numeric(X, var_names=edata.var_names, vars=vars)
+
+        if changed:
+            logger.debug(
+                f"ed.harmonize_missing_values treats the implicit zero fill value as missing in sparse.COO layer {'X' if layer is None else layer}."
+            )
+            if layer is None:
+                edata.X = harmonized
+            else:
+                edata.layers[layer] = harmonized
+        else:
+            logger.debug(
+                f"ed.harmonize_missing_values does not affect numeric layer {'X' if layer is None else layer}."
+            )
+
         return edata if copy else None
 
     df = pd.DataFrame(X.reshape(-1, edata.shape[1]), columns=edata.var_names)
@@ -335,12 +566,17 @@ def harmonize_missing_values(
 
 
 def _harmonize_on_read(edata: EHRData) -> None:
-    if edata.X is not None:
-        harmonize_missing_values(edata)
-        logger.info("Harmonizing missing values of X")
-
-    for key in edata.layers:
-        if key is None:  # anndata 0.13: the unified `.X` slot, already harmonized above
-            continue
-        harmonize_missing_values(edata, layer=key)
-        logger.info(f"Harmonizing missing values of layer {key}")
+    layers = [None] if edata.X is not None else []
+    # anndata 0.13's unified `.X` slot shows up as key None in edata.layers too; skip the duplicate
+    layers += [key for key in edata.layers if key is not None]
+
+    for layer in layers:
+        label = "X" if layer is None else f"layer {layer}"
+        try:
+            harmonize_missing_values(edata, layer=layer)
+            logger.info(f"Harmonizing missing values of {label}")
+        except NotImplementedError:
+            # scipy.sparse (CSR/CSC) can't be harmonized without densifying
+            logger.debug(
+                f"Skipping missing-value harmonization of {label}: not supported for scipy.sparse without densifying."
+            )

tests/io/test_h5ed.py

@@ -230,7 +230,7 @@ def test_write_read_h5ed_sparse_coo_3d(slot, tmp_path):
         assert {"indices_0", "indices_1", "indices_2", "values", "fill_value"} <= set(group.keys())
         assert group["indices_0"].dtype == np.int64  # indices are always written as int64 on disk
 
-    edata_read = read_h5ed(path)
+    edata_read = read_h5ed(path, harmonize_missing_values=False)
     restored = edata_read.X if slot == "X" else edata_read.layers["tem_data"]
     assert isinstance(restored, sparse.COO)
     assert np.array_equal(restored.todense(), dense)

tests/io/test_zarr.py

@@ -248,7 +248,7 @@ def test_write_read_zarr_sparse_coo_3d(slot, chunks, tmp_path):
     assert {"indices_0", "indices_1", "indices_2", "values", "fill_value"} <= set(group.keys())
     assert group["indices_0"].dtype == np.int64  # indices are always written as int64 on disk
 
-    edata_read = read_zarr(path)
+    edata_read = read_zarr(path, harmonize_missing_values=False)
     restored = edata_read.X if slot == "X" else edata_read.layers["tem_data"]
     assert isinstance(restored, sparse.COO)
     assert np.array_equal(restored.todense(), dense)