fix curvy bugs

desc/coils.py

@@ -883,21 +883,22 @@ def linspaced_angular(
         axis : array-like, shape(3,)
             axis to rotate about
         angle : float
-            total rotational extend of coil set.
+            total rotational extent of coil set
         n : int
             number of copies of original coil
         endpoint : bool
             whether to include a coil at final angle
+
         """
         assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
         if current is None:
             current = coil.current
         currents = jnp.broadcast_to(current, (n,))
+        phi = jnp.linspace(0, angle, n, endpoint=endpoint)
         coils = []
-        phis = jnp.linspace(0, angle, n, endpoint=endpoint)
         for i in range(n):
             coili = coil.copy()
-            coili.rotate(axis, angle=phis[i])
+            coili.rotate(axis=axis, angle=phi[i])
             coili.current = currents[i]
             coils.append(coili)
         return cls(*coils)
@@ -920,14 +921,15 @@ def linspaced_linear(
             number of copies of original coil
         endpoint : bool
             whether to include a coil at final point
+
         """
         assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
         if current is None:
             current = coil.current
         currents = jnp.broadcast_to(current, (n,))
         displacement = jnp.asarray(displacement)
-        coils = []
         a = jnp.linspace(0, 1, n, endpoint=endpoint)
+        coils = []
         for i in range(n):
             coili = coil.copy()
             coili.translate(a[i] * displacement)
@@ -953,6 +955,7 @@ def from_symmetry(cls, coils, NFP, sym=False):
             number of field periods
         sym : bool
             whether coils should be stellarator symmetric
+
         """
         if not isinstance(coils, CoilSet):
             coils = CoilSet(coils)

desc/compute/_curve.py

@@ -3,14 +3,8 @@
 from desc.backend import jnp
 
 from .data_index import register_compute_fun
-from .geom_utils import (
-    _rotation_matrix_from_normal,
-    rpz2xyz,
-    rpz2xyz_vec,
-    xyz2rpz,
-    xyz2rpz_vec,
-)
-from .utils import cross, dot
+from .geom_utils import rotation_matrix, rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
+from .utils import cross, dot, safenormalize
 
 
 @register_compute_fun(
@@ -184,16 +178,19 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs):
     basis="basis",
 )
 def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
-    # create planar curve at z==0
+    # create planar curve at Z==0
     r = transforms["r"].transform(params["r_n"], dz=0)
     Z = jnp.zeros_like(r)
     X = r * jnp.cos(data["s"])
     Y = r * jnp.sin(data["s"])
     coords = jnp.array([X, Y, Z]).T
     # rotate into place
-    R = _rotation_matrix_from_normal(params["normal"])
-    coords = jnp.matmul(coords, R.T) + params["center"]
-    coords = jnp.matmul(coords, params["rotmat"].T) + params["shift"]
+    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"])))
+    A = rotation_matrix(axis=axis, angle=angle)
+    coords = jnp.matmul(coords, A.T) + params["center"]
+    coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz(coords)
     data["x"] = coords
@@ -224,16 +221,22 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     dY = dr * jnp.sin(data["s"]) + r * jnp.cos(data["s"])
     dZ = jnp.zeros_like(dX)
     coords = jnp.array([dX, dY, dZ]).T
-    A = _rotation_matrix_from_normal(params["normal"])
+    # 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"])))
+    A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
-    coords = jnp.matmul(coords, params["rotmat"].T)
+    coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     if kwargs.get("basis", "rpz").lower() == "rpz":
         X = r * jnp.cos(data["s"])
         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, params["rotmat"].T) + params["shift"]
+        xyzcoords = (
+            jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
+        )
         x, y, z = xyzcoords.T
         coords = xyz2rpz_vec(coords, x=x, y=y)
     data["x_s"] = coords
@@ -269,16 +272,22 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     )
     d2Z = jnp.zeros_like(d2X)
     coords = jnp.array([d2X, d2Y, d2Z]).T
-    A = _rotation_matrix_from_normal(params["normal"])
+    # 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"])))
+    A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
-    coords = jnp.matmul(coords, params["rotmat"].T)
+    coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     if kwargs.get("basis", "rpz").lower() == "rpz":
         X = r * jnp.cos(data["s"])
         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, params["rotmat"].T) + params["shift"]
+        xyzcoords = (
+            jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
+        )
         x, y, z = xyzcoords.T
         coords = xyz2rpz_vec(coords, x=x, y=y)
     data["x_ss"] = coords
@@ -321,16 +330,22 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     )
     d3Z = jnp.zeros_like(d3X)
     coords = jnp.array([d3X, d3Y, d3Z]).T
-    A = _rotation_matrix_from_normal(params["normal"])
+    # 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"])))
+    A = rotation_matrix(axis=axis, angle=angle)
     coords = jnp.matmul(coords, A.T)
-    coords = jnp.matmul(coords, params["rotmat"].T)
+    coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     if kwargs.get("basis", "rpz").lower() == "rpz":
         X = r * jnp.cos(data["s"])
         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, params["rotmat"].T) + params["shift"]
+        xyzcoords = (
+            jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
+        )
         x, y, z = xyzcoords.T
         coords = xyz2rpz_vec(coords, x=x, y=y)
     data["x_sss"] = coords
@@ -363,7 +378,9 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.stack([R, phi, Z], axis=1)
     # convert to xyz for displacement and rotation
     coords = rpz2xyz(coords)
-    coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+    coords = (
+        coords @ params["rotmat"].reshape((3, 3)).T + params["shift"][jnp.newaxis, :]
+    )
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz(coords)
     data["x"] = coords
@@ -397,7 +414,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.stack([dR, dphi, dZ], axis=1)
     # convert to xyz for displacement and rotation
     coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2])
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2])
     data["x_s"] = coords
@@ -435,7 +452,7 @@ def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.stack([R, phi, Z], axis=1)
     # convert to xyz for displacement and rotation
     coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2])
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2])
     data["x_ss"] = coords
@@ -473,7 +490,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coords = jnp.stack([R, phi, Z], axis=1)
     # convert to xyz for displacement and rotation
     coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2])
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2])
     data["x_sss"] = coords
@@ -504,7 +521,9 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     Y = transforms["Y"].transform(params["Y_n"], dz=0)
     Z = transforms["Z"].transform(params["Z_n"], dz=0)
     coords = jnp.stack([X, Y, Z], axis=1)
-    coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+    coords = (
+        coords @ params["rotmat"].reshape((3, 3)).T + params["shift"][jnp.newaxis, :]
+    )
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz(coords)
     data["x"] = coords
@@ -535,7 +554,7 @@ def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     dY = transforms["Y"].transform(params["Y_n"], dz=1)
     dZ = transforms["Z"].transform(params["Z_n"], dz=1)
     coords = jnp.stack([dX, dY, dZ], axis=1)
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(
             coords,
@@ -570,7 +589,7 @@ def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     d2Y = transforms["Y"].transform(params["Y_n"], dz=2)
     d2Z = transforms["Z"].transform(params["Z_n"], dz=2)
     coords = jnp.stack([d2X, d2Y, d2Z], axis=1)
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(
             coords,
@@ -605,7 +624,7 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     d3Y = transforms["Y"].transform(params["Y_n"], dz=3)
     d3Z = transforms["Z"].transform(params["Z_n"], dz=3)
     coords = jnp.stack([d3X, d3Y, d3Z], axis=1)
-    coords = coords @ params["rotmat"].T
+    coords = coords @ params["rotmat"].reshape((3, 3)).T
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz_vec(
             coords,
@@ -662,7 +681,9 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     )
 
     coords = jnp.stack([Xq, Yq, Zq], axis=1)
-    coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+    coords = (
+        coords @ params["rotmat"].reshape((3, 3)).T + params["shift"][jnp.newaxis, :]
+    )
     if kwargs.get("basis", "rpz").lower() == "rpz":
         coords = xyz2rpz(coords)
     data["x"] = coords
@@ -715,7 +736,7 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     )
 
     coords_s = jnp.stack([dXq, dYq, dZq], axis=1)
-    coords_s = coords_s @ params["rotmat"].T
+    coords_s = coords_s @ params["rotmat"].reshape((3, 3)).T
 
     if kwargs.get("basis", "rpz").lower() == "rpz":
         # calculate the xy coordinates to rotate to rpz
@@ -745,7 +766,10 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
         )
 
         coords = jnp.stack([Xq, Yq, Zq], axis=1)
-        coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+        coords = (
+            coords @ params["rotmat"].reshape((3, 3)).T
+            + params["shift"][jnp.newaxis, :]
+        )
 
         coords_s = xyz2rpz_vec(coords_s, x=coords[:, 0], y=coords[:, 1])
     data["x_s"] = coords_s
@@ -798,7 +822,7 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     )
 
     coords_ss = jnp.stack([d2Xq, d2Yq, d2Zq], axis=1)
-    coords_ss = coords_ss @ params["rotmat"].T
+    coords_ss = coords_ss @ params["rotmat"].reshape((3, 3)).T
 
     if kwargs.get("basis", "rpz").lower() == "rpz":
         # calculate the xy coordinates to rotate to rpz
@@ -827,7 +851,10 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
             period=2 * jnp.pi,
         )
         coords = jnp.stack([Xq, Yq, Zq], axis=1)
-        coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+        coords = (
+            coords @ params["rotmat"].reshape((3, 3)).T
+            + params["shift"][jnp.newaxis, :]
+        )
 
         coords_ss = xyz2rpz_vec(coords_ss, x=coords[:, 0], y=coords[:, 1])
     data["x_ss"] = coords_ss
@@ -880,7 +907,7 @@ def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     )
 
     coords_sss = jnp.stack([d3Xq, d3Yq, d3Zq], axis=1)
-    coords_sss = coords_sss @ params["rotmat"].T
+    coords_sss = coords_sss @ params["rotmat"].reshape((3, 3)).T
 
     if kwargs.get("basis", "rpz").lower() == "rpz":
         # calculate the xy coordinates to rotate to rpz
@@ -909,7 +936,10 @@ def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
             period=2 * jnp.pi,
         )
         coords = jnp.stack([Xq, Yq, Zq], axis=1)
-        coords = coords @ params["rotmat"].T + params["shift"][jnp.newaxis, :]
+        coords = (
+            coords @ params["rotmat"].reshape((3, 3)).T
+            + params["shift"][jnp.newaxis, :]
+        )
 
         coords_sss = xyz2rpz_vec(coords_sss, x=coords[:, 0], y=coords[:, 1])
     data["x_sss"] = coords_sss
@@ -976,7 +1006,7 @@ def _frenet_normal(params, transforms, profiles, data, **kwargs):
 def _frenet_binormal(params, transforms, profiles, data, **kwargs):
     data["frenet_binormal"] = cross(
         data["frenet_tangent"], data["frenet_normal"]
-    ) * jnp.linalg.det(params["rotmat"])
+    ) * jnp.linalg.det(params["rotmat"].reshape((3, 3)))
     return data
 
 

desc/compute/geom_utils.py

@@ -4,6 +4,8 @@
 
 from desc.backend import jnp
 
+from .utils import safenorm, safenormalize
+
 
 def reflection_matrix(normal):
     """Matrix to reflect points across plane through origin with specified normal.
@@ -35,17 +37,20 @@ def rotation_matrix(axis, angle=None):
 
     Returns
     -------
-    rot : ndarray, shape(3,3)
-        Matrix to rotate points in cartesian (X,Y,Z) coordinates
+    rotmat : ndarray, shape(3,3)
+        Matrix to rotate points in cartesian (X,Y,Z) coordinates.
+
     """
     axis = jnp.asarray(axis)
+    norm = safenorm(axis)
+    axis = safenormalize(axis)
     if angle is None:
-        angle = jnp.linalg.norm(axis)
-    axis = axis / jnp.linalg.norm(axis)
+        angle = norm
+    eps = 1e2 * jnp.finfo(axis.dtype).eps
     R1 = jnp.cos(angle) * jnp.eye(3)
     R2 = jnp.sin(angle) * jnp.cross(axis, jnp.identity(axis.shape[0]) * -1)
     R3 = (1 - jnp.cos(angle)) * jnp.outer(axis, axis)
-    return R1 + R2 + R3
+    return jnp.where(norm < eps, jnp.eye(3), R1 + R2 + R3)  # if axis=0, no rotation
 
 
 def xyz2rpz(pts):
@@ -152,17 +157,3 @@ def inner(vec, phi):
         return cart
 
     return inner(vec, phi)
-
-
-def _rotation_matrix_from_normal(normal):
-    nx, ny, nz = normal
-    nxny = jnp.sqrt(nx**2 + ny**2)
-    R = jnp.array(
-        [
-            [ny / nxny, -nx / nxny, 0],
-            [nx * nx / nxny, ny * nz / nxny, -nxny],
-            [nx, ny, nz],
-        ]
-    ).T
-    R = jnp.where(nxny == 0, jnp.eye(3), R)
-    return R