changed handling of FIM in numpy

Author: davidsebfischer

batchglm/train/numpy/base_glm/model.py

@@ -71,15 +71,15 @@ def ll_byfeature_j(self, j) -> np.ndarray:
         return np.sum(self.ll_j(j=j), axis=0)
 
     @abc.abstractmethod
-    def fim_weight(self) -> np.ndarray:
+    def fim_weight_aa(self) -> np.ndarray:
         pass
 
     @abc.abstractmethod
     def ybar(self) -> np.ndarray:
         pass
 
     @abc.abstractmethod
-    def fim_weight_j(self, j) -> np.ndarray:
+    def fim_weight_aa_j(self, j) -> np.ndarray:
         pass
 
     @abc.abstractmethod
@@ -95,12 +95,13 @@ def ybar_j(self, j) -> np.ndarray:
         pass
 
     @property
-    def fim(self) -> np.ndarray:
+    def fim_aa(self) -> np.ndarray:
         """
+        Location-location coefficient block of FIM
 
         :return: (features x inferred param x inferred param)
         """
-        w = self.fim_weight  # (observations x features)
+        w = self.fim_weight_aa  # (observations x features)
         # constraints: (observed param x inferred param)
         # design: (observations x observed param)
         # w: (observations x features)
@@ -112,6 +113,30 @@ def fim(self) -> np.ndarray:
             xh
         )
 
+    @abc.abstractmethod
+    def fim_ab(self) -> np.ndarray:
+        pass
+
+    @property
+    def fim_bb(self) -> np.ndarray:
+        pass
+
+    @property
+    def fim(self) -> np.ndarray:
+        """
+        Full FIM
+
+        :return: (features x inferred param x inferred param)
+        """
+        fim_aa = self.fim_aa
+        fim_bb = self.fim_bb
+        fim_ab = self.fim_ab
+        fim_ba = np.transpose(fim_ab, axes=[0, 2, 1])
+        return - np.concatenate([
+            np.concatenate([fim_aa, fim_ab], axis=2),
+            np.concatenate([fim_ba, fim_bb], axis=2)
+        ], axis=1)
+
     @abc.abstractmethod
     def hessian_weight_aa(self) -> np.ndarray:
         pass
@@ -190,7 +215,7 @@ def jac_a(self) -> np.ndarray:
 
         :return: (features x inferred param)
         """
-        w = self.fim_weight  # (observations x features)
+        w = self.fim_weight_aa  # (observations x features)
         ybar = self.ybar  # (observations x features)
         xh = np.matmul(self.design_loc, self.constraints_loc)  # (observations x inferred param)
         return np.einsum(
@@ -207,7 +232,7 @@ def jac_a_j(self, j) -> np.ndarray:
         # Make sure that dimensionality of sliced array is kept:
         if isinstance(j, int) or isinstance(j, np.int32) or isinstance(j, np.int64):
             j = [j]
-        w = self.fim_weight_j(j=j)  # (observations x features)
+        w = self.fim_weight_aa_j(j=j)  # (observations x features)
         ybar = self.ybar_j(j=j)  # (observations x features)
         xh = np.matmul(self.design_loc, self.constraints_loc)  # (observations x inferred param)
         return np.einsum(

batchglm/train/numpy/glm_nb/model.py

@@ -36,7 +36,7 @@ def __init__(
         )
 
     @property
-    def fim_weight(self):
+    def fim_weight_aa(self):
         """
 
         :return: observations x features
@@ -51,7 +51,7 @@ def ybar(self) -> np.ndarray:
         """
         return np.asarray(self.x - self.location) / self.location
 
-    def fim_weight_j(self, j):
+    def fim_weight_aa_j(self, j):
         """
 
         :return: observations x features
@@ -113,6 +113,38 @@ 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 fim_ab(self) -> np.ndarray:
+        """
+        Location-scale coefficient block of FIM
+
+        The negative binomial model is not fit as whole with IRLS but only the location model.
+        The location model is conditioned on the scale model estimates, which is why we only
+        supply the FIM of the location model and return an empty FIM for scale model components.
+        Note that there is also no closed form FIM for the scale-scale block. Returning a zero-array
+        here leads to singular matrices for the whole location-scale FIM in some cases that throw
+        linear algebra errors when inverted.
+
+        :return: (features x inferred param x inferred param)
+        """
+        return np.zeros([self.b_var.shape[1], 0, 0])
+
+    @property
+    def fim_bb(self) -> np.ndarray:
+        """
+        Scale-scale coefficient block of FIM
+
+        The negative binomial model is not fit as whole with IRLS but only the location model.
+        The location model is conditioned on the scale model estimates, which is why we only
+        supply the FIM of the location model and return an empty FIM for scale model components.
+        Note that there is also no closed form FIM for the scale-scale block. Returning a zero-array
+        here leads to singular matrices for the whole location-scale FIM in some cases that throw
+        linear algebra errors when inverted.
+
+        :return: (features x inferred param x inferred param)
+        """
+        return np.zeros([self.b_var.shape[1], 0, 0])
+
     @property
     def hessian_weight_ab(self):
         scale = self.scale
@@ -209,3 +241,18 @@ def fun(x, eta_loc, b_var, xh_scale):
                 raise ValueError("type x %s not supported" % type(x))
             return self.np_clip_param(ll, "ll")
         return fun
+
+    def jac_b_handle(self):
+        def fun(x, eta_loc, b_var, xh_scale):
+            scale = np.exp(b_var)
+            loc = np.exp(eta_loc)
+            scale_plus_x = scale + x
+            r_plus_mu = scale + loc
+
+            # Define graphs for individual terms of constant term of hessian:
+            const1 = scipy.special.digamma(scale_plus_x) - scipy.special.digamma(scale)
+            const2 = - scale_plus_x / r_plus_mu
+            const3 = np.log(scale) + np.ones_like(scale) - np.log(r_plus_mu)
+            return scale * (const1 + const2 + const3)
+
+        return fun