Merge pull request #61 from theislab/trustregion_plus_sparse

Trustregion NR and IRLS and full sparse data support

Author: davidsebfischer

batchglm/data.py

@@ -13,6 +13,8 @@
 import dask
 import dask.array
 
+from .external import SparseXArrayDataArray, SparseXArrayDataSet
+
 try:
     import anndata
 except ImportError:
@@ -31,17 +33,18 @@ def fetch_X(idx):
         if idx.size == 1:
             retval = np.squeeze(retval, axis=0)
 
-        return retval.astype(np.float32)
+        return retval.astype(np.float64)
 
     delayed_fetch = dask.delayed(fetch_X, pure=True)
     X = [
         dask.array.from_delayed(
             delayed_fetch(idx),
             shape=(num_features,),
-            dtype=np.float32
+            dtype=np.float64
         ) for idx in range(num_observations)
     ]
-    X = xr.DataArray(dask.array.stack(X), dims=dims)
+
+    X = data
 
     # currently broken:
     # X = data.X
@@ -53,9 +56,9 @@ def fetch_X(idx):
 
 
 def xarray_from_data(
-        data: Union[anndata.AnnData, xr.DataArray, xr.Dataset, np.ndarray],
+        data: Union[anndata.AnnData, anndata.base.Raw, xr.DataArray, xr.Dataset, np.ndarray, scipy.sparse.csr_matrix],
         dims: Union[Tuple, List] = ("observations", "features")
-) -> xr.DataArray:
+):
     """
     Parse any array-like object, xr.DataArray, xr.Dataset or anndata.Anndata and return a xarray containing
     the observations.
@@ -64,26 +67,52 @@ def xarray_from_data(
     :param dims: tuple or list with two strings. Specifies the names of the xarray dimensions.
     :return: xr.DataArray of shape `dims`
     """
-    if anndata is not None and isinstance(data, anndata.AnnData):
+    if anndata is not None and (isinstance(data, anndata.AnnData) or isinstance(data, anndata.base.Raw)):
+        # Anndata.raw does not have obs_names.
+        if isinstance(data, anndata.AnnData):
+            obs_names = np.asarray(data.obs_names)
+        else:
+            obs_names = ["obs_" + str(i) for i in range(data.X.shape[0])]
+
         if scipy.sparse.issparse(data.X):
-            X = _sparse_to_xarray(data.X, dims=dims)
-            X.coords[dims[0]] = np.asarray(data.obs_names)
-            X.coords[dims[1]] = np.asarray(data.var_names)
+            # X = _sparse_to_xarray(data.X, dims=dims)
+            # X.coords[dims[0]] = np.asarray(data.obs_names)
+            # X.coords[dims[1]] = np.asarray(data.var_names)
+            X = SparseXArrayDataSet(
+                X=data.X,
+                obs_names=np.asarray(obs_names),
+                feature_names=np.asarray(data.var_names),
+                dims=dims
+            )
         else:
-            X = data.X
-            X = xr.DataArray(X, dims=dims, coords={
-                dims[0]: np.asarray(data.obs_names),
-                dims[1]: np.asarray(data.var_names),
-            })
+            X = xr.DataArray(
+                data.X,
+                dims=dims,
+                coords={
+                    dims[0]: np.asarray(obs_names),
+                    dims[1]: np.asarray(data.var_names),
+                }
+            )
     elif isinstance(data, xr.Dataset):
         X: xr.DataArray = data["X"]
     elif isinstance(data, xr.DataArray):
         X = data
+    elif isinstance(data, SparseXArrayDataSet):
+        X = data
+    elif scipy.sparse.issparse(data):
+        # X = _sparse_to_xarray(data, dims=dims)
+        # X.coords[dims[0]] = np.asarray(data.obs_names)
+        # X.coords[dims[1]] = np.asarray(data.var_names)
+        X = SparseXArrayDataSet(
+            X=data,
+            obs_names=None,
+            feature_names=None,
+            dims=dims
+        )
+    elif isinstance(data, np.ndarray):
+        X = xr.DataArray(data, dims=dims)
     else:
-        if scipy.sparse.issparse(data):
-            X = _sparse_to_xarray(data, dims=dims)
-        else:
-            X = xr.DataArray(data, dims=dims)
+        raise ValueError("batchglm data parsing: data format %s not recognized" % type(data))
 
     return X
 
