Add eq_fixed option to BoundaryError

CHANGELOG.md

@@ -25,6 +25,7 @@ New Features
 - Adds an optional attribute `ion_density` to the `Equilibrium` class, to allow the ion density profile to be set independently of the electron density and effective atomic number. Also adds compute functions for ``"ni_rr"`` and ``"Zeff_rr"``.
 - Modernizes dependencies to use [``nvidia-ml-py``](https://pypi.org/project/nvidia-ml-py/) in place of [``nvgpu``](https://github.com/rossumai/nvgpu).
   If you are updating an existing software environment uninstall ``pynvml`` first and then reinstall the dependencies to correctly get ``nvidia-ml-py``.
+- Option for ``eq_fixed`` added to ``BoundaryError`` to remove equilibrium as a degree of freedom from boundary error objective. 
 
 Bug Fixes
 

desc/objectives/_free_boundary.py

@@ -447,6 +447,7 @@
         "_B_plasma_chunk_size",
         "_bs_chunk_size",
         "_eq_data_keys",
+        "_eq_fixed",
         "_field_fixed",
         "_q",
         "_sheet_current",
@@ -478,6 +479,7 @@
         eval_grid=None,
         field_grid=None,
         field_fixed=False,
+        eq_fixed=False,
         name="Boundary error",
         jac_chunk_size=None,
         *,
@@ -487,12 +489,20 @@
     ):
         if target is None and bounds is None:
             target = 0
+        self._eq = eq
         self._source_grid = source_grid
         self._eval_grid = eval_grid
         self._st, self._sz = s if isinstance(s, (tuple, list)) else (s, s)
         self._q = q
         self._field = [field] if not isinstance(field, list) else field
         self._field_grid = field_grid
+        errorif(
+            field_fixed and eq_fixed,
+            ValueError,
+            "At least one of eq_fixed or field_fixed must be false.",
+        )
+        self._field_fixed = field_fixed
+        self._eq_fixed = eq_fixed
         self._bs_chunk_size = bs_chunk_size
         B_plasma_chunk_size = parse_argname_change(
             B_plasma_chunk_size, kwargs, "loop", "B_plasma_chunk_size"
@@ -501,7 +511,9 @@
             B_plasma_chunk_size = None
         self._B_plasma_chunk_size = B_plasma_chunk_size
         self._sheet_current = hasattr(eq.surface, "Phi_mn")
-        things = [eq]
+        things = []
+        if not eq_fixed:
+            things.append(self._eq)
         if not field_fixed:
             things.append(self._field)
         super().__init__(
@@ -530,7 +542,7 @@
         """
         from desc.magnetic_fields import SumMagneticField
 
-        eq = self.things[0]
+        eq = self._eq
 
         if self._source_grid is None:
             # for axisymmetry we still need to know about toroidal effects, so its
@@ -628,6 +640,9 @@
 
         neq = 3 if self._sheet_current else 2  # number of equations we're using
 
+        if self._eq_fixed:
+            eq = self._eq
+
         self._constants = {
             "eval_transforms": eval_transforms,
             "eval_profiles": eval_profiles,
@@ -651,6 +666,74 @@
                 )
             )
 
+        if self._eq_fixed:
+            source_data = compute_fun(
+                "desc.equilibrium.equilibrium.Equilibrium",
+                self._eq_data_keys,
+                params=self._eq.params_dict,
+                transforms=source_transforms,
+                profiles=source_profiles,
+            )
+            eval_data = (
+                source_data
+                if self._use_same_grid
+                else compute_fun(
+                    "desc.equilibrium.equilibrium.Equilibrium",
+                    self._eq_data_keys,
+                    params=self._eq.params_dict,
+                    transforms=eval_transforms,
+                    profiles=eval_profiles,
+                )
+            )
+
+            Bplasma = virtual_casing_biot_savart(
+                eval_data,
+                source_data,
+                interpolator,
+                chunk_size=self._B_plasma_chunk_size,
+            )
+            # need extra factor of B/2 bc we're evaluating on plasma surface
+            Bplasma = Bplasma + eval_data["B"] / 2
+
+            self._constants["source_data"] = source_data
+            self._constants["eval_data"] = eval_data
+            self._constants["Bplasma"] = Bplasma
+
+            # sheet current stuff
+            if self._sheet_current:
+                p = (
+                    "desc.magnetic_fields._current_potential."
+                    "FourierCurrentPotentialField"
+                )
+                sheet_params = {
+                    "R_lmn": self._eq.params_dict["Rb_lmn"],
+                    "Z_lmn": self._eq.params_dict["Zb_lmn"],
+                    "I": self._eq.params_dict["I"],
+                    "G": self._eq.params_dict["G"],
+                    "Phi_mn": self._eq.params_dict["Phi_mn"],
+                }
+                sheet_source_data = compute_fun(
+                    p,
+                    self._sheet_data_keys,
+                    params=sheet_params,
+                    transforms=self._constants["sheet_source_transforms"],
+                    profiles={},
+                )
+                sheet_eval_data = (
+                    sheet_source_data
+                    if self._use_same_grid
+                    else compute_fun(
+                        p,
+                        self._sheet_data_keys,
+                        params=sheet_params,
+                        transforms=self._constants["sheet_eval_transforms"],
+                        profiles={},
+                    )
+                )
+
+                self._constants["sheet_eval_data"] = sheet_eval_data
+                source_data["K_vc"] += sheet_source_data["K"]
+
         timer.stop("Precomputing transforms")
         if verbose > 1:
             timer.disp("Precomputing transforms")
@@ -674,15 +757,19 @@
 
         super().build(use_jit=use_jit, verbose=verbose)
 
-    def compute(self, eq_params, *field_params, constants=None):
+    def compute(self, *params, constants=None):
         """Compute boundary force error.
 
         Parameters
         ----------
-        eq_params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        field_params : dict
-            Dictionary of field parameters, if field is not fixed.
+        params_1 : dict
+            Dictionary of equilibrium degrees of freedom, eg
+            ``Equilibrium.params_dict`` if self._eq_fixed is False, else is the field
+            degrees of freedom
+        params_2 : dict
+            Dictionary of field degrees of freedom,
+            eg ``CoilSet.params_dict``
+            Only required if ``self._field_fixed = False``.
         constants : dict
             Dictionary of constant data, eg transforms, profiles etc. Defaults to
             self.constants. (Deprecated)
@@ -695,64 +782,84 @@
             √g||μ₀𝐊 − 𝐧 × [𝐁]|| in T*m^2
 
         """
-        if field_params == ():  # common case for field_fixed=True
+        if self._eq_fixed:
+            field_params = params
+            eq_params = None
+        elif self._field_fixed:
+            eq_params = params[0]
             field_params = None
+        else:
+            eq_params = params[0]
+            field_params = params[1:]
         constants = self._get_deprecated_constants(constants)
-        source_data = compute_fun(
-            "desc.equilibrium.equilibrium.Equilibrium",
-            self._eq_data_keys,
-            params=eq_params,
-            transforms=constants["source_transforms"],
-            profiles=constants["source_profiles"],
-        )
-        eval_data = (
-            source_data
-            if self._use_same_grid
-            else compute_fun(
+
+        if self._eq_fixed:
+            source_data = constants["source_data"]
+            eval_data = constants["eval_data"]
+            Bplasma = constants["Bplasma"]
+            if self._sheet_current:
+                sheet_eval_data = constants["sheet_eval_data"]
+        else:
+            source_data = compute_fun(
                 "desc.equilibrium.equilibrium.Equilibrium",
                 self._eq_data_keys,
                 params=eq_params,
-                transforms=constants["eval_transforms"],
-                profiles=constants["eval_profiles"],
+                transforms=constants["source_transforms"],
+                profiles=constants["source_profiles"],
             )
-        )
-        if self._sheet_current:
-            p = "desc.magnetic_fields._current_potential.FourierCurrentPotentialField"
-            sheet_params = {
-                "R_lmn": eq_params["Rb_lmn"],
-                "Z_lmn": eq_params["Zb_lmn"],
-                "I": eq_params["I"],
-                "G": eq_params["G"],
-                "Phi_mn": eq_params["Phi_mn"],
-            }
-            sheet_source_data = compute_fun(
-                p,
-                self._sheet_data_keys,
-                params=sheet_params,
-                transforms=constants["sheet_source_transforms"],
-                profiles={},
-            )
-            sheet_eval_data = (
-                sheet_source_data
+            eval_data = (
+                source_data
                 if self._use_same_grid
                 else compute_fun(
+                    "desc.equilibrium.equilibrium.Equilibrium",
+                    self._eq_data_keys,
+                    params=eq_params,
+                    transforms=constants["eval_transforms"],
+                    profiles=constants["eval_profiles"],
+                )
+            )
+
+            Bplasma = virtual_casing_biot_savart(
+                eval_data,
+                source_data,
+                constants["interpolator"],
+                chunk_size=self._B_plasma_chunk_size,
+            )
+            # need extra factor of B/2 bc we're evaluating on plasma surface
+            Bplasma = Bplasma + eval_data["B"] / 2
+
+            if self._sheet_current:
+                p = (
+                    "desc.magnetic_fields._current_potential."
+                    "FourierCurrentPotentialField"
+                )
+                sheet_params = {
+                    "R_lmn": eq_params["Rb_lmn"],
+                    "Z_lmn": eq_params["Zb_lmn"],
+                    "I": eq_params["I"],
+                    "G": eq_params["G"],
+                    "Phi_mn": eq_params["Phi_mn"],
+                }
+                sheet_source_data = compute_fun(
                     p,
                     self._sheet_data_keys,
                     params=sheet_params,
-                    transforms=constants["sheet_eval_transforms"],
+                    transforms=constants["sheet_source_transforms"],
                     profiles={},
                 )
-            )
-            source_data["K_vc"] += sheet_source_data["K"]
+                sheet_eval_data = (
+                    sheet_source_data
+                    if self._use_same_grid
+                    else compute_fun(
+                        p,
+                        self._sheet_data_keys,
+                        params=sheet_params,
+                        transforms=constants["sheet_eval_transforms"],
+                        profiles={},
+                    )
+                )
+                source_data["K_vc"] += sheet_source_data["K"]
 
-        Bplasma = virtual_casing_biot_savart(
-            eval_data,
-            source_data,
-            constants["interpolator"],
-            chunk_size=self._B_plasma_chunk_size,
-        )
-        # need extra factor of B/2 bc we're evaluating on plasma surface
-        Bplasma = Bplasma + eval_data["B"] / 2
         x = jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T
         # can always pass in field params. If they're None, it just uses the
         # defaults for the given field.

tests/test_objective_funs.py

@@ -610,6 +610,43 @@ def test_boundary_error_biest(self):
         # next n should be B^2 errors
         np.testing.assert_allclose(f[n : 2 * n], 0, atol=5e-2)
 
+    @pytest.mark.unit
+    def test_boundary_error_things_fixed(self):
+        """Test that BoundaryError runs correctly for eq_fixed/field_fixed combos."""
+        eq = Equilibrium(L=3, M=3, N=3, Psi=np.pi)
+        eq.solve()
+        coil = FourierXYZCoil(5e5)
+        coilset = CoilSet.linspaced_angular(coil, n=100, check_intersection=False)
+        field=[coilset, ToroidalMagneticField(B0=0, R0=1)]
+
+        def test(eq_fixed=False, field_fixed=False):
+            obj = BoundaryError(
+                eq,
+                field,
+                eq_fixed=eq_fixed,
+                field_fixed=field_fixed,
+            )
+            obj.build()
+            if eq_fixed:
+                f = obj.compute(*[ff.params_dict for ff in field])
+            elif field_fixed:
+                f = obj.compute(eq.params_dict)
+            else:
+                f = obj.compute(eq.params_dict, *[ff.params_dict for ff in field])
+            f=obj.compute_scaled_error(*obj.xs())
+            n = len(f) // 2
+            # first n should be B*n errors
+            np.testing.assert_allclose(f[:n], 0, atol=1e-4)
+            # next n should be B^2 errors
+            np.testing.assert_allclose(f[n : 2 * n], 0, atol=5e-2)
+
+        test(eq_fixed=False, field_fixed=False)
+        test(eq_fixed=True)
+        test(field_fixed=True)
+
+        with pytest.raises(ValueError, match="At least one"):
+            BoundaryError(eq, field, eq_fixed=True, field_fixed=True)
+
     @pytest.mark.unit
     def test_boundary_error_vacuum(self):
         """Test calculation of vacuum boundary error."""