From 90296ea83fd4a8fbefba8badfd175fd8c4157b88 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 17 Jun 2024 18:20:49 -0400
Subject: [PATCH] update test with new UI

---
 desc/objectives/__init__.py  |  7 ++++++-
 tests/test_examples.py       | 35 +++++++++++++++++++++++++++++++++--
 tests/test_objective_funs.py |  3 ---
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
index 29d875b7ed..05a44a63d4 100644
--- 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
index 5ee219d2fb..b523250282 100644
--- 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
index 1fcbc95291..9dd4ec5fee 100644
--- 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,