update test with new UI

PlasmaControl · Jun 17, 2024 · 90296ea · 90296ea
1 parent 30aeea4
commit 90296ea
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 6 deletions.
diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
@@ -18,7 +18,12 @@
     RadialForceBalance,
 )
 from ._free_boundary import BoundaryError, VacuumBoundaryError
-from ._generic import GenericObjective, LinearObjectiveFromUser, ObjectiveFromUser
+from ._generic import (
+    GenericObjective,
+    LinearObjectiveFromUser,
+    ObjectiveFromUser,
+    _ExternalObjective,
+)
 from ._geometry import (
     AspectRatio,
     BScaleLength,

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -31,7 +31,6 @@
     CoilLength,
     CoilTorsion,
     CurrentDensity,
-    ExternalObjective,
     FixBoundaryR,
     FixBoundaryZ,
     FixCurrent,
@@ -56,6 +55,7 @@
     QuasisymmetryTwoTerm,
     VacuumBoundaryError,
     Volume,
+    _ExternalObjective,
     get_fixed_boundary_constraints,
     get_NAE_constraints,
 )
@@ -1319,6 +1319,7 @@ def test_optimize_with_fourier_planar_coil():
 
 
 @pytest.mark.unit
+@pytest.mark.slow
 def test_external_vs_generic_objectives(tmpdir_factory):
     """Test ExternalObjective compared to GenericObjective."""
     target = np.array([6.2e-3, 1.1e-1, 6.5e-3, 0])  # values at p_l = [2e2, -2e2]
@@ -1333,6 +1334,36 @@ def data_from_vmec(eq, path=""):
         file.close()
         return np.atleast_1d([betatot, betapol, betator, presf1])
 
+    class TestExternalObjective(_ExternalObjective):
+
+        def __init__(
+            self,
+            eq,
+            target=None,
+            bounds=None,
+            weight=1,
+            normalize=False,
+            normalize_target=False,
+            loss_function=None,
+            path="",
+            name="external",
+        ):
+            super().__init__(
+                eq=eq,
+                fun=data_from_vmec,
+                dim_f=4,
+                target=target,
+                bounds=bounds,
+                weight=weight,
+                normalize=normalize,
+                normalize_target=normalize_target,
+                loss_function=loss_function,
+                fd_step=1e-4,
+                vectorized=False,
+                name=name,
+                path=path,
+            )
+
     eq0 = get("SOLOVEV")
     optimizer = Optimizer("lsq-exact")
 
@@ -1369,7 +1400,7 @@ def data_from_vmec(eq, path=""):
     dir = tmpdir_factory.mktemp("results")
     path = dir.join("wout_result.nc")
     objective = ObjectiveFunction(
-        ExternalObjective(eq=eq0, fun=data_from_vmec, dim_f=4, target=target, path=path)
+        TestExternalObjective(eq=eq0, target=target, path=path)
     )
     constraints = FixParameters(
         eq0,

diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py
@@ -39,7 +39,6 @@
     CoilTorsion,
     Elongation,
     Energy,
-    ExternalObjective,
     ForceBalance,
     ForceBalanceAnisotropic,
     GenericObjective,
@@ -1755,7 +1754,6 @@ class TestComputeScalarResolution:
         # need to avoid blowup near the axis
         MercierStability,
         # don't test these since they depend on what user wants
-        ExternalObjective,
         LinearObjectiveFromUser,
         ObjectiveFromUser,
     ]
@@ -2077,7 +2075,6 @@ class TestObjectiveNaNGrad:
         QuadraticFlux,
         ToroidalFlux,
         # we don't test these since they depend too much on what exactly the user wants
-        ExternalObjective,
         GenericObjective,
         LinearObjectiveFromUser,
         ObjectiveFromUser,