added linalg error propagation as nan which is then separately dealt with as singular matrix fix

Author: davidsebfischer

batchglm/train/numpy/base_glm/estimator.py

@@ -3,7 +3,6 @@
 import logging
 import multiprocessing
 import numpy as np
-import pprint
 import scipy
 import scipy.sparse
 import scipy.optimize
@@ -266,20 +265,25 @@ def iwls_step(
         if isinstance(delta_theta, dask.array.core.Array):
             delta_theta = delta_theta.compute()
 
-        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.
-            if a.shape[0] != 1:
-                delta_theta[:, idx_update] = dask.array.map_blocks(
-                    np.linalg.solve, a, b[:, :, None], chunks=b[:, :, None].shape
-                ).squeeze().T.compute()
+        with np.errstate(linalg="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.
+                if a.shape[0] != 1:
+                    delta_theta[:, idx_update] = dask.array.map_blocks(
+                        np.linalg.solve, a, b[:, :, None], chunks=b[:, :, None].shape
+                    ).squeeze().T.compute()
+                else:
+                    delta_theta[:, idx_update] = np.expand_dims(
+                        np.linalg.solve(a[0], b[0]).compute(),
+                        axis=-1
+                    )
             else:
-                delta_theta[:, idx_update] = np.expand_dims(
-                    np.linalg.solve(a[0], b[0]).compute(),
-                    axis=-1
-                )
-        else:
-            delta_theta[:, idx_update] = np.linalg.solve(a, b).T
+                delta_theta[:, idx_update] = np.linalg.solve(a, b).T
+        linalg_errors = np.isnan(delta_theta[:, 0, 0])
+        if np.any(linalg_errors):
+            print("caught %i linalg errors" % np.sum(linalg_errors))
+            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)
@@ -512,7 +516,8 @@ def finalize(self):
         """
         # Read from numpy-IRLS estimator specific model:
         self._hessian = - self.model.fim.compute()
-        self._fisher_inv = np.linalg.inv(- self._hessian)
+        with np.errstate(linalg="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()
         self._loss = np.sum(self._log_likelihood)