Use qr_multiply to allow better reuse of factorisation

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.

desc/backend.py

@@ -87,6 +87,23 @@
     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 @@
         cho_factor,
         cho_solve,
         qr,
+        qr_multiply,
         solve_triangular,
     )
     from scipy.special import gammaln  # noqa: F401

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 @@
             # 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 @@
                 )
             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.

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.
 

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