Always re-solve eq in ProximalProjection before first computation of value and jacobian

CHANGELOG.md

@@ -9,6 +9,11 @@ Bug Fixes
 ---------
 - [Fixes straight field line equilibrium conversion](https://github.com/PlasmaControl/DESC/pull/1880).
 
+Backend
+
+- When using any of the ``"proximal-"`` optimization methods, the equilbrium is now always solved before beginning optimization to the specified tolerance (as determined, for example, by ``options={"solve_options":{"ftol"...}}`` passed to the ``desc.optimize.Optimizer.optimize`` call). This ensures the assumptions of the proximal projection method are enforced starting from the first step of the optimization.
+
+
 v0.15.0
 -------
 
@@ -37,6 +42,7 @@ Bug Fixes
 - Fixes bug in ``desc.geometry.curve.FourierRZCurve.from_values`` when numpy array is passed in for the ``coords`` argument.
 
 Backend
+
 - Significant changes to how DESC handles static attributes during JIT compilation. Going forward if any class/object has attributes that should be treated as static by `jax.jit`, these should be declared at the class level like `_static_attrs = ["foo", "bar"]`. Generally, non-arraylike attributes such as functions, strings etc should be marked static, as well as any attributes used for control flow. Previously this was done automatically, but in a way that caused a lot of performance bugs and unnecessary recompilation. These changes have been implemented for all classes in the `desc` repository, but if you have custom objectives or other local objects that subclass from `desc` you may need to add this yourself. JAX error messages usually do a good job of alerting you to things that need to be static, and feel free to open an issue with `desc` if you have any questions.
 - No longer closes over the field in ``desc.magnetic_fields._core.field_line_integrate``, which can dramatically reduce compile times when the field being traced has large size attributes (for example, when using a ``desc.magnetic_fields._core.SplineMagneticField`` object).
 

desc/optimize/_constraint_wrappers.py

@@ -609,6 +609,9 @@ def __init__(
         self._objective = objective
         self._constraint = constraint
         solve_options = {} if solve_options is None else solve_options
+        self._solve_during_proximal_build = solve_options.pop(
+            "solve_during_proximal_build", True
+        )  # If user does not want the solve during build, mainly for debug purposes
         perturb_options = {} if perturb_options is None else perturb_options
         perturb_options.setdefault("verbose", 0)
         perturb_options.setdefault("include_f", False)
@@ -771,7 +774,17 @@ def build(self, use_jit=None, verbose=1):  # noqa: C901
             [np.atleast_2d(foo) for foo in self._feasible_tangents], axis=-1
         )
 
-        # history and caching
+        ## history and caching
+        # first, ensure equilibrium is solved to the
+        # specified tolerances, necessary as we assume
+        # eq is solved when taking the derivatives later
+        if self._solve_during_proximal_build:
+            self._eq.solve(
+                objective=self._eq_solve_objective,
+                constraints=None,
+                **self._solve_options,
+            )
+        # then store the now-solved eq state as the initial state
         self._x_old = self.x(self.things)
         self._allx = [self._x_old]
         self._allxopt = [self._objective.x(*self.things)]

tests/benchmarks/benchmark_cpu_small.py

@@ -333,7 +333,9 @@ def test_proximal_jac_atf(benchmark):
     grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.linspace(0.1, 1, 10))
     objective = ObjectiveFunction(QuasisymmetryTwoTerm(eq, grid=grid))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     prox.build()
     x = prox.x(eq)
     prox.jac_scaled_error(x, prox.constants).block_until_ready()
@@ -361,7 +363,11 @@ def test_proximal_jac_atf_with_eq_update(benchmark):
         constraint,
         eq,
         perturb_options={"verbose": 3},
-        solve_options={"verbose": 3, "maxiter": 0},
+        solve_options={
+            "verbose": 3,
+            "maxiter": 0,
+            "solve_during_proximal_build": False,
+        },
     )
     prox.build(verbose=3)
     x = prox.x(eq)
@@ -388,7 +394,9 @@ def test_proximal_freeb_compute(benchmark):
     field = ToroidalMagneticField(1.0, 1.0)  # just a dummy field for benchmarking
     objective = ObjectiveFunction(BoundaryError(eq, field=field))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -412,7 +420,9 @@ def test_proximal_freeb_jac(benchmark):
     field = ToroidalMagneticField(1.0, 1.0)  # just a dummy field for benchmarking
     objective = ObjectiveFunction(BoundaryError(eq, field=field))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -532,7 +542,9 @@ def _test_objective_ripple(benchmark, spline, method):
         ]
     )
     constraint = ObjectiveFunction([ForceBalance(eq)])
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     prox.build()
     x = prox.x(eq)
     _ = getattr(prox, method)(x, prox.constants).block_until_ready()

tests/benchmarks/benchmark_gpu_small.py

@@ -336,7 +336,9 @@ def test_proximal_jac_atf(benchmark):
     grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.linspace(0.1, 1, 10))
     objective = ObjectiveFunction(QuasisymmetryTwoTerm(eq, grid=grid))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     prox.build()
     x = prox.x(eq)
     prox.jac_scaled_error(x, prox.constants).block_until_ready()
@@ -364,7 +366,11 @@ def test_proximal_jac_atf_with_eq_update(benchmark):
         constraint,
         eq,
         perturb_options={"verbose": 3},
-        solve_options={"verbose": 3, "maxiter": 0},
+        solve_options={
+            "verbose": 3,
+            "maxiter": 0,
+            "solve_during_proximal_build": False,
+        },
     )
     prox.build(verbose=3)
     x = prox.x(eq)
@@ -391,7 +397,9 @@ def test_proximal_freeb_compute(benchmark):
     field = ToroidalMagneticField(1.0, 1.0)  # just a dummy field for benchmarking
     objective = ObjectiveFunction(BoundaryError(eq, field=field))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -415,7 +423,9 @@ def test_proximal_freeb_jac(benchmark):
     field = ToroidalMagneticField(1.0, 1.0)  # just a dummy field for benchmarking
     objective = ObjectiveFunction(BoundaryError(eq, field=field))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -535,7 +545,9 @@ def _test_objective_ripple(benchmark, spline, method):
         ]
     )
     constraint = ObjectiveFunction([ForceBalance(eq)])
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     prox.build()
     x = prox.x(eq)
     _ = getattr(prox, method)(x, prox.constants).block_until_ready()

tests/benchmarks/memory_funcs.py

@@ -73,7 +73,11 @@ def test_proximal_jac_w7x_with_eq_update():
         constraint,
         eq,
         perturb_options={"verbose": 0},
-        solve_options={"verbose": 0, "maxiter": 0},
+        solve_options={
+            "verbose": 0,
+            "maxiter": 0,
+            "solve_during_proximal_build": False,
+        },
     )
     prox.build(verbose=0)
     x = prox.x(eq)
@@ -97,7 +101,9 @@ def test_proximal_freeb_jac():
     field = ToroidalMagneticField(1.0, 1.0)  # just a dummy field for benchmarking
     objective = ObjectiveFunction(BoundaryError(eq, field=field))
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -125,7 +131,9 @@ def test_proximal_freeb_jac_batched():
         jac_chunk_size=100,
     )
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -152,7 +160,9 @@ def test_proximal_freeb_jac_blocked():
         deriv_mode="blocked",
     )
     constraint = ObjectiveFunction(ForceBalance(eq))
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     obj = LinearConstraintProjection(
         prox, ObjectiveFunction((FixCurrent(eq), FixPressure(eq), FixPsi(eq)))
     )
@@ -196,7 +206,9 @@ def _test_proximal_ripple(spline, method):
         ]
     )
     constraint = ObjectiveFunction([ForceBalance(eq)])
-    prox = ProximalProjection(objective, constraint, eq)
+    prox = ProximalProjection(
+        objective, constraint, eq, solve_options={"solve_during_proximal_build": False}
+    )
     prox.build(verbose=0)
     x = prox.x(eq)
     for _ in range(3):

tests/test_optimizer.py

@@ -409,7 +409,8 @@ def compute(self, params, constants=None):
     np.random.seed(0)
     objective = ObjectiveFunction(DummyObjective(things=eq), use_jit=False)
     # make gradient super noisy so it stalls
-    objective.jac_scaled_error = lambda x, *args: objective._jac_scaled_error(
+    objective.build()
+    objective.jac_scaled_error = lambda x, *args: objective.jac_scaled_error(
         x
     ) + 1e2 * (np.random.random((objective._dim_f, x.size)) - 0.5)
 
@@ -447,7 +448,13 @@ def compute(self, params, constants=None):
         options={
             "initial_trust_radius": 0.5,
             "perturb_options": {"verbose": 0, "order": 1},
-            "solve_options": {"verbose": 0, "maxiter": 2},
+            "solve_options": {
+                "verbose": 0,
+                "maxiter": 2,
+                # Hidden kwarg just for debug/tests, to not solve
+                # during build
+                "solve_during_proximal_build": False,
+            },
         },
     )