fixing tests

PlasmaControl · Jun 28, 2024 · 989ad9f · 989ad9f
1 parent 9c60daa
commit 989ad9f
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 11 deletions.
diff --git a/desc/coils.py b/desc/coils.py
@@ -224,7 +224,7 @@ def compute_magnetic_field(
             current = params.pop("current", self.current)
 
         data = self.compute(
-            ["x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz"
+            ["phi", "x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz"
         )
         B = biot_savart_quad(
             coords, data["x"], data["x_s"] * data["ds"][:, None], current

diff --git a/desc/compute/utils.py b/desc/compute/utils.py
@@ -9,6 +9,7 @@
 
 from desc.backend import cond, fori_loop, jnp, put
 from desc.grid import ConcentricGrid, Grid, LinearGrid
+from desc.utils import errorif
 
 from .data_index import allowed_kwargs, data_index
 
@@ -105,7 +106,12 @@ def compute(parameterization, names, params, transforms, profiles, data=None, **
             if name == "x":
                 data[name] = rpz2xyz(data[name])
             else:
-                # TODO: check if phi in data?
+                errorif(
+                    "phi" not in data.keys(),
+                    ValueError,
+                    "'phi' must be included in the compute data "
+                    + "to convert to 'xyz' basis.",
+                )
                 data[name] = rpz2xyz_vec(data[name], phi=data["phi"])
 
     return data

diff --git a/tests/test_curves.py b/tests/test_curves.py
@@ -55,19 +55,15 @@ def test_torsion(self):
     def test_frenet(self):
         """Test frenet-serret frame of circular curve."""
         c = FourierRZCurve()
-        data = c.compute(
-            ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", grid=0
-        )
+        data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0)
         T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"]
         np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12)
         np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12)
         np.testing.assert_allclose(B, np.array([[0, 0, 1]]), atol=1e-12)
         c.rotate(angle=np.pi)
         c.flip([0, 1, 0])
         c.translate([1, 1, 1])
-        data = c.compute(
-            ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0
-        )
+        data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0)
         T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"]
         np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12)
         np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12)
@@ -470,9 +466,9 @@ def test_rotation(self):
         datax = cx.compute("x", grid=20, basis="xyz")
         datay = cy.compute("x", grid=20, basis="xyz")
         dataz = cz.compute("x", grid=20, basis="xyz")
-        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=1e-16)  # only in Y-Z plane
-        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=1e-16)  # only in X-Z plane
-        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=1e-16)  # only in X-Y plane
+        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=2e-16)  # only in Y-Z plane
+        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=2e-16)  # only in X-Z plane
+        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=2e-16)  # only in X-Y plane
 
     @pytest.mark.unit
     def test_length(self):