Patch to work around discretization error impeding optimizer

CHANGELOG.md

@@ -3,6 +3,7 @@ Changelog
 
 New Features
 ------------
+- Inverse method free surface optimization.
 - Automatically-differentiable, non-singular Laplace BIE solver.
 - Improved performance and accuracy of FFT interpolation in singular integrals
   ([1](https://github.com/f0uriest/interpax/pull/116), [2](https://github.com/f0uriest/interpax/pull/117)).

desc/objectives/_free_boundary.py

@@ -6,6 +6,7 @@
 from desc.backend import jnp
 from desc.compute import get_profiles, get_transforms
 from desc.compute.utils import _compute as compute_fun
+from desc.compute.utils import _compute_RpZ_data as compute_fun_rpz
 from desc.grid import LinearGrid
 from desc.integrals import get_interpolator, virtual_casing_biot_savart
 from desc.nestor import Nestor
@@ -1023,10 +1024,24 @@ def __init__(
         self._chunk_size = chunk_size
         self._B_coil_chunk_size = B_coil_chunk_size
         self._grad_keys = ["grad(theta)", "grad(zeta)", "n_rho"]
-        self._inner_keys = ["|B|^2", "p", "I", "|e_theta x e_zeta|"] + self._grad_keys
+        self._inner_keys = [
+            "|B|^2",
+            "p",
+            "I",
+            "|e_theta x e_zeta|",
+            "R",
+            "phi",
+            "zeta",
+            "omega",
+            "Z",
+        ] + self._grad_keys
         self._reuseable_keys = [
             "0",
             "R",
+            "Z",
+            "phi",
+            "zeta",
+            "omega",
             "R_t",
             "R_z",
             "Z_t",
@@ -1210,7 +1225,11 @@ def compute(self, params, constants=None):
 
         outer = compute_fun(
             self._field,
-            ["K_vc", "n_rho x B_coil"] if self._is_neumann else "|K_vc|^2",
+            (
+                ["K_vc", "n_rho x B_coil", "K_vc (periodic)"]
+                if self._is_neumann
+                else "|K_vc|^2"
+            ),
             field_params,
             constants["eval_transforms"],
             constants["profiles"],
@@ -1223,11 +1242,34 @@ def compute(self, params, constants=None):
             field_grid=self._coil_grid,
             **problem,
         )
+
+        outer_rpz = compute_fun_rpz(
+            self._field,
+            ["B"],
+            field_params,
+            constants["eval_transforms"],
+            constants["profiles"],
+            data=outer,
+            maxiter=self._maxiter,
+            chunk_size=self._chunk_size,
+            B_coil_chunk_size=self._B_coil_chunk_size,
+            B_coil=self._B_coil,
+            field_grid=self._coil_grid,
+            RpZ_data={"R": outer["R"], "Z": outer["Z"], "phi": outer["phi"]},
+            RpZ_grid=constants["eval_transforms"]["grid"],
+            eval_interpolator=outer["interpolator"],
+            B0=self._field._B0,
+            on_boundary=True,
+            **problem,
+        )
+
         if self._is_neumann:
-            outer["K_vc"] -= outer["n_rho x B_coil"]
-            outer["|K_vc|^2"] = dot(outer["K_vc"], outer["K_vc"])
+            # Compute from equation 4.26 instead of the equation
+            # sandwhiched between 4.23 and 4.24 due to discretization error
+            # in quadrature messing up the optimizer.
+            outer["|B|^2"] = dot(outer_rpz["B"], outer_rpz["B"])
 
-        return (outer["|K_vc|^2"] - inner["|B|^2"] - 2 * mu_0 * inner["p"]) * inner[
+        return (outer["|B|^2"] - inner["|B|^2"] - 2 * mu_0 * inner["p"]) * inner[
             "|e_theta x e_zeta|"
         ]