Fix tests for older dependencies

CHANGELOG.md

@@ -14,6 +14,8 @@ Performance Improvements
 Bug Fixes
 
 - Fixes bug in ``auglag`` optimizers which prevented them from accepting solver hyperparameters.
+- Improves the handling of failed Cholesky factorizations in the ``"cho"`` trust-region method.
+- Adds a warning when sub-objectives with ``bounds`` can make the Jacobian rank-deficient, since the default ``"qr"`` trust-region method may then fail to solve the subproblem, suggesting ``options={"tr_method": "svd"}`` instead.
 - Fixes bug in modified Cholesky factorization used by the trust-region
   subproblems when the Gershgorin lower bound of the Hessian was exactly zero
   (e.g. a Hessian with an all-zero row), producing NaN steps in ``fmintr`` and

desc/optimize/_desc_wrappers.py

@@ -1,6 +1,7 @@
 from scipy.optimize import NonlinearConstraint
 
 from desc.backend import jnp
+from desc.utils import warnif
 
 from .aug_lagrangian import fmin_auglag
 from .aug_lagrangian_ls import lsq_auglag
@@ -45,6 +46,36 @@
 ]
 
 
+def _warn_if_bounds(objective, constraint, x0, options):
+    """Warn if bounds on sub-objectives can make the Jacobian rank-deficient.
+
+    A sub-objective with bounds contributes zero residuals, and hence zero Jacobian
+    rows, whenever it is inside its bounds. The full rank of the (m, n) Jacobian is
+    min(m, n), so those rows can only cost rank if the ones that always remain are
+    fewer than that. Only ``"svd"`` handles the deficient case reliably, so warn if
+    the user hasn't picked a method themselves.
+    """
+    if "tr_method" in options:  # user picked a method, don't second guess it
+        return
+    sub = objective
+    while hasattr(sub, "_objective"):  # unwrap Proximal/LinearConstraintProjection
+        sub = sub._objective
+    bounded = [obj for obj in sub.objectives if obj.bounds is not None]
+    dim_f_bounded = sum(obj.dim_f for obj in bounded)
+    # constraint rows never vanish, bounds there become slack variables instead
+    m = objective.dim_f + (0 if constraint is None else constraint.dim_f)
+    n = x0.size
+    warnif(
+        m - dim_f_bounded < min(m, n),
+        UserWarning,
+        f"Objectives {[obj.name for obj in bounded]} use bounds instead of target, so "
+        + f"they can zero out {dim_f_bounded} of the {m} rows of the ({m}, {n}) "
+        + f"Jacobian and drop its rank below {min(m, n)}. The default 'qr' trust "
+        + "region method may then fail to solve the subproblem, in that case pass "
+        + "options={'tr_method': 'svd'}.",
+    )
+
+
 @register_optimizer(
     name=["fmin-auglag", "fmin-auglag-bfgs"],
     description=[
@@ -206,6 +237,7 @@
     if not isinstance(x_scale, str) and jnp.allclose(x_scale, 1):
         options.setdefault("initial_trust_radius", 1e-3)
         options.setdefault("max_trust_radius", 1.0)
+    _warn_if_bounds(objective, constraint, x0, options)
     options["max_nfev"] = stoptol["max_nfev"]
 
     if constraint is not None:
@@ -300,6 +332,7 @@
         options.setdefault("max_trust_radius", 1.0)
     elif options.get("initial_trust_radius", "scipy") == "scipy":
         options.setdefault("initial_trust_ratio", 0.1)
+    _warn_if_bounds(objective, constraint, x0, options)
     options["max_nfev"] = stoptol["max_nfev"]
 
     result = lsqtr(

desc/optimize/aug_lagrangian_ls.py

@@ -173,7 +173,9 @@ def lsq_auglag(  # noqa: C901
           Cholesky factorizations (generally 2-3), while ``"svd"`` uses one singular
           value decomposition. ``"cho"`` is generally the fastest for large systems,
           especially on GPU, but may be less accurate for badly scaled systems.
-          ``"svd"`` is the most accurate but significantly slower. Default ``"qr"``.
+          ``"svd"`` is the most accurate but significantly slower. If any of the
+          sub-objective includes bounds, the ``'svd'`` is recommended since the linear
+          system has a chance to be rank-deficient. Default ``"qr"``.
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 

desc/optimize/least_squares.py

@@ -139,7 +139,9 @@ def lsqtr(  # noqa: C901
           Cholesky factorizations (generally 2-3), while ``"svd"`` uses one singular
           value decomposition. ``"cho"`` is generally the fastest for large systems,
           especially on GPU, but may be less accurate for badly scaled systems.
-          ``"svd"`` is the most accurate but significantly slower. Default ``"qr"``.
+          ``"svd"`` is the most accurate but significantly slower. If any of the
+          sub-objective includes bounds, the ``'svd'`` is recommended since the linear
+          system has a chance to be rank-deficient. Default ``"qr"``.
         - ``"scaled_termination"`` : Whether to evaluate termination criteria for
           ``xtol`` and ``gtol`` in scaled / normalized units (default) or base units.
 

desc/optimize/utils.py

@@ -175,6 +175,20 @@ def gershgorin_bounds(H):
     return lb, ub
 
 
+@jit
+def _chol_failed(L):
+    """Whether a Cholesky factor is unusable because the matrix was not positive.
+
+    ``potrf`` only errors out, and so returns nan, when a pivot comes out
+    non-positive. A pivot at the noise level instead completes but gives a factor that
+    is useless to solve with, and which of the two a marginal matrix gets depends on
+    the LAPACK build, so treat both as a failure.
+    """
+    d = jnp.diag(L)
+    tol = L.shape[-1] * jnp.finfo(L.dtype).eps
+    return jnp.any(jnp.isnan(L)) | (jnp.min(d) ** 2 <= tol * jnp.max(d) ** 2)
+
+
 @jit
 def _cholmod(A, maxiter=4):
     """Modified Cholesky factorization of indefinite matrix.
@@ -241,21 +255,21 @@ def _cholmod(A, maxiter=4):
     alphas = jnp.logspace(lb, ub, k)
     kbest = k // 2
     klow = 0
-    khigh = k
+    khigh = k - 1
     # first we try alpha = max, which we know will succeed by gershgorin bounds
     # but might be too big a correction, so then we try to reduce it while keeping
     # A + alpha*I positive definite
-    Lbest = jnp.linalg.cholesky(A + alphas[k] * eye)
+    Lbest = jnp.linalg.cholesky(A + alphas[-1] * eye)
     for i in range(maxiter):
         L = jnp.linalg.cholesky(A + alphas[kbest] * eye)
         # check if it succeeded
-        isnan = jnp.any(jnp.isnan(L))
+        failed = _chol_failed(L)
         # adjust bounds for correction
-        klow = isnan * kbest + (1 - isnan) * klow
-        khigh = isnan * khigh + (1 - isnan) * kbest
+        klow = failed * kbest + (1 - failed) * klow
+        khigh = failed * khigh + (1 - failed) * kbest
         kbest = (klow + khigh) // 2
         # if it succeeded, mark it as the best so far
-        Lbest = cond(isnan, lambda _: Lbest, lambda _: L, None)
+        Lbest = cond(failed, lambda _: Lbest, lambda _: L, None)
     return Lbest
 
 
@@ -279,7 +293,7 @@ def chol(A):
 
     """
     L = jnp.linalg.cholesky(A)
-    L = cond(jnp.any(jnp.isnan(L)), lambda A: _cholmod(A), lambda A: L, A)
+    L = cond(_chol_failed(L), lambda A: _cholmod(A), lambda A: L, A)
     return L
 
 

tests/test_compute_everything.py

@@ -49,13 +49,16 @@ def _compare_against_master(
             else:
                 mean = np.mean(np.atleast_1d(np.abs(master_data[p][name])))
             try:
+                rtol = 1e-5 if "Gamma_" in name and OLD_FINUFFT else 1e-8
+                atol = 1e-4 if "Gamma_" in name and OLD_FINUFFT else 1e-8
+                atol = atol * mean + 1e-9  # add 1e-9 for basically-zero things
                 err_msg = f"Parameterization: {p}. Name: {name}."
                 assert np.isfinite(mean).all(), err_msg
                 np.testing.assert_allclose(
                     actual=data[p][name],
                     desired=master_data[p][name],
-                    atol=1e-8 * mean + 1e-9,  # add 1e-9 for basically-zero things
-                    rtol=1e-8,
+                    atol=atol,
+                    rtol=rtol,
                     err_msg=err_msg,
                 )
             except AssertionError as e:

tests/test_examples.py

@@ -174,7 +174,15 @@ def test_solve_bounds():
     obj = ObjectiveFunction(
         ForceBalance(normalize=False, normalize_target=False, bounds=(-3e3, 3e3), eq=eq)
     )
-    eq.solve(objective=obj, ftol=1e-16, xtol=1e-16, maxiter=200, verbose=3)
+    # solve with bounds creates singular Jacobian which QR cannot handle
+    eq.solve(
+        objective=obj,
+        ftol=1e-16,
+        xtol=1e-16,
+        maxiter=200,
+        verbose=3,
+        options={"tr_method": "svd"},
+    )
 
     # check that all errors are nearly 0, since residual values are within target bounds
     f = obj.compute_scaled_error(obj.x(eq))
@@ -1171,7 +1179,8 @@ def test_omnigenity_proximal():
         FixPsi(eq=eq),
     )
     optimizer = Optimizer("proximal-lsq-exact")
-    [eq], _ = optimizer.optimize(eq, objective, constraints, maxiter=2, verbose=3)
+    with pytest.warns(UserWarning, match="use bounds instead of target"):
+        [eq], _ = optimizer.optimize(eq, objective, constraints, maxiter=2, verbose=3)
 
     # second, test optimizing both the equilibrium and the field simultaneously
     objective = ObjectiveFunction(
@@ -1188,9 +1197,10 @@ def test_omnigenity_proximal():
         FixPsi(eq=eq),
     )
     optimizer = Optimizer("proximal-lsq-exact")
-    (eq, field), _ = optimizer.optimize(
-        (eq, field), objective, constraints, maxiter=2, verbose=3
-    )
+    with pytest.warns(UserWarning, match="use bounds instead of target"):
+        (eq, field), _ = optimizer.optimize(
+            (eq, field), objective, constraints, maxiter=2, verbose=3
+        )
 
 
 @pytest.mark.unit
@@ -2408,7 +2418,9 @@ def test_ballooning_stability_opt():
         gtol=1e-6,
         maxiter=2,  # increase maxiter to 50 for a better result
         verbose=3,
-        options={"initial_trust_ratio": 2e-3},
+        # Jacobian has only 2 rows and 1 of them can be full of 0s
+        # default QR can fail, choose SVD instead
+        options={"initial_trust_ratio": 2e-3, "tr_method": "svd"},
     )
     data = eq.compute("ideal ballooning lambda", grid=grid)
     lam2_optimized = data["ideal ballooning lambda"].max((-1, -2, -3))