fix curvy bugs

desc/compute/_curve.py

@@ -3,13 +3,7 @@
 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 .geom_utils import rotation_matrix, rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
 from .utils import cross, dot
 
 
@@ -184,16 +178,19 @@
     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, 1])  # 2D curve in X-Y plane has normal = +Z axis
+    axis = cross(Zaxis, params["normal"])
+    angle = jnp.arccos(dot(Zaxis, params["normal"]) / jnp.linalg.norm(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,21 @@
     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"])
+    Zaxis = jnp.array([0, 0, 1])  # 2D curve in X-Y plane has normal = +Z axis
+    axis = cross(Zaxis, params["normal"])
+    angle = jnp.arccos(dot(Zaxis, params["normal"]) / jnp.linalg.norm(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 +271,21 @@
     )
     d2Z = jnp.zeros_like(d2X)
     coords = jnp.array([d2X, d2Y, d2Z]).T
-    A = _rotation_matrix_from_normal(params["normal"])
+    Zaxis = jnp.array([0, 0, 1])  # 2D curve in X-Y plane has normal = +Z axis
+    axis = cross(Zaxis, params["normal"])
+    angle = jnp.arccos(dot(Zaxis, params["normal"]) / jnp.linalg.norm(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 +328,21 @@
     )
     d3Z = jnp.zeros_like(d3X)
     coords = jnp.array([d3X, d3Y, d3Z]).T
-    A = _rotation_matrix_from_normal(params["normal"])
+    Zaxis = jnp.array([0, 0, 1])  # 2D curve in X-Y plane has normal = +Z axis
+    axis = cross(Zaxis, params["normal"])
+    angle = jnp.arccos(dot(Zaxis, params["normal"]) / jnp.linalg.norm(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 +375,9 @@
     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 +411,7 @@
     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 +449,7 @@
     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 +487,7 @@
     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 +518,9 @@
     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 +551,7 @@
     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 +586,7 @@
     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 +621,7 @@
     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 +678,9 @@
     )
 
     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 +733,7 @@
     )
 
     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 +763,10 @@
         )
 
         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 +819,7 @@
     )
 
     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 +848,10 @@
             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 +904,7 @@
     )
 
     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 +933,10 @@
             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 +1003,7 @@
 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

@@ -37,11 +37,14 @@
     -------
     rot : ndarray, shape(3,3)
         Matrix to rotate points in cartesian (X,Y,Z) coordinates
+
     """
+    # FIXME: does not work with AD
     axis = jnp.asarray(axis)
+    norm = jnp.linalg.norm(axis)
     if angle is None:
-        angle = jnp.linalg.norm(axis)
-    axis = axis / jnp.linalg.norm(axis)
+        angle = norm
+    axis = jnp.nan_to_num(axis / norm)
     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)
@@ -152,17 +155,3 @@
         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

desc/geometry/core.py

@@ -5,7 +5,6 @@
 
 import numpy as np
 
-from desc.backend import jnp
 from desc.compute import compute as compute_fun
 from desc.compute import data_index
 from desc.compute.geom_utils import reflection_matrix, rotation_matrix
@@ -17,18 +16,44 @@
 )
 from desc.grid import LinearGrid, QuadratureGrid, _Grid
 from desc.io import IOAble
-from desc.optimizable import Optimizable
+from desc.optimizable import Optimizable, optimizable_parameter
 
 
 class Curve(IOAble, Optimizable, ABC):
     """Abstract base class for 1D curves in 3D space."""
 
-    _io_attrs_ = ["_name", "shift", "rotmat"]
+    _io_attrs_ = ["_name", "_shift", "_rotmat"]
 
     def __init__(self, name=""):
-        self.shift = jnp.array([0, 0, 0]).astype(float)
-        self.rotmat = jnp.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).astype(float)
-        self.name = name
+        self._shift = np.array([0, 0, 0])
+        self._rotmat = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).flatten()
+        self._name = name
+
+    @optimizable_parameter
+    @property
+    def shift(self):
+        """Displacement of curve in X, Y, Z."""
+        return self._shift
+
+    @shift.setter
+    def shift(self, new):
+        if len(new) == 3:
+            self._shift = np.asarray(new)
+        else:
+            raise ValueError("shift should be a 3 element vector, got {}".format(new))
+
+    @optimizable_parameter
+    @property
+    def rotmat(self):
+        """Rotation matrix of curve in X, Y, Z."""
+        return self._rotmat
+
+    @rotmat.setter
+    def rotmat(self, new):
+        if len(new) == 9:
+            self._rotmat = np.asarray(new)
+        else:
+            self._rotmat = np.asarray(new.flatten())
 
     @property
     def name(self):
@@ -144,7 +169,7 @@
 
     def translate(self, displacement=[0, 0, 0]):
         """Translate the curve by a rigid displacement in x, y, z."""
-        self.shift += jnp.asarray(displacement)
+        self.shift += np.asarray(displacement)
 
     def rotate(self, axis=[0, 0, 1], angle=0):
         """Rotate the curve by a fixed angle about axis in xyz coordinates."""

tests/test_coils.py

@@ -489,8 +489,8 @@ def test_save_and_load_makegrid_coils_rotated(tmpdir_factory):
         Z2 = coords2[:, 2]
 
         np.testing.assert_allclose(c1.current, c2.current, err_msg=f"Coil {i}")
-        np.testing.assert_allclose(X1, X2, err_msg=f"Coil {i}")
-        np.testing.assert_allclose(Y1, Y2, err_msg=f"Coil {i}")
+        np.testing.assert_allclose(X1, X2, err_msg=f"Coil {i}", atol=1e-16)
+        np.testing.assert_allclose(Y1, Y2, err_msg=f"Coil {i}", atol=1e-16)
         np.testing.assert_allclose(Z1, Z2, atol=2e-7, err_msg=f"Coil {i}")
 
     # check Bnormal on torus and ensure is near zero
@@ -564,8 +564,8 @@ def test_save_and_load_makegrid_coils_rotated_int_grid(tmpdir_factory):
         Z2 = coords2[:, 2]
 
         np.testing.assert_allclose(c1.current, c2.current, err_msg=f"Coil {i}")
-        np.testing.assert_allclose(X1, X2, err_msg=f"Coil {i}")
-        np.testing.assert_allclose(Y1, Y2, err_msg=f"Coil {i}")
+        np.testing.assert_allclose(X1, X2, err_msg=f"Coil {i}", atol=1e-16)
+        np.testing.assert_allclose(Y1, Y2, err_msg=f"Coil {i}", atol=1e-16)
         np.testing.assert_allclose(Z1, Z2, atol=2e-7, err_msg=f"Coil {i}")
 
     # check Bnormal on torus and ensure is near zero