@@ -537,7 +566,7 @@ def parse_constraints(
     Parse constraint matrix into xarray.
 
     :param dmat: Design matrix.
-    :param a constraint matrix
+    :param constraints: a constraint matrix
     :return: constraint matrix in xarray format
     """
     constraints_ar = xr.DataArray(

batchglm/external.py

@@ -0,0 +1 @@
+from batchglm.xarray_sparse import SparseXArrayDataArray, SparseXArrayDataSet

batchglm/models/base/__init__.py

@@ -1,4 +1,4 @@
-from .input import _InputData_Base, INPUT_DATA_PARAMS
+from .input import _InputData_Base, INPUT_DATA_PARAMS, SparseXArrayDataSet, SparseXArrayDataArray
 from .estimator import _Estimator_Base, _EstimatorStore_XArray_Base
 from .model import _Model_Base, _Model_XArray_Base
 from .simulator import _Simulator_Base

batchglm/models/base/external.py

@@ -1,2 +1,3 @@
 import batchglm.pkg_constants as pkg_constants
 import batchglm.data as data_utils
+from batchglm.xarray_sparse import SparseXArrayDataSet, SparseXArrayDataArray

batchglm/models/base/input.py

@@ -1,15 +1,19 @@
 import abc
 import os
 import logging
+from typing import Union
 
+import numpy as np
+import scipy
+import scipy.sparse
 import xarray as xr
 
 try:
     import anndata
 except ImportError:
     anndata = None
 
-from .external import pkg_constants, data_utils
+from .external import pkg_constants, data_utils, SparseXArrayDataSet, SparseXArrayDataArray
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +25,7 @@ class _InputData_Base:
     """
     Base class for all input data types.
     """
-    data: xr.Dataset
+    data: Union[xr.Dataset, SparseXArrayDataSet]
 
     @classmethod
     @abc.abstractmethod
@@ -57,20 +61,40 @@ def new(cls, data, observation_names=None, feature_names=None, cast_dtype=None):
             X = X.astype(cast_dtype)
             # X = X.chunk({"observations": 1})
 
-        retval = cls(xr.Dataset({
-            "X": X,
-        }, coords={
-            "feature_allzero": ~X.any(dim="observations")
-        }))
-        if observation_names is not None:
-            retval.observations = observation_names
-        elif "observations" not in retval.data.coords:
-            retval.observations = retval.data.coords["observations"]
-
-        if feature_names is not None:
-            retval.features = feature_names
-        elif "features" not in retval.data.coords:
-            retval.features = retval.data.coords["features"]
+        if scipy.sparse.issparse(X):
+            retval = cls(SparseXArrayDataSet(
+                X=X,
+                obs_names=observation_names,
+                feature_names=feature_names
+            ))
+        elif isinstance(X, SparseXArrayDataArray):
+            retval = cls(SparseXArrayDataSet(
+                X=X.X,
+                obs_names=X.coords[X.dims[0]] if observation_names is None else observation_names,
+                feature_names=X.coords[X.dims[1]] if feature_names is None else feature_names,
+                dims=X.dims
+            ))
+        elif isinstance(X, SparseXArrayDataSet):
+            retval = cls(X)
+            if observation_names is not None:
+                retval.observations = observation_names
+            if feature_names is not None:
+                retval.features = feature_names
+        else:
+            retval = cls(xr.Dataset({
+                "X": X,
+            }, coords={
+                "feature_allzero": ~X.any(dim="observations")
+            }))
+            if observation_names is not None:
+                retval.observations = observation_names
+            elif "observations" not in retval.data.coords:
+                retval.observations = retval.data.coords["observations"]
+
+            if feature_names is not None:
+                retval.features = feature_names
+            elif "features" not in retval.data.coords:
+                retval.features = retval.data.coords["features"]
 
         return retval
 
@@ -117,7 +141,7 @@ def save(self, path, group="", append=False):
         )
 
     @property
-    def X(self) -> xr.DataArray:
+    def X(self):
         return self.data.X
 
     @X.setter
@@ -156,9 +180,24 @@ def feature_isnonzero(self):
     def feature_isallzero(self):
         return self.data.coords["feature_allzero"]
 
-    def fetch_X(self, idx):
+    def fetch_X_dense(self, idx):
         return self.X[idx].values
 
+    def fetch_X_sparse(self, idx):
+        assert isinstance(self.X.X, scipy.sparse.csr_matrix), "tried to fetch sparse from non csr matrix"
+
+        data = self.X.X[idx]
+
+        data_idx = np.asarray(np.vstack(data.nonzero()).T, np.int64)
+        data_val = np.asarray(data.data, np.float64)
+        data_shape = np.asarray(data.shape, np.int64)
+
+        if idx.shape[0] == 1:
+            data_val = np.squeeze(data_val, axis=0)
+            data_idx = np.squeeze(data_idx, axis=0)
+
+        return data_idx, data_val, data_shape
+
     def set_chunk_size(self, cs: int):
         self.X = self.X.chunk({"observations": cs})
 

batchglm/models/base/model.py

@@ -62,7 +62,7 @@ def to_xarray(self, parm: Union[str, list], coords=None):
             output = xr.Dataset(output)
             if coords is not None:
                 for i in output.dims:
-                    if i in coords:
+                    if i in coords.coords:
                         output.coords[i] = coords[i]
 
         return output

batchglm/models/base_glm/external.py

@@ -3,6 +3,7 @@
 from batchglm.models.base import _Model_Base, _Model_XArray_Base
 from batchglm.models.base import _Simulator_Base
 from batchglm.models.base import INPUT_DATA_PARAMS
+from batchglm.models.base import SparseXArrayDataArray, SparseXArrayDataSet
 
 import batchglm.data as data_utils
 from batchglm.utils.linalg import groupwise_solve_lm

batchglm/models/base_glm/input.py

@@ -6,10 +6,11 @@
     anndata = None
 import xarray as xr
 import numpy as np
+import scipy.sparse
 import pandas as pd
 
 from .utils import parse_constraints, parse_design
-from .external import _InputData_Base, INPUT_DATA_PARAMS
+from .external import _InputData_Base, INPUT_DATA_PARAMS, SparseXArrayDataSet, SparseXArrayDataArray
 
 import patsy
 
@@ -33,7 +34,7 @@ def param_shapes(cls) -> dict:
     @classmethod
     def new(
             cls,
-            data: Union[np.ndarray, anndata.AnnData, xr.DataArray, xr.Dataset],
+            data: Union[np.ndarray, anndata.AnnData, xr.DataArray, xr.Dataset, scipy.sparse.csr_matrix],
             design_loc: Union[np.ndarray, pd.DataFrame, patsy.design_info.DesignMatrix, xr.DataArray] = None,
             design_loc_names: Union[list, np.ndarray, xr.DataArray] = None,
             design_scale: Union[np.ndarray, pd.DataFrame, patsy.design_info.DesignMatrix, xr.DataArray] = None,
@@ -216,12 +217,16 @@ def size_factors(self):
     def size_factors(self, data):
         if data is None and "size_factors" in self.data.coords:
             del self.data.coords["size_factors"]
-        else:
-            dims = self.param_shapes()["size_factors"]
-            self.data.coords["size_factors"] = xr.DataArray(
+
+        dims = self.param_shapes()["size_factors"]
+        sf = xr.DataArray(
                 dims=dims,
                 data=np.broadcast_to(data, [self.data.dims[d] for d in dims])
             )
+        if isinstance(self.data, SparseXArrayDataSet):
+            self.data.size_factors = sf
+        else:
+            self.data.coords["size_factors"] = sf
 
     @property
     def num_design_loc_params(self):

batchglm/models/base_glm/simulator.py

@@ -109,8 +109,14 @@ def generate_params(
             default: rand_fn = lambda shape: np.random.uniform(0.5, 2, shape)
         :param rand_fn_loc: random function taking one argument `shape`.
             If not provided, will use `rand_fn` instead.
+            This function generates location model parameters in inverse linker space,
+            ie. these parameter will be log transformed if a log linker function is used!
+            Values below 1e-08 will be set to 1e-08 to map them into the positive support.
         :param rand_fn_scale: random function taking one argument `shape`.
             If not provided, will use `rand_fn` instead.
+            This function generates scale model parameters in inverse linker space,
+            ie. these parameter will be log transformed if a log linker function is used!
+            Values below 1e-08 will be set to 1e-08 to map them into the positive support.
         """
         if rand_fn_loc is None:
             rand_fn_loc = rand_fn
@@ -157,8 +163,9 @@ def generate_params(
             dims=self.param_shapes()["a_var"],
             data=np.log(
                 np.concatenate([
-                    np.expand_dims(rand_fn_ave(self.num_features), axis=0),  # intercept
-                    rand_fn_loc((self.data.design_loc.shape[1] - 1, self.num_features))
+                    np.expand_dims(rand_fn_ave([self.num_features]), axis=0),  # intercept
+                    np.maximum(rand_fn_loc((self.data.design_loc.shape[1] - 1, self.num_features)),
+                               np.zeros([self.data.design_loc.shape[1] - 1, self.num_features]) + 1e-08)
                 ], axis=0)
             ),
             coords={"loc_params": self.data.loc_params}
@@ -167,7 +174,8 @@ def generate_params(
             dims=self.param_shapes()["b_var"],
             data=np.log(
                 np.concatenate([
-                    rand_fn_scale((self.data.design_scale.shape[1], self.num_features))
+                    np.maximum(rand_fn_scale((self.data.design_scale.shape[1], self.num_features)),
+                               np.zeros([self.data.design_scale.shape[1], self.num_features]) + 1e-08)
                 ], axis=0)
             ),
             coords={"scale_params": self.data.scale_params}
@@ -217,12 +225,12 @@ def size_factors(self):
     def size_factors(self, data):
         if data is None and "size_factors" in self.data.coords:
             del self.data.coords["size_factors"]
-        else:
-            dims = self.param_shapes()["size_factors"]
-            self.data.coords["size_factors"] = xr.DataArray(
-                dims=dims,
-                data=np.broadcast_to(data, [self.data.dims[d] for d in dims])
-            )
+
+        dims = self.param_shapes()["size_factors"]
+        self.data.coords["size_factors"] = xr.DataArray(
+            dims=dims,
+            data=np.broadcast_to(data, [self.data.dims[d] for d in dims])
+        )
 
     @property
     def a_var(self):