batchglm numpy backend support functional from diffxpy

Author: davidsebfischer

batchglm/train/numpy/base_glm/estimator.py

@@ -1,11 +1,15 @@
 import abc
+import logging
 import numpy as np
+import pprint
 import scipy
 import scipy.optimize
 
 from .external import _EstimatorGLM, pkg_constants
 from .training_strategies import TrainingStrategies
 
+logger = logging.getLogger("batchglm")
+
 
 class EstimatorGlm(_EstimatorGLM, metaclass=abc.ABCMeta):
     """
@@ -18,6 +22,8 @@ def __init__(
             input_data,
             dtype,
     ):
+        if input_data.design_scale.shape[1] != 1:
+            raise ValueError("cannot model more than one scale parameter with numpy backend right now.")
         _EstimatorGLM.__init__(
             self=self,
             model=model,
@@ -32,9 +38,17 @@ def __init__(
     def initialize(self):
         pass
 
-    def train_sequence(self, training_strategy: str):
+    def train_sequence(
+            self,
+            training_strategy: str = "DEFAULT"
+    ):
         if isinstance(training_strategy, str):
             training_strategy = self.TrainingStrategies[training_strategy].value[0]
+
+        if training_strategy is None:
+            training_strategy = self.TrainingStrategies.DEFAULT.value
+
+        logging.getLogger("batchglm").info("training strategy:\n%s", pprint.pformat(training_strategy))
         self.train(**training_strategy)
 
     def train(
@@ -44,17 +58,17 @@ def train(
     ):
         # Iterate until conditions are fulfilled.
         train_step = 0
-        delayed_b_converged = np.tile(False, self.model.model_vars.n_features)
+        delayed_converged = np.tile(False, self.model.model_vars.n_features)
 
         ll_current = - self.model.ll_byfeature
-        print("iter %i: ll=%f" % (0, np.sum(ll_current)))
-        while (np.any(np.logical_not(self.model.converged)) or np.any(np.logical_not(delayed_b_converged))) and \
+        logging.getLogger("batchglm").debug("iter %i: ll=%f" % (0, np.sum(ll_current)))
+        while np.any(np.logical_not(delayed_converged)) and \
                 train_step < max_steps:
             # Update parameters:
             # Line search step for scale model:
             if train_step % update_b_freq == 0 and train_step > 0:
                 b_var_cache = self.model.b_var.copy()
-                self.model.b_var = self.b_step(idx=np.where(np.logical_not(delayed_b_converged))[0])
+                self.model.b_var = self.b_step(idx=np.where(np.logical_not(delayed_converged))[0])
                 # Reverse update by feature if update leads to worse loss:
                 ll_proposal = - self.model.ll_byfeature
                 b_var_new = self.model.b_var.copy()
@@ -71,10 +85,14 @@ def train(
             # Location model convergence status has to be updated if b model was updated
             if train_step % update_b_freq == 0 and train_step > 0:
                 self.model.converged = converged_f
+                delayed_converged = converged_f
             else:
                 self.model.converged = np.logical_or(self.model.converged, converged_f)
             train_step += 1
-            print("iter %i: ll=%f, converged: %i" % (train_step, np.sum(ll_current), np.sum(self.model.converged)))
+            logging.getLogger("batchglm").debug(
+                "iter %i: ll=%f, converged: %i" %
+                (train_step, np.sum(ll_current), np.sum(self.model.converged))
+            )
             self.lls.append(ll_current)
 
     def iwls_step(self) -> np.ndarray:
@@ -167,10 +185,11 @@ def finalize(self):
         transfers relevant attributes.
         """
         # Read from numpy-IRLS estimator specific model:
-        self._hessian = - self.model.fim
+
+        self._hessian = self.model.hessian
         self._fisher_inv = np.linalg.inv(- self._hessian)
-        self._jacobian = self.model.jac
-        self._log_likelihood = self.model.ll
+        self._jacobian = np.sum(np.abs(self.model.jac / self.model.x.shape[0]), axis=1)
+        self._log_likelihood = self.model.ll_byfeature
         self._loss = np.sum(self._log_likelihood)
 
     @abc.abstractmethod

batchglm/train/numpy/base_glm/model.py

@@ -109,6 +109,74 @@ def fim(self) -> np.ndarray:
             xh
         )
 
+    @abc.abstractmethod
+    def hessian_weight_aa(self) -> np.ndarray:
+        pass
+
+    @property
+    def hessian_aa(self) -> np.ndarray:
+        """
+
+        :return: (features x inferred param x inferred param)
+        """
+        w = self.hessian_weight_aa
+        xh = np.matmul(self.design_loc, self.constraints_loc)
+        return np.einsum(
+            'fob,oc->fbc',
+            np.einsum('ob,of->fob', xh, w),
+            xh
+        )
+
+    @abc.abstractmethod
+    def hessian_weight_ab(self) -> np.ndarray:
+        pass
+
+    @property
+    def hessian_ab(self) -> np.ndarray:
+        """
+
+        :return: (features x inferred param x inferred param)
+        """
+        w = self.hessian_weight_ab
+        return np.einsum(
+            'fob,oc->fbc',
+            np.einsum('ob,of->fob', np.matmul(self.design_loc, self.constraints_loc), w),
+            np.matmul(self.design_scale, self.constraints_scale)
+        )
+
+    @abc.abstractmethod
+    def hessian_weight_bb(self) -> np.ndarray:
+        pass
+
+    @property
+    def hessian_bb(self) -> np.ndarray:
+        """
+
+        :return: (features x inferred param x inferred param)
+        """
+        w = self.hessian_weight_bb
+        xh = np.matmul(self.design_scale, self.constraints_scale)
+        return np.einsum(
+            'fob,oc->fbc',
+            np.einsum('ob,of->fob', xh, w),
+            xh
+        )
+
+    @property
+    def hessian(self) -> np.ndarray:
+        """
+
+        :return: (features x inferred param x inferred param)
+        """
+        h_aa = self.hessian_aa
+        h_bb = self.hessian_bb
+        h_ab = self.hessian_ab
+        h_ba = np.transpose(h_ab, axes=[0, 2, 1])
+        return np.concatenate([
+            np.concatenate([h_aa, h_ab], axis=2),
+            np.concatenate([h_ba, h_bb], axis=2)
+        ], axis=1)
+
     @property
     def jac(self) -> np.ndarray:
         return np.concatenate([self.jac_a, self.jac_b], axis=-1)

batchglm/train/numpy/glm_nb/model.py

@@ -101,6 +101,39 @@ def jac_weight_b_j(self, j):
         const3 = np.log(scale) + np.ones_like(scale) - np.log(r_plus_mu)
         return scale * (const1 + const2 + const3)
 
+    @property
+    def hessian_weight_ab(self):
+        scale = self.scale
+        loc = self.location
+        return np.multiply(
+            loc * scale,
+            np.asarray(self.x - loc) / np.square(loc + scale)
+        )
+
+    @property
+    def hessian_weight_aa(self):
+        scale = self.scale
+        loc = self.location
+        if isinstance(self.x, np.ndarray):
+            x_by_scale_plus_one = self.x / scale + np.ones_like(scale)
+        else:
+            x_by_scale_plus_one = np.asarray(self.x.divide(scale) + np.ones_like(scale))
+
+        return - loc * x_by_scale_plus_one / np.square((loc / scale) + np.ones_like(loc))
+
+    @property
+    def hessian_weight_bb(self):
+        scale = self.scale
+        loc = self.location
+        scale_plus_x = np.asarray(self.x + scale)
+        scale_plus_loc = scale + loc
+        # Define graphs for individual terms of constant term of hessian:
+        const1 = scipy.special.digamma(scale_plus_x) + scale * scipy.special.polygamma(n=1, x=scale_plus_x)
+        const2 = - scipy.special.digamma(scale) + scale * scipy.special.polygamma(n=1, x=scale)
+        const3 = - loc * scale_plus_x + np.ones_like(scale) * 2. * scale * scale_plus_loc / np.square(scale_plus_loc)
+        const4 = np.log(scale) + np.ones_like(scale) * 2. - np.log(scale_plus_loc)
+        return scale * (const1 + const2 + const3 + const4)
+
     @property
     def ll(self):
         scale = self.scale