fixed some bugs in data parsing

davidsebfischer · davidsebfischer · commit f2aff559c841 · 2019-01-25T13:13:03.000+01:00
diff --git a/batchglm/data.py b/batchglm/data.py
@@ -56,7 +56,7 @@ 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")
 ):
     """
@@ -85,27 +85,34 @@ def xarray_from_data(
                 dims=dims
             )
         else:
-            X = xr.DataArray(data.X, dims=dims, coords={
-                dims[0]: np.asarray(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)
-            # 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
-            )
-        else:
-            X = xr.DataArray(data, dims=dims)
+        raise ValueError("batchglm data parsing: data format %s not recognized" % type(data))
 
     return X
 
diff --git a/batchglm/train/tf/glm_nb/estimator.py b/batchglm/train/tf/glm_nb/estimator.py
@@ -106,134 +106,134 @@ def init_par(
                 shape=[input_data.num_observations, input_data.num_features]
             )
 
-        groupwise_means = None
-        init_a_str = None
-        if isinstance(init_a, str):
-            init_a_str = init_a.lower()
-            # Chose option if auto was chosen
-            if init_a.lower() == "auto":
-                init_a = "closed_form"
-
-            if init_a.lower() == "closed_form":
-                #try:
-                groupwise_means, init_a, rmsd_a = closedform_nb_glm_logmu(
-                    X=input_data.X,
-                    design_loc=input_data.design_loc,
-                    constraints_loc=input_data.constraints_loc.values,
-                    size_factors=size_factors_init,
-                    link_fn=lambda mu: np.log(self.np_clip_param(mu, "mu"))
-                )
-
-                # train mu, if the closed-form solution is inaccurate
-                self._train_loc = not np.all(rmsd_a == 0)
-
-                if input_data.size_factors is not None:
-                    if np.any(input_data.size_factors != 1):
-                        self._train_loc = True
-
-                logger.debug("Using closed-form MLE initialization for mean")
-                logger.debug("Should train mu: %s", self._train_loc)
-                #except np.linalg.LinAlgError:
-                #    logger.warning("Closed form initialization failed!")
-            elif init_a.lower() == "standard":
-                if isinstance(input_data.X, SparseXArrayDataArray):
-                    overall_means = input_data.X.mean(dim="observations")
-                else:
-                    overall_means = input_data.X.mean(dim="observations").values  # directly calculate the mean
-                overall_means = self.np_clip_param(overall_means, "mu")
-
-                init_a = np.zeros([input_data.num_loc_params, input_data.num_features])
-                init_a[0, :] = np.log(overall_means)
-                self._train_loc = True
-
-                logger.debug("Using standard initialization for mean")
-                logger.debug("Should train mu: %s", self._train_loc)
-            elif init_a.lower() == "all_zero":
-                init_a = np.zeros([input_data.num_loc_params, input_data.num_features])
-                self._train_loc = True
-
-                logger.debug("Using all_zero initialization for mean")
-                logger.debug("Should train mu: %s", self._train_loc)
-            else:
-                raise ValueError("init_a string %s not recognized" % init_a)
-
-        if isinstance(init_b, str):
-            if init_b.lower() == "auto":
-                init_b = "standard"
-
-            if init_b.lower() == "closed_form" or init_b.lower() == "standard":
-                #try:
-                # Check whether it is necessary to recompute group-wise means.
-                dmats_unequal = False
-                if input_data.design_loc.shape[1] == input_data.design_scale.shape[1]:
-                    if np.any(input_data.design_loc.values != input_data.design_scale.values):
-                        dmats_unequal = True
-
-                inits_unequal = False
-                if init_a_str is not None:
-                    if init_a_str != init_b:
-                        inits_unequal = True
-
-                if inits_unequal or dmats_unequal:
-                    groupwise_means = None
-
-                # Watch out: init_mu is full obs x features matrix and is very large in many cases.
-                if inits_unequal or dmats_unequal:
+        if init_model is None:
+            groupwise_means = None
+            init_a_str = None
+            if isinstance(init_a, str):
+                init_a_str = init_a.lower()
+                # Chose option if auto was chosen
+                if init_a.lower() == "auto":
+                    init_a = "closed_form"
+
+                if init_a.lower() == "closed_form":
+                    #try:
+                    groupwise_means, init_a, rmsd_a = closedform_nb_glm_logmu(
+                        X=input_data.X,
+                        design_loc=input_data.design_loc,
+                        constraints_loc=input_data.constraints_loc.values,
+                        size_factors=size_factors_init,
+                        link_fn=lambda mu: np.log(self.np_clip_param(mu, "mu"))
+                    )
+
+                    # train mu, if the closed-form solution is inaccurate
+                    self._train_loc = not np.all(rmsd_a == 0)
+
+                    if input_data.size_factors is not None:
+                        if np.any(input_data.size_factors != 1):
+                            self._train_loc = True
+
+                    logger.debug("Using closed-form MLE initialization for mean")
+                    logger.debug("Should train mu: %s", self._train_loc)
+                    #except np.linalg.LinAlgError:
+                    #    logger.warning("Closed form initialization failed!")
+                elif init_a.lower() == "standard":
                     if isinstance(input_data.X, SparseXArrayDataArray):
-                        init_mu = np.matmul(
-                                input_data.design_loc.values,
-                                np.matmul(input_data.constraints_loc.values, init_a)
-                        )
+                        overall_means = input_data.X.mean(dim="observations")
                     else:
-                        init_a_xr = data_utils.xarray_from_data(init_a, dims=("loc_params", "features"))
-                        init_a_xr.coords["loc_params"] = input_data.constraints_loc.coords["loc_params"]
-                        init_mu = input_data.design_loc.dot(input_data.constraints_loc.dot(init_a_xr))
+                        overall_means = input_data.X.mean(dim="observations").values  # directly calculate the mean
+                    overall_means = self.np_clip_param(overall_means, "mu")
 
-                    if size_factors_init is not None:
-                        init_mu = init_mu + np.log(size_factors_init)
-                    init_mu = np.exp(init_mu)
-                else:
-                    init_mu = None
+                    init_a = np.zeros([input_data.num_loc_params, input_data.num_features])
+                    init_a[0, :] = np.log(overall_means)
+                    self._train_loc = True
 
-                if init_b.lower() == "closed_form":
-                    groupwise_scales, init_b, rmsd_b = closedform_nb_glm_logphi(
-                        X=input_data.X,
-                        mu=init_mu,
-                        design_scale=input_data.design_scale,
-                        constraints=input_data.constraints_scale.values,
-                        size_factors=size_factors_init,
-                        groupwise_means=groupwise_means,
-                        link_fn=lambda r: np.log(self.np_clip_param(r, "r"))
-                    )
+                    logger.debug("Using standard initialization for mean")
+                    logger.debug("Should train mu: %s", self._train_loc)
+                elif init_a.lower() == "all_zero":
+                    init_a = np.zeros([input_data.num_loc_params, input_data.num_features])
+                    self._train_loc = True
 
-                    logger.debug("Using closed-form MME initialization for dispersion")
-                    logger.debug("Should train r: %s", self._train_scale)
-                elif init_b.lower() == "standard":
-                    groupwise_scales, init_b_intercept, rmsd_b = closedform_nb_glm_logphi(
-                        X=input_data.X,
-                        mu=init_mu,
-                        design_scale=input_data.design_scale[:,[0]],
-                        constraints=input_data.constraints_scale[[0], [0]].values,
-                        size_factors=size_factors_init,
-                        groupwise_means=None,
-                        link_fn=lambda r: np.log(self.np_clip_param(r, "r"))
-                    )
+                    logger.debug("Using all_zero initialization for mean")
+                    logger.debug("Should train mu: %s", self._train_loc)
+                else:
+                    raise ValueError("init_a string %s not recognized" % init_a)
+
+            if isinstance(init_b, str):
+                if init_b.lower() == "auto":
+                    init_b = "standard"
+
+                if init_b.lower() == "closed_form" or init_b.lower() == "standard":
+                    #try:
+                    # Check whether it is necessary to recompute group-wise means.
+                    dmats_unequal = False
+                    if input_data.design_loc.shape[1] == input_data.design_scale.shape[1]:
+                        if np.any(input_data.design_loc.values != input_data.design_scale.values):
+                            dmats_unequal = True
+
+                    inits_unequal = False
+                    if init_a_str is not None:
+                        if init_a_str != init_b:
+                            inits_unequal = True
+
+                    if inits_unequal or dmats_unequal:
+                        groupwise_means = None
+
+                    # Watch out: init_mu is full obs x features matrix and is very large in many cases.
+                    if inits_unequal or dmats_unequal:
+                        if isinstance(input_data.X, SparseXArrayDataArray):
+                            init_mu = np.matmul(
+                                    input_data.design_loc.values,
+                                    np.matmul(input_data.constraints_loc.values, init_a)
+                            )
+                        else:
+                            init_a_xr = data_utils.xarray_from_data(init_a, dims=("loc_params", "features"))
+                            init_a_xr.coords["loc_params"] = input_data.constraints_loc.coords["loc_params"]
+                            init_mu = input_data.design_loc.dot(input_data.constraints_loc.dot(init_a_xr))
+
+                        if size_factors_init is not None:
+                            init_mu = init_mu + np.log(size_factors_init)
+                        init_mu = np.exp(init_mu)
+                    else:
+                        init_mu = None
+
+                    if init_b.lower() == "closed_form":
+                        groupwise_scales, init_b, rmsd_b = closedform_nb_glm_logphi(
+                            X=input_data.X,
+                            mu=init_mu,
+                            design_scale=input_data.design_scale,
+                            constraints=input_data.constraints_scale.values,
+                            size_factors=size_factors_init,
+                            groupwise_means=groupwise_means,
+                            link_fn=lambda r: np.log(self.np_clip_param(r, "r"))
+                        )
+
+                        logger.debug("Using closed-form MME initialization for dispersion")
+                        logger.debug("Should train r: %s", self._train_scale)
+                    elif init_b.lower() == "standard":
+                        groupwise_scales, init_b_intercept, rmsd_b = closedform_nb_glm_logphi(
+                            X=input_data.X,
+                            mu=init_mu,
+                            design_scale=input_data.design_scale[:,[0]],
+                            constraints=input_data.constraints_scale[[0], [0]].values,
+                            size_factors=size_factors_init,
+                            groupwise_means=None,
+                            link_fn=lambda r: np.log(self.np_clip_param(r, "r"))
+                        )
+                        init_b = np.zeros([input_data.num_scale_params, input_data.X.shape[1]])
+                        init_b[0, :] = init_b_intercept
+
+                        logger.debug("Using closed-form MME initialization for dispersion")
+                        logger.debug("Should train r: %s", self._train_scale)
+                    #except np.linalg.LinAlgError:
+                    #    logger.warning("Closed form initialization failed!")
+                elif init_b.lower() == "all_zero":
                     init_b = np.zeros([input_data.num_scale_params, input_data.X.shape[1]])
-                    init_b[0, :] = init_b_intercept
 
-                    logger.debug("Using closed-form MME initialization for dispersion")
+                    logger.debug("Using standard initialization for dispersion")
                     logger.debug("Should train r: %s", self._train_scale)
-                #except np.linalg.LinAlgError:
-                #    logger.warning("Closed form initialization failed!")
-            elif init_b.lower() == "all_zero":
-                init_b = np.zeros([input_data.num_scale_params, input_data.X.shape[1]])
-
-                logger.debug("Using standard initialization for dispersion")
-                logger.debug("Should train r: %s", self._train_scale)
-            else:
-                raise ValueError("init_b string %s not recognized" % init_b)
-
-        if init_model is not None:
+                else:
+                    raise ValueError("init_b string %s not recognized" % init_b)
+        else:
             # Locations model:
             if isinstance(init_a, str) and (init_a.lower() == "auto" or init_a.lower() == "init_model"):
                 my_loc_names = set(input_data.design_loc_names.values)
diff --git a/batchglm/unit_test/base_glm/test_data_types_glm.py b/batchglm/unit_test/base_glm/test_data_types_glm.py
@@ -92,8 +92,9 @@ def basic_test(
         return estimator.test_estimation()
 
     def _test_numpy(self, sparse):
+        X = self.sim.X
         if sparse:
-            X = scipy.sparse.csr_matrix(self.sim.X)
+            X = scipy.sparse.csr_matrix(X)
 
         success = self.basic_test(
             data=X,
diff --git a/batchglm/xarray_sparse/base.py b/batchglm/xarray_sparse/base.py
@@ -116,23 +116,27 @@ def square(self, copy=False):
         else:
             self.X = new_x
 
-    def mean(self, dim=None):
+    def mean(self, dim: str = None, axis: int = None):
+        assert not (dim is not None and axis is not None), "only supply dim or axis"
         if dim is not None:
             assert dim in self.dims, "dim not recognized"
             axis = self.dims.index(dim)
             return np.asarray(self.X.mean(axis=axis)).flatten()
+        elif axis is not None:
+            assert axis < len(self.X.shape), "axis index out of range"
+            return np.asarray(self.X.mean(axis=axis)).flatten()
         else:
             return np.asarray(self.X.mean()).flatten()
 
-    def var(self, dim):
+    def var(self, dim: str):
         assert dim in self.dims, "dim not recognized"
         axis = self.dims.index(dim)
         Xsq = self.square(copy=True)
         expect_x_sq = np.square(self.mean(dim=dim))
         expect_xsq = np.mean(Xsq, axis=axis)
         return np.asarray(expect_xsq - expect_x_sq).flatten()
 
-    def std(self, dim):
+    def std(self, dim: str):
         return np.sqrt(self.var(dim=dim))
 
     def groupby(self, key):
@@ -215,6 +219,10 @@ def __init__(
     def ndim(self):
         return len(self.dims)
 
+    @property
+    def shape(self):
+        return self.X.shape
+
     @property
     def feature_allzero(self):
         return self.X.coords["feature_allzero"]
