backport cleaned-up closed form initialization; fix some unit test bugs

Hoeze · Hoeze · commit 1d0abc0daab1 · 2018-11-21T12:03:04.000+01:00
diff --git a/batchglm/api/utils/__init__.py b/batchglm/api/utils/__init__.py
@@ -1,2 +1,4 @@
 from . import stats
 from . import random
+from . import numeric
+from . import linalg
diff --git a/batchglm/api/utils/linalg.py b/batchglm/api/utils/linalg.py
@@ -0,0 +1 @@
+from batchglm.utils.linalg import stacked_lstsq, groupwise_solve_lm
diff --git a/batchglm/api/utils/numeric.py b/batchglm/api/utils/numeric.py
@@ -0,0 +1 @@
+from batchglm.utils.numeric import combine_matrices, softmax, weighted_mean, weighted_variance
diff --git a/batchglm/models/nb_glm/utils.py b/batchglm/models/nb_glm/utils.py
@@ -0,0 +1,139 @@
+from typing import Union
+
+import numpy as np
+import xarray as xr
+
+from batchglm.utils.linalg import groupwise_solve_lm
+from batchglm.utils.numeric import weighted_mean
+from batchglm.models.glm import closedform_glm_mean
+
+
+def closedform_nb_glm_logmu(
+        X: xr.DataArray,
+        design_loc,
+        constraints=None,
+        size_factors=None,
+        weights=None,
+        link_fn=np.log
+):
+    r"""
+    Calculates a closed-form solution for the `mu` parameters of negative-binomial GLMs.
+
+    :param X: The sample data
+    :param design_loc: design matrix for location
+    :param constraints: some design constraints
+    :param size_factors: size factors for X
+    :param weights: the weights of the arrays' elements; if `none` it will be ignored.
+    :return: tuple: (groupwise_means, mu, rmsd)
+    """
+    return closedform_glm_mean(
+        X=X,
+        dmat=design_loc,
+        constraints=constraints,
+        size_factors=size_factors,
+        weights=weights,
+        link_fn=link_fn
+    )
+
+
+def closedform_nb_glm_logphi(
+        X: xr.DataArray,
+        design_scale: xr.DataArray,
+        constraints=None,
+        size_factors=None,
+        weights: Union[np.ndarray, xr.DataArray] = None,
+        mu=None,
+        groupwise_means=None,
+        link_fn=np.log
+):
+    r"""
+    Calculates a closed-form solution for the log-scale parameters of negative-binomial GLMs.
+    Based on the Method-of-Moments estimator.
+
+    :param X: The sample data
+    :param design_scale: design matrix for scale
+    :param constraints: some design constraints
+    :param size_factors: size factors for X
+    :param weights: the weights of the arrays' elements; if `none` it will be ignored.
+    :param mu: optional, if there are for example different mu's per observation.
+
+        Used to calculate `Xdiff = X - mu`.
+    :param groupwise_means: optional, in case if already computed this can be specified to spare double-calculation
+    :return: tuple (groupwise_scales, logphi, rmsd)
+    """
+    if size_factors is not None:
+        X = np.divide(X, size_factors)
+
+    # to circumvent nonlocal error
+    provided_groupwise_means = groupwise_means
+    provided_weights = weights
+    provided_mu = mu
+
+    def apply_fun(grouping):
+        grouped_X = X.assign_coords(group=((X.dims[0],), grouping))
+
+        # convert weights into a xr.DataArray
+        if provided_weights is not None:
+            weights = xr.DataArray(
+                data=provided_weights,
+                dims=(X.dims[0],),
+                coords={
+                    "group": ((X.dims[0],), grouping),
+                }
+            )
+        else:
+            weights = None
+
+        # calculate group-wise means if necessary
+        if provided_groupwise_means is None:
+            if weights is None:
+                groupwise_means = grouped_X.mean(X.dims[0]).values
+            else:
+                # for each group: calculate weighted mean
+                groupwise_means: xr.DataArray = xr.concat([
+                    weighted_mean(d, w, axis=0) for (g, d), (g, w) in zip(
+                        grouped_X.groupby("group"),
+                        weights.groupby("group"))
+                ], dim="group")
+        else:
+            groupwise_means = provided_groupwise_means
+
+        # calculated (x - mean) depending on whether `mu` was specified
+        if provided_mu is None:
+            Xdiff = grouped_X - groupwise_means
+        else:
+            Xdiff = grouped_X - provided_mu
+
+        if weights is None:
+            # for each group:
+            #   calculate mean of (X - mean)^2
+            variance = np.square(Xdiff).groupby("group").mean(X.dims[0])
+        else:
+            # for each group:
+            #   calculate weighted mean of (X - mean)^2
+            variance: xr.DataArray = xr.concat([
+                weighted_mean(d, w, axis=0) for (g, d), (g, w) in zip(
+                    np.square(Xdiff).groupby("group"),
+                    weights.groupby("group")
+                )
+            ], dim="group")
+
+        denominator = np.fmax(variance - groupwise_means, np.sqrt(np.nextafter(0, 1, dtype=variance.dtype)))
+        groupwise_scales = np.square(groupwise_means) / denominator
+
+        # # clipping
+        # # r = np_clip_param(r, "r")
+        # groupwise_scales = np.nextafter(0, 1, out=groupwise_scales,
+        #                                 where=groupwise_scales == 0,
+        #                                 dtype=groupwise_scales.dtype)
+        # groupwise_scales = np.fmin(groupwise_scales, np.finfo(groupwise_scales.dtype).max)
+
+        return link_fn(groupwise_scales)
+
+    groupwise_scales, logphi, rmsd, rank, _ = groupwise_solve_lm(
+        dmat=design_scale,
+        apply_fun=apply_fun,
+        constraints=constraints
+    )
+
+    return groupwise_scales, logphi, rmsd
diff --git a/batchglm/train/tf/nb_glm/estimator.py b/batchglm/train/tf/nb_glm/estimator.py
@@ -18,7 +18,7 @@
 from .base import param_bounds, tf_clip_param, np_clip_param, apply_constraints
 
 from .external import AbstractEstimator, XArrayEstimatorStore, InputData, Model, MonitoredTFEstimator, TFEstimatorGraph
-from .external import nb_utils, train_utils, op_utils, rand_utils, data_utils
+from .external import nb_utils, train_utils, op_utils, rand_utils, data_utils, nb_glm_utils
 from .external import pkg_constants
 from .hessians import Hessians
 from .jacobians import Jacobians
@@ -759,55 +759,23 @@ def __init__(
 
                 if init_a.lower() == "closed_form":
                     try:
-                        unique_design_loc, inverse_idx = np.unique(input_data.design_loc, axis=0, return_inverse=True)
-                        if input_data.constraints_loc is not None:
-                            unique_design_loc_constraints = input_data.constraints_loc.copy()
-                            # -1 in the constraint matrix is used to indicate which variable
-                            # is made dependent so that the constrained is fullfilled.
-                            # This has to be rewritten here so that the design matrix is full rank
-                            # which is necessary so that it can be inverted for parameter
-                            # initialisation.
-                            unique_design_loc_constraints[unique_design_loc_constraints == -1] = 1
-                            # Add constraints into design matrix to remove structural unidentifiability.
-                            unique_design_loc = np.vstack([unique_design_loc, unique_design_loc_constraints])
-
-                        if unique_design_loc.shape[1] > np.linalg.matrix_rank(unique_design_loc):
-                            logger.warning("Location model is not full rank!")
-                        X = input_data.X.assign_coords(group=(("observations",), inverse_idx))
-                        if size_factors_init is not None:
-                            X = np.divide(X, size_factors_init)
-
-                        groupwise_means = X.groupby("group").mean(dim="observations").values
-                        # clipping
-                        groupwise_means = np_clip_param(groupwise_means, "mu")
-                        # mean = np.nextafter(0, 1, out=mean.values, where=mean == 0, dtype=mean.dtype)
-
-                        a = np.log(groupwise_means)
-                        if input_data.constraints_loc is not None:
-                            a_constraints = np.zeros([input_data.constraints_loc.shape[0], a.shape[1]])
-                            # Add constraints (sum to zero) to value vector to remove structural unidentifiability.
-                            a = np.vstack([a, a_constraints])
-
-                        # inv_design = np.linalg.pinv(unique_design_loc) # NOTE: this is numerically inaccurate!
-                        # inv_design = np.linalg.inv(unique_design_loc) # NOTE: this is exact if full rank!
-                        # init_a = np.matmul(inv_design, a)
-                        #
-                        # Use least-squares solver to calculate a':
-                        # This is faster and more accurate than using matrix inversion.
-                        logger.debug(" ** Solve lstsq problem")
-                        a_prime = np.linalg.lstsq(unique_design_loc, a, rcond=None)
-                        init_a = a_prime[0]
-                        # stat_utils.rmsd(np.exp(unique_design_loc @ init_a), mean)
+                        groupwise_means, init_a, rmsd_a = nb_glm_utils.closedform_nb_glm_logmu(
+                            X=input_data.X,
+                            design_loc=input_data.design_loc,
+                            constraints=input_data.constraints_loc,
+                            size_factors=size_factors_init,
+                            link_fn=lambda mu: np.log(np_clip_param(mu, "mu"))
+                        )
 
                         # train mu, if the closed-form solution is inaccurate
-                        self._train_mu = not np.all(a_prime[1] == 0)
+                        self._train_mu = not np.all(rmsd_a == 0)
 
                         # Temporal fix: train mu if size factors are given as closed form may be different:
                         if input_data.size_factors is not None:
                             self._train_mu = True
 
                         logger.info("Using closed-form MLE initialization for mean")
-                        logger.debug("RMSE of closed-form mean:\n%s", a_prime[1])
+                        logger.debug("RMSE of closed-form mean:\n%s", rmsd_a)
                         logger.info("Should train mu: %s", self._train_mu)
                     except np.linalg.LinAlgError:
                         logger.warning("Closed form initialization failed!")
@@ -831,63 +799,22 @@ def __init__(
 
                 if init_b.lower() == "closed_form":
                     try:
-                        unique_design_scale, inverse_idx = np.unique(input_data.design_scale, axis=0,
-                                                                     return_inverse=True)
-                        if input_data.constraints_scale is not None:
-                            unique_design_scale_constraints = input_data.constraints_scale.copy()
-                            # -1 in the constraint matrix is used to indicate which variable
-                            # is made dependent so that the constrained is fullfilled.
-                            # This has to be rewritten here so that the design matrix is full rank
-                            # which is necessary so that it can be inverted for parameter
-                            # initialisation.
-                            unique_design_scale_constraints[unique_design_scale_constraints == -1] = 1
-                            # Add constraints into design matrix to remove structural unidentifiability.
-                            unique_design_scale = np.vstack([unique_design_scale, unique_design_scale_constraints])
-
-                        if unique_design_scale.shape[1] > np.linalg.matrix_rank(unique_design_scale):
-                            logger.warning("Scale model is not full rank!")
-
-                        X = input_data.X.assign_coords(group=(("observations",), inverse_idx))
-                        if input_data.size_factors is not None:
-                            X = np.divide(X, size_factors_init)
-
-                        # Xdiff = X - np.exp(input_data.design_loc @ init_a)
-                        # Define xarray version of init so that Xdiff can be evaluated lazy by dask.
                         init_a_xr = data_utils.xarray_from_data(init_a, dims=("design_loc_params", "features"))
                         init_a_xr.coords["design_loc_params"] = input_data.design_loc.coords["design_loc_params"]
-                        logger.debug(" ** Define Xdiff")
-                        Xdiff = X - np.exp(input_data.design_loc.dot(init_a_xr))
-                        variance = np.square(Xdiff).groupby("group").mean(dim="observations")
-
-                        if groupwise_means is None:
-                            groupwise_means = X.groupby("group").mean(dim="observations")
-                        denominator = np.fmax(variance - groupwise_means, 0)
-                        denominator = np.nextafter(0, 1, out=denominator.values, where=denominator == 0,
-                                                   dtype=denominator.dtype)
-                        r = np.asarray(np.square(groupwise_means) / denominator)
-                        # clipping
-                        r = np_clip_param(r, "r")
-                        # r = np.nextafter(0, 1, out=r.values, where=r == 0, dtype=r.dtype)
-                        # r = np.fmin(r, np.finfo(r.dtype).max)
-
-                        b = np.log(r)
-                        if input_data.constraints_scale is not None:
-                            b_constraints = np.zeros([input_data.constraints_scale.shape[0], b.shape[1]])
-                            # Add constraints (sum to zero) to value vector to remove structural unidentifiability.
-                            b = np.vstack([b, b_constraints])
-
-                        # inv_design = np.linalg.pinv(unique_design_scale) # NOTE: this is numerically inaccurate!
-                        # inv_design = np.linalg.inv(unique_design_scale) # NOTE: this is exact if full rank!
-                        # init_b = np.matmul(inv_design, b)
-                        #
-                        # Use least-squares solver to calculate a':
-                        # This is faster and more accurate than using matrix inversion.
-                        logger.debug(" ** Solve lstsq problem")
-                        b_prime = np.linalg.lstsq(unique_design_scale, b, rcond=None)
-                        init_b = b_prime[0]
+                        init_mu = np.exp(input_data.design_loc.dot(init_a_xr))
+
+                        groupwise_scales, init_b, rmsd_b = nb_glm_utils.closedform_nb_glm_logphi(
+                            X=input_data.X,
+                            mu=init_mu,
+                            design_scale=input_data.design_scale,
+                            constraints=input_data.constraints_scale,
+                            size_factors=size_factors_init,
+                            groupwise_means=groupwise_means,
+                            link_fn=lambda r: np.log(np_clip_param(r, "r"))
+                        )
 
                         logger.info("Using closed-form MME initialization for dispersion")
-                        logger.debug("RMSE of closed-form dispersion:\n%s", b_prime[1])
+                        logger.debug("RMSE of closed-form dispersion:\n%s", rmsd_b)
                         logger.info("Should train r: %s", self._train_r)
                     except np.linalg.LinAlgError:
                         logger.warning("Closed form initialization failed!")
@@ -903,8 +830,11 @@ def __init__(
                     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)
 
-                    # Initialize new parameters to zero:
-                    init_loc = np.zeros(shape=(input_data.num_design_loc_params, input_data.num_features))
+                    init_loc = np.random.uniform(
+                        low=np.nextafter(0, 1, dtype=input_data.X.dtype),
+                        high=np.sqrt(np.nextafter(0, 1, dtype=input_data.X.dtype)),
+                        size=(input_data.num_design_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)
@@ -917,8 +847,11 @@ def __init__(
                     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)
 
-                    # Initialize new parameters to zero:
-                    init_scale = np.zeros(shape=(input_data.num_design_scale_params, input_data.num_features))
+                    init_scale = np.random.uniform(
+                        low=np.nextafter(0, 1, dtype=input_data.X.dtype),
+                        high=np.sqrt(np.nextafter(0, 1, dtype=input_data.X.dtype)),
+                        size=(input_data.num_design_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)
diff --git a/batchglm/train/tf/nb_glm/external.py b/batchglm/train/tf/nb_glm/external.py
@@ -10,5 +10,6 @@
 
 # from train.tf.nb import EstimatorGraph as NegativeBinomialEstimatorGraph
 
+import batchglm.models.nb_glm.utils as nb_glm_utils
 import batchglm.utils.random as rand_utils
 from batchglm import pkg_constants
diff --git a/batchglm/unit_test/test_nb_glm.py b/batchglm/unit_test/test_nb_glm.py
@@ -100,9 +100,9 @@ def test_nonconfounded_fit(self):
         sim.generate_sample_description(num_conditions=0, num_batches=4)
         sim.generate()
 
-        sample_description = data_utils.sample_description_from_xarray(sim.data, dim="observations")
-        design_loc = data_utils.design_matrix(sample_description, formula="~ 1 - 1 + batch")
-        design_scale = data_utils.design_matrix(sample_description, formula="~ 1 - 1 + batch")
+        sample_description = glm.data.sample_description_from_xarray(sim.data, dim="observations")
+        design_loc = glm.data.design_matrix(sample_description, formula="~ 1 - 1 + batch")
+        design_scale = glm.data.design_matrix(sample_description, formula="~ 1 - 1 + batch")
 
         input_data = InputData.new(sim.X, design_loc=design_loc, design_scale=design_scale)
 
@@ -127,7 +127,13 @@ def test_nonconfounded_fit(self):
 
     def test_anndata(self):
         adata = self.sim.data_to_anndata()
-        idata = InputData.new(adata)
+        design_loc = self.sim.design_loc
+        design_scale = self.sim.design_scale
+        idata = InputData.new(
+            data=adata,
+            design_loc=design_loc,
+            design_scale=design_scale,
+        )
 
         wd = os.path.join(self.working_dir.name, "anndata")
         os.makedirs(wd, exist_ok=True)
@@ -141,7 +147,13 @@ def test_anndata(self):
     def test_anndata_sparse(self):
         adata = self.sim.data_to_anndata()
         adata.X = scipy.sparse.csr_matrix(adata.X)
-        idata = InputData.new(adata)
+        design_loc = self.sim.design_loc
+        design_scale = self.sim.design_scale
+        idata = InputData.new(
+            data=adata,
+            design_loc=design_loc,
+            design_scale=design_scale,
+        )
 
         wd = os.path.join(self.working_dir.name, "anndata")
         os.makedirs(wd, exist_ok=True)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from batchglm.utils.linalg import stacked_lstsq, groupwise_solve_lm`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from batchglm.utils.numeric import combine_matrices, softmax, weighted_mean, weighted_variance`