Merge pull request #19 from ihmeuw-msca/feat/quit-early-when-nans

Quit early when nans encountered

Author: kels271828

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "msca"
-version = "0.3.1"
+version = "0.3.2"
 description = "Mathematical sciences and computational algorithms"
 readme = "README.md"
 requires-python = ">=3.11,<3.13"

src/msca/optim/solver/ntcgsolver.py

@@ -144,16 +144,19 @@ def get_cg_maxiter(niter: int) -> int | None:
         step = 1.0
         niter = 0
         success = False
+        failure = False
 
         x_pair = deque([x], maxlen=2)
         g_pair = deque([g], maxlen=2)
 
         if verbose:
             fun = self.fun(x)
             print(f"{type(self).__name__}:")
-            print(f"{niter=:3d}, {fun=:.2e}, {gnorm=:.2e}, {xdiff=:.2e}, {step=:.2e}")
+            print(
+                f"{niter=:3d}, {fun=:.2e}, {gnorm=:.2e}, {xdiff=:.2e}, {step=:.2e}"
+            )
 
-        while (not success) and (niter < maxiter):
+        while (not success) and (not failure) and (niter < maxiter):
             niter += 1
 
             # compute all directions
@@ -185,13 +188,19 @@ def cg_iter_counter(xk, cg_info):
             x_pair.append(x)
             g_pair.append(g)
 
+            fun = self.fun(x)
             if verbose:
-                fun = self.fun(x)
                 print(
                     f"{niter=:3d}, {fun=:.2e}, {gnorm=:.2e}, {xdiff=:.2e}, "
                     f"{step=:.2e}, cg_iter={cg_info['iter']}"
                 )
             success = gnorm <= gtol or xdiff <= xtol
+            failure = not (
+                np.isfinite(fun)
+                and np.isfinite(gnorm)
+                and np.isfinite(xdiff)
+                and np.isfinite(step)
+            )
 
         result = NTCGResult(
             x=x,