Nonlinear max steps reached error on implicit solvers with PID step controller

[zhou-haoyang]

Hi!

I was using implicit solvers with a PID step size controller for solving stiff ODEs and I got the error of equinox.EquinoxRuntimeError: The maximum number of steps was reached in the nonlinear solver. The problem may not be solveable (e.g., a root-find on a function that has no roots), or you may need to increase max_steps. If I understand correctly, the PID controller should instead reject the step and try with a smaller step in this case.

Checking the code shown below, the integrator doesn't seem to check the solver_result returned by the solver, and the PID controller relies on y_error=inf as a sign of a failed implicit step. In my case, the root finder returns max_steps_reached but with a non-infinite y_error. This is probably the reason why the error got propagated to the final state.

diffrax/diffrax/_integrate.py

Lines 373 to 405 in 7f0001c

	(y, y_error, dense_info, solver_state, solver_result) = solver.step(
	terms,
	state.tprev,
	state.tnext,
	state.y,
	args,
	state.solver_state,
	state.made_jump,
	)
	# e.g. if someone has a sqrt(y) in the vector field, and dt0 is so large that
	# we get a negative value for y, and then get a NaN vector field. (And then
	# everything breaks.) See #143.
	y_error = jtu.tree_map(lambda x: jnp.where(jnp.isnan(x), jnp.inf, x), y_error)
	error_order = solver.error_order(terms)
	(
	keep_step,
	tprev,
	tnext,
	made_jump,
	controller_state,
	stepsize_controller_result,
	) = stepsize_controller.adapt_step_size(
	state.tprev,
	state.tnext,
	state.y,
	y,
	args,
	y_error,
	error_order,
	state.controller_state,
	)

diffrax/diffrax/_step_size_controller/pid.py

Lines 533 to 538 in 7f0001c

	# E.g. we failed an implicit step, so y_error=inf, so inv_scaled_error=0,
	# so factor=factormin, and we shrunk our step.
	# If we're using a PI or PID controller we shouldn't then force shrinking on
	# the next or next two steps as well!
	pred = (inv_scaled_error == 0) | jnp.isinf(inv_scaled_error)
	inv_scaled_error = jnp.where(pred, 1, inv_scaled_error)

In my case, applying this simple fix after solver.step seems to solve the problem, but I'm not sure if this is as intended.

y_error = jtu.tree_map(lambda x: jnp.where(is_okay(solver_result), x, jnp.inf), y_error)

[patrick-kidger]

So at least for Runge-Kutta methods then exactly the fix you have in mind is performed here:

diffrax/diffrax/_solver/runge_kutta.py

Lines 1194 to 1197 in 7f0001c

	y_error = jtu.tree_map(
	lambda _y_error: jnp.where(is_okay(result), _y_error, jnp.inf),
	y_error,
	) # i.e. an implicit step failed to converge

Which solver are you using? It might be that we left this out somewhere else!

[zhou-haoyang]

Ah that explains it. I was using implicit Euler, and it's missing there. I didn't run into this issue with the Runge-Kutta methods.

I can open a PR with a small patch if that sounds good to you :)

[patrick-kidger]

Yes please! :)

[zhou-haoyang]

Fixed in #753