Flexibility in how solves are terminated

[michael-0brien]

Hello again! First of all, amazing package. The modularity here is really fantastic.

For me, I think a key missing feature is some kind of flexibility in how to terminate solves. There are two aspects of this I think

1. I sometimes struggle with the interpretability of the cauchy termination criteria. My parameters all have different scales, and the meaning of “atol” and “rtol” applied to all parameters is not always meaningful. I would find it useful if there were 1) a way to converge just on function tolerance or 2) to have pytree-valued rtol and atol applied to parameters.

2. It is very cool that OptaxMinimiser is provided. For me, the ability to do the converse—wrap an optimistix minimizer into an optax-like style API—would be very useful. For my work, I don’t always want to use the cauchy termination criteria at all—I may want to just step through a solve and inspect the results. I also may want to vmap over solves, and I don't think it is always wise to use while loop logic in this case. I wrote my own API to give this behavior using AbstractMinimiser.init and AbstractMinimiser.step, and this was fairly straight forward to do! But I definitely needed the source code to do so and some prior Equinox expertise. Also, related to the above, rtol and atol I noticed were required parameters.

I realize what I’m describing is not a small feature request by any means and there is no immediate need for this on my end, but I thought it would be interesting to have some discussion about this. Let me know what you think!

[johannahaffner]

Hi again :)

1. I think that different termination criteria could in principle be supported - this is a use case @bagibence is also interested in, if I recall correctly? This should be fairly simple to do, we could create a callable solver attribute that defaults to cauchy_termination and may be overridden by an interested user by providing any other callable with the right signature when instantiating the solver. On our end, this would only require a minor break (adding a solver attribute), and it would give power users a lot more flexibility.
   Termination conditions could potentially be upstreamed in the future, although this would not be directly necessary.

2. For the second request - interactive solving is already supported: https://docs.kidger.site/optimistix/examples/interactive/, as you note. If you have an Optax-style API that works for you, that sounds like a good candidate for the Awesome List, exactly the kind of thing it is meant for - a small tool that accomplishes a specific, ecosystem-relevant task.

[michael-0brien]

1. I think that different termination criteria could in principle be supported - this is a use case @bagibence is also interested in, if I recall correctly? This should be fairly simple to do, we could create a callable solver attribute that defaults to cauchy_termination and may be overridden by an interested user by providing any other callable with the right signature when instantiating the solver. On our end, this would only require a minor break (adding a solver attribute), and it would give power users a lot more flexibility.

Cool! This would solve things in a pretty straight forward way. Rather than having a pytree-valued rtol and atol, I believe this means I could do appropriate rescalings in this function.

2. For the second request - interactive solving is already supported: https://docs.kidger.site/optimistix/examples/interactive/, as you note. If you have an Optax-style API that works for you, that sounds like a good candidate for the Awesome List, exactly the kind of thing it is meant for - a small tool that accomplishes a specific, ecosystem-relevant task.

This makes sense. I did know about this tutorial; I think my comment was more about API friendliness for 1) Instantiating f_struct and aux_struct and 2) Cases where we don’t want to use rtol and atol at all.

Though looking again at this tutorial, point 1 really isn’t bad at all. I was biased by trying to do it in the general case via eqx.filter_closure_convert. For point 2, I suppose API design shouldn’t be too informed by edge cases; it’s really no problem to just set these to dummy values if I don’t want to use them.

[johannahaffner]

Do you want to try your hand at a PR for custom termination criteria? This requires mostly small tweaks, but in lots of places. The attribute should be added here

optimistix/optimistix/_iterate.py

Line 26 in 5b1fa30

class AbstractIterativeSolver(eqx.Module, Generic[Y, Out, Aux, SolverState]):

and then all abstract solvers should get the appropriate AbstractVar, and all concrete solvers should have this attribute set in their __init__.

There is a root finder that currently uses a boolean attribute to select termination criteria, and we should think about how to best update it:

optimistix/optimistix/_solver/newton_chord.py

Line 65 in 5b1fa30

cauchy_termination: bool = True

Most importantly, we should think about the signature of termination criteria. Currently, cauchy_termination has the signature Callable[[Float, Float, Callable[[PyTree], Scalar], Y, Y, Out, Out], Bool], for arguments rtol, atol, norm, y, y_diff, f, f_diff, where f is typically either a scalar or a PyTree of residuals, and y is a PyTree of arrays.

Enabling custom termination criteria means exposing this API, and perhaps the current signature is too restrictive (e.g. if gradient values are supposed to be considered). On the other hand, it should be clear what may go into this function.

For the tutorial / interactive solving: glad you like it! Feel free to suggest tweaks that would have been helpful to you.

[michael-0brien]

On the termination criteria, one thought is that norm is somewhat superfluous with having a function for the criteria itself. The issue though would be that it is far simpler to change a criteria based on norm, and possibly too cumbersome to change termination function for non-advanced users.

Another thought is that rtol, atol, and norm are not runtime arguments, so in principle an AbstractTerminationCriteria could be created with a concrete __init__ for non-runtime arguments, but an abstract __call__ for runtime arguments. This would be instantiated in the AbstractMinimiser, etc __init__. This option would be a breaking change if anyone has code that operates on rtol, atol, and norm as pytree leaves. This is a little overengineered I think.

Another option would be to let users fully define and pass an AbstractTerminationCriteria and rtol, atol, and norm are only in the cauchy concrete implementation. While this would make instantiation more cumbersome, it might be the most appealing option in theory. In practice, it would be totally API breaking as it would change the AbstractMinimiser.__init__.

None of these options sound too appealing. I think it would be fine to keep the signature the same and add the termination function to __init__s; it would be a feature for advanced users anyway!

Unfortunately things are kind of crazy for me right now I don’t think I could get to this for some time!