From 1a272b38d82afffe4ad02387c87d456df0f484e7 Mon Sep 17 00:00:00 2001 From: dpanici Date: Wed, 29 Jul 2026 13:05:00 -0400 Subject: [PATCH 01/12] attempt fix for nan --- desc/utils.py | 9 +++++++-- tests/test_objective_funs.py | 3 ++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/desc/utils.py b/desc/utils.py index 34088ebb01..df5995fd9f 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -998,14 +998,19 @@ def rotation_matrix(axis, angle=None): """ axis = jnp.asarray(axis) norm = safenorm(axis) - axis = safenormalize(axis) if angle is None: angle = norm eps = 1e2 * jnp.finfo(axis.dtype).eps + no_rotation = norm < eps + # near-zero axis = no rotation; sanitize axis+angle so the unused where-branch + # never forms cos(inf)/0-norm and can't leak a nan gradient (forward or reverse) + axis = jnp.where(no_rotation, jnp.array([1.0, 0.0, 0.0], dtype=axis.dtype), axis) + angle = jnp.where(no_rotation, 0.0, angle) + axis = safenormalize(axis) 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 + return jnp.where(no_rotation, jnp.eye(3), R1 + R2 + R3) # if axis=0, no rotation def xyz2rpz(pts): diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 6ba349333e..5e8900c2f2 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -4077,7 +4077,8 @@ def test(normal): test([0, 0, -1]) # antiparallel test([0, 1e-4, 1]) # nearly parallel test([0, 1e-4, -1]) # nearly antiparallel - + test([0, 1e-8, 1]) # nearly parallel + test([0, 1e-8, -1]) # nearly antiparallel # have use_jit=True here to check that runs correctly with spline # coils, see PR #1656 obj = ObjectiveFunction(SurfaceQuadraticFlux(surf, ext_field), use_jit=True) From 0b7f12d391b1a80eaeeae6770cbcc6679dd79e2b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 19 Aug 2026 12:18:53 -0400 Subject: [PATCH 02/12] change eps to be sqrt eps to fix NaN gradient issue --- desc/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/utils.py b/desc/utils.py index df5995fd9f..1c903469c9 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -1000,7 +1000,7 @@ def rotation_matrix(axis, angle=None): norm = safenorm(axis) if angle is None: angle = norm - eps = 1e2 * jnp.finfo(axis.dtype).eps + eps = jnp.sqrt(jnp.finfo(axis.dtype).eps) no_rotation = norm < eps # near-zero axis = no rotation; sanitize axis+angle so the unused where-branch # never forms cos(inf)/0-norm and can't leak a nan gradient (forward or reverse) From 8052752898300d76b06a8dee29f2467cae9b507a Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 19 Aug 2026 16:17:15 -0400 Subject: [PATCH 03/12] add rotate_vector_to_vector function that can handle (without NaN or surprise 0 gradients) rotations from one vector to another --- desc/compute/_curve.py | 83 +++++++--------------------------- desc/utils.py | 86 ++++++++++++++++++++++++++++++++++++ tests/test_objective_funs.py | 4 +- 3 files changed, 104 insertions(+), 69 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index ef2c5b3aca..bf283386eb 100644 --- a/desc/compute/_curve.py +++ b/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 diff --git a/desc/utils.py b/desc/utils.py index 1c903469c9..e4066d2885 100644 --- a/desc/utils.py +++ b/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 axis when the angle is nearly zero (specifically, when + the angle Date: Wed, 19 Aug 2026 16:23:22 -0400 Subject: [PATCH 04/12] update changelog --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 389fc8d989..db6ad8b6de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,8 @@ Bug Fixes - Bug (#2120) in ``desc.magnetic_fields.OmnigenousField`` computing NaNs when the ``B_lm`` corresponded to flat magnetic wells fixed by ``interpax`` ``v0.3.14``, updated tests to exercise this. - Fixes bug in ``reactor_QA.py`` script where the current profile was allowed to have a nonzero rho^1 component, which resulted in an unphysical profile near-axis. - Updates ``"reactor_QA"`` in ``desc.examples`` to fix this. Note that if using ``"reactor_QA"`` example from ``v0.16.0`` until this fix, the current profile in that example has this issue. +- 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. + Breaking Changes From f769aef1755b730372dbc88f13db9d29e1cadbe4 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 19 Aug 2026 16:26:41 -0400 Subject: [PATCH 05/12] clean up comments --- desc/utils.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/desc/utils.py b/desc/utils.py index e4066d2885..65684bc0e0 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -984,7 +984,7 @@ 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 axis when the angle is nearly zero (specifically, when + the angle when the angle is nearly zero (specifically, when the angle Date: Wed, 19 Aug 2026 16:36:24 -0400 Subject: [PATCH 06/12] add test --- tests/test_geometry.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 952b80b160..c38c6d673d 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -3,7 +3,15 @@ import numpy as np import pytest -from desc.utils import rotation_matrix, rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.utils import ( + rotate_vector_to_vector, + rotation_matrix, + rpz2xyz, + rpz2xyz_vec, + safenormalize, + xyz2rpz, + xyz2rpz_vec, +) @pytest.mark.unit @@ -14,6 +22,23 @@ def test_rotation_matrix(): np.testing.assert_allclose(A, At, atol=1e-10) +@pytest.mark.unit +def test_rotate_vector_to_vector(): + """Test calculation of rotation matrices given two vectors.""" + u = safenormalize(np.random.rand(3)) + v = safenormalize(np.random.rand(3) + np.array([0.1, 0.2, 0.3])) + A = rotate_vector_to_vector(u, v) + np.testing.assert_allclose(v, u @ A.T) + # edge case: vectors parallel + v = u + A = rotate_vector_to_vector(u, v) + np.testing.assert_allclose(v, u @ A.T) + # edge case: vectors antiparallel + v = -u + A = rotate_vector_to_vector(u, v) + np.testing.assert_allclose(v, u @ A.T) + + @pytest.mark.unit def test_xyz2rpz(): """Test converting between cartesian and polar coordinates.""" From e82024042a691a498d846651c65ccd607a16683b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 20 Aug 2026 14:06:42 -0400 Subject: [PATCH 07/12] fix error --- desc/utils.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/desc/utils.py b/desc/utils.py index 65684bc0e0..98a289800c 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -1007,19 +1007,14 @@ def rotation_matrix(axis, angle=None): """ axis = jnp.asarray(axis) norm = safenorm(axis) + axis = safenormalize(axis) if angle is None: angle = norm eps = jnp.sqrt(jnp.finfo(axis.dtype).eps) - no_rotation = norm < eps - # near-zero axis = no rotation; sanitize axis+angle so the unused where-branch - # never forms cos(inf)/0-norm and can't leak a nan gradient (forward or reverse) - axis = jnp.where(no_rotation, jnp.array([1.0, 0.0, 0.0], dtype=axis.dtype), axis) - angle = jnp.where(no_rotation, 0.0, angle) - axis = safenormalize(axis) 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(no_rotation, jnp.eye(3), R1 + R2 + R3) # if axis=0, no rotation + return jnp.where(norm < eps, jnp.eye(3), R1 + R2 + R3) # if axis=0, no rotation def rotate_vector_to_vector(u, v): From f41f74d43014f2ce27169c81d3c48a246dfc80c6 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 21 Aug 2026 16:48:47 -0400 Subject: [PATCH 08/12] fix test which was due to too tight precision tol in compute_hess_scale --- desc/geometry/core.py | 2 +- desc/geometry/curve.py | 4 ++-- desc/optimize/utils.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 38c70f764b..0d414c3cd8 100644 --- a/desc/geometry/core.py +++ b/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 diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 54532ff336..79b9f36056 100644 --- a/desc/geometry/curve.py +++ b/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 " diff --git a/desc/optimize/utils.py b/desc/optimize/utils.py index 41845ad392..729d7cf4fd 100644 --- a/desc/optimize/utils.py +++ b/desc/optimize/utils.py @@ -502,7 +502,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: From 0ee1273e008ebf84cd07a9364bc687775f578d50 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 21 Aug 2026 16:53:40 -0400 Subject: [PATCH 09/12] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 96d576acd9..057f7862d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ 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. v0.17.3 ------- From f8eceacfc82df544f9c34541d53e1e4b6150659c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 21 Aug 2026 20:38:34 -0400 Subject: [PATCH 10/12] relax test tol which was too tight before --- tests/test_curves.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_curves.py b/tests/test_curves.py index 5c089b1ed6..fdd4663f40 100644 --- a/tests/test_curves.py +++ b/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): From 389f4c6589cce411c12f73a7ec4588f92c3104da Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 23 Aug 2026 11:48:47 -0400 Subject: [PATCH 11/12] reduce test tol --- tests/test_coils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_coils.py b/tests/test_coils.py index 9eab47d8ba..5a3d3db392 100644 --- a/tests/test_coils.py +++ b/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] From de0392352aa89f1077453c498193a5152627af32 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 3 Sep 2026 14:08:56 -0400 Subject: [PATCH 12/12] tighten allclose tol, add tests on gradient against finite diffs --- desc/utils.py | 2 +- tests/test_geometry.py | 81 ++++++++++++++++++++++++++++++++++++ tests/test_objective_funs.py | 4 ++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/desc/utils.py b/desc/utils.py index 98a289800c..24fa5862cf 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -1087,7 +1087,7 @@ def rotate_vector_to_vector(u, v): ] ) # where to return antiparallel R if needed, else normal R - return jnp.where(jnp.allclose(dot, -1.0), R_antiparallel, R) + return jnp.where(jnp.allclose(dot, -1.0, atol=1e-8, rtol=1e-8), R_antiparallel, R) def xyz2rpz(pts): diff --git a/tests/test_geometry.py b/tests/test_geometry.py index c38c6d673d..0025da7705 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -2,7 +2,9 @@ import numpy as np import pytest +from jax.test_util import check_grads +from desc.backend import jnp from desc.utils import ( rotate_vector_to_vector, rotation_matrix, @@ -39,6 +41,85 @@ def test_rotate_vector_to_vector(): np.testing.assert_allclose(v, u @ A.T) +@pytest.mark.unit +def test_rotate_vector_to_vector_gradients(): + """Test grad of rotate_vector_to_vector against finite diffs.""" + # Standard non-aligned case + u = jnp.array([1.0, 2.0, 3.0]) + v = jnp.array([4.0, 1.0, -2.0]) + + # Validates order 1 derivatives wrt positional args (0: u, 1: v) + # using forward-mode ('fwd') and reverse-mode ('rev') AD against finite differences + check_grads( + rotate_vector_to_vector, + args=(u, v), + order=1, + modes=["fwd", "rev"], + eps=1e-4, + atol=1e-3, + rtol=1e-3, + ) + + # Parallel case + u_par = jnp.array([1.0, 0.0, 0.0]) + v_par = jnp.array([2.0, 0.0, 0.0]) + check_grads( + rotate_vector_to_vector, + args=(u_par, v_par), + order=1, + modes=["fwd", "rev"], + eps=1e-4, + atol=1e-3, + ) + # Nearly Parallel case + u_par = jnp.array([1.0, 1e-5, 0.0]) + v_par = jnp.array([2.0, 0.0, 0.0]) + check_grads( + rotate_vector_to_vector, + args=(u_par, v_par), + order=1, + modes=["fwd", "rev"], + eps=1e-4, + atol=1e-3, + ) + + # Nearly Antiparallel case + u_anti = jnp.array([1.0, 1e-2, 0.0]) + v_anti = jnp.array([-1.0, 0.0, 0.0]) + check_grads( + rotate_vector_to_vector, + args=(u_anti, v_anti), + order=1, + modes=["fwd", "rev"], + eps=1e-6, + atol=1e-6, + ) + + # Nearly Antiparallel case + u_anti = jnp.array([1.0, 1e-5, 0.0]) + v_anti = jnp.array([-1.0, 0.0, 0.0]) + check_grads( + rotate_vector_to_vector, + args=(u_anti, v_anti), + order=1, + modes=["fwd", "rev"], + eps=1e-6, + atol=1e-6, + ) + + # Antiparallel case + u_anti = jnp.array([1.0, 0.0, 0.0]) + v_anti = jnp.array([-1.0, 0.0, 0.0]) + check_grads( + rotate_vector_to_vector, + args=(u_anti, v_anti), + order=1, + modes=["fwd", "rev"], + eps=1e-6, + atol=1e-6, + ) + + @pytest.mark.unit def test_xyz2rpz(): """Test converting between cartesian and polar coordinates.""" diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 7feb8ea42c..14d639fb2f 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -4111,6 +4111,10 @@ def test(normal): test([0, 0, -1]) # antiparallel test([0, 1e-4, 1]) # nearly parallel test([0, 1e-4, -1]) # nearly antiparallel + test([0, 1e-6, 1]) # nearly parallel + test([0, 1e-6, -1]) # nearly antiparallel + test([0, 1e-8, 1]) # nearly parallel + test([0, 1e-8, -1]) # nearly antiparallel test([0, 1e-12, 1]) # nearly parallel test([0, 1e-12, -1]) # nearly antiparallel # have use_jit=True here to check that runs correctly with spline