fixed issues with previous linalg fix

Author: davidsebfischer

batchglm/models/base_glm/simulator.py

@@ -171,3 +171,16 @@ def constraints_loc(self):
     @property
     def constraints_scale(self):
         return np.identity(n=self.b_var.shape[0])
+
+    def np_clip_param(
+            self,
+            param,
+            name
+    ):
+        # TODO: inherit this from somewhere?
+        bounds_min, bounds_max = self.param_bounds(param.dtype)
+        return np.clip(
+            param,
+            bounds_min[name],
+            bounds_max[name]
+        )

batchglm/models/glm_nb/external.py

@@ -5,4 +5,6 @@
 from batchglm.models.base_glm import closedform_glm_mean, closedform_glm_scale
 
 import batchglm.data as data_utils
-from batchglm.utils.linalg import groupwise_solve_lm
+from batchglm.utils.linalg import groupwise_solve_lm
+
+from batchglm import pkg_constants

batchglm/models/glm_nb/simulator.py

@@ -2,6 +2,7 @@
 
 from .model import Model
 from .external import _SimulatorGLM, InputDataGLM
+from .external import pkg_constants
 
 
 class Simulator(_SimulatorGLM, Model):
@@ -58,3 +59,35 @@ def generate_data(self):
             design_scale_names=None
         )
 
+    def param_bounds(
+            self,
+            dtype
+    ):
+        # TODO: inherit this from somewhere?
+        dtype = np.dtype(dtype)
+        dmin = np.finfo(dtype).min
+        dmax = np.finfo(dtype).max
+        dtype = dtype.type
+
+        sf = dtype(pkg_constants.ACCURACY_MARGIN_RELATIVE_TO_LIMIT)
+        bounds_min = {
+            "a_var": np.log(np.nextafter(0, np.inf, dtype=dtype)) / sf,
+            "b_var": np.log(np.nextafter(0, np.inf, dtype=dtype)) / sf,
+            "eta_loc": np.log(np.nextafter(0, np.inf, dtype=dtype)) / sf,
+            "eta_scale": np.log(np.nextafter(0, np.inf, dtype=dtype)) / sf,
+            "loc": np.nextafter(0, np.inf, dtype=dtype),
+            "scale": np.nextafter(0, np.inf, dtype=dtype),
+            "likelihood": dtype(0),
+            "ll": np.log(np.nextafter(0, np.inf, dtype=dtype)),
+        }
+        bounds_max = {
+            "a_var": np.nextafter(np.log(dmax), -np.inf, dtype=dtype) / sf,
+            "b_var": np.nextafter(np.log(dmax), -np.inf, dtype=dtype) / sf,
+            "eta_loc": np.nextafter(np.log(dmax), -np.inf, dtype=dtype) / sf,
+            "eta_scale": np.nextafter(np.log(dmax), -np.inf, dtype=dtype) / sf,
+            "loc": np.nextafter(dmax, -np.inf, dtype=dtype) / sf,
+            "scale": np.nextafter(dmax, -np.inf, dtype=dtype) / sf,
+            "likelihood": dtype(1),
+            "ll": dtype(0),
+        }
+        return bounds_min, bounds_max

batchglm/train/numpy/base_glm/estimator.py

@@ -265,7 +265,7 @@ def iwls_step(
         if isinstance(delta_theta, dask.array.core.Array):
             delta_theta = delta_theta.compute()
 
-        with np.errstate(linalg="ignore"):
+        with np.errstate(all="ignore"):
             if isinstance(a, dask.array.core.Array):
                 # Have to use a workaround to solve problems in parallel in dask here. This workaround does
                 # not work if there is only a single problem, ie. if the first dimension of a and b has length 1.
@@ -280,10 +280,10 @@ def iwls_step(
                     )
             else:
                 delta_theta[:, idx_update] = np.linalg.solve(a, b).T
-        linalg_errors = np.isnan(delta_theta[:, 0, 0])
+        linalg_errors = np.isnan(delta_theta[0, :])
         if np.any(linalg_errors):
             print("caught %i linalg errors" % np.sum(linalg_errors))
-            delta_theta[linalg_errors, :, :] = 0.
+            delta_theta[:, linalg_errors] = 0.
         # Via np.linalg.lsts:
         #delta_theta[:, idx_update] = np.concatenate([
         #    np.expand_dims(np.linalg.lstsq(a[i, :, :], b[i, :])[0], axis=-1)
@@ -516,7 +516,7 @@ def finalize(self):
         """
         # Read from numpy-IRLS estimator specific model:
         self._hessian = - self.model.fim.compute()
-        with np.errstate(linalg="ignore"):
+        with np.errstate(all="ignore"):
             self._fisher_inv = np.linalg.inv(- self._hessian)
         self._jacobian = np.sum(np.abs(self.model.jac.compute() / self.model.x.shape[0]), axis=1)
         self._log_likelihood = self.model.ll_byfeature.compute()