Add array type handling for normalization

.github/workflows/run_notebooks.yml

@@ -1,7 +1,6 @@
 name: Run Notebooks
 
 on:
-    - push
     - pull_request
 
 jobs:

ehrapy/_compat.py

@@ -1,7 +1,6 @@
 # Since we might check whether an object is an instance of dask.array.Array
 # without requiring dask installed in the environment.
-# This would become obsolete should dask become a requirement for ehrapy
-
+from collections.abc import Callable
 
 try:
     import dask.array as da
@@ -11,6 +10,12 @@
     DASK_AVAILABLE = False
 
 
+def _raise_array_type_not_implemented(func: Callable, type_: type) -> NotImplementedError:
+    raise NotImplementedError(
+        f"{func.__name__} does not support array type {type_}. Must be of type {func.registry.keys()}."  # type: ignore
+    )
+
+
 def is_dask_array(array):
     if DASK_AVAILABLE:
         return isinstance(array, da.Array)

ehrapy/preprocessing/_normalization.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
+from functools import singledispatch
 from typing import TYPE_CHECKING
 
 import numpy as np
 import sklearn.preprocessing as sklearn_pp
 
-from ehrapy._compat import is_dask_array
+from ehrapy._compat import _raise_array_type_not_implemented
 
 try:
+    import dask.array as da
     import dask_ml.preprocessing as daskml_pp
+
+    DASK_AVAILABLE = True
 except ImportError:
     daskml_pp = None
+    DASK_AVAILABLE = False
+
 
 from ehrapy.anndata.anndata_ext import (
     assert_numeric_vars,
@@ -69,6 +75,23 @@ def _scale_func_group(
         return None
 
 
+@singledispatch
+def _scale_norm_function(arr):
+    _raise_array_type_not_implemented(_scale_norm_function, type(arr))
+
+
+@_scale_norm_function.register
+def _(arr: np.ndarray, **kwargs):
+    return sklearn_pp.StandardScaler(**kwargs).fit_transform
+
+
+if DASK_AVAILABLE:
+
+    @_scale_norm_function.register
+    def _(arr: da.Array, **kwargs):
+        return daskml_pp.StandardScaler(**kwargs).fit_transform
+
+
 def scale_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -98,10 +121,7 @@ def scale_norm(
         >>> adata_norm = ep.pp.scale_norm(adata, copy=True)
     """
 
-    if is_dask_array(adata.X):
-        scale_func = daskml_pp.StandardScaler(**kwargs).fit_transform
-    else:
-        scale_func = sklearn_pp.StandardScaler(**kwargs).fit_transform
+    scale_func = _scale_norm_function(adata.X, **kwargs)
 
     return _scale_func_group(
         adata=adata,
@@ -113,6 +133,23 @@ def scale_norm(
     )
 
 
+@singledispatch
+def _minmax_norm_function(arr):
+    _raise_array_type_not_implemented(_minmax_norm_function, type(arr))
+
+
+@_minmax_norm_function.register
+def _(arr: np.ndarray, **kwargs):
+    return sklearn_pp.MinMaxScaler(**kwargs).fit_transform
+
+
+if DASK_AVAILABLE:
+
+    @_minmax_norm_function.register
+    def _(arr: da.Array, **kwargs):
+        return daskml_pp.MinMaxScaler(**kwargs).fit_transform
+
+
 def minmax_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -143,10 +180,7 @@ def minmax_norm(
         >>> adata_norm = ep.pp.minmax_norm(adata, copy=True)
     """
 
-    if is_dask_array(adata.X):
-        scale_func = daskml_pp.MinMaxScaler(**kwargs).fit_transform
-    else:
-        scale_func = sklearn_pp.MinMaxScaler(**kwargs).fit_transform
+    scale_func = _minmax_norm_function(adata.X, **kwargs)
 
     return _scale_func_group(
         adata=adata,
@@ -158,6 +192,16 @@ def minmax_norm(
     )
 
 
+@singledispatch
+def _maxabs_norm_function(arr):
+    _raise_array_type_not_implemented(_scale_norm_function, type(arr))
+
+
+@_maxabs_norm_function.register
+def _(arr: np.ndarray):
+    return sklearn_pp.MaxAbsScaler().fit_transform
+
+
 def maxabs_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -184,10 +228,8 @@ def maxabs_norm(
         >>> adata = ep.dt.mimic_2(encoded=True)
         >>> adata_norm = ep.pp.maxabs_norm(adata, copy=True)
     """
-    if is_dask_array(adata.X):
-        raise NotImplementedError("MaxAbsScaler is not implemented in dask_ml.")
-    else:
-        scale_func = sklearn_pp.MaxAbsScaler().fit_transform
+
+    scale_func = _maxabs_norm_function(adata.X)
 
     return _scale_func_group(
         adata=adata,
@@ -199,6 +241,23 @@ def maxabs_norm(
     )
 
 
+@singledispatch
+def _robust_scale_norm_function(arr, **kwargs):
+    _raise_array_type_not_implemented(_robust_scale_norm_function, type(arr))
+
+
+@_robust_scale_norm_function.register
+def _(arr: np.ndarray, **kwargs):
+    return sklearn_pp.RobustScaler(**kwargs).fit_transform
+
+
+if DASK_AVAILABLE:
+
+    @_robust_scale_norm_function.register
+    def _(arr: da.Array, **kwargs):
+        return daskml_pp.RobustScaler(**kwargs).fit_transform
+
+
 def robust_scale_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -229,10 +288,8 @@ def robust_scale_norm(
         >>> adata = ep.dt.mimic_2(encoded=True)
         >>> adata_norm = ep.pp.robust_scale_norm(adata, copy=True)
     """
-    if is_dask_array(adata.X):
-        scale_func = daskml_pp.RobustScaler(**kwargs).fit_transform
-    else:
-        scale_func = sklearn_pp.RobustScaler(**kwargs).fit_transform
+
+    scale_func = _robust_scale_norm_function(adata.X, **kwargs)
 
     return _scale_func_group(
         adata=adata,
@@ -244,6 +301,23 @@ def robust_scale_norm(
     )
 
 
+@singledispatch
+def _quantile_norm_function(arr):
+    _raise_array_type_not_implemented(_quantile_norm_function, type(arr))
+
+
+@_quantile_norm_function.register
+def _(arr: np.ndarray, **kwargs):
+    return sklearn_pp.QuantileTransformer(**kwargs).fit_transform
+
+
+if DASK_AVAILABLE:
+
+    @_quantile_norm_function.register
+    def _(arr: da.Array, **kwargs):
+        return daskml_pp.QuantileTransformer(**kwargs).fit_transform
+
+
 def quantile_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -273,10 +347,8 @@ def quantile_norm(
         >>> adata = ep.dt.mimic_2(encoded=True)
         >>> adata_norm = ep.pp.quantile_norm(adata, copy=True)
     """
-    if is_dask_array(adata.X):
-        scale_func = daskml_pp.QuantileTransformer(**kwargs).fit_transform
-    else:
-        scale_func = sklearn_pp.QuantileTransformer(**kwargs).fit_transform
+
+    scale_func = _quantile_norm_function(adata.X, **kwargs)
 
     return _scale_func_group(
         adata=adata,
@@ -288,6 +360,16 @@ def quantile_norm(
     )
 
 
+@singledispatch
+def _power_norm_function(arr, **kwargs):
+    _raise_array_type_not_implemented(_power_norm_function, type(arr))
+
+
+@_power_norm_function.register
+def _(arr: np.ndarray, **kwargs):
+    return sklearn_pp.PowerTransformer(**kwargs).fit_transform
+
+
 def power_norm(
     adata: AnnData,
     vars: str | Sequence[str] | None = None,
@@ -317,10 +399,8 @@ def power_norm(
         >>> adata = ep.dt.mimic_2(encoded=True)
         >>> adata_norm = ep.pp.power_norm(adata, copy=True)
     """
-    if is_dask_array(adata.X):
-        raise NotImplementedError("dask-ml has no PowerTransformer, this is only available in scikit-learn")
-    else:
-        scale_func = sklearn_pp.PowerTransformer(**kwargs).fit_transform
+
+    scale_func = _power_norm_function(adata.X, **kwargs)
 
     return _scale_func_group(
         adata=adata,

pyproject.toml

@@ -72,7 +72,7 @@ medcat = [
     "medcat",
 ]
 dask = [
-    "dask",
+    "anndata[dask]",
     "dask-ml",
 ]
 dev = [
@@ -136,7 +136,8 @@ filterwarnings = [
     "ignore:`flavor='seurat_v3'` expects raw count data, but non-integers were found:UserWarning",
     "ignore:All-NaN slice encountered:RuntimeWarning",
     "ignore:Observation names are not unique. To make them unique, call `.obs_names_make_unique`.:UserWarning",
-    "ignore:Trying to modify attribute .var of view"
+    "ignore:Trying to modify attribute `.var` of view, initializing view as actual.:anndata.ImplicitModificationWarning",
+    "ignore:Transforming to str index.:anndata.ImplicitModificationWarning:"
 ]
 minversion = 6.0
 norecursedirs = [ '.*', 'build', 'dist', '*.egg', 'data', '__pycache__']