fix type J must be J_a

PlasmaControl · Aug 15, 2024 · dc1c6aa · dc1c6aa
1 parent 5e30fdc
commit dc1c6aa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/desc/optimize/aug_lagrangian_ls.py b/desc/optimize/aug_lagrangian_ls.py
@@ -371,12 +371,12 @@ def lagjac(z, y, mu, *args):
             B_h = jnp.dot(J_a.T, J_a)
         elif tr_method == "qr":
             # try full newton step
-            tall = J.shape[0] >= J.shape[1]
+            tall = J_a.shape[0] >= J_a.shape[1]
             if tall:
-                Q, R = qr(J, mode="economic")
+                Q, R = qr(J_a, mode="economic")
                 p_newton = solve_triangular_regularized(R, -Q.T @ f)
             else:
-                Q, R = qr(J.T, mode="economic")
+                Q, R = qr(J_a.T, mode="economic")
                 p_newton = Q @ solve_triangular_regularized(R.T, -f, lower=True)
 
         actual_reduction = -1

diff --git a/desc/optimize/least_squares.py b/desc/optimize/least_squares.py
@@ -271,12 +271,12 @@ def lsqtr(  # noqa: C901 - FIXME: simplify this
             B_h = jnp.dot(J_a.T, J_a)
         elif tr_method == "qr":
             # try full newton step
-            tall = J.shape[0] >= J.shape[1]
+            tall = J_a.shape[0] >= J_a.shape[1]
             if tall:
-                Q, R = qr(J, mode="economic")
+                Q, R = qr(J_a, mode="economic")
                 p_newton = solve_triangular_regularized(R, -Q.T @ f)
             else:
-                Q, R = qr(J.T, mode="economic")
+                Q, R = qr(J_a.T, mode="economic")
                 p_newton = Q @ solve_triangular_regularized(R.T, -f, lower=True)
 
         actual_reduction = -1