Skip to content

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

Open
dpanici wants to merge 17 commits into
masterfrom
dp/fix-nan-planarcoil
Open

Fix for NaN in reverse mode gradient of FourierPlanarCoil and Hessian Scaling Issue for small problems#2277
dpanici wants to merge 17 commits into
masterfrom
dp/fix-nan-planarcoil

Conversation

@dpanici

@dpanici dpanici commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

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)

Fixes

  • NaN gradient occuring for horizontal coils
  • zero gradient occuring for horizontal coils which arrests the optimization there incorrectly
  • Hessian scaling issue for tiny problems (small # DOFs) which popped up in the tests
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

dpanici commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author
    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

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

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

github-actions Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.44 %    |     4.224e+03      |     4.242e+03      |    18.55     |       30.80        |       27.04        |
  test_proximal_jac_w7x_with_eq_update   |    0.27 %    |     6.807e+03      |     6.826e+03      |    18.62     |       132.22       |       132.26       |
  test_proximal_freeb_jac                |   -0.04 %    |     1.354e+04      |     1.353e+04      |    -5.12     |       81.53        |       81.91        |
  test_proximal_freeb_jac_blocked        |   -0.16 %    |     7.868e+03      |     7.855e+03      |    -12.39    |       74.32        |       73.90        |
  test_proximal_freeb_jac_batched        |   -0.44 %    |     7.893e+03      |     7.858e+03      |    -34.99    |       74.81        |       73.61        |
  test_proximal_jac_ripple               |   -1.86 %    |     3.800e+03      |     3.730e+03      |    -70.56    |       50.91        |       49.39        |
  test_proximal_jac_ripple_bounce1d      |    0.03 %    |     3.782e+02      |     3.783e+02      |     0.11     |        3.23        |        3.20        |
  test_eq_solve                          |   -0.09 %    |     1.827e+03      |     1.826e+03      |    -1.73     |       45.82        |       46.29        |
  test_objective_quadratic_flux_jac      |    0.25 %    |     1.891e+03      |     1.896e+03      |     4.76     |       33.33        |       32.89        |

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

@dpanici

dpanici commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator Author

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)

@dpanici

dpanici commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator Author

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

@dpanici
dpanici marked this pull request as ready for review August 19, 2026 20:20
@dpanici
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

dpanici commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator Author

The test failing is not due to this PR really, as the main issue is that the rotation matrix derivative AD wrt normal vector z component can be really tiny/near zero in certain cases. This leads then to the Hessian scale (~1/diag of hess) being massive. Our epsilon check when Hess size is small (like in the test, only 8x8) allows things like 1/1e-14 to get thru which then completely ruins the scale of the problem.

This PR addresses that as well by just slightly increasing the epsilon under which scale_inv is converted to 1 (multiplies by a factor of 10). This was enough to overcome the numerical noise, it seems

Using this code:

scales=[]
scales_inv=[]

hess_mats=[]
for normal_comp in np.logspace(-16,-1,16):
    c = FourierPlanarCoil(normal=[normal_comp,1.0,0.0])

    angle = np.arccos(np.dot(np.array([0,0,1]), c.normal))
    print(f"{angle=:1.2e}")
    
    obj = ObjectiveFunction(CoilLength(c))
    obj.build(verbose=0)
    h=obj.hess(obj.x(c))
    hess_mats.append(h)
    scale,scale_inv = compute_hess_scale(h)
    print(scale)
    scales.append(scale.copy())
    scales_inv.append(scale_inv.copy())
    
import matplotlib.pyplot as plt
plt.figure()
plt.scatter(np.logspace(-16,-1,16), np.max(scales,axis=1))
plt.xlabel("off-normal component of normal vector")
plt.ylabel("Max xscale from hessian")
plt.xscale("log")
plt.yscale("log")

With the default epsilon we had before in compute_hess_scale, the scale was just noise (1/epsilon basically). It was proportional to the normal vector component which made no sense

image

Increasing epsilon by a factor of 10 in this PR at least makes it make sense, this now is proportional to one of the rotmat things which does not really matter anyways as that is always fixed in optimization. There is no more random numerical noise here too

image

Comment thread tests/test_curves.py
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

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

relaxed this bc 2e-16 was a too tight tolerance. The new method for rotations has a few more arithmetic operations than the old which just pushed it past the 2e-16 threshold (it was failing at like 2.4e-16 which is well within machine precision).

@codecov

codecov Bot commented Aug 23, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.34%. Comparing base (ad105c5) to head (b541eb0).

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #2277      +/-   ##
==========================================
- Coverage   94.35%   94.34%   -0.02%     
==========================================
  Files         101      101              
  Lines       29092    29082      -10     
==========================================
- Hits        27451    27438      -13     
- Misses       1641     1644       +3     
Files with missing lines Coverage Δ
desc/compute/_curve.py 100.00% <100.00%> (ø)
desc/geometry/core.py 95.30% <100.00%> (ø)
desc/geometry/curve.py 96.44% <100.00%> (ø)
desc/optimize/utils.py 95.54% <ø> (ø)
desc/utils.py 92.02% <100.00%> (-0.48%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@dpanici dpanici changed the title Fix for NaN in reverse mode gradient of FourierPlanarCoil Fix for NaN in reverse mode gradient of FourierPlanarCoil and Hessian Scaling Issue for small problems Aug 27, 2026

@YigitElma YigitElma left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couple nitpicks. It looks good but I am not sure if the edge cases still exist.

Comment thread desc/utils.py
Parameters
----------
u : array-like, shape(3,)
first vector, to be rotated onto second vector v

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These need to be in cartesian coordinates, right?

Comment thread desc/utils.py
u = safenormalize(u)
v = safenormalize(v)

# 2. Compute the dot product

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# 2. Compute the dot product
# Compute the dot product

Comment thread desc/optimize/utils.py
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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if we have something like,

eps = jnp.maximum(jnp.finfo(H.dtype).eps * max(H.shape), 1e-12)
scale_inv = jnp.where(scale_inv < eps, 1, scale_inv)

1e-12 or something else can be chosen as the minimum threshold.

Comment thread desc/geometry/core.py
Comment thread desc/utils.py Outdated
]
)
# where to return antiparallel R if needed, else normal R
return jnp.where(jnp.allclose(dot, -1.0), R_antiparallel, R)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Default rtol=1e-5 and atol=1e-8 here. Do the tests hit somewhere in between?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

FourierPlanarCoil still gives NaN gradient when axis is vertical

2 participants