Merge branch 'master' into rc/interpolate

PlasmaControl · Nov 14, 2023 · 6b18da4 · 6b18da4
2 parents afa2ae1 + 3eb55a0
commit 6b18da4
Show file tree

Hide file tree

Showing 19 changed files with 108 additions and 220 deletions.
diff --git a/desc/compat.py b/desc/compat.py
@@ -204,4 +204,8 @@ def rescale(eq, L=("R0", None), B=("B0", None), verbose=0):
     if eq.current is not None:
         eq.c_l *= cL * cB
 
+    # boundary & axis
+    eq.axis = eq.get_axis()
+    eq.surface = eq.get_surface_at(rho=1)
+
     return eq
diff --git a/desc/continuation.py b/desc/continuation.py
@@ -98,11 +98,7 @@ def _solve_axisym(
             deltas = get_deltas({"surface": surf_i}, {"surface": surf_i2})
             surf_i = surf_i2
 
-        constraints_i = get_fixed_boundary_constraints(
-            eq=eqi,
-            iota=eqi.iota,
-            kinetic=eqi.electron_temperature,
-        )
+        constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
 
         if verbose:
@@ -227,11 +223,7 @@ def _add_pressure(
         deltas["p_l"] *= pres_step
         pres_ratio += pres_step
 
-        constraints_i = get_fixed_boundary_constraints(
-            eq=eqi,
-            iota=eqi.iota,
-            kinetic=eqi.electron_temperature,
-        )
+        constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
 
         if verbose:
@@ -360,11 +352,7 @@ def _add_shaping(
             deltas["Zb_lmn"] *= bdry_step
         bdry_ratio += bdry_step
 
-        constraints_i = get_fixed_boundary_constraints(
-            eq=eqi,
-            iota=eqi.iota,
-            kinetic=eqi.electron_temperature,
-        )
+        constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
 
         if verbose:
@@ -662,11 +650,7 @@ def solve_continuation(  # noqa: C901
     if not isinstance(optimizer, Optimizer):
         optimizer = Optimizer(optimizer)
     objective_i = get_equilibrium_objective(eq=eqfam[0], mode=objective)
-    constraints_i = get_fixed_boundary_constraints(
-        eq=eqfam[0],
-        iota=objective != "vacuum" and eqfam[0].iota is not None,
-        kinetic=eqfam[0].electron_temperature is not None,
-    )
+    constraints_i = get_fixed_boundary_constraints(eq=eqfam[0])
 
     ii = 0
     nn = len(eqfam)
@@ -713,11 +697,7 @@ def solve_continuation(  # noqa: C901
             # TODO: pass Jx if available
             eqp = eqfam[ii - 1].copy()
             objective_i = get_equilibrium_objective(eq=eqp, mode=objective)
-            constraints_i = get_fixed_boundary_constraints(
-                eq=eqp,
-                iota=eqp.iota,
-                kinetic=eqp.electron_temperature,
-            )
+            constraints_i = get_fixed_boundary_constraints(eq=eqp)
             eqp.change_resolution(**eqi.resolution)
             eqp.perturb(
                 objective=objective_i,
@@ -737,11 +717,7 @@ def solve_continuation(  # noqa: C901
         if not stop:
             # TODO: add ability to rebind objectives
             objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
-            constraints_i = get_fixed_boundary_constraints(
-                eq=eqi,
-                iota=eqi.iota,
-                kinetic=eqi.electron_temperature,
-            )
+            constraints_i = get_fixed_boundary_constraints(eq=eqi)
             eqi.solve(
                 optimizer=optimizer,
                 objective=objective_i,

diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py
@@ -1704,11 +1704,7 @@ def solve(
 
         """
         if constraints is None:
-            constraints = get_fixed_boundary_constraints(
-                eq=self,
-                iota=objective != "vacuum" and self.iota is not None,
-                kinetic=self.electron_temperature is not None,
-            )
+            constraints = get_fixed_boundary_constraints(eq=self)
         if not isinstance(objective, ObjectiveFunction):
             objective = get_equilibrium_objective(eq=self, mode=objective)
         if not isinstance(optimizer, Optimizer):
@@ -1809,11 +1805,7 @@ def optimize(
         if not isinstance(optimizer, Optimizer):
             optimizer = Optimizer(optimizer)
         if constraints is None:
-            constraints = get_fixed_boundary_constraints(
-                eq=self,
-                iota=self.iota is not None,
-                kinetic=self.electron_temperature is not None,
-            )
+            constraints = get_fixed_boundary_constraints(eq=self)
             constraints = (ForceBalance(eq=self), *constraints)
         if not isinstance(constraints, (list, tuple)):
             constraints = tuple([constraints])
@@ -2047,17 +2039,9 @@ def perturb(
             objective = get_equilibrium_objective(eq=self)
         if constraints is None:
             if "Ra_n" in deltas or "Za_n" in deltas:
-                constraints = get_fixed_axis_constraints(
-                    eq=self,
-                    iota=self.iota is not None,
-                    kinetic=self.electron_temperature is not None,
-                )
+                constraints = get_fixed_axis_constraints(eq=self)
             else:
-                constraints = get_fixed_boundary_constraints(
-                    eq=self,
-                    iota=self.iota is not None,
-                    kinetic=self.electron_temperature is not None,
-                )
+                constraints = get_fixed_boundary_constraints(eq=self)
 
         eq = perturb(
             self,

diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py
@@ -15,13 +15,13 @@ class ForceBalance(_Objective):
 
     Given force densities:
 
-    Fᵨ = √g (B^ζ J^θ - B^θ J^ζ) - ∇ p
+    Fᵨ = √g (J^θ B^ζ - J^ζ B^θ) - ∇ p
 
     Fₕₑₗᵢ √g J^ρ
 
     and helical basis vector:
 
-    𝐞ʰᵉˡⁱ = −B^ζ ∇ θ + B^θ ∇ ζ
+    𝐞ʰᵉˡⁱ = B^ζ ∇ θ - B^θ ∇ ζ
 
     Minimizes the magnitude of the forces:
 

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
@@ -30,6 +30,17 @@
 from .nae_utils import calc_zeroth_order_lambda, make_RZ_cons_1st_order
 from .objective_funs import ObjectiveFunction
 
+_PROFILE_CONSTRAINTS = {
+    "pressure": FixPressure,
+    "iota": FixIota,
+    "current": FixCurrent,
+    "electron_density": FixElectronDensity,
+    "electron_temperature": FixElectronTemperature,
+    "ion_temperature": FixIonTemperature,
+    "atomic_number": FixAtomicNumber,
+    "anisotropy": FixAnisotropy,
+}
+
 
 def get_equilibrium_objective(eq, mode="force", normalize=True):
     """Get the objective function for a typical force balance equilibrium problem.
@@ -70,23 +81,15 @@ def get_equilibrium_objective(eq, mode="force", normalize=True):
     return ObjectiveFunction(objectives)
 
 
-def get_fixed_axis_constraints(
-    eq, profiles=True, iota=True, kinetic=False, anisotropy=False, normalize=True
-):
+def get_fixed_axis_constraints(eq, profiles=True, normalize=True):
     """Get the constraints necessary for a fixed-axis equilibrium problem.
 
     Parameters
     ----------
     eq : Equilibrium
-        Equilibrium being constrained.
+        Equilibrium to constrain.
     profiles : bool
-        Whether to also return constraints to fix input profiles.
-    iota : bool
-        Whether to add FixIota or FixCurrent as a constraint.
-    kinetic : bool
-        Whether to add constraints to fix kinetic profiles or pressure
-    anisotropy : bool
-        Whether to add constraint to fix anisotropic pressure.
+        If True, also include constraints to fix all profiles assigned to equilibrium.
     normalize : bool
         Whether to apply constraints in normalized units.
 
@@ -102,58 +105,24 @@ def get_fixed_axis_constraints(
         FixPsi(eq=eq, normalize=normalize, normalize_target=normalize),
     )
     if profiles:
-        if kinetic:
-            constraints += (
-                FixElectronDensity(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixElectronTemperature(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixIonTemperature(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixAtomicNumber(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-        else:
-            constraints += (
-                FixPressure(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-            if anisotropy:
+        for name, con in _PROFILE_CONSTRAINTS.items():
+            if getattr(eq, name) is not None:
                 constraints += (
-                    FixAnisotropy(
-                        eq=eq, normalize=normalize, normalize_target=normalize
-                    ),
+                    con(eq=eq, normalize=normalize, normalize_target=normalize),
                 )
 
-        if iota:
-            constraints += (
-                FixIota(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-        else:
-            constraints += (
-                FixCurrent(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
     return constraints
 
 
-def get_fixed_boundary_constraints(
-    eq, profiles=True, iota=True, kinetic=False, anisotropy=False, normalize=True
-):
+def get_fixed_boundary_constraints(eq, profiles=True, normalize=True):
     """Get the constraints necessary for a typical fixed-boundary equilibrium problem.
 
     Parameters
     ----------
     eq : Equilibrium
-        Equilibrium to constraint.
+        Equilibrium to constrain.
     profiles : bool
-        Whether to also return constraints to fix input profiles.
-    iota : bool
-        Whether to add FixIota or FixCurrent as a constraint.
-    kinetic : bool
-        Whether to also fix kinetic profiles.
-    anisotropy : bool
-        Whether to add constraint to fix anisotropic pressure.
+        If True, also include constraints to fix all profiles assigned to equilibrium.
     normalize : bool
         Whether to apply constraints in normalized units.
 
@@ -169,37 +138,12 @@ def get_fixed_boundary_constraints(
         FixPsi(eq=eq, normalize=normalize, normalize_target=normalize),
     )
     if profiles:
-        if kinetic:
-            constraints += (
-                FixElectronDensity(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixElectronTemperature(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixIonTemperature(
-                    eq=eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixAtomicNumber(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-        else:
-            constraints += (
-                FixPressure(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-            if anisotropy:
+        for name, con in _PROFILE_CONSTRAINTS.items():
+            if getattr(eq, name) is not None:
                 constraints += (
-                    FixAnisotropy(
-                        eq=eq, normalize=normalize, normalize_target=normalize
-                    ),
+                    con(eq=eq, normalize=normalize, normalize_target=normalize),
                 )
-        if iota:
-            constraints += (
-                FixIota(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
-        else:
-            constraints += (
-                FixCurrent(eq=eq, normalize=normalize, normalize_target=normalize),
-            )
+
     return constraints
 
 
@@ -208,9 +152,6 @@ def get_NAE_constraints(
     qsc_eq,
     order=1,
     profiles=True,
-    iota=False,
-    kinetic=False,
-    anisotropy=False,
     normalize=True,
     N=None,
     fix_lambda=False,
@@ -227,13 +168,7 @@ def get_NAE_constraints(
     order : int
         order (in rho) of near-axis behavior to constrain
     profiles : bool
-        Whether to also return constraints to fix input profiles.
-    iota : bool
-        Whether to add `FixIota` or `FixCurrent` as a constraint.
-    kinetic : bool
-        Whether to also fix kinetic profiles.
-    anisotropy : bool
-        Whether to add constraint to fix anisotropic pressure.
+        If True, also include constraints to fix all profiles assigned to equilibrium.
     normalize : bool
         Whether to apply constraints in normalized units.
     N : int
@@ -256,43 +191,14 @@ def get_NAE_constraints(
         FixAxisZ(eq=desc_eq, normalize=normalize, normalize_target=normalize),
         FixPsi(eq=desc_eq, normalize=normalize, normalize_target=normalize),
     )
+
     if profiles:
-        if kinetic:
-            constraints += (
-                FixElectronDensity(
-                    eq=desc_eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixElectronTemperature(
-                    eq=desc_eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixIonTemperature(
-                    eq=desc_eq, normalize=normalize, normalize_target=normalize
-                ),
-                FixAtomicNumber(
-                    eq=desc_eq, normalize=normalize, normalize_target=normalize
-                ),
-            )
-        else:
-            constraints += (
-                FixPressure(
-                    eq=desc_eq, normalize=normalize, normalize_target=normalize
-                ),
-            )
-            if anisotropy:
+        for name, con in _PROFILE_CONSTRAINTS.items():
+            if getattr(desc_eq, name) is not None:
                 constraints += (
-                    FixAnisotropy(
-                        eq=desc_eq, normalize=normalize, normalize_target=normalize
-                    ),
+                    con(eq=desc_eq, normalize=normalize, normalize_target=normalize),
                 )
 
-        if iota:
-            constraints += (
-                FixIota(eq=desc_eq, normalize=normalize, normalize_target=normalize),
-            )
-        else:
-            constraints += (
-                FixCurrent(eq=desc_eq, normalize=normalize, normalize_target=normalize),
-            )
     if fix_lambda or (fix_lambda >= 0 and type(fix_lambda) is int):
         L_axis_constraints, _, _ = calc_zeroth_order_lambda(
             qsc=qsc_eq, desc_eq=desc_eq, N=N

diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py
@@ -157,8 +157,8 @@ def recover(x_reduced):
         np.testing.assert_allclose(
             y1,
             y2,
-            atol=1e-14,
-            rtol=1e-14,
+            atol=2e-14,
+            rtol=5e-14,
             err_msg="Incompatible constraints detected, cannot satisfy "
             + f"constraint {con}",
         )

diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py
@@ -473,11 +473,7 @@ def build(self, eq=None, use_jit=None, verbose=1):  # noqa: C901
         timer.start("Proximal projection build")
 
         self._eq = eq
-        self._linear_constraints = get_fixed_boundary_constraints(
-            eq=eq,
-            iota=self._eq.iota,
-            kinetic=self._eq.electron_temperature,
-        )
+        self._linear_constraints = get_fixed_boundary_constraints(eq=eq)
         self._linear_constraints = maybe_add_self_consistency(
             self._eq, self._linear_constraints
         )