From 317208119a00774503f1052740b1a47b0e791ec9 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 10 Aug 2026 12:36:24 -0400 Subject: [PATCH 01/10] fix jax-finufft related error in test_compute_everything --- tests/test_compute_everything.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index 0889b46bfc..46471b7456 100644 --- a/tests/test_compute_everything.py +++ b/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: From a584f175ffc75cec7fd50cdb5a80d3231aabf8e1 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 10 Aug 2026 12:43:50 -0400 Subject: [PATCH 02/10] temporarily change all tests to use jax-finufft v1.2.0 --- .github/workflows/regression_tests.yml | 2 +- .github/workflows/unit_tests.yml | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml index 750f9e05ae..87d40161b2 100644 --- a/.github/workflows/regression_tests.yml +++ b/.github/workflows/regression_tests.yml @@ -22,7 +22,7 @@ jobs: GH_TOKEN: ${{ github.token }} strategy: matrix: - python-version: ['3.11'] + python-version: ['3.10'] group: [1, 2, 3, 4, 5, 6] steps: diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index f7e63729ec..65984386a4 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -23,13 +23,13 @@ jobs: strategy: matrix: combos: [{group: 1, python_version: '3.10'}, - {group: 2, python_version: '3.11'}, - {group: 3, python_version: '3.11'}, - {group: 4, python_version: '3.12'}, - {group: 5, python_version: '3.12'}, - {group: 6, python_version: '3.13'}, - {group: 7, python_version: '3.13'}, - {group: 8, python_version: '3.14'}] + {group: 2, python_version: '3.10'}, + {group: 3, python_version: '3.10'}, + {group: 4, python_version: '3.10'}, + {group: 5, python_version: '3.10'}, + {group: 6, python_version: '3.10'}, + {group: 7, python_version: '3.10'}, + {group: 8, python_version: '3.10'}] steps: - uses: actions/checkout@v7 From 3eade7ce08ca7d3eec849ab67e24a0bc17fe7f8c Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 10 Aug 2026 18:19:09 -0400 Subject: [PATCH 03/10] update test, update docs --- desc/optimize/aug_lagrangian_ls.py | 4 +++- desc/optimize/least_squares.py | 4 +++- tests/test_examples.py | 10 +++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py index 5badddc61f..9d3636a468 100644 --- a/desc/optimize/aug_lagrangian_ls.py +++ b/desc/optimize/aug_lagrangian_ls.py @@ -172,7 +172,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 high 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. diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py index 125247f1e6..ac5516526a 100644 --- a/desc/optimize/least_squares.py +++ b/desc/optimize/least_squares.py @@ -138,7 +138,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 high 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. diff --git a/tests/test_examples.py b/tests/test_examples.py index 81d8f660c7..cde947b4c1 100644 --- a/tests/test_examples.py +++ b/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)) From ae24b1b3da1077927ca15e2fe3e2a0d104a9f126 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 10 Aug 2026 19:19:52 -0400 Subject: [PATCH 04/10] add a warning for possible rank deficient case --- desc/optimize/_desc_wrappers.py | 33 ++++++++++++++++++++++++++++++ desc/optimize/aug_lagrangian_ls.py | 2 +- desc/optimize/least_squares.py | 2 +- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/desc/optimize/_desc_wrappers.py b/desc/optimize/_desc_wrappers.py index da648ac7d1..0ed3702a27 100644 --- a/desc/optimize/_desc_wrappers.py +++ b/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 @@ def _optimize_desc_aug_lagrangian_least_squares( 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 @@ def _optimize_desc_least_squares( 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( diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py index 3fd9094e67..5be0c312e3 100644 --- a/desc/optimize/aug_lagrangian_ls.py +++ b/desc/optimize/aug_lagrangian_ls.py @@ -175,7 +175,7 @@ def lsq_auglag( # noqa: C901 especially on GPU, but may be less accurate for badly scaled systems. ``"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 high chance to be rank-deficient. Default ``"qr"``. + 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. diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py index c3eddba0c0..472baf4e77 100644 --- a/desc/optimize/least_squares.py +++ b/desc/optimize/least_squares.py @@ -141,7 +141,7 @@ def lsqtr( # noqa: C901 especially on GPU, but may be less accurate for badly scaled systems. ``"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 high chance to be rank-deficient. Default ``"qr"``. + 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. From bc43dfce8aab02093e2326a7f46daba4ac797270 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 10 Aug 2026 21:20:04 -0400 Subject: [PATCH 05/10] update tests, they were not doing actual optimization, but still fix them --- tests/test_examples.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index cde947b4c1..b8ea8f0255 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1179,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( @@ -1196,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 @@ -2416,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)) From 1c11a7bbe20a6982a8c75c65f63eee94bdacfe6d Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 11 Aug 2026 12:36:14 -0400 Subject: [PATCH 06/10] revert back to original python versions --- .github/workflows/regression_tests.yml | 2 +- .github/workflows/unit_tests.yml | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml index 87d40161b2..750f9e05ae 100644 --- a/.github/workflows/regression_tests.yml +++ b/.github/workflows/regression_tests.yml @@ -22,7 +22,7 @@ jobs: GH_TOKEN: ${{ github.token }} strategy: matrix: - python-version: ['3.10'] + python-version: ['3.11'] group: [1, 2, 3, 4, 5, 6] steps: diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 65984386a4..f7e63729ec 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -23,13 +23,13 @@ jobs: strategy: matrix: combos: [{group: 1, python_version: '3.10'}, - {group: 2, python_version: '3.10'}, - {group: 3, python_version: '3.10'}, - {group: 4, python_version: '3.10'}, - {group: 5, python_version: '3.10'}, - {group: 6, python_version: '3.10'}, - {group: 7, python_version: '3.10'}, - {group: 8, python_version: '3.10'}] + {group: 2, python_version: '3.11'}, + {group: 3, python_version: '3.11'}, + {group: 4, python_version: '3.12'}, + {group: 5, python_version: '3.12'}, + {group: 6, python_version: '3.13'}, + {group: 7, python_version: '3.13'}, + {group: 8, python_version: '3.14'}] steps: - uses: actions/checkout@v7 From c569650ee8b0b9616f104995f72f7ce0a4ffb10e Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sat, 15 Aug 2026 02:32:24 -0400 Subject: [PATCH 07/10] solve the cho failed bug --- desc/optimize/utils.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index 41845ad392..1f0ce11c80 100644 --- a/desc/optimize/utils.py +++ b/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. @@ -233,13 +247,13 @@ def _cholmod(A, maxiter=4): 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 @@ -263,7 +277,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 From 5b0c38e8c9ec2a02a9eae32de4df4452e225241c Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 18 Aug 2026 00:15:32 -0400 Subject: [PATCH 08/10] fix small out of bounds indexing, update changelog --- CHANGELOG.md | 3 +++ desc/optimize/utils.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b209ad29dc..9e435a47bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,9 @@ 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. + v0.17.3 ------- diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index 1f0ce11c80..605211e989 100644 --- a/desc/optimize/utils.py +++ b/desc/optimize/utils.py @@ -239,11 +239,11 @@ 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 From c849dc9612e6357d894d0e28c332c90c167de3f2 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 18 Aug 2026 00:23:18 -0400 Subject: [PATCH 09/10] add todo --- desc/optimize/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index 605211e989..700386ee0f 100644 --- a/desc/optimize/utils.py +++ b/desc/optimize/utils.py @@ -229,6 +229,7 @@ def _cholmod(A, maxiter=4): # upper bound on log(alpha) such that A + alpha*I > 0, ie we know alpha < ub # lower bound on eig(A) = upper bound on alpha, +1 in log scale to make sure # that it's actually greater than the maximum alpha + # TODO: add special handling for lb = 0, i.e. 0 row of A, rest diag dominant ub = jnp.log10(jnp.abs(lb)) + 1 # we know alpha > 0 because otherwise initial factorization would have succeeded # but we'd like to be a bit better (in log scaling). This is just a heuristic but From cd40f7d3b06777c5228b026aa083daaf7f2808d9 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 18 Aug 2026 00:26:40 -0400 Subject: [PATCH 10/10] add issue number --- desc/optimize/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index 700386ee0f..3cf8ebdc0e 100644 --- a/desc/optimize/utils.py +++ b/desc/optimize/utils.py @@ -229,7 +229,7 @@ def _cholmod(A, maxiter=4): # upper bound on log(alpha) such that A + alpha*I > 0, ie we know alpha < ub # lower bound on eig(A) = upper bound on alpha, +1 in log scale to make sure # that it's actually greater than the maximum alpha - # TODO: add special handling for lb = 0, i.e. 0 row of A, rest diag dominant + # TODO(#2296): add special handling for lb = 0, i.e. 0 row of A, rest diag dominant ub = jnp.log10(jnp.abs(lb)) + 1 # we know alpha > 0 because otherwise initial factorization would have succeeded # but we'd like to be a bit better (in log scaling). This is just a heuristic but