Fix for NaN in reverse mode gradient of FourierPlanarCoil by dpanici · Pull Request #2277 · PlasmaControl/DESC

dpanici · 2026-07-29T17:10:28Z

Resolves #2276

Fix was mainly in switching from axis-angle to quaternion for general case, and being sure that the antiparllel edge-case is handled by making a rotation of 180 degrees around a perpendicular vector composed of the input vector (ensuring that the derivative does not arbitrarily go to exactly zero wrt normal at that edge case, which could stall optijmizations even if the deriv was not NaN there)

from desc.coils import FourierPlanarCoil
from desc.objectives import ObjectiveFunction, CoilLength
from desc.optimize import Optimizer
import numpy as np

def test(normal):
    coil = FourierPlanarCoil(normal=normal)
    opt = Optimizer("lsq-exact")

    obj = ObjectiveFunction(CoilLength(coil))
    obj.build(verbose=0)
    g = obj.grad(obj.x(coil))
    return np.any(np.isnan(g))
for comp in np.concatenate([np.array([0.0]), np.logspace(-16,-6,11)]):
    was_nan = test([comp,comp,1.0])
    print(f"Did Normal of [{comp:1.2e}, {comp:1.2e},1.0] result in nan gradient?  {was_nan}")

On master:

Did Normal of [0.00e+00, 0.00e+00,1.0] result in nan gradient?  False
Did Normal of [1.00e-16, 1.00e-16,1.0] result in nan gradient?  True
Did Normal of [1.00e-15, 1.00e-15,1.0] result in nan gradient?  True
Did Normal of [1.00e-14, 1.00e-14,1.0] result in nan gradient?  True
Did Normal of [1.00e-13, 1.00e-13,1.0] result in nan gradient?  True
Did Normal of [1.00e-12, 1.00e-12,1.0] result in nan gradient?  True
Did Normal of [1.00e-11, 1.00e-11,1.0] result in nan gradient?  True
Did Normal of [1.00e-10, 1.00e-10,1.0] result in nan gradient?  True
Did Normal of [1.00e-09, 1.00e-09,1.0] result in nan gradient?  True
Did Normal of [1.00e-08, 1.00e-08,1.0] result in nan gradient?  True
Did Normal of [1.00e-07, 1.00e-07,1.0] result in nan gradient?  False
Did Normal of [1.00e-06, 1.00e-06,1.0] result in nan gradient?  False

On this PR:

Did Normal of [0.00e+00, 0.00e+00,1.0] result in nan gradient?  False
Did Normal of [1.00e-16, 1.00e-16,1.0] result in nan gradient?  False
Did Normal of [1.00e-15, 1.00e-15,1.0] result in nan gradient?  False
Did Normal of [1.00e-14, 1.00e-14,1.0] result in nan gradient?  False
Did Normal of [1.00e-13, 1.00e-13,1.0] result in nan gradient?  False
Did Normal of [1.00e-12, 1.00e-12,1.0] result in nan gradient?  False
Did Normal of [1.00e-11, 1.00e-11,1.0] result in nan gradient?  False
Did Normal of [1.00e-10, 1.00e-10,1.0] result in nan gradient?  False
Did Normal of [1.00e-09, 1.00e-09,1.0] result in nan gradient?  False
Did Normal of [1.00e-08, 1.00e-08,1.0] result in nan gradient?  False
Did Normal of [1.00e-07, 1.00e-07,1.0] result in nan gradient?  False
Did Normal of [1.00e-06, 1.00e-06,1.0] result in nan gradient?  False

dpanici · 2026-07-29T17:24:23Z

    axis = jnp.asarray(axis)
    norm = safenorm(axis)
    if angle is None:
        angle = norm
    eps = 1e3 * jnp.finfo(axis.dtype).eps # 1e3 instead of 1e2
    no_rotation = norm < eps

If the eps is made an order of magnitude larger, the 1e-13 results in no NaN.

So the issue comes when the axis magnitude from the cross of the normal and the Z-axis is just large enough that the norm becomes greater than our threshold for zero, but is small enough still that we don't trigger some other conditional on the magnitude of some intermediate quantity.

So I guess there is some inconsistency in thresholding somewhere in this.

I think it is that safearccos will return an infinite for the angle when their dotprod evaluates to 1 s.t. abs(dotprod)==1 is True

DESC/desc/utils.py

Line 945 in c2dfad1

safe_x = jnp.where(jnp.abs(x) == 1, 0, x)

and we then rely on that only happening when the Zaxis and normal are nearly parallel, so their cross product's norm is smaller than 1e2*eps in rotation_matrix and thus we never take cos(inf) but instead just return the identity matrix

DESC/desc/utils.py

Line 1008 in c2dfad1

    
           return jnp.where(norm < eps, jnp.eye(3), R1 + R2 + R3)  # if axis=0, no rotation

or, if the axis is parallel to -Z, the correct reflection

DESC/desc/compute/_curve.py

Lines 228 to 232 in c2dfad1

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

I know the answer is in the following data, just need to wrap my head around it:

from desc.coils import FourierPlanarCoil
from desc.objectives import ObjectiveFunction, CoilLength
from desc.optimize import Optimizer
import numpy as np
from desc.utils import dot
from desc.utils import safearccos,safenormalize
from desc.backend import jnp
zaxis = jnp.array([0,0,1])


def test(normal):
    coil = FourierPlanarCoil(normal=normal,basis="xyz")
    print(f"norm of normal, pre-normalization: [{np.linalg.norm(coil.normal):1.16e}]")
    print(f"coil normal after normalization: [{coil.normal[0]:1.16e},{coil.normal[1]:1.16e},{coil.normal[2]:1.16e}]")

    opt = Optimizer("lsq-exact")

    obj = ObjectiveFunction(CoilLength(coil))
    obj.build(verbose=0)
    g = obj.grad(obj.x(coil))
    return np.any(np.isnan(g)), coil
for comp in np.concatenate([np.array([0.0]), np.logspace(-16,-6,11)]):
    was_nan, coil = test([comp,comp,1.0])
    print(f"Did Normal of [{comp:1.2e}, {comp:1.2e},1.0] result in nan gradient?  {was_nan}")
    dotprod = dot(zaxis,coil.normal)
    print(f"dot of normal and z axis == 1: {jnp.abs(dotprod) == 1}")
    print(f"arccos of dotprod: {safearccos(dotprod)}")

norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [0.0000000000000000e+00,0.0000000000000000e+00,1.0000000000000000e+00]
Did Normal of [0.00e+00, 0.00e+00,1.0] result in nan gradient?  False
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [9.9999999999999998e-17,9.9999999999999998e-17,1.0000000000000000e+00]
Did Normal of [1.00e-16, 1.00e-16,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000001e-15,1.0000000000000001e-15,1.0000000000000000e+00]
Did Normal of [1.00e-15, 1.00e-15,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000000e-14,1.0000000000000000e-14,1.0000000000000000e+00]
Did Normal of [1.00e-14, 1.00e-14,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000000e-13,1.0000000000000000e-13,1.0000000000000000e+00]
Did Normal of [1.00e-13, 1.00e-13,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [9.9999999999999998e-13,9.9999999999999998e-13,1.0000000000000000e+00]
Did Normal of [1.00e-12, 1.00e-12,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [9.9999999999999994e-12,9.9999999999999994e-12,1.0000000000000000e+00]
Did Normal of [1.00e-11, 1.00e-11,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000000e-10,1.0000000000000000e-10,1.0000000000000000e+00]
Did Normal of [1.00e-10, 1.00e-10,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000001e-09,1.0000000000000001e-09,1.0000000000000000e+00]
Did Normal of [1.00e-09, 1.00e-09,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [1.0000000000000000e-08,1.0000000000000000e-08,1.0000000000000000e+00]
Did Normal of [1.00e-08, 1.00e-08,1.0] result in nan gradient?  True
dot of normal and z axis == 1: True
arccos of dotprod: inf
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [9.9999999999999003e-08,9.9999999999999003e-08,9.9999999999999001e-01]
Did Normal of [1.00e-07, 1.00e-07,1.0] result in nan gradient?  False
dot of normal and z axis == 1: False
arccos of dotprod: 1.4136482746161737e-07
norm of normal, pre-normalization: [1.0000000000000000e+00]
coil normal after normalization: [9.9999999999900003e-07,9.9999999999900003e-07,9.9999999999900013e-01]
Did Normal of [1.00e-06, 1.00e-06,1.0] result in nan gradient?  False
dot of normal and z axis == 1: False
arccos of dotprod: 1.4141194121979817e-06

OK, so issue is that when a normal is almost aligned with the z-axis (with a z-component of 1 but some additional other tiny but nonzero component), but the misalignment is so small that the norm of the vector is 1 to machine precision, then when we go to normalize the vector when it is set in the FourierPlanarCoil, the normalization does not change the z-axis of the normal.
So, that sets the scene for the bug: when the z-axis component of a FourierPlanarCoil's normal is exactly 1.0, but it has nonzero off-z-axis components to it (tiny enough that the norm of the vector is 1.0 to machine precision, thus the normalization of the normal does not reduce the z-axis of it to below 1.0)

Then, we can have this inconsistency occur where we in one part of the compute logic, (the safearccos), the checked quantity (dotprod of normal and Zaxis) is 1.0 to machine precision and thus we trigger the safe route there and return inf
but then inside of rotation_matrix, the norm of the axis (cross product of zaxis and normal) is NOT small enough to trigger the safe route there, and we end up returning things that result in nans (I THINK, I am not sure)

github-actions · 2026-07-29T17:33:07Z

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.06 %    |     4.237e+03      |     4.240e+03      |     2.67     |       29.91        |       28.47        |
  test_proximal_jac_w7x_with_eq_update   |    0.18 %    |     6.807e+03      |     6.819e+03      |    12.02     |       153.25       |       153.64       |
  test_proximal_freeb_jac                |    0.02 %    |     1.355e+04      |     1.355e+04      |     2.49     |       80.34        |       78.61        |
  test_proximal_freeb_jac_blocked        |    0.06 %    |     7.878e+03      |     7.883e+03      |     4.45     |       69.00        |       67.56        |
  test_proximal_freeb_jac_batched        |    0.31 %    |     7.864e+03      |     7.888e+03      |    24.24     |       68.14        |       68.04        |
  test_proximal_jac_ripple               |    0.30 %    |     3.811e+03      |     3.823e+03      |    11.37     |       53.20        |       54.69        |
  test_proximal_jac_ripple_bounce1d      |   -0.87 %    |     4.020e+03      |     3.985e+03      |    -34.84    |       67.47        |       69.18        |
  test_eq_solve                          |   -0.16 %    |     1.832e+03      |     1.829e+03      |    -3.00     |       52.20        |       52.76        |
  test_objective_quadratic_flux_jac      |    0.14 %    |     1.898e+03      |     1.901e+03      |     2.73     |       34.10        |       35.18        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

dpanici · 2026-08-19T18:54:44Z

Updated test and solution

from desc.coils import FourierPlanarCoil
from desc.objectives import ObjectiveFunction, CoilLength
from desc.optimize import Optimizer
import numpy as np
from desc.utils import dot, cross
from desc.utils import safearccos,safenormalize, safenorm
from desc.backend import jnp
zaxis = jnp.array([0,0,1])


def test(normal):
    coil = FourierPlanarCoil(normal=normal,basis="xyz")
    print(f"norm of normal, pre-normalization: [{np.linalg.norm(coil.normal):1.16e}]")
    print(f"coil normal after normalization: [{coil.normal[0]:1.16e},{coil.normal[1]:1.16e},{coil.normal[2]:1.16e}]")

    opt = Optimizer("lsq-exact")

    obj = ObjectiveFunction(CoilLength(coil))
    obj.build(verbose=0)
    g = obj.grad(obj.x(coil))
    return np.any(np.isnan(g)), coil
for comp in np.concatenate([np.array([0.0]), np.logspace(-16,-6,11)]):
    normal = np.array([comp,comp,1.0])
    was_nan, coil = test(normal)
    axis = cross(zaxis, normal)
    norm = safenorm(axis)

    eps = 1e8 * jnp.finfo(axis.dtype).eps
    no_rotation = norm < eps
    print(f"Did Normal of [{comp:1.2e}, {comp:1.2e},1.0] result in nan gradient?  {was_nan}")
    dotprod = dot(zaxis,coil.normal)
    
    print(f"arccos of dotprod: {safearccos(dotprod)}")
    print(f"zaxis x normal = rotation axis: {axis}")
    print(f"norm of rotation axis: {norm}")
    print(f"dot of normal and z axis == 1: {jnp.abs(dotprod) == 1}")
    print(f"is norm of rotation axis < eps?: {no_rotation}")
    print("#"*15)

We need a consistency between the two "safeness" logic checks: is the dot of the normal and zaxiz ==1 (or -1), and is the norm of the rotation axis < eps

Because the axis could contain two components which are < sqrt(eps) ~ 1e-8, when we take the norm of axis we square it and they become comparable to zero.

Gist of it is, we should use sqrt(eps) instead of eps when we try to determine if an axis' norm is zero or not. Making that change results in no NaN.

Outstanding issue: the gradient wrt the normal vector before was ZERO for normal=[eps,eps,1] with eps<1e-15 and then NaN for 1e-15>eps>1e-8. With the change in this PR, the gradient is now just ZERO for normal=[eps,eps,1] with eps<1e-8. This can cause issues in optimization if your coil is initialized as horizontal

normal = [0,0,1.0]

coil = FourierPlanarCoil(center = [10,1e-1,-3], r_n=0.5,normal=normal,basis="xyz", current=1e6)

from desc.objectives import SurfaceQuadraticFlux, FixCoilCurrent, FixParameters
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
surf = FourierRZToroidalSurface() # should be R=10 r=1 surface

# opt problem  fix the geometry of coil, the current, and its center location of the coil
# to minimize  Bn, should rotate from horizontal to vertical

cons = (FixCoilCurrent(coil), FixParameters(coil, {"r_n":True,"center":True}), FixParameters(surf))
obj = ObjectiveFunction(SurfaceQuadraticFlux(surf,coil,eval_grid=LinearGrid(N=10,M=10), field_grid=LinearGrid(N=10)))
opt = Optimizer("lsq-exact")
opt.optimize((surf,coil),objective=obj, constraints=cons,verbose=3,ftol=0,gtol=0);

This stalls with 0 gradient initially

normal = [1e-7,0,1.0]

coil = FourierPlanarCoil(center = [10,1e-1,-3], r_n=0.5,normal=normal,basis="xyz", current=1e6)

from desc.objectives import SurfaceQuadraticFlux, FixCoilCurrent, FixParameters
from desc.geometry import FourierRZToroidalSurface
from desc.grid import LinearGrid
surf = FourierRZToroidalSurface() # should be R=10 r=1 surface

# opt problem  fix the geometry of coil, the current, and its center location of the coil
# to minimize  Bn, should rotate from horizontal to vertical

cons = (FixCoilCurrent(coil), FixParameters(coil, {"r_n":True,"center":True}), FixParameters(surf))
obj = ObjectiveFunction(SurfaceQuadraticFlux(surf,coil,eval_grid=LinearGrid(N=10,M=10), field_grid=LinearGrid(N=10)))
opt = Optimizer("lsq-exact")
opt.optimize((surf,coil),objective=obj, constraints=cons,verbose=3,ftol=0,gtol=0);

This does not stall out and correctly moves the coil normal so that the dipole moment is parallel to the surface (minimizing Bn)

…surprise 0 gradients) rotations from one vector to another

dpanici · 2026-08-19T20:20:34Z

OK this works: switched to using quaternions for general rotation from one vector onto another, and for the antiparallel case made sure to use the vector that we want the deriv wrt (normal vector) to be nonzero so that case is covered too. This works for both parallel to z and antiparallel to z without nan OR zero gradient (for cases where e.g. coil is horizontal but we would expect the gradient to be nonzero

… into dp/fix-nan-planarcoil

attempt fix for nan

1a272b3

change eps to be sqrt eps to fix NaN gradient issue

0b7f12d

dpanici and others added 2 commits August 19, 2026 16:17

add rotate_vector_to_vector function that can handle (without NaN or …

8052752

…surprise 0 gradients) rotations from one vector to another

Merge branch 'master' into dp/fix-nan-planarcoil

ea858c8

dpanici marked this pull request as ready for review August 19, 2026 20:20

dpanici added 2 commits August 19, 2026 16:23

update changelog

b712fe8

Merge branch 'dp/fix-nan-planarcoil' of github.com:PlasmaControl/DESC…

27feda3

… into dp/fix-nan-planarcoil

dpanici requested review from a team, YigitElma, ddudt, f0uriest, lkadz, rahulgaur104, singh-jaydeep and unalmis and removed request for a team August 19, 2026 20:24

dpanici added 2 commits August 19, 2026 16:26

clean up comments

f769aef

add test

daf4e45

Navigation Menu

Uh oh!

Fix for NaN in reverse mode gradient of FourierPlanarCoil - #2277

Fix for NaN in reverse mode gradient of FourierPlanarCoil#2277
dpanici wants to merge 8 commits into
masterfrom
dp/fix-nan-planarcoil

dpanici commented Jul 29, 2026 •

edited

Loading

Uh oh!

dpanici commented Jul 29, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 29, 2026 •

edited

Loading

Uh oh!

dpanici commented Aug 19, 2026

Uh oh!

dpanici commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

dpanici commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dpanici commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Memory benchmark result

Uh oh!

dpanici commented Aug 19, 2026

Uh oh!

dpanici commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

dpanici commented Jul 29, 2026 •

edited

Loading

dpanici commented Jul 29, 2026 •

edited

Loading

github-actions Bot commented Jul 29, 2026 •

edited

Loading