fixed some functions for interface to diffxpy

Author: davidsebfischer

batchglm/data.py

@@ -566,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/models/base_glm/utils.py

@@ -77,7 +77,7 @@ def parse_constraints(
         dmat: xr.Dataset,
         dims,
         constraints: np.ndarray = None,
-        constraint_par_names: list = None,
+        constraint_par_names: list = None
 ) -> xr.DataArray:
     r"""
     Parser for constraint matrices.
@@ -90,7 +90,11 @@ def parse_constraints(
     """
     if constraints is None:
         constraints = np.identity(n=dmat.shape[1])
+        # Use given parameter names if constraint matrix is identity.
+        par_names = dmat.coords[dims[0]]
     else:
+        # Cannot use given parameter names if constraint matrix is not identity: Make up new ones.
+        par_names = ["var_"+str(x) for x in range(constraints.shape[1])]
         assert constraints.shape[0] == dmat.shape[1], "constraint dimension mismatch"
 
     constraints_mat = xr.DataArray(
@@ -99,7 +103,7 @@ def parse_constraints(
     )
     constraints_mat.coords[dims[0]] = dmat.coords[dims[0]]
     if constraint_par_names is None:
-        constraint_par_names = ["var_"+str(x) for x in range(constraints_mat.shape[1])]
+        constraint_par_names = par_names
 
     constraints_mat.coords[dims[1]] = constraint_par_names
 

batchglm/train/tf/glm_nb/estimator.py

@@ -187,7 +187,7 @@ def init_par(
                             )
                         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_a_xr.coords["loc_params"] = input_data.constraints_loc.coords["loc_params"].values
                             init_mu = input_data.design_loc.dot(input_data.constraints_loc.dot(init_a_xr))
 
                         if size_factors_init is not None:
@@ -236,28 +236,28 @@ def init_par(
         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)
-                my_loc_names = my_loc_names.intersection(init_model.input_data.design_loc_names.values)
+                my_loc_names = set(input_data.loc_names.values)
+                my_loc_names = my_loc_names.intersection(set(init_model.input_data.loc_names.values))
 
                 init_loc = np.zeros([input_data.num_loc_params, input_data.num_features])
                 for parm in my_loc_names:
-                    init_idx = np.where(init_model.input_data.design_loc_names == parm)
-                    my_idx = np.where(input_data.design_loc_names == parm)
-                    init_loc[my_idx] = init_model.par_link_loc[init_idx]
+                    init_idx = np.where(init_model.input_data.loc_names == parm)[0]
+                    my_idx = np.where(input_data.loc_names == parm)[0]
+                    init_loc[my_idx] = init_model.a_var[init_idx]
 
                 init_a = init_loc
                 logger.debug("Using initialization based on input model for mean")
 
             # Scale model:
             if isinstance(init_b, str) and (init_b.lower() == "auto" or init_b.lower() == "init_model"):
-                my_scale_names = set(input_data.design_scale_names.values)
-                my_scale_names = my_scale_names.intersection(init_model.input_data.design_scale_names.values)
+                my_scale_names = set(input_data.scale_names.values)
+                my_scale_names = my_scale_names.intersection(init_model.input_data.scale_names.values)
 
                 init_scale = np.zeros([input_data.num_scale_params, input_data.num_features])
                 for parm in my_scale_names:
-                    init_idx = np.where(init_model.input_data.design_scale_names == parm)
-                    my_idx = np.where(input_data.design_scale_names == parm)
-                    init_scale[my_idx] = init_model.par_link_scale[init_idx]
+                    init_idx = np.where(init_model.input_data.scale_names == parm)[0]
+                    my_idx = np.where(input_data.scale_names == parm)[0]
+                    init_scale[my_idx] = init_model.b_var[init_idx]
 
                 init_b = init_scale
                 logger.debug("Using initialization based on input model for dispersion")

batchglm/xarray_sparse/base.py

@@ -62,6 +62,9 @@ def new_from_x(self, x):
     def dtype(self):
         return self.X.dtype
 
+    def astype(self, dtype):
+        return self.new_from_x(self.X.astype(dtype))
+
     @property
     def shape(self):
         return self.X.shape
@@ -128,12 +131,19 @@ def mean(self, dim: str = None, axis: int = None):
         else:
             return np.asarray(self.X.mean()).flatten()
 
-    def var(self, dim: str):
-        assert dim in self.dims, "dim not recognized"
-        axis = self.dims.index(dim)
+    def var(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)
+        elif axis is not None:
+            assert axis < len(self.X.shape), "axis index out of range"
+        else:
+            assert False, "supply either dim or axis"
+
         Xsq = self.square(copy=True)
-        expect_x_sq = np.square(self.mean(dim=dim))
-        expect_xsq = np.mean(Xsq, axis=axis)
+        expect_x_sq = np.square(self.mean(axis=axis))
+        expect_xsq = Xsq.mean(axis=axis)
         return np.asarray(expect_xsq - expect_x_sq).flatten()
 
     def std(self, dim: str):
@@ -177,6 +187,14 @@ def group_vars(self, dim):
     def __copy__(self):
         return type(self)(self.X)
 
+    def __getitem__(self, key):
+        if isinstance(key, np.ndarray) or isinstance(key, slice):  # This is an observation wise slice!
+            return self.new_from_x(x=self.X[key])
+        elif key in self.coords:
+            return self.coords[key]
+        else:
+            return self.__getattribute__(key)
+
 
 class SparseXArrayDataSet:
     """