From 5ba60045f372800864c96331e8ccfb5308a746e0 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 13 May 2024 15:36:24 -0400
Subject: [PATCH 1/9] FourierPlanarCurve basis option

---
 desc/coils.py          | 15 +++++----
 desc/compute/_curve.py | 52 ++++++++++++++++++++++-------
 desc/geometry/core.py  |  6 ++--
 desc/geometry/curve.py | 74 +++++++++++++++++++++++++++++++++++-------
 tests/test_curves.py   | 24 ++++++++++++++
 5 files changed, 139 insertions(+), 32 deletions(-)

diff --git a/desc/coils.py b/desc/coils.py
index 8fe01baf41..743b2160c7 100644
--- a/desc/coils.py
+++ b/desc/coils.py
@@ -458,17 +458,19 @@ class FourierPlanarCoil(_Coil, FourierPlanarCurve):
     Parameters
     ----------
     current : float
-        current through the coil, in Amperes
+        Current through the coil, in Amperes.
     center : array-like, shape(3,)
-        x,y,z coordinates of center of coil
+        Coordinates of center of curve, in system determined by basis.
     normal : array-like, shape(3,)
-        x,y,z components of normal vector to planar surface
+        Components of normal vector to planar surface, in system determined by basis.
     r_n : array-like
-        fourier coefficients for radius from center as function of polar angle
+        Fourier coefficients for radius from center as function of polar angle
     modes : array-like
         mode numbers associated with r_n
+    basis : {'xyz', 'rpz'}
+        Coordinate system for center and normal vectors. Default = 'xyz'.
     name : str
-        name for this coil
+        Name for this coil.
 
     Examples
     --------
@@ -514,9 +516,10 @@ def __init__(
         normal=[0, 1, 0],
         r_n=2,
         modes=None,
+        basis="xyz",
         name="",
     ):
-        super().__init__(current, center, normal, r_n, modes, name)
+        super().__init__(current, center, normal, r_n, modes, basis, name)
 
 
 class SplineXYZCoil(_Coil, SplineXYZCurve):
diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py
index 85d03ca31f..7271095b1e 100644
--- a/desc/compute/_curve.py
+++ b/desc/compute/_curve.py
@@ -176,8 +176,15 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs):
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+        center = rpz2xyz(params["center"])
+        normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
+    else:
+        center = params["center"]
+        normal = params["normal"]
     # create planar curve at Z==0
     r = transforms["r"].transform(params["r_n"], dz=0)
     Z = jnp.zeros_like(r)
@@ -186,10 +193,10 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.array([X, Y, Z]).T
     # rotate into place
     Zaxis = jnp.array([0.0, 0.0, 1.0])  # 2D curve in X-Y plane has normal = +Z axis
-    axis = cross(Zaxis, params["normal"])
-    angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"])))
+    axis = cross(Zaxis, normal)
+    angle = jnp.arccos(dot(Zaxis, safenormalize(normal)))
     A = rotation_matrix(axis=axis, angle=angle)
-    coords = jnp.matmul(coords, A.T) + params["center"]
+    coords = jnp.matmul(coords, A.T) + center
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz(coords)
@@ -213,8 +220,15 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+        center = rpz2xyz(params["center"])
+        normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
+    else:
+        center = params["center"]
+        normal = params["normal"]
     r = transforms["r"].transform(params["r_n"], dz=0)
     dr = transforms["r"].transform(params["r_n"], dz=1)
     dX = dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"])
@@ -223,8 +237,8 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.array([dX, dY, dZ]).T
     # rotate into place
     Zaxis = jnp.array([0.0, 0.0, 1.0])  # 2D curve in X-Y plane has normal = +Z axis
-    axis = cross(Zaxis, params["normal"])
-    angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"])))
+    axis = cross(Zaxis, normal)
+    angle = jnp.arccos(dot(Zaxis, safenormalize(normal)))
     A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
@@ -233,7 +247,7 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
         Y = r * jnp.sin(data["s"])
         Z = jnp.zeros_like(X)
         xyzcoords = jnp.array([X, Y, Z]).T
-        xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"]
+        xyzcoords = jnp.matmul(xyzcoords, A.T) + center
         xyzcoords = (
             jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
         )
@@ -259,8 +273,15 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+        center = rpz2xyz(params["center"])
+        normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
+    else:
+        center = params["center"]
+        normal = params["normal"]
     r = transforms["r"].transform(params["r_n"], dz=0)
     dr = transforms["r"].transform(params["r_n"], dz=1)
     d2r = transforms["r"].transform(params["r_n"], dz=2)
@@ -274,8 +295,8 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.array([d2X, d2Y, d2Z]).T
     # rotate into place
     Zaxis = jnp.array([0.0, 0.0, 1.0])  # 2D curve in X-Y plane has normal = +Z axis
-    axis = cross(Zaxis, params["normal"])
-    angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"])))
+    axis = cross(Zaxis, normal)
+    angle = jnp.arccos(dot(Zaxis, safenormalize(normal)))
     A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
@@ -284,7 +305,7 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
         Y = r * jnp.sin(data["s"])
         Z = jnp.zeros_like(X)
         xyzcoords = jnp.array([X, Y, Z]).T
-        xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"]
+        xyzcoords = jnp.matmul(xyzcoords, A.T) + center
         xyzcoords = (
             jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
         )
@@ -310,8 +331,15 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+        center = rpz2xyz(params["center"])
+        normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
+    else:
+        center = params["center"]
+        normal = params["normal"]
     r = transforms["r"].transform(params["r_n"], dz=0)
     dr = transforms["r"].transform(params["r_n"], dz=1)
     d2r = transforms["r"].transform(params["r_n"], dz=2)
@@ -332,8 +360,8 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.array([d3X, d3Y, d3Z]).T
     # rotate into place
     Zaxis = jnp.array([0.0, 0.0, 1.0])  # 2D curve in X-Y plane has normal = +Z axis
-    axis = cross(Zaxis, params["normal"])
-    angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"])))
+    axis = cross(Zaxis, normal)
+    angle = jnp.arccos(dot(Zaxis, safenormalize(normal)))
     A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
@@ -342,7 +370,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
         Y = r * jnp.sin(data["s"])
         Z = jnp.zeros_like(X)
         xyzcoords = jnp.array([X, Y, Z]).T
-        xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"]
+        xyzcoords = jnp.matmul(xyzcoords, A.T) + center
         xyzcoords = (
             jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
         )
diff --git a/desc/geometry/core.py b/desc/geometry/core.py
index dfc3b53ba1..a637ff529e 100644
--- a/desc/geometry/core.py
+++ b/desc/geometry/core.py
@@ -178,17 +178,17 @@ def compute(
         return data
 
     def translate(self, displacement=[0, 0, 0]):
-        """Translate the curve by a rigid displacement in X, Y, Z."""
+        """Translate the curve by a rigid displacement in X,Y,Z coordinates."""
         self.shift = self.shift + jnp.asarray(displacement)
 
     def rotate(self, axis=[0, 0, 1], angle=0):
-        """Rotate the curve by a fixed angle about axis in X, Y, Z coordinates."""
+        """Rotate the curve by a fixed angle about axis in X,Y,Z coordinates."""
         R = rotation_matrix(axis=axis, angle=angle)
         self.rotmat = (R @ self.rotmat.reshape(3, 3)).flatten()
         self.shift = self.shift @ R.T
 
     def flip(self, normal=[0, 0, 1]):
-        """Flip the curve about the plane with specified normal."""
+        """Flip the curve about the plane with specified normal in X,Y,Z coordinates."""
         F = reflection_matrix(normal)
         self.rotmat = (F @ self.rotmat.reshape(3, 3)).flatten()
         self.shift = self.shift @ F.T
diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py
index 0c4edbb90d..3f670656e7 100644
--- a/desc/geometry/curve.py
+++ b/desc/geometry/curve.py
@@ -556,24 +556,21 @@ class FourierPlanarCurve(Curve):
     Parameters
     ----------
     center : array-like, shape(3,)
-        x,y,z coordinates of center of curve
+        Coordinates of center of curve, in system determined by basis.
     normal : array-like, shape(3,)
-        x,y,z components of normal vector to planar surface
+        Components of normal vector to planar surface, in system determined by basis.
     r_n : array-like
         Fourier coefficients for radius from center as function of polar angle
     modes : array-like
         mode numbers associated with r_n
+    basis : {'xyz', 'rpz'}
+        Coordinate system for center and normal vectors. Default = 'xyz'.
     name : str
-        name for this curve
+        Name for this curve.
 
     """
 
-    _io_attrs_ = Curve._io_attrs_ + [
-        "_r_n",
-        "_center",
-        "_normal",
-        "_r_basis",
-    ]
+    _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "_basis"]
 
     # Reference frame is centered at the origin with normal in the +Z direction.
     # The curve is computed in this frame and then shifted/rotated to the correct frame.
@@ -583,6 +580,7 @@ def __init__(
         normal=[0, 1, 0],
         r_n=2,
         modes=None,
+        basis="xyz",
         name="",
     ):
         super().__init__(name)
@@ -593,6 +591,7 @@ def __init__(
             modes = np.asarray(modes)
         assert issubclass(modes.dtype.type, np.integer)
         assert r_n.size == modes.size, "r_n size and modes must be the same size"
+        assert basis.lower() in ["xyz", "rpz"]
 
         N = np.max(abs(modes))
         self._r_basis = FourierSeries(N, NFP=1, sym=False)
@@ -600,6 +599,7 @@ def __init__(
 
         self.normal = normal
         self.center = center
+        self._basis = basis
 
     @property
     def r_basis(self):
@@ -632,7 +632,9 @@ def center(self, new):
             self._center = np.asarray(new)
         else:
             raise ValueError(
-                "center should be a 3 element vector [cx, cy, cz], got {}".format(new)
+                "center should be a 3 element vector in "
+                + self._basis
+                + " coordinates, got {}".format(new)
             )
 
     @optimizable_parameter
@@ -647,7 +649,9 @@ def normal(self, new):
             self._normal = np.asarray(new) / np.linalg.norm(new)
         else:
             raise ValueError(
-                "normal should be a 3 element vector [nx, ny, nz], got {}".format(new)
+                "normal should be a 3 element vector in "
+                + self._basis
+                + " coordinates, got {}".format(new)
             )
 
     @optimizable_parameter
@@ -685,6 +689,54 @@ def set_coeffs(self, n, r=None):
             if rr is not None:
                 self.r_n = put(self.r_n, idx, rr)
 
+    def compute(
+        self,
+        names,
+        grid=None,
+        params=None,
+        transforms=None,
+        data=None,
+        override_grid=True,
+        **kwargs,
+    ):
+        """Compute the quantity given by name on grid.
+
+        Parameters
+        ----------
+        names : str or array-like of str
+            Name(s) of the quantity(s) to compute.
+        grid : Grid or int, optional
+            Grid of coordinates to evaluate at. Defaults to a Linear grid.
+            If an integer, uses that many equally spaced points.
+        params : dict of ndarray
+            Parameters from the equilibrium. Defaults to attributes of self.
+        transforms : dict of Transform
+            Transforms for R, Z, lambda, etc. Default is to build from grid
+        data : dict of ndarray
+            Data computed so far, generally output from other compute functions
+        override_grid : bool
+            If True, override the user supplied grid if necessary and use a full
+            resolution grid to compute quantities and then downsample to user requested
+            grid. If False, uses only the user specified grid, which may lead to
+            inaccurate values for surface or volume averages.
+
+        Returns
+        -------
+        data : dict of ndarray
+            Computed quantity and intermediate variables.
+
+        """
+        return super().compute(
+            names=names,
+            grid=grid,
+            params=params,
+            transforms=transforms,
+            data=data,
+            override_grid=override_grid,
+            basis_in=self._basis,
+            **kwargs,
+        )
+
 
 class SplineXYZCurve(Curve):
     """Curve parameterized by spline knots in X,Y,Z.
diff --git a/tests/test_curves.py b/tests/test_curves.py
index ce1a04e3a9..66c4a12edf 100644
--- a/tests/test_curves.py
+++ b/tests/test_curves.py
@@ -554,6 +554,30 @@ def test_coords(self):
         np.testing.assert_allclose(y, -11)
         np.testing.assert_allclose(z, 1)
 
+    @pytest.mark.unit
+    def test_basis(self):
+        """Test xyz vs rpz basis."""
+        cxyz = FourierPlanarCurve(center=[1, 1, 0], normal=[-1, 1, 0], basis="xyz")
+        crpz = FourierPlanarCurve(
+            center=[np.sqrt(2), np.pi / 4, 0], normal=[0, 1, 0], basis="rpz"
+        )
+
+        x_xyz = cxyz.compute("x")["x"]
+        x_rpz = crpz.compute("x")["x"]
+        np.testing.assert_allclose(x_xyz, x_rpz)
+
+        xs_xyz = cxyz.compute("x_s")["x_s"]
+        xs_rpz = crpz.compute("x_s")["x_s"]
+        np.testing.assert_allclose(xs_xyz, xs_rpz, atol=2e-15)
+
+        xss_xyz = cxyz.compute("x_ss")["x_ss"]
+        xss_rpz = crpz.compute("x_ss")["x_ss"]
+        np.testing.assert_allclose(xss_xyz, xss_rpz, atol=2e-15)
+
+        xsss_xyz = cxyz.compute("x_sss")["x_sss"]
+        xsss_rpz = crpz.compute("x_sss")["x_sss"]
+        np.testing.assert_allclose(xsss_xyz, xsss_rpz, atol=2e-15)
+
     @pytest.mark.unit
     def test_misc(self):
         """Test getting/setting misc attributes of FourierPlanarCurve."""

From 28cc1926db8a063bfd81ad0a6fcc5a76b80735a1 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 14 May 2024 11:55:31 -0400
Subject: [PATCH 2/9] move kwarg to transforms

---
 desc/compute/_curve.py | 68 ++++++++++--------------------------------
 desc/geometry/curve.py | 54 ++-------------------------------
 2 files changed, 19 insertions(+), 103 deletions(-)

diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py
index 7271095b1e..d63d3e678c 100644
--- a/desc/compute/_curve.py
+++ b/desc/compute/_curve.py
@@ -168,18 +168,15 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={
-        "r": [[0, 0, 0]],
-    },
+    transforms={"r": [[0, 0, 0]], "basis": []},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
-    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+    if transforms["basis"].lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -212,18 +209,15 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, first derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={
-        "r": [[0, 0, 0], [0, 0, 1]],
-    },
+    transforms={"r": [[0, 0, 0], [0, 0, 1]], "basis": []},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
-    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+    if transforms["basis"].lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -265,18 +259,15 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, second derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={
-        "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]],
-    },
+    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], "basis": []},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
-    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+    if transforms["basis"].lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -323,18 +314,15 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, third derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={
-        "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]],
-    },
+    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], "basis": []},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
-    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if kwargs.get("basis_in", "xyz").lower() == "rpz":
+    if transforms["basis"].lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -388,11 +376,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve",
     dim=3,
     params=["R_n", "Z_n", "rotmat", "shift"],
-    transforms={
-        "R": [[0, 0, 0]],
-        "Z": [[0, 0, 0]],
-        "grid": [],
-    },
+    transforms={"R": [[0, 0, 0]], "Z": [[0, 0, 0]], "grid": []},
     profiles=[],
     coordinates="s",
     data=[],
@@ -423,11 +407,7 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, first derivative",
     dim=3,
     params=["R_n", "Z_n", "rotmat"],
-    transforms={
-        "R": [[0, 0, 0], [0, 0, 1]],
-        "Z": [[0, 0, 1]],
-        "grid": [],
-    },
+    transforms={"R": [[0, 0, 0], [0, 0, 1]], "Z": [[0, 0, 1]], "grid": []},
     profiles=[],
     coordinates="s",
     data=[],
@@ -457,11 +437,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, second derivative",
     dim=3,
     params=["R_n", "Z_n", "rotmat"],
-    transforms={
-        "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]],
-        "Z": [[0, 0, 2]],
-        "grid": [],
-    },
+    transforms={"R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], "Z": [[0, 0, 2]], "grid": []},
     profiles=[],
     coordinates="s",
     data=[],
@@ -533,11 +509,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve",
     dim=3,
     params=["X_n", "Y_n", "Z_n", "rotmat", "shift"],
-    transforms={
-        "X": [[0, 0, 0]],
-        "Y": [[0, 0, 0]],
-        "Z": [[0, 0, 0]],
-    },
+    transforms={"X": [[0, 0, 0]], "Y": [[0, 0, 0]], "Z": [[0, 0, 0]]},
     profiles=[],
     coordinates="s",
     data=[],
@@ -671,9 +643,7 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve",
     dim=3,
     params=["X", "Y", "Z", "knots", "rotmat", "shift"],
-    transforms={
-        "method": [],
-    },
+    transforms={"method": []},
     profiles=[],
     coordinates="s",
     data=["s"],
@@ -726,9 +696,7 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, first derivative",
     dim=3,
     params=["X", "Y", "Z", "knots", "rotmat", "shift"],
-    transforms={
-        "method": [],
-    },
+    transforms={"method": []},
     profiles=[],
     coordinates="s",
     data=["s"],
@@ -812,9 +780,7 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, second derivative",
     dim=3,
     params=["X", "Y", "Z", "knots", "rotmat", "shift"],
-    transforms={
-        "method": [],
-    },
+    transforms={"method": []},
     profiles=[],
     coordinates="s",
     data=["s"],
@@ -897,9 +863,7 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, third derivative",
     dim=3,
     params=["X", "Y", "Z", "knots", "rotmat", "shift"],
-    transforms={
-        "method": [],
-    },
+    transforms={"method": []},
     profiles=[],
     coordinates="s",
     data=["s"],
diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py
index 3f670656e7..8b649e9d34 100644
--- a/desc/geometry/curve.py
+++ b/desc/geometry/curve.py
@@ -599,7 +599,7 @@ def __init__(
 
         self.normal = normal
         self.center = center
-        self._basis = basis
+        self.basis = basis
 
     @property
     def r_basis(self):
@@ -633,7 +633,7 @@ def center(self, new):
         else:
             raise ValueError(
                 "center should be a 3 element vector in "
-                + self._basis
+                + self.basis
                 + " coordinates, got {}".format(new)
             )
 
@@ -650,7 +650,7 @@ def normal(self, new):
         else:
             raise ValueError(
                 "normal should be a 3 element vector in "
-                + self._basis
+                + self.basis
                 + " coordinates, got {}".format(new)
             )
 
@@ -689,54 +689,6 @@ def set_coeffs(self, n, r=None):
             if rr is not None:
                 self.r_n = put(self.r_n, idx, rr)
 
-    def compute(
-        self,
-        names,
-        grid=None,
-        params=None,
-        transforms=None,
-        data=None,
-        override_grid=True,
-        **kwargs,
-    ):
-        """Compute the quantity given by name on grid.
-
-        Parameters
-        ----------
-        names : str or array-like of str
-            Name(s) of the quantity(s) to compute.
-        grid : Grid or int, optional
-            Grid of coordinates to evaluate at. Defaults to a Linear grid.
-            If an integer, uses that many equally spaced points.
-        params : dict of ndarray
-            Parameters from the equilibrium. Defaults to attributes of self.
-        transforms : dict of Transform
-            Transforms for R, Z, lambda, etc. Default is to build from grid
-        data : dict of ndarray
-            Data computed so far, generally output from other compute functions
-        override_grid : bool
-            If True, override the user supplied grid if necessary and use a full
-            resolution grid to compute quantities and then downsample to user requested
-            grid. If False, uses only the user specified grid, which may lead to
-            inaccurate values for surface or volume averages.
-
-        Returns
-        -------
-        data : dict of ndarray
-            Computed quantity and intermediate variables.
-
-        """
-        return super().compute(
-            names=names,
-            grid=grid,
-            params=params,
-            transforms=transforms,
-            data=data,
-            override_grid=override_grid,
-            basis_in=self._basis,
-            **kwargs,
-        )
-
 
 class SplineXYZCurve(Curve):
     """Curve parameterized by spline knots in X,Y,Z.

From 9320bd1087936ae4f83dd91a9f55e9365ac350f6 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 14 May 2024 11:59:17 -0400
Subject: [PATCH 3/9] fix io_attrs for basis

---
 desc/geometry/curve.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py
index 8b649e9d34..4d9d7a50a7 100644
--- a/desc/geometry/curve.py
+++ b/desc/geometry/curve.py
@@ -570,7 +570,7 @@ class FourierPlanarCurve(Curve):
 
     """
 
-    _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "_basis"]
+    _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "basis"]
 
     # Reference frame is centered at the origin with normal in the +Z direction.
     # The curve is computed in this frame and then shifted/rotated to the correct frame.

From f7757d1251e3e0ab8f2393e543936a68f3dd3c9b Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 15 May 2024 11:40:44 -0400
Subject: [PATCH 4/9] add failing test

---
 tests/test_examples.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index adad7d7104..ba72a4ee18 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -14,7 +14,7 @@
 from desc.continuation import solve_continuation_automatic
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.examples import get
-from desc.geometry import FourierRZToroidalSurface
+from desc.geometry import FourierPlanarCurve, FourierRZToroidalSurface
 from desc.grid import LinearGrid
 from desc.io import load
 from desc.magnetic_fields import (
@@ -1294,3 +1294,13 @@ def test_second_stage_optimization():
     np.testing.assert_allclose(field[0].R0, 3.5)  # this value was fixed
     np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (no change)
     np.testing.assert_allclose(field[1].B0, 0, atol=1e-12)  # vertical field (vanishes)
+
+
+@pytest.mark.unit
+def test_optimize_with_fourier_planar_curve():
+    """Test optimizing a FourierPlanarCurve."""
+    c = FourierPlanarCurve()
+    objective = ObjectiveFunction(CoilLength(c, target=11))
+    optimizer = Optimizer("fmintr")
+    (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
+    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)

From 13db5e8f3a7cb8f759a4ea1fd9816335bb3839a2 Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 15 May 2024 11:47:23 -0400
Subject: [PATCH 5/9] add test that triggers jax issue

---
 tests/test_examples.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index ba72a4ee18..86ce42a7f9 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -10,7 +10,7 @@
 from qsc import Qsc
 
 from desc.backend import jnp
-from desc.coils import FourierRZCoil
+from desc.coils import FourierPlanarCoil, FourierRZCoil
 from desc.continuation import solve_continuation_automatic
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.examples import get
@@ -1304,3 +1304,13 @@ def test_optimize_with_fourier_planar_curve():
     optimizer = Optimizer("fmintr")
     (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
     np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)
+
+
+@pytest.mark.unit
+def test_optimize_with_fourier_planar_coil():
+    """Test optimizing a FourierPlanarCoil."""
+    c = FourierPlanarCoil()
+    objective = ObjectiveFunction(CoilLength(c, target=11))
+    optimizer = Optimizer("fmintr")
+    (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
+    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)

From b3832f28f85edf1f36840eaa26e208319e88602f Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 15 May 2024 11:49:11 -0400
Subject: [PATCH 6/9] remove unneeded test

---
 tests/test_examples.py | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index 86ce42a7f9..6e8f801415 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -14,7 +14,7 @@
 from desc.continuation import solve_continuation_automatic
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.examples import get
-from desc.geometry import FourierPlanarCurve, FourierRZToroidalSurface
+from desc.geometry import FourierRZToroidalSurface
 from desc.grid import LinearGrid
 from desc.io import load
 from desc.magnetic_fields import (
@@ -1296,16 +1296,6 @@ def test_second_stage_optimization():
     np.testing.assert_allclose(field[1].B0, 0, atol=1e-12)  # vertical field (vanishes)
 
 
-@pytest.mark.unit
-def test_optimize_with_fourier_planar_curve():
-    """Test optimizing a FourierPlanarCurve."""
-    c = FourierPlanarCurve()
-    objective = ObjectiveFunction(CoilLength(c, target=11))
-    optimizer = Optimizer("fmintr")
-    (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
-    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)
-
-
 @pytest.mark.unit
 def test_optimize_with_fourier_planar_coil():
     """Test optimizing a FourierPlanarCoil."""

From e6065f904fea0ec968125551180630038ac72954 Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 15 May 2024 12:22:16 -0400
Subject: [PATCH 7/9] add test with mixed coilset

---
 tests/test_examples.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index 6e8f801415..ee9c6b7d61 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -10,7 +10,7 @@
 from qsc import Qsc
 
 from desc.backend import jnp
-from desc.coils import FourierPlanarCoil, FourierRZCoil
+from desc.coils import FourierPlanarCoil, FourierRZCoil, MixedCoilSet
 from desc.continuation import solve_continuation_automatic
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.examples import get
@@ -1304,3 +1304,13 @@ def test_optimize_with_fourier_planar_coil():
     optimizer = Optimizer("fmintr")
     (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
     np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)
+
+
+@pytest.mark.unit
+def test_optimize_with_fourier_planar_coil_in_MixedCoilSet():
+    """Test optimizing a FourierPlanarCoil."""
+    c = MixedCoilSet(FourierRZCoil(), FourierPlanarCoil())
+    objective = ObjectiveFunction(CoilLength(c, target=11))
+    optimizer = Optimizer("fmintr")
+    (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
+    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)

From 0497627d23dc751408c6c0a3427058e8e18db03d Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Wed, 15 May 2024 17:21:19 -0400
Subject: [PATCH 8/9] revert to original method

---
 desc/compute/_curve.py | 20 +++++++++------
 desc/geometry/curve.py | 56 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py
index d63d3e678c..40d77174b3 100644
--- a/desc/compute/_curve.py
+++ b/desc/compute/_curve.py
@@ -168,15 +168,16 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={"r": [[0, 0, 0]], "basis": []},
+    transforms={"r": [[0, 0, 0]]},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if transforms["basis"].lower() == "rpz":
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -209,15 +210,16 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, first derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={"r": [[0, 0, 0], [0, 0, 1]], "basis": []},
+    transforms={"r": [[0, 0, 0], [0, 0, 1]]},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if transforms["basis"].lower() == "rpz":
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -259,15 +261,16 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, second derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], "basis": []},
+    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]]},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if transforms["basis"].lower() == "rpz":
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
@@ -314,15 +317,16 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     description="Position vector along curve, third derivative",
     dim=3,
     params=["r_n", "center", "normal", "rotmat", "shift"],
-    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], "basis": []},
+    transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]},
     profiles=[],
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
     basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
+    basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'",
 )
 def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    if transforms["basis"].lower() == "rpz":
+    if kwargs.get("basis_in", "xyz").lower() == "rpz":
         center = rpz2xyz(params["center"])
         normal = rpz2xyz_vec(params["normal"], phi=params["center"][1])
     else:
diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py
index 4d9d7a50a7..3f670656e7 100644
--- a/desc/geometry/curve.py
+++ b/desc/geometry/curve.py
@@ -570,7 +570,7 @@ class FourierPlanarCurve(Curve):
 
     """
 
-    _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "basis"]
+    _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "_basis"]
 
     # Reference frame is centered at the origin with normal in the +Z direction.
     # The curve is computed in this frame and then shifted/rotated to the correct frame.
@@ -599,7 +599,7 @@ def __init__(
 
         self.normal = normal
         self.center = center
-        self.basis = basis
+        self._basis = basis
 
     @property
     def r_basis(self):
@@ -633,7 +633,7 @@ def center(self, new):
         else:
             raise ValueError(
                 "center should be a 3 element vector in "
-                + self.basis
+                + self._basis
                 + " coordinates, got {}".format(new)
             )
 
@@ -650,7 +650,7 @@ def normal(self, new):
         else:
             raise ValueError(
                 "normal should be a 3 element vector in "
-                + self.basis
+                + self._basis
                 + " coordinates, got {}".format(new)
             )
 
@@ -689,6 +689,54 @@ def set_coeffs(self, n, r=None):
             if rr is not None:
                 self.r_n = put(self.r_n, idx, rr)
 
+    def compute(
+        self,
+        names,
+        grid=None,
+        params=None,
+        transforms=None,
+        data=None,
+        override_grid=True,
+        **kwargs,
+    ):
+        """Compute the quantity given by name on grid.
+
+        Parameters
+        ----------
+        names : str or array-like of str
+            Name(s) of the quantity(s) to compute.
+        grid : Grid or int, optional
+            Grid of coordinates to evaluate at. Defaults to a Linear grid.
+            If an integer, uses that many equally spaced points.
+        params : dict of ndarray
+            Parameters from the equilibrium. Defaults to attributes of self.
+        transforms : dict of Transform
+            Transforms for R, Z, lambda, etc. Default is to build from grid
+        data : dict of ndarray
+            Data computed so far, generally output from other compute functions
+        override_grid : bool
+            If True, override the user supplied grid if necessary and use a full
+            resolution grid to compute quantities and then downsample to user requested
+            grid. If False, uses only the user specified grid, which may lead to
+            inaccurate values for surface or volume averages.
+
+        Returns
+        -------
+        data : dict of ndarray
+            Computed quantity and intermediate variables.
+
+        """
+        return super().compute(
+            names=names,
+            grid=grid,
+            params=params,
+            transforms=transforms,
+            data=data,
+            override_grid=override_grid,
+            basis_in=self._basis,
+            **kwargs,
+        )
+
 
 class SplineXYZCurve(Curve):
     """Curve parameterized by spline knots in X,Y,Z.

From 0f9a6fe57fae62c56384081bff22a47c46996d79 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Wed, 15 May 2024 17:24:21 -0400
Subject: [PATCH 9/9] consolidate tests

---
 tests/test_examples.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index ee9c6b7d61..b84cd967b1 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1296,21 +1296,20 @@ def test_second_stage_optimization():
     np.testing.assert_allclose(field[1].B0, 0, atol=1e-12)  # vertical field (vanishes)
 
 
+# TODO: replace this with the solution to Issue #1021
 @pytest.mark.unit
 def test_optimize_with_fourier_planar_coil():
     """Test optimizing a FourierPlanarCoil."""
+    # single coil
     c = FourierPlanarCoil()
     objective = ObjectiveFunction(CoilLength(c, target=11))
     optimizer = Optimizer("fmintr")
     (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
     np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)
 
-
-@pytest.mark.unit
-def test_optimize_with_fourier_planar_coil_in_MixedCoilSet():
-    """Test optimizing a FourierPlanarCoil."""
+    # in MixedCoilSet
     c = MixedCoilSet(FourierRZCoil(), FourierPlanarCoil())
     objective = ObjectiveFunction(CoilLength(c, target=11))
     optimizer = Optimizer("fmintr")
     (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0)
-    np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3)
+    np.testing.assert_allclose(c.compute("length")[1]["length"], 11, atol=1e-3)