diff --git a/CHANGELOG.md b/CHANGELOG.md index f1a2a5651e..2449a4d868 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,15 @@ Changelog ========= + +Performance Improvements + +- Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep. On ``jax >= 0.10.0`` this uses ``qr_multiply`` to additionally avoid forming ``Q`` explicitly; on older versions a fallback preserves the same results. + + +v0.17.2 +------- + New Features - Adds ``desc.objectives.DeflationOperator``, a new objective class which can be used to apply deflation techniques to equilibrium and optimization problems to find multiple local minima or multiple solutions from a single initial point, either by wrapping an existing ``desc.objectives._Objective`` object or by including as an additional penalty or constraint. Also adds a tutorial showing this functionality. diff --git a/desc/backend.py b/desc/backend.py index a5dbd9a17b..e893d540f5 100644 --- a/desc/backend.py +++ b/desc/backend.py @@ -87,6 +87,23 @@ def _diag_to_full(d, e): from jax.numpy.fft import ifft, irfft, irfft2, rfft, rfft2 from jax.scipy.fft import dct, dctn, idct, idctn from jax.scipy.linalg import block_diag, cho_factor, cho_solve, qr, solve_triangular + + # TODO: remove fallback once JAX min version >= 0.10.0 + if Version(jax.__version__) >= Version("0.10.0"): + from jax.scipy.linalg import qr_multiply + else: + + def qr_multiply(a, c, mode="right"): + """Fallback for ``jax.scipy.linalg.qr_multiply`` (added in JAX 0.10.0).""" + Q, R = qr(a, mode="economic") + if mode == "right": + # 1-D c (all DESC uses) matches the old Q.T @ c; c @ Q keeps + # higher-dim c consistent with qr_multiply rather than silently wrong + cq = Q.T @ c if c.ndim == 1 else c @ Q + else: + cq = Q @ c + return cq, R + from jax.scipy.special import gammaln from jax.tree_util import ( register_pytree_node, @@ -539,6 +556,7 @@ def bodyfun(state): cho_factor, cho_solve, qr, + qr_multiply, solve_triangular, ) from scipy.special import gammaln # noqa: F401 diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py index 5dfd9fe97b..5badddc61f 100644 --- a/desc/optimize/aug_lagrangian_ls.py +++ b/desc/optimize/aug_lagrangian_ls.py @@ -2,7 +2,7 @@ from scipy.optimize import NonlinearConstraint, OptimizeResult -from desc.backend import jnp, qr +from desc.backend import jnp, qr, qr_multiply from desc.utils import errorif, safediv, setdefault from .bound_utils import ( @@ -408,15 +408,15 @@ def lagjac(z, y, mu, *args): # try full newton step tall = J_a.shape[0] >= J_a.shape[1] if tall: - Q, R = qr(J_a, mode="economic") - p_newton = solve_triangular_regularized(R, -Q.T @ L_a) + Qt_La, R = qr_multiply(J_a, L_a, mode="right") + p_newton = solve_triangular_regularized(R, -Qt_La) else: - Q, R = qr(J_a.T, mode="economic") - p_newton = Q @ solve_triangular_regularized(R.T, -L_a, lower=True) - # We don't need the Q and R matrices anymore - # Trust region solver will solve the augmented system - # with a new Q and R - del Q, R + # min-norm Newton step uses the QR of J_a.T + Q, Rt = qr(J_a.T, mode="economic") + p_newton = Q @ solve_triangular_regularized(Rt.T, -L_a, lower=True) + del Q, Rt + # the tr subproblem still needs the QR of J_a itself + Qt_La, R = qr_multiply(J_a, L_a, mode="right") actual_reduction = -1 Lactual_reduction = -1 @@ -439,7 +439,7 @@ def lagjac(z, y, mu, *args): ) elif tr_method == "qr": step_h, hits_boundary, alpha = trust_region_step_exact_qr( - p_newton, L_a, J_a, trust_radius, alpha + p_newton, Qt_La, R, trust_radius, alpha ) step = d * step_h # Trust-region solution in the original space. diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py index 089871a5a8..125247f1e6 100644 --- a/desc/optimize/least_squares.py +++ b/desc/optimize/least_squares.py @@ -2,7 +2,7 @@ from scipy.optimize import OptimizeResult -from desc.backend import jnp, qr +from desc.backend import jnp, qr, qr_multiply from desc.utils import errorif, safediv, setdefault from .bound_utils import ( @@ -302,15 +302,15 @@ def lsqtr( # noqa: C901 # try full newton step tall = J_a.shape[0] >= J_a.shape[1] if tall: - Q, R = qr(J_a, mode="economic") - p_newton = solve_triangular_regularized(R, -Q.T @ f_a) + Qt_fa, R = qr_multiply(J_a, f_a, mode="right") + p_newton = solve_triangular_regularized(R, -Qt_fa) else: - Q, R = qr(J_a.T, mode="economic") - p_newton = Q @ solve_triangular_regularized(R.T, -f_a, lower=True) - # We don't need the Q and R matrices anymore - # Trust region solver will solve the augmented system - # with a new Q and R - del Q, R + # min-norm Newton step uses the QR of J_a.T + Q, Rt = qr(J_a.T, mode="economic") + p_newton = Q @ solve_triangular_regularized(Rt.T, -f_a, lower=True) + del Q, Rt + # the tr subproblem still needs the QR of J_a itself + Qt_fa, R = qr_multiply(J_a, f_a, mode="right") actual_reduction = -1 @@ -332,7 +332,7 @@ def lsqtr( # noqa: C901 ) elif tr_method == "qr": step_h, hits_boundary, alpha = trust_region_step_exact_qr( - p_newton, f_a, J_a, trust_radius, alpha + p_newton, Qt_fa, R, trust_radius, alpha ) step = d * step_h # Trust-region solution in the original space. diff --git a/desc/optimize/tr_subproblems.py b/desc/optimize/tr_subproblems.py index 0d592fec84..9335015a3e 100644 --- a/desc/optimize/tr_subproblems.py +++ b/desc/optimize/tr_subproblems.py @@ -9,6 +9,7 @@ jit, jnp, qr, + qr_multiply, solve_triangular, while_loop, ) @@ -357,7 +358,7 @@ def loop_body(state): @jit def trust_region_step_exact_qr( - p_newton, f, J, trust_radius, initial_alpha=0.0, rtol=0.01, max_iter=10 + p_newton, z, R, trust_radius, initial_alpha=0.0, rtol=0.01, max_iter=10 ): """Solve a trust-region problem using a semi-exact method. @@ -373,12 +374,19 @@ def trust_region_step_exact_qr( which is equivalent to || [J; sqrt(alpha)*I].Tp - [f; 0].T ||^2 + The caller supplies the factorization ``J = Q1@R`` (and ``z = Q1.T@f``), so + the alpha-loop only retriangularizes the small reduced system + ``[R; sqrt(alpha)*I]`` instead of refactorizing ``J`` each iteration. + Parameters ---------- - f : ndarray - Vector of residuals. - J : ndarray - Jacobian matrix. + p_newton : ndarray + The full (unregularized) Newton step, returned as-is if it lies within + the trust region. + z : ndarray + ``Q1.T@f``, where ``J = Q1@R`` is the (economic) QR factorization of J. + R : ndarray + The R factor of J, as returned by ``qr_multiply(J, f, mode="right")``. trust_radius : float Radius of a trust region. initial_alpha : float, optional @@ -407,13 +415,15 @@ def truefun(*_): return p_newton, False, 0.0 def falsefun(*_): - alpha_upper = jnp.linalg.norm(J.T @ f) / trust_radius + # J.T@f == R.T@z, so we never need J or f here + alpha_upper = jnp.linalg.norm(R.T @ z) / trust_radius alpha_lower = 0.0 alpha = initial_alpha alpha = jnp.clip(alpha, alpha_lower, alpha_upper) k = 0 - fp = jnp.pad(f, (0, J.shape[1])) + n = R.shape[1] + zp = jnp.concatenate([z, jnp.zeros(n)]) def loop_cond(state): p, alpha, alpha_lower, alpha_upper, phi, k = state @@ -422,17 +432,16 @@ def loop_cond(state): def loop_body(state): p, alpha, alpha_lower, alpha_upper, phi, k = state - Ji = jnp.vstack([J, jnp.sqrt(alpha) * jnp.eye(J.shape[1])]) - # Ji is always tall since its padded by alpha*I - Q, R = qr(Ji, mode="economic") + A = jnp.vstack([R, jnp.sqrt(alpha) * jnp.eye(n)]) + Qtz, Rtil = qr_multiply(A, zp, mode="right") - p = solve_triangular_regularized(R, -Q.T @ fp) + p = solve_triangular_regularized(Rtil, -Qtz) p_norm = jnp.linalg.norm(p) phi = p_norm - trust_radius alpha_upper = jnp.where(phi < 0, alpha, alpha_upper) alpha_lower = jnp.where(phi > 0, alpha, alpha_lower) - q = solve_triangular_regularized(R.T, p, lower=True) + q = solve_triangular_regularized(Rtil.T, p, lower=True) q_norm = jnp.linalg.norm(q) alpha += (p_norm / q_norm) ** 2 * phi / trust_radius