fix self consistency stuff better

PlasmaControl · Dec 14, 2024 · 59fb333 · 59fb333
1 parent 654d3b3
commit 59fb333
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 31 deletions.
diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py
@@ -2275,22 +2275,11 @@ def solve(
             options=options,
             copy=copy,
         )
-        from desc.objectives import (
-            AxisRSelfConsistency,
-            BoundaryRSelfConsistency,
-            SectionRSelfConsistency,
-        )
+        from desc.objectives import SectionRSelfConsistency
         from desc.objectives.getters import maybe_add_self_consistency
 
         constraints = maybe_add_self_consistency(self, constraints)
         # TODO: make this more general and probably put it somewhere else
-        if not (
-            is_any_instance(constraints, BoundaryRSelfConsistency)
-            and is_any_instance(constraints, AxisRSelfConsistency)
-        ):
-            things[0].surface = things[0].get_surface_at(rho=1.0)
-        if not is_any_instance(constraints, AxisRSelfConsistency):
-            things[0].axis = things[0].get_axis()
         if not is_any_instance(constraints, SectionRSelfConsistency):
             things[0].xsection = things[0].get_surface_at(zeta=0)
 
@@ -2376,22 +2365,11 @@ def optimize(
             options=options,
             copy=copy,
         )
-        from desc.objectives import (
-            AxisRSelfConsistency,
-            BoundaryRSelfConsistency,
-            SectionRSelfConsistency,
-        )
+        from desc.objectives import SectionRSelfConsistency
         from desc.objectives.getters import maybe_add_self_consistency
 
         constraints = maybe_add_self_consistency(self, constraints)
         # TODO: make this more general and probably put it somewhere else
-        if not (
-            is_any_instance(constraints, BoundaryRSelfConsistency)
-            and is_any_instance(constraints, AxisRSelfConsistency)
-        ):
-            things[0].surface = things[0].get_surface_at(rho=1.0)
-        if not is_any_instance(constraints, AxisRSelfConsistency):
-            things[0].axis = things[0].get_axis()
         if not is_any_instance(constraints, SectionRSelfConsistency):
             things[0].xsection = things[0].get_surface_at(zeta=0)
 

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
@@ -370,6 +370,9 @@ def maybe_add_self_consistency(thing, constraints):  # noqa: C901
     # For example, if we are fixing the boundary R, then we add a dummy fixed
     # parameter for the section R. This is not ideal, but it is a workaround.
     # This helps with the factorization of the linear constraints matrix.
+    # Extra conditions are added to avoid edge cases like calling this function
+    # twice. The first turn will change the constraints, and the second turn
+    # should account for that.
     if {"R_lmn", "Rb_lmn"} <= params and not is_any_instance(
         constraints, BoundaryRSelfConsistency
     ):
@@ -393,6 +396,7 @@ def maybe_add_self_consistency(thing, constraints):  # noqa: C901
     if (
         {"R_lmn", "Rp_lmn"} <= params
         and not is_any_instance(constraints, SectionRSelfConsistency)
+        and not is_any_instance(constraints, FixBoundaryR)
         and is_any_instance(constraints, FixSectionR)
     ):
         constraints += (SectionRSelfConsistency(eq=thing),)
@@ -403,6 +407,7 @@ def maybe_add_self_consistency(thing, constraints):  # noqa: C901
     if (
         {"Z_lmn", "Zp_lmn"} <= params
         and not is_any_instance(constraints, SectionZSelfConsistency)
+        and not is_any_instance(constraints, FixBoundaryZ)
         and is_any_instance(constraints, FixSectionZ)
     ):
         constraints += (SectionZSelfConsistency(eq=thing),)
@@ -413,6 +418,7 @@ def maybe_add_self_consistency(thing, constraints):  # noqa: C901
     if (
         {"L_lmn", "Lp_lmn"} <= params
         and not is_any_instance(constraints, SectionLambdaSelfConsistency)
+        and not is_any_instance(constraints, FixBoundaryR)
         and is_any_instance(constraints, FixSectionLambda)
     ):
         constraints += (SectionLambdaSelfConsistency(eq=thing),)

diff --git a/desc/perturbations.py b/desc/perturbations.py
@@ -446,13 +446,6 @@ def perturb(  # noqa: C901
         timer.disp("Total perturbation")
 
     # TODO: make this more general and probably put it somewhere else
-    if not (
-        is_any_instance(constraints, BoundaryRSelfConsistency)
-        and is_any_instance(constraints, AxisRSelfConsistency)
-    ):
-        eq_new.surface = eq_new.get_surface_at(rho=1.0)
-    if not is_any_instance(constraints, AxisRSelfConsistency):
-        eq_new.axis = eq_new.get_axis()
     if not is_any_instance(constraints, SectionRSelfConsistency):
         eq_new.xsection = eq_new.get_surface_at(zeta=0)