What does "parameter splitting" mean?

[Datseris]

Question❓

I am updating my codebase to MTK v9 (congrats on the release btw). I am trying to understand what "parameter splitting" means, as this is one of the largest entries in the changelog.

In the docs (at least the current docs) I couldn't find a section on this. Can someone explain to me what this splitting means? Does it simply mean that instead of the parameter container being a vector, it is a custom struct?

If that is the case, what's the reason to have a custom struct...?

[ChrisRackauckas]

What is Parameter Splitting?

With the MTK v9 we now fully support the Symbolics type information. For example @parameters a::Int b::Bool for integer and boolean parameters. In particular, this is required for mixed discrete-continuous models to be optimal, instead of requiring that everything is a floating point number, allowing for some boolean variables and integer variables then lets you have pieces like a state machine optimally controlling a continuous system.

However, if we were to then make p = Vector{Any} our performance would be terrible. If p is just a straight tuple you get type-instability in the indexing for generic loop usage and it doesn't scale to large numbers of parameters.

The optimal choice is to then make p a tuple (Vector{Float64}, Vector{Int}, Vector{Bool}, Vector{T}, ...) for every type that is used. That's what we call the split form. This existed in v8 because it's required for many of the complex scenarios described before.

Connection Between MTKParameters and SciMLStructures

However with the v9, we now wrap this object in the MTKParameters structure and endow it with extra properties. In particular, it satisfies the https://github.com/SciML/SciMLStructures.jl interface, which means that you can query it for subsections that are potentially useful, and the structure is somewhat more elaborate. So for example, you may have some tunable parameters and some non-tunable parameters. This gives a structure (Vector{Float64}, Vector{Float64}, Vector{Int}, Vector{Bool}, Vector{T}, ...) where you have a separate vector for the tunable and non-tuanbles, then canonicalize(Tunable(), prob) gives a vector of the tunable parameters (non-allocating by aliasing the tunable subset), and repack(new_tunable) updates them. This is then what's being setup downstream SciML/SciMLSensitivity.jl#1010 so that our packages for example automatically know to only differentiate the tunable subset of the parameters, and optimally filter out the size of the adjoint equation to just the subset that is required.

There's some other goodies in there as well to be added soon. So for example if you are updating those Bools and Ints as control parameters within callbacks, MTK will soon declare those as Discrete values, so canonicalize(Discrete(), p) will give the (Vector{Int}, Vector{Bool}) which are modified/structurally modifiable (i.e. can be non-constant in a simulation given some input, not necessarily that they are modified for the default ones). The solvers are then being updated so that you can turn on/off saving of the Discrete canonicalization and then it can properly recreate the full time series of any process with modified parameters. This will fix a few lazy interpolations in some edge cases (I always like bug fixes), but also this is a useful feature for plotting, visualization, and other analyses on discrete-continuous systems.

Additionally, the split form can have for example Vector{Vector{Float64}} which for example can hold cache variables which are used to write non-allocating code, which are somewhat hidden from the user but make the in-place form fully non-allocating when using symbolic array functions. For example, A * B .+ c should have a cache variable for the A*B result that is then stored in the MTKParameters so that the operations can be fully non-allocating. Right now the codegen does not do this, but this is planned as a way to then allow for fully optimal array structures code (no one is planning to do it for Symbolics.jl codegen right now, anyone is welcome to take this up, but the JuliaSimCompiler downstream will make use of it). Then inside of SciML tools which need to know cache values for correctness they can query for it. For example, Ensemble solving in multithreading requires that caches are copied so they don't collide, so there will then be a check for if cache is empty, safetycopy, and then in the future make that even simpler where if it's an MTK generated system, then we can guarantee the only possible shared memory is in the canonicalize(Caches(), prob), in which case we only safety copy these to limit memory usage at the multithreading points. This can also be used to give better error messages when building GPU kernels.

Now one thing this is really alluding to is that "parameters" is a bit overloaded. Indeed, in SciML we have used the parameters for caches, things changed in callbacks, extra functions, and other wacky things since the dawn of time. This is now just formalizing the idea so that downstream code that needs to know this kind of information can now make use of it. Before for example, SciMLSensitivity would require that your parameters were an AbstractArray. Other parameter structures would work, but it would fail if you tried to use an adjoint method on it. But now you will get a message that you need to define it as a SciMLStructure and it will then use the Tunable interface to automatically pull the right subset of the parameters from your arbitrary user structure that may also contain other caches, dataframes, or whatever the heck you're doing.

Some of that is "coming soon", we didn't want to wait any longer on the v9 since adding those things is non-breaking, so for now we have just done the breaking part of making the MTKParameters struct and setting up the core canonicalizations required for "most" (all?) packages to get an easy update, but throughout the v9 timeframe you can see this gives us lots of space to make more goodies.

Indexing

Now one obvious loss from all of this is the indexing p[i] isn't going to work. p[i][j] in theory, but now it's a bit obscure where the parameters are. So the SymbolicIndexingInterface is even more important for the v9 parameters, hence all of the work just prior to the v9 to make it very solid. Now parameter indexes are Cartesian to handle the splitting, but we think that the getp/setp interface has been refined enough to be optimal in performance and also be quite ergonomic (with some slower more ergonomic methods) and so we'd highly recommend using the direct interfaces and not using the indices of the parameters.

[AayushSabharwal]

So for example if you are updating those Bools and Ints as control parameters within callbacks, MTK will soon declare those as Discrete values, so canonicalize(Discrete(), p) will give the (Vector{Int}, Vector{Bool}) which are modified/structurally modifiable

This already works 😄 It just requires that for FunctionAffects, users declare the modified parameters explicitly. Docs will soon be updated to reflect this.

[ChrisRackauckas]

Awesome, I didn't know that was already done. I'll try and see if I can update the solvers over the ASKEM week.