How to pass kwargs to `root_find`

[jpbrodrick89]

I haven't trawled through the source code but couldn't find this documented. It would be nice to be able to pass kwargs such as max_steps and tags to root_find when doing an implicit solve but not sure how this would work as max_steps coincides with a diffrax kwargs and you would often want very different values for the two.

[johannahaffner]

The implicit solvers have a root_find_max_steps attribute, which is then passed to Optimistix e.g. here:

https://github.com/patrick-kidger/diffrax/blob/d3c1430b76e7d0158536823724371d7b26f5db8f/diffrax/_solver/implicit_euler.py#L85

I imagine that tags could equally be a solver attribute added to their base class here

https://github.com/patrick-kidger/diffrax/blob/d3c1430b76e7d0158536823724371d7b26f5db8f/diffrax/_solver/base.py#L207-L215

and then implemented in the subclasses, defaulting to an empty frozenset.

[jpbrodrick89]

Thanks @johannahaffner don't know how I missed that in the docs, it was plastered everywhere 🤦‍♂️I'm going to pretend I must have been looking at an old version 🙄happy to submit a PR on the tags front. 🙂👍

[patrick-kidger]

Do note that such a tag describes the structure of the particular root find problem that's being tackled, which is kind of complicated.

That aside, it might be easier for you to directly pick a linear solver instead, that just quietly assumes the particular structure.

[jpbrodrick89]

I'm looking to solve a nonlinear tridiagonal system (diffusion equation with nonlinear diffusion coefficient), so would need to use a root finder of some sort. However, when I tried to set the linear_solver kwarg in optx.Newton I get an error if I don't add tags as discussed in https://github.com/patrick-kidger/lineax/issues/149

[patrick-kidger]

Right, but - at least in principle - one could provide a solver that perform a tridiagonal solve whilst completely ignoring tags.

I'm just concious that the tags here would refer to the structure of the Jacobian of any of the following functions:

https://github.com/patrick-kidger/diffrax/blob/d3c1430b76e7d0158536823724371d7b26f5db8f/diffrax/_solver/implicit_euler.py#L19

https://github.com/patrick-kidger/diffrax/blob/d3c1430b76e7d0158536823724371d7b26f5db8f/diffrax/_solver/runge_kutta.py#L250

https://github.com/patrick-kidger/diffrax/blob/d3c1430b76e7d0158536823724371d7b26f5db8f/diffrax/_solver/runge_kutta.py#L271

and exposing this to a user would definitely be a footgun. Hence why I'm thinking that to support the sufficiently-bold power user (you) it might be better to take the solver-based workaround instead.

[jpbrodrick89]

Ah, yes, I see, exposing these for general IRK solvers could get confusing and error-prone fast. If you try and implicitly solve both stages of a RK2 at once the tridiagonal operator would become pentadiagonal and your promise gets broken. I think using a custom diffrax solver that accepts tags is probably what I'll try for now, but it may be fairly safe to expose this to ImplicitEuler only?