SciML / SciMLBase.jl

The Base interface of the SciML ecosystem
https://docs.sciml.ai/SciMLBase/stable
MIT License
118 stars 91 forks source link

Consider splitting index provider from `(f::ODEFunction).sys` #699

Closed hexaeder closed 1 month ago

hexaeder commented 1 month ago

Is your feature request related to a problem? Please describe.

I am working on a package which constructs a rhs function of an ODE and I am using the SymbolicIndexingInterface for symbolic indexing. All the machinery for symbolic indexing basically starts at odefunction.sys to get the index provider.

My RHS object also acts as the index provider. So currently, I define

SciMLBase.__has_sys(rhs::MyRHSwhichIsIndexProvider) = true
Base.getproperty(rhs::MyRHSwhichIsIndexProvider, s::Symbol) = s===:sys ? rhs : getfield(rhs, s)

to hook into that machinery. Feels a bit hacky, because I am basically providing a "system" and I don't know which other interfaces I am thus violating.

Describe the solution you’d like It would be cool if the index provider functionality would be decoupled from the sys field. So maybe not implementing the index provider interface for AbstractODEFunction at all but relying on the default definition

symbolic_container(fn::AbstractODEFunction) = has_sys(fn) ? fn.sys : SymbolCache()

which could be overloaded by packages. This kinda works already but AbstractODESolution for example forwards the calls to the interface implemented by ODEFunction instead of respecting the symbolic_container.

AayushSabharwal commented 1 month ago

So you're suggesting this method:

SymbolicIndexingInterface.symbolic_container(A::AbstractSolution) = A.prob.f

should be replaced by

SymbolicIndexingInterface.symbolic_container(A::AbstractSolution) = A.prob

?

I'm not sure why symbolic_container skips prob completely. Although it's valid from the standpoint of SII, it's definitely better behavior for SciMLBase to forward it to A.prob.

Or have I misunderstood the issue?

hexaeder commented 1 month ago

Hm did a bit of digging, the problem seems to be different from what I described initially, sorry. Lets say I define

SII.symbolic_container(odef::ODEFunction{<:Any,<:Any,<:MyTypeWhichIsIndexProvier}) = odef.f

then a call SII.default_values(odef) works because I directly overloaded that method. A call to the problem SII.default_values(odeproblem) also works, because the dispatch is still in place. However after the solving, the call SII.default_values(sol) does not work anymore, because sol.prob.f changed the type of f such that MyType... is now wrapped in multiple levels of FunctionWrappersWrappers.

It all works as expected once I provide the indp as the sys filed of the ODEFunction...

AayushSabharwal commented 1 month ago

I think that's more a side effect of how specialization of these SciMLFunctions works. They're wrapped in FunctionWrappersWrappers to avoid overspecialization (and I believe reduce REPL latency by avoiding repeated recompilation). Also, the type parameters of ODEFunction (and related functions) is not public API, so that dispatch is not guaranteed to be stable anyway.

What if you put rhs in the sys field as well? That way f.sys would always be your index provider. Basically, ODEFunction(; f = rhs, sys = rhs, ...)

hexaeder commented 1 month ago

Providing sys = rhs is essentially what I do right now and everything works, so this is not a pressing issue. However feels a bit hacky/internal because I don't know if the sys is used for purposes other than symbolic indexing.

If the goal is to make SII part of the public interface, there should be clear way of constructing an ODEFunction and supplying it with both a rhs and a symbolic index provider, decoupled from the sys property

f = ODEFUnction(myrhs, symbols=SymbolCache([:u1, :u2, :u3])
prob = ODEProblem(f,...)

A long time ago, something like this was basically possible by providing a vector of symbols to the syms keyword to the ODEFunction. The alternative would be to explicitly use sys for that, but then it would be nice to document what sys actually means outside the MTK context and whether its fine to use it like that.

AayushSabharwal commented 1 month ago

sys is just for symbolic indexing

AayushSabharwal commented 1 month ago

The previous behavior using syms = is now replaced by providing a SymbolCache as the sys keyword argument.

hexaeder commented 1 month ago

Allright thanks for clarification! To me the name implied, that there's more to it, but then I'll close here.