State and Parameter vectors for SciML applications

New State & Parameter Vector Types

I’m hoping for others’ thoughts on this idea, including suggestions, alternatives, rotten fruit, etc. All input is welcome.

The “Problem”

Users have tried to add unit handling to their numerical integration functions. It’s usually a bit tricky. I haven’t seen a perfect solution yet.

Motivation

Many users describe physical systems with ModelingToolkit or self-written ODEFunctions. Some try to add unit handling, and initialize the dynamics with Unitful quantities; at this point, they usually see some error. Several Discourse posts and GitHub issues have been created for this over the years. including the two linked below.

There are a few issues with Unitful quantities in ODEFunctions today. Some math operators aren’t defined for Unitful.Unitlike or similar types. As an anecdotal example, I have encountered an error like “no method defined for one(::Unitlike)”. Generally, math with Unitful quantities is very good. That’s a testament to the Unitful developers, and Julia’s composability. Still, I know there’s lots of numerics magic going on in numerical integrators these days, so defining every possible required method for Unitful types seems like a game of whack-a-mole.

Current Solutions

I’ve seen the following solutions so far.

Stripping Units

If we strip units before integrating, then all’s fine! This can be a bit clunky, though.

Defining all Methods for Unitful Types

We could play that game of whack-a-mole, and define missing methods for Unitful types as we find them. That’s all fine, and Unitful vectors may even work with DifferentialEquations out of the box today (they did a few months ago). Still, then you run into an issue where your state vector has non-concrete eltypes.

If you pass a ComponentArray or a LabelledArray with Unitful element types, then you could make sure the formal eltype is Union{<:Unitful.LengthUnit, <:Unitful.VelocityUnit} (as an example for a system with position and velocity states), but this is still less efficient than simply passing a Vector{Float64} or similar.

Proposed Solution

What if we parameterized our state vectors and parameter vectors by their units. I’ve been using this for some time with GeneralAstrodynamics. The relevant source code is here. Right now, that functionality is pulling a lot of code from LabelledArrays, hence the ParameterizedLabelledArrays filename linked there.

As an example, we could define the following types.


"""
A vector with physical elements (aka `Unitful`)
"""
abstract type PhysicalVector{F, MU, LU, TU, AU,L} <: AbstractVector{F} end

"""
A representation for all state vectors for physical systems.
The parameters are...

- F  => Floating-point type
- MU => Mass unit
- LU => Length unit
- TU => Time unit
- AU => Angular unit 
- L => Labels
"""
abstract type StateVector{F, MU, LU, TU, AU,L} <: PhysicalVector{F, MU, LU, TU, AU,L} end

"""
A representation for all parameter vectors for physical systems.
The parameters are...

- F  => Floating-point type
- MU => Mass unit
- LU => Length unit
- TU => Time unit
- AU => Angular unit 
- L => Labels
"""
abstract type ParameterVector{F, MU, LU, TU, AU,L} <: PhysicalVector{F, MU, LU, TU, AU,L} end

Then, we could define a CartesianState like so…

struct CartesianState{F, MU, LU, TU, AU} <: StateVector{F, MU, LU, TU, AU,(x = 1, y = 2, z = 3, ẋ = 4, ẏ = 5, ż = 6, r = 1:3, v = 4:6)}

  data::MVector{6,F}

end

And we could dispatch on functions like…

Base.@pure lengthunit(::CartesianState{F,MU,LU}) = LU
position(state::CartesianState) = state[1:3] * LU

Conclusion

Anyone have any thoughts on this? I think the struct definition is rough, but I do think there’s promise in parameterizing our state and parameter vectors by Unitful units. Maybe instead of carrying around units for all dimensions (like in the MU, LU, TU, AU above), we have a NamedTuple struct parameter which tracks units in the same way that LabelledArrays (and the copycat examples above) track labels.

Two things. One, Unitful is the wrong idea for a units package and I would want to write a different one. But two, ModelingToolkit shouldn’t have any of these issues. Did you try doing it directly in MTK? It eliminates units after code construction and gets type stability.

Just curious, what about the Unitful design would you change? Would you prefer for units and values to be separated? In other words, is the "combine units and values into a new Number type, Quantity" idea the one you would change?

As for MTK – this is what I’m trying to do. Note that the u0 type is not concrete (unless I’m misunderstanding Julia’s types). These kinds of complications are why I’ve started parameterizing the state and parameter vectors by units. The information is still there, but the numerical integrators don’t have to care about it.

julia> using AstrodynamicalModels, ModelingToolkit, Unitful

julia> model = R2BP()
Model R2BP with 6 equations
States (6):
  x(t)
  y(t)
  z(t)
  ẋ(t)
  ẏ(t)
  ż(t)
Parameters (1):
  μ
Incidence matrix:
 ×  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ×  ⋅  ⋅
 ⋅  ×  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ×  ⋅
 ⋅  ⋅  ×  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ⋅  ×
 ⋅  ⋅  ⋅  ×  ⋅  ⋅  ×  ×  ×  ⋅  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  ×  ⋅  ×  ×  ×  ⋅  ⋅  ⋅
 ⋅  ⋅  ⋅  ⋅  ⋅  ×  ×  ×  ×  ⋅  ⋅  ⋅

julia> ICs = @nonamespace [
           model.x => 1e3u"km",
           model.y => 0u"km",
           model.z => 11e3u"km",
           model.ẋ => -2u"km/s",
           model.ẏ => 1u"km/s",
           model.ż => -4u"km/s",
       ];

julia> params = @nonamespace [model.μ => 5e4u"km^3/s^2"];

julia> tspan = (0.0u"s", 1e4u"s");

julia> problem = ODEProblem(model, ICs, tspan, params)
ODEProblem with uType Vector{Quantity{Float64, D, U} where {D, U}} and tType Quantity{Float64, 𝐓, Unitful.FreeUnits{(s,), 𝐓, nothing}}. In-place: true
timespan: (0.0 s, 10000.0 s)
u0: 6-element Vector{Quantity{Float64, D, U} where {D, U}}:
   1000.0 km
      0.0 km
  11000.0 km
 -2.0 km s⁻¹
  1.0 km s⁻¹
 -4.0 km s⁻¹

The units should be in the variables. Model Validation and Units · ModelingToolkit.jl

1 Like