fix zeta phi problem causing nans

PlasmaControl · Oct 20, 2024 · b7675e6 · b7675e6
1 parent c40eaf1
commit b7675e6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py
@@ -103,7 +103,7 @@ def test_map_coordinates_derivative():
     def foo(params, in_coords):
         out = eq.map_coordinates(
             in_coords,
-            ("rho", "alpha", "zeta"),  # for this test, zeta==phi
+            ("rho", "alpha", "phi"),
             ("rho", "theta_PEST", "zeta"),
             np.array([rho, theta, zeta]).T,
             params,
@@ -123,9 +123,9 @@ def foo(params, in_coords):
     # this time _map_clebsch_coordinates is called inside map_coordinates
     @jax.jit
     def foo(params, in_coords):
-        out, _ = eq.map_coordinates(
+        out, (_, _) = eq.map_coordinates(
             in_coords,
-            ("rho", "alpha", "zeta"),  # for this test, zeta==phi
+            ("rho", "alpha", "phi"),
             ("rho", "theta", "zeta"),
             np.array([rho, theta, zeta]).T,
             params,