Propagate `kwargs` through `update_coefficients!`

[gaurav-arya]

In some rare cases, an operator depends on more state than just u, p, and t. For example, the W operator in OrdinaryDiffEq depends on dtgamma.

If kwargs could be propagated through update_coefficients!, we'd be able to take advantage of the recursive infrastructure in these cases too. The call might look something like update_coefficients!(W, u, p, t; dtgamma): operators which do not use dtgamma would just ignore the extra info, while an operator which would like to use it catches it and uses it. Of course, if an operator requires dtgamma (or contains an operator that does), it could not be used as a regular SciMLOperator for the usual use cases, but it wouldn't make sense for that to be possible in any case if the operator needs extra state.

For this to work, it seems like all that would be needed to be done is:

• Recursively propagate kwargs in update_coefficients! for ComposedOperator's, etc.
• User-provided update_func signatures which do not catch kwargs should be detected and modified to catch kwargs (and not use them), so that the usual use case of update_coefficients! still works seamlessly. (Something like https://stackoverflow.com/a/42579941?)

In general, something like this just seems like good design: while the positional args u, p, t capture 99% of use cases, there are definitely cases (like the W operator) which could naturally be expressed as a SciMLOperator if we just allowed an extra state dependence. Happy to take a stab at this if it's deemed a good idea.

[vpuri3]

that is easy to implement.

@ChrisRackauckas thoughts on the design change? and do we have any updates on the sciml parameter interface? SciML/DifferentialEquations.jl#881

[vpuri3]

@ChrisRackauckas thoughts?

[gaurav-arya]

Personally I'm in favor of something like this to easily support cases where the state isn't easily shoehorned into u, p, or t (e.g. the WOperator depends on dtgamma in OrdinaryDiffEq). The technical challenge I hit, though, was with:

User-provided update_func signatures which do not catch kwargs should be detected and modified to
catch kwargs (and not use them), so that the usual use case of update_coefficients! still works
seamlessly. (Something like https://stackoverflow.com/a/42579941?)

The idea would be that if the user provides update_func(A,u,p,t), we'd convert it into something with the effect of update_func(A,u,p,t; kwargs...) with the kwargs dropped, if the user provides update_func(A,u,p,t; dtgamma) it would convert into update_func(A,u,p,t; dtgamma, kwargs...), etc. Handling this in generality seemed a bit non-trivial to me, not sure if it's done elsewhere in SciML.

[gaurav-arya]

An easier alternative would be to not allow arbitrary kwargs but just an extra_info object as a fifth argument, either positional or keyword, which again would automatically be strapped on if the user-provided function doesn't accept it. The docs would be something like "Some operators are used in a context where they expect extra info to be provided to them in a particular format. If you are implementing such an operator, use the signature update_func(A, u, p, t, extra_info)" (or update_func(A, u, p, t; extra_info) if we go with keyword).

This makes it easier to do the automatic handling than with kwargs, but is perhaps a bit clunkier. For comparison, the docs with the kwargs approach might read "Some operators are used in a context where they expect specific extra info to be provided to them as keyword arguments in update_coefficients!. If you are implementing such an operator, simply capture these keyword arguments in the provided update_func signature, and update_coefficients! will automatically work with the keyword arguments.

(The fundamental problem is I don't think this extra info can be put into either u, p, or t, because all three have a pre-fixed meaning based on the problem definition.)

[vpuri3]

Just remembered - this won't be necessary. For the WOperator, W= 1/\gamma * J + M (or whatever), we plan to make J, M subtypes of AbstractSciMLOperator, and \gamma a ScalarOperator. Therefore W will be a ComposedOperator, and we can just write a simple update rule for \gamma.

[gaurav-arya]

"simple update rule for gamma" -- can you elaborate? Do you mean you'll manually update it, separately from running update coefficients?

The issue I see is that if you run update_coefficients on the WOperator, and it'll recurse onto the ScalarOperator gamma, but gamma needs the dtgamma to update.

[vpuri3]

gamma = ScalarOperator(0.0, update_func = (dtold, u, p, t) -> integrator.dt)
W = 1/gamma * M + J
update_coefficients!(W, u, p, t)

[gaurav-arya]

So update_func would be a closure over the integrator? That feels a bit strange to me, feels like that's suggesting that the dtgamma should really be being propagated through update_coefficients...

[vpuri3]

^that should work for now. Though your point is valid. FEniCS has a mature UFL which just lets you say gamma.assign(1.0) and the value of gamma is updated everywhere. Let's think of an API for SciMLOperators that lets us do just this

[vpuri3]

something like

gamma = ScalarOperator(0.0, update_func = (oldval, u, p, t) -> 1/integrator.dt))
W = 1/gamma * M + J
integrator.dt = dtnew

# option 1 - use update_coefficients!
W(du, u, p, t)

# option 2 - override update_coefficients!
update_coefficients!(W, u, p, t)
set!(gamma, 1/integrator.dt)
mul!(du, W, u)

and then we can write set! for every SciMLOperator. I think regardless of the update_coefficients! mechanism, the latter method would be a very nice convenience.

[gaurav-arya]

Right, that works.

I still feel like update_coefficients! with recursively prop'd extra state (like kwargs) is the right way to go here. set! and update_coefficients! are very similar in spirit, so it doesn't feel right to separate the two just because some state doesn't fit into u/p/t. I think I know how to implement this now, how about I put together a prototype tomorrow and we can go from there?

[gaurav-arya]

Implemented a prototype at #143, let me know your thoughts!

[vpuri3]

hold off on implementation till we get some input from Chris. after all it has to fit with the rest of the ecosystem