From 0d4c33143d67da5841ea2e15c8217e6cb83e6802 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 2 Dec 2024 17:13:24 -0500
Subject: [PATCH] bug fix for proximal scalar optimization

---
 desc/optimize/_constraint_wrappers.py | 19 ++++++++++++++
 tests/test_optimizer.py               | 38 +++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py
index dc2e3033c7..7695a671b4 100644
--- a/desc/optimize/_constraint_wrappers.py
+++ b/desc/optimize/_constraint_wrappers.py
@@ -836,6 +836,25 @@ def compute_scaled_error(self, x, constants=None):
         xopt, _ = self._update_equilibrium(x, store=False)
         return self._objective.compute_scaled_error(xopt, constants[0])
 
+    def compute_scalar(self, x, constants=None):
+        """Compute the sum of squares error.
+
+        Parameters
+        ----------
+        x : ndarray
+            State vector.
+        constants : list
+            Constant parameters passed to sub-objectives.
+
+        Returns
+        -------
+        f : float
+            Objective function scalar value.
+
+        """
+        f = jnp.sum(self.compute_scaled_error(x, constants=constants) ** 2) / 2
+        return f
+
     def compute_unscaled(self, x, constants=None):
         """Compute the raw value of the objective function.
 
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 4f0cce0e0b..9a4aac89bf 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -324,6 +324,44 @@ def test_no_iterations():
     np.testing.assert_allclose(x0, out2["x"])
 
 
+@pytest.mark.regression
+@pytest.mark.optimize
+def test_proximal_scalar():
+    """Test that proximal scalar optimization works."""
+    # test fix for GH issue #1403
+
+    # optimize to reduce DSHAPE volume from 100 m^3 to 90 m^3
+    eq = desc.examples.get("DSHAPE")
+    optimizer = Optimizer("proximal-fmintr")  # proximal scalar optimizer
+    R_modes = np.vstack(
+        (
+            [0, 0, 0],
+            eq.surface.R_basis.modes[
+                np.max(np.abs(eq.surface.R_basis.modes), 1) > 1, :
+            ],
+        )
+    )
+    Z_modes = eq.surface.Z_basis.modes[
+        np.max(np.abs(eq.surface.Z_basis.modes), 1) > 1, :
+    ]
+    objective = ObjectiveFunction(Volume(eq=eq, target=90))  # scalar objective function
+    constraints = (
+        FixBoundaryR(eq=eq, modes=R_modes),
+        FixBoundaryZ(eq=eq, modes=Z_modes),
+        FixIota(eq=eq),
+        FixPressure(eq=eq),
+        FixPsi(eq=eq),
+        ForceBalance(eq=eq),  # force balance constraint for proximal projection
+    )
+    [eq], _ = optimizer.optimize(
+        things=eq,
+        objective=objective,
+        constraints=constraints,
+        verbose=3,
+    )
+    np.testing.assert_allclose(eq.compute("V")["V"], 90)
+
+
 @pytest.mark.regression
 @pytest.mark.slow
 @pytest.mark.optimize