Fix for NaN in reverse mode gradient of FourierPlanarCoil and Hessian Scaling Issue for small problems

CHANGELOG.md

@@ -21,6 +21,8 @@ Breaking Changes and Deprecations
 Bug Fixes
 
 - Fixes bug in ``auglag`` optimizers which prevented them from accepting solver hyperparameters.
+- Improves planar coil representation (``desc.coils.FourierPlanarCoil`` and ``desc.coils.FourierXYCoil``) internal rotation methods to avoid potential NaNs which could occur when the normal is parallel or antiparallel to Z-axis( to within machine epsilon), and also ensure the gradient at those edge cases is not only not NaN but also non-zero to avoid optimizer stalls at those cases.
+- Fixes potential scaling-based issue in DESC-based optimization methods which use adaptive Hessian scaling (e.g. ``"fmintr"``) that could occur when the problem size was small and there were directions of near-zero derivative in the problem.
 - Fixes bug in modified Cholesky factorization used by the trust-region
   subproblems when the Gershgorin lower bound of the Hessian was exactly zero
   (e.g. a Hessian with an all-zero row), producing NaN steps in ``fmintr`` and

desc/compute/_curve.py

@@ -5,11 +5,9 @@
 from ..utils import (
     cross,
     dot,
-    rotation_matrix,
+    rotate_vector_to_vector,
     rpz2xyz,
     rpz2xyz_vec,
-    safearccos,
-    safenormalize,
     xyz2rpz,
     xyz2rpz_vec,
 )
@@ -222,14 +220,7 @@ 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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
     coords = jnp.matmul(coords, A.T) + center
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
     # convert back to rpz
@@ -267,14 +258,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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz
@@ -317,14 +302,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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz
@@ -374,14 +353,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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz
@@ -422,14 +395,8 @@ def _x_FourierXYCurve(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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T) + center
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"]
     # convert back to rpz
@@ -465,14 +432,8 @@ def _x_s_FourierXYCurve(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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz
@@ -508,14 +469,8 @@ def _x_ss_FourierXYCurve(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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz
@@ -551,14 +506,8 @@ def _x_sss_FourierXYCurve(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, normal)
-    dotprod = dot(Zaxis, safenormalize(normal))
-    angle = safearccos(dotprod)
-    A = jnp.where(  # handle the case where normal is aligned with the -Z axis
-        jnp.allclose(dotprod, -1.0),
-        jnp.diag(jnp.array([1.0, -1.0, -1.0])),
-        rotation_matrix(axis, angle),
-    )
+    A = rotate_vector_to_vector(Zaxis, normal)
+
     coords = jnp.matmul(coords, A.T)
     coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T)
     # convert back to rpz

desc/geometry/core.py

@@ -190,7 +190,7 @@ def translate(self, displacement=[0, 0, 0]):
 
     def rotate(self, axis=[0, 0, 1], angle=0):
         """Rotate the curve by a fixed angle about axis in X,Y,Z coordinates."""
-        R = rotation_matrix(axis=axis, angle=angle)
+        R = rotation_matrix(axis=np.asarray(axis).astype(float), angle=angle)
         self.rotmat = (R @ self.rotmat.reshape(3, 3)).flatten()
         self.shift = self.shift @ R.T
 

desc/geometry/curve.py

@@ -767,7 +767,7 @@ def normal(self):
     @normal.setter
     def normal(self, new):
         if len(np.asarray(new)) == 3:
-            self._normal = np.asarray(new) / np.linalg.norm(new)
+            self._normal = np.asarray(new).astype(float) / np.linalg.norm(new)
         else:
             raise ValueError(
                 "normal should be a 3 element vector in "
@@ -1121,7 +1121,7 @@ def normal(self):
     @normal.setter
     def normal(self, new):
         if len(np.asarray(new)) == 3:
-            self._normal = np.asarray(new) / np.linalg.norm(new)
+            self._normal = np.asarray(new).astype(float) / np.linalg.norm(new)
         else:
             raise ValueError(
                 "normal should be a 3 element vector in "

desc/optimize/utils.py

@@ -518,7 +518,7 @@ def compute_hess_scale(H, prev_scale_inv=None):
     """Compute scaling factors based on diagonal of Hessian matrix."""
     scale_inv = jnp.abs(jnp.diag(H))
     scale_inv = jnp.where(
-        scale_inv < jnp.finfo(H.dtype).eps * max(H.shape), 1, scale_inv
+        scale_inv < jnp.finfo(H.dtype).eps * max(H.shape) * 1e1, 1, scale_inv
     )
 
     if prev_scale_inv is not None:

desc/utils.py

@@ -983,6 +983,15 @@ def reflection_matrix(normal):
 def rotation_matrix(axis, angle=None):
     """Matrix to rotate points about axis by given angle.
 
+    NOTE: This function works but will have zero gradient w.r.t.
+    the angle when the angle is nearly zero (specifically, when
+    the angle<sqrt(epsilon)) as there is a manual replacement of the
+    result for this case to avoid a NaN leak due to 1/0 division.
+
+    If seeking to rotate one vector onto another, it is recommended
+    to use rotate_vector_to_vector instead which uses quaternions
+    to avoid this issue
+
     Parameters
     ----------
     axis : array-like, shape(3,)
@@ -1001,13 +1010,86 @@ def rotation_matrix(axis, angle=None):
     axis = safenormalize(axis)
     if angle is None:
         angle = norm
-    eps = 1e2 * jnp.finfo(axis.dtype).eps
+    eps = jnp.sqrt(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 jnp.where(norm < eps, jnp.eye(3), R1 + R2 + R3)  # if axis=0, no rotation
 
 
+def rotate_vector_to_vector(u, v):
+    """
+    Computes a rotation matrix that rotates vector u onto vector v.
+
+    Numerically stable when vectors are parallel or anti-parallel, and avoids
+    both NaN gradient when parallel/antiparallel, and also avoids zero
+    gradient (wrt v)
+
+    Parameters
+    ----------
+    u : array-like, shape(3,)
+        first vector, to be rotated onto second vector v
+    v :  array-like, shape(3,)
+        vector to rotate u onto
+
+    Returns
+    -------
+    rotmat : ndarray, shape(3,3)
+        Matrix to rotate points in cartesian (X,Y,Z) coordinates.
+
+    """
+    u = safenormalize(u)
+    v = safenormalize(v)
+
+    # 2. Compute the dot product
+    dot = jnp.dot(u, v)
+
+    # things needed for near antiparallel case
+    axis = jnp.where(
+        jnp.all(jnp.less(jnp.abs(v[0]), jnp.abs(v[1]))),
+        jnp.array([0, -v[2], v[1]]),
+        jnp.array([-v[2], 0, v[0]]),
+    )
+    axis = safenormalize(axis)
+    K = jnp.array(
+        [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
+    )
+    R_antiparallel = jnp.eye(3) + 2 * jnp.dot(K, K)
+
+    # Handle general cases using the half-vector (bisector) method
+    # This prevents the division-by-zero errors common in standard cross-product methods
+    h = u + v
+    h = safenormalize(h)
+
+    # Calculate quaternion components q = [w, x, y, z]
+    # Equivalent to rotating around (u x h) by the half-angle
+    qw = jnp.dot(u, h)
+    qx, qy, qz = jnp.cross(u, h)
+
+    # Convert quaternion to a standard 3x3 rotation matrix
+    R = jnp.array(
+        [
+            [
+                1 - 2 * (qy**2 + qz**2),
+                2 * (qx * qy - qw * qz),
+                2 * (qx * qz + qw * qy),
+            ],
+            [
+                2 * (qx * qy + qw * qz),
+                1 - 2 * (qx**2 + qz**2),
+                2 * (qy * qz - qw * qx),
+            ],
+            [
+                2 * (qx * qz - qw * qy),
+                2 * (qy * qz + qw * qx),
+                1 - 2 * (qx**2 + qy**2),
+            ],
+        ]
+    )
+    # where to return antiparallel R if needed, else normal R
+    return jnp.where(jnp.allclose(dot, -1.0, atol=1e-8, rtol=1e-8), R_antiparallel, R)
+
+
 def xyz2rpz(pts):
     """Transform points from cartesian (X,Y,Z) to polar (R,phi,Z) form.
 

tests/test_coils.py

@@ -672,7 +672,7 @@ def test_properties(self):
             ).reshape((4, 1, 3)),
         )
         np.testing.assert_allclose([dat["curvature"] for dat in data], 1 / 2)
-        np.testing.assert_allclose([dat["torsion"] for dat in data], 0)
+        np.testing.assert_allclose([dat["torsion"] for dat in data], 0, atol=2e-16)
         T = [dat["frenet_tangent"] for dat in data]
         N = [dat["frenet_normal"] for dat in data]
         B = [dat["frenet_binormal"] for dat in data]

tests/test_curves.py

@@ -552,9 +552,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=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
+        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=5e-16)  # only in Y-Z plane
+        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=5e-16)  # only in X-Z plane
+        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=5e-16)  # only in X-Y plane
 
     @pytest.mark.unit
     def test_length(self):
@@ -763,9 +763,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=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
+        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=5e-16)  # only in Y-Z plane
+        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=5e-16)  # only in X-Z plane
+        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=5e-16)  # only in X-Y plane
 
     @pytest.mark.unit
     def test_length(self):