diffrax
diffrax copied to clipboard
NewtonNonlinearSolver returns nan when initialized with root value
It seems like the Newton solver fails when initialized with the root value at which also $f'(x) = 0$. A minimal working example:
from diffrax import NewtonNonlinearSolver
def fun(x, args): return x**3
solver = NewtonNonlinearSolver(rtol=1e-3, atol=1e-6)
sol = solver(fun, 0., None)
print(sol.root, fun(sol.root, None)) # nan nan
It looks like the conditions here would catch it, where it not for at_least_two
iterations. Of course, Newtons method can't do much in case of $f'(x) = 0$, but I would expect the solver to stop when it figures out it is already at the root.
Yeah, this is due to the linear solve Ax=b
basically turning into 0 @ x = 0
, which isn't defined.
The current nonlinear code is due to get overhauled fairly shortly! I'll mark this as a bug until then.