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__(
@@ -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")
@@ -680,7 +763,8 @@
         Parameters
         ----------
         eq_params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
+            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict,
+            if equilibrium is not fixed
         field_params : dict
             Dictionary of field parameters, if field is not fixed.
         constants : dict
@@ -697,62 +781,78 @@
         """
         if field_params == ():  # common case for field_fixed=True
             field_params = None
+        if eq_params == ():
+            eq_params = None
+
         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.