Propagate kwargs through update_coefficients!

[gaurav-arya]

Resolves #125. For example usage, see the added test for a scalar operator. I left function operator mostly alone because I don't understand some aspects of it, and from some of the other issues it seems it still might be in flux. More comprehensive tests can be added in a follow up PR, but note that this change is totally backwards compatible. (Addressed)

For developers of this library, this is what you should know about the new design:

• Operator constructors called by the user should accept update_func and accepted_kwargs arguments, and preprocess them using the preprocess_update_func utility before making the struct, as done here.
• For internal calls to constructors where the update_func has already been processed, accepted_kwargs should be set to NoKwargFilter(), which indicates that no kwarg filtering should be done.

[codecov]

Codecov Report

Merging #143 (3a665ca) into master (ac683c3) will increase coverage by 0.73%.
The diff coverage is 84.55%.

@@            Coverage Diff             @@
##           master     #143      +/-   ##
==========================================
+ Coverage   71.85%   72.58%   +0.73%     
==========================================
  Files          11       11              
  Lines        1407     1430      +23     
==========================================
+ Hits         1011     1038      +27     
+ Misses        396      392       -4     

Impacted Files	Coverage Δ
src/batch.jl	74.19% <56.25%> (+2.76%)	⬆️
src/scalar.jl	62.40% <81.81%> (-2.21%)	⬇️
src/matrix.jl	77.10% <86.36%> (+2.73%)	⬆️
src/func.jl	82.02% <89.28%> (+1.01%)	⬆️
src/interface.jl	58.40% <94.44%> (+3.06%)	⬆️
src/utils.jl	80.00% <100.00%> (+13.33%)	⬆️

... and 1 file with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vpuri3]

I like the idea. We can simply pass AD-able parameters in p, which can be array-like, and operator relevant context or struct objects in kwargs. The latter may contain mesh information, numerical parameters, discretization information. This solving the primary problem in SciML/DifferentialEquations.jl#881.

[vpuri3]

You can go ahead and modify FunctionOperator. I'll hold off on any changes for the meantime.

[vpuri3]

Would be good if you can modify the operator algebra test in test/total.jl and have D1, D2, T1, T2 take in kwargs.

[gaurav-arya]

Sure, working on it

[gaurav-arya]

Regarding SciML/DifferentialEquations.jl#881: there's still more design work to do in terms of how users can specify extra parameters, but yes hopefully this comes in handy. The nice thing about this PR is that it's very general. All it is doing is making SciMLOperators more flexible: if you want to use a SciMLOperator in a context where something else is expected beyond u/p/t, you can do it.

It's still up to upstream projects how to use this flexibility: as one possibility, SciML-verse problems could take in an operator_state keyword argument for non-differentiable, fixed, parameters. If the user were to make use of this keyword argument, they would be expected to write SciMLOps / DiffEq funcs that accept the operator_state kwarg, and SciML would indeed provide it to them.

[vpuri3]

When you do modify FunctionOperator, can you propagate the kwargs down to L.op, L.op_adjoint?

function update_coeff...
for op in getops(L)
  if !isnothing(op)
    update_coefficients!(op, u, p, t; kw...)
  end
end

[gaurav-arya]

check now:)

The construction is a bit awkward, because the FunctionOperator already takes some kwargs. I did something reasonable for now, but happy to work with you on improving it in the future when #162 is addessed

[gaurav-arya]

@ChrisRackauckas ready for merge. The PR body, added docs+example, and #143 (comment) should give the relevant context if you want to give it a look over

[gaurav-arya]

@ChrisRackauckas bumping this since it's a bit of a blocker for work on other operators (because the constructor style is changed slightly here), and for integration into OrdinaryDiffEq.jl.

[vpuri3]

@ChrisRackauckas ping :)

[gaurav-arya]

@vpuri3 I decided I want to finish and merge SciML/OrdinaryDiffEq.jl#1917 before this one, if that sounds good? So I can make sure I get the design right here 🙂

[gaurav-arya]

I want to revise the design update_coefficients(u, p, t; kwargs...) to update_coefficients(u, p, t; extra_state=...) where extra_state can be any user specified structure (e.g. a named tuple). While the former is more elegant, I think it could hurt us down the road to take up the entire kwarg space as part of the API. The latter should also reduce the implementation complexity a lot.

I'll make the change in the next day or two!

[vpuri3]

I guess I'm not getting this. Can't you pass arbitrary structs into kwargs anyway? What would the extra_state be exactly? Here's a use case I'm playing with right now. Can you illustrate your point with this eg?

update_func(diag, u, p, t; NN = NN, x = x) = NN(p)(x)
L = DiagonalOperator(zeros(N); update_func = update_func)

# with this branch
L(u, p, t; x = x1, NN = NN1) # == NN1(x1) .* u
L(u, p, t; x = x2, NN = NN2) # == NN2(x2) .* u

[gaurav-arya]

The extra_state option would be less permissive, not more permissive. So we'd have to stuff NN and x into e.g. extra_state=(; x = x1, NN = N1) and each op extracts what it needs from that.

But actually, I think I'm OK with the current approach in this PR. It does have its own benefits too, like doing the keyword filtering automatically for each op rather than having each op extract what it needs from the extra_state. It's just a hard decision:)

I am a little bit swamped with work right now but I've added you as a contributor on the branch if you'd like to resolve the conflicts before I get back to this

[ChrisRackauckas]

This looks reasonable enough to me.

[vpuri3]

great. we just gotta clean up FunctionOperator and this this is good to go.

[vpuri3]

@ChrisRackauckas this is done.