FourierPlanarCurve basis option (#1017)

Adds an option for the `center` and `normal` vectors of
`FourierPlanarCurve` to be specified in $(R,\phi,Z)$ coordinates,
instead of the previous limitation to Cartesian coordinates.

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):

desc/compute/_curve.py

@@ -168,16 +168,21 @@ 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]]},
     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":
+        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 +191,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)
@@ -205,16 +210,21 @@ 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]]},
     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":
+        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 +233,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 +243,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"]
         )
@@ -251,16 +261,21 @@ 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]]},
     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":
+        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 +289,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 +299,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"]
         )
@@ -302,16 +317,21 @@ 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]]},
     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":
+        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 +352,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 +362,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"]
         )
@@ -360,11 +380,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=[],
@@ -395,11 +411,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=[],
@@ -429,11 +441,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=[],
@@ -505,11 +513,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=[],
@@ -643,9 +647,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"],
@@ -698,9 +700,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"],
@@ -784,9 +784,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"],
@@ -869,9 +867,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"],

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