Default MTK problems to AutoDespecialize

Project.toml

@@ -74,7 +74,7 @@ ConstructionBase = "1.5.8"
 DataInterpolations = "9"
 DataStructures = "0.18, 0.19"
 DelayDiffEq = "6"
-DiffEqBase = "7.12"
+DiffEqBase = "7.14"
 DiffEqNoiseProcess = "5.25.0"
 DifferentiationInterface = "0.7.13"
 DocStringExtensions = "0.9"
@@ -87,7 +87,7 @@ Libdl = "1"
 LinearAlgebra = "1"
 LinearSolve = "5.4"
 Logging = "1"
-ModelingToolkitBase = "1.63"
+ModelingToolkitBase = "1.64"
 ModelingToolkitStandardLibrary = "2.20"
 ModelingToolkitTearing = "1.19.2"
 Moshi = "0.3.6"
@@ -112,7 +112,7 @@ REPL = "1"
 Reexport = "1"
 RuntimeGeneratedFunctions = "0.5.12"
 SCCNonlinearSolve = "1.13"
-SciMLBase = "3.19"
+SciMLBase = "3.46"
 SciMLPublic = "1.0.0"
 Serialization = "1"
 Setfield = "1"
@@ -139,7 +139,7 @@ RecursiveArrayTools = "4"
 ReferenceTests = "0.10"
 SafeTestsets = "0.1"
 SciCompDSL = "1"
-SciMLStructures = "1"
+SciMLStructures = "1.10.2"
 SciMLTesting = "2.4"
 SpecialFunctions = "2"
 StableRNGs = "1"

docs/src/API/problems.md

@@ -89,10 +89,16 @@ ModelingToolkit.get_u0
 ModelingToolkit.varmap_to_vars
 ```
 
-By default, the parameters of the system are stored in a custom data structure called
-`MTKParameters`. The internals of this data structure are undocumented, and it should
-only be interacted with through defined public API. SymbolicIndexingInterface.jl contains
-functionality useful for this purpose.
+The parameters of a split system are stored in a custom data structure called
+`MTKParameters`. ModelingToolkit problem constructors use
+[`SciMLBase.AutoDespecialize`](https://docs.sciml.ai/SciMLBase/stable/interfaces/Problems/)
+by default. Solvers that support this policy wrap the parameters in
+[`SciMLBase.DespecializedParameters`](https://docs.sciml.ai/SciMLBase/stable/interfaces/Problems/)
+at solve time so compiled code
+can be reused across parameter-buffer layouts. Explicit `AutoSpecialize` and
+`FullSpecialize` problems retain their existing behavior. These objects should only be
+interacted with through their defined public API.
+SymbolicIndexingInterface.jl contains functionality useful for this purpose.
 
 ```@docs
 MTKParameters

docs/src/basics/FAQ.md

@@ -3,9 +3,14 @@
 ## Why are my parameters some obscure object?
 
 In ModelingToolkit.jl version 9, the parameter vector was replaced with a custom
-`MTKParameters` object, whose internals are intentionally undocumented and subject
-to change without a breaking release. This enables us to efficiently store and generate
-code for parameters of multiple types. To obtain parameter values use
+`MTKParameters` object. ModelingToolkit problems use
+[`SciMLBase.AutoDespecialize`](https://docs.sciml.ai/SciMLBase/stable/interfaces/Problems/)
+by default. Supporting solvers expose a
+[`SciMLBase.DespecializedParameters`](https://docs.sciml.ai/SciMLBase/stable/interfaces/Problems/) wrapper
+around the `MTKParameters` object while solving so compiled code can be reused across
+parameter layouts. The internals of `MTKParameters` are intentionally undocumented and
+subject to change without a breaking release. This representation enables us to efficiently
+store and generate code for parameters of multiple types. To obtain parameter values use
 [SymbolicIndexingInterface.jl](https://github.com/SciML/SymbolicIndexingInterface.jl/) or
 [SciMLStructures.jl](https://github.com/SciML/SciMLStructures.jl/). For example:
 
@@ -16,24 +21,25 @@ getβ(sol) # can be used on any object that is based off of the same system
 getβ(prob)
 ```
 
-Indexes into the `MTKParameters` object take the form of `ParameterIndex` objects, which
-are similarly undocumented. Following is the list of behaviors that should be relied on for
-`MTKParameters`:
+Indexes into these parameter objects take the form of `ParameterIndex` objects, which are
+similarly undocumented. The following behaviors can be relied on for both `MTKParameters`
+and `SciMLBase.DespecializedParameters`:
 
   - It implements the SciMLStructures interface.
   - It can be queried for parameters using functions returned from
     `SymbolicIndexingInterface.getp`.
-  - `getindex(::MTKParameters, ::ParameterIndex)` can be used to obtain the value of a
-    parameter with the given index.
-  - `setindex!(::MTKParameters, value, ::ParameterIndex)` can be used to set the value of a
-    parameter with the given index.
+  - `getindex` with a `ParameterIndex` can be used to obtain the value of a parameter with
+    the given index.
+  - `setindex!` with a `ParameterIndex` can be used to set the value of a parameter with the
+    given index.
+  - `SciMLBase.unwrap_parameters` recovers the wrapped `MTKParameters` object from an
+    `SciMLBase.DespecializedParameters` object.
   - `parameter_index(sys, sym)` will return a `ParameterIndex` object if `sys` has been
     `complete`d (through `mtkcompile`, `complete` or `@mtkcompile`).
-  - `copy(::MTKParameters)` is defined and duplicates the parameter object, including the
-    memory used by the underlying buffers.
+  - `copy` duplicates the parameter object, including the memory used by the underlying
+    buffers.
 
-Any other behavior of `MTKParameters` (other `getindex`/`setindex!` methods, etc.) is an
-undocumented internal and should not be relied upon.
+Any other behavior of these types is an undocumented internal and should not be relied upon.
 
 ## How do I use non-numeric/array-valued parameters?
 
@@ -68,10 +74,10 @@ The same principle applies to any parameter type that is not `Float64`.
 
 ## Getting the index for a symbol
 
-Ordering of symbols is not guaranteed after symbolic transformations, and parameters
-are now stored in a custom `MTKParameters` object instead of a vector. Thus, values
-should be referred to by their name. For example `sol[lorenz.x]`. To obtain the index,
-use the following functions from
+Ordering of symbols is not guaranteed after symbolic transformations, and parameters use
+the custom parameter representation described above instead of a vector. Thus, values should
+be referred to by their name. For example `sol[lorenz.x]`. To obtain the index, use the
+following functions from
 [SymbolicIndexingInterface.jl](https://github.com/SciML/SymbolicIndexingInterface.jl/):
 
 ```julia

lib/ModelingToolkitBase/Project.toml

@@ -112,7 +112,7 @@ ControlSystemsBase = "1"
 DataInterpolations = "9"
 DataStructures = "0.18, 0.19"
 DelayDiffEq = "6"
-DiffEqBase = "7.12"
+DiffEqBase = "7.14"
 DiffEqCallbacks = "4"
 DiffEqNoiseProcess = "5"
 DiffRules = "1"
@@ -176,7 +176,7 @@ SCCNonlinearSolve = "1.13"
 SafeTestsets = "0.1"
 SciMLBase = "3.48"
 SciMLPublic = "1.0.0"
-SciMLStructures = "1.7"
+SciMLStructures = "1.10.2"
 SciMLTesting = "2.4"
 Serialization = "1"
 Setfield = "1"

lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl

@@ -219,6 +219,9 @@ function define_params(p::MTKParameters, t, names = nothing)
     end
 end
 
+define_params(p::SciMLBase.DespecializedParameters, t, names = nothing) =
+    define_params(SciMLBase.unwrap_parameters(p), t, names)
+
 """
     $(TYPEDSIGNATURES)
 
@@ -233,6 +236,9 @@ function to_paramvec(p::MTKParameters)
     return reduce(vcat, collect(p); init = [])
 end
 
+to_paramvec(p::SciMLBase.DespecializedParameters) =
+    to_paramvec(SciMLBase.unwrap_parameters(p))
+
 """
     $(TYPEDSIGNATURES)
 
@@ -384,8 +390,19 @@ Obtain default values for unknowns `vars` and parameters `paramvec`
 given the problem `prob` and symbolic parameter object `paramobj`.
 """
 function defaults_from_u0_p(prob, vars, paramobj, paramvec)
+    return defaults_from_u0_p(prob, vars, paramobj, paramvec, parameter_values(prob))
+end
+
+function defaults_from_u0_p(
+        prob, vars, paramobj, paramvec, p::SciMLBase.DespecializedParameters
+    )
+    return defaults_from_u0_p(
+        prob, vars, paramobj, paramvec, SciMLBase.unwrap_parameters(p)
+    )
+end
+
+function defaults_from_u0_p(prob, vars, paramobj, paramvec, p)
     u0 = state_values(prob)
-    p = parameter_values(prob)
     defaults = Dict{Any, Any}(vec(vars) .=> vec(collect(u0)))
     if !(p isa Union{SciMLBase.NullParameters, Nothing})
         if p isa Union{NamedTuple, AbstractDict}

lib/ModelingToolkitBase/src/precompile.jl

@@ -111,19 +111,15 @@ PrecompileTools.@compile_workload begin
     v = [p]
     isempty(v)
     # mtkcompile(sys)
-    # This ODEProblem construction now goes through AutoSpecialize (the new default)
-    # instead of FullSpecialize via the @fallback_iip_specialize macro.
-    prob_precompile = ODEProblem(mtkcompile(System([ModelingToolkitBase.D_nounits(x) ~ 2x + 1], ModelingToolkitBase.t_nounits; name = :a)), [x => 1], (0.0, 1.0))
-
-    # Precompile the FunctionWrappersWrapper wrapping that DiffEqBase.promote_f
-    # performs at solve time for AutoSpecialize ODEProblems. This warms up the
-    # FunctionWrappersWrapper construction path so that the first `solve()` call
-    # does not pay this cost.
-    _f_unwrapped = prob_precompile.f.f
-    _u0 = prob_precompile.u0
-    _p = prob_precompile.p
-    _t = prob_precompile.tspan[1]
-    DiffEqBase.wrapfun_iip(_f_unwrapped, (_u0, _u0, _p, _t))
+    ODEProblem(
+        mtkcompile(
+            System(
+                [ModelingToolkitBase.D_nounits(x) ~ 2x + 1],
+                ModelingToolkitBase.t_nounits; name = :a
+            )
+        ),
+        [x => 1], (0.0, 1.0)
+    )
 end
 
 precompile(Tuple{typeof(SymbolicUtils.isequal_somescalar), Float64, Float64})

lib/ModelingToolkitBase/src/problems/initializationproblem.jl

@@ -40,7 +40,7 @@ as an explicit keyword.
 function InitializationProblem{iip}(
         sys::AbstractSystem, t, op, opts::SciMLProblemOptions; kwargs...
     ) where {iip}
-    return InitializationProblem{iip, SciMLBase.AutoSpecialize}(sys, t, op, opts; kwargs...)
+    return InitializationProblem{iip, SciMLBase.AutoDespecialize}(sys, t, op, opts; kwargs...)
 end
 
 function InitializationProblem{iip, specialize}(
@@ -168,7 +168,13 @@ function InitializationProblem{iip, specialize}(
     # Only forward `check_length` when the caller explicitly set it; otherwise let the
     # underlying problem type apply its own default (see the keyword's definition above).
     check_length_kw = check_length === nothing ? (;) : (; check_length)
-    return TProb{_iip}(
+    problem_constructor = if TProb === LinearInitializationProblem ||
+            TProb === SCCNonlinearProblem
+        TProb{_iip}
+    else
+        TProb{_iip, specialize}
+    end
+    return problem_constructor(
         isys, op; kwargs..., check_length_kw...,
         u0_constructor, p_constructor, missing_guess_value,
         eval_expression, eval_module, warn_cyclic_dependency,

lib/ModelingToolkitBase/src/problems/jumpproblem.jl

@@ -188,6 +188,13 @@ function updateparams!(
     return nothing
 end
 
+function updateparams!(
+        ratemap::JumpSysMajParamMapper{U, V, W},
+        params::SciMLBase.DespecializedParameters
+    ) where {U <: AbstractArray, V <: AbstractArray, W}
+    return updateparams!(ratemap, SciMLBase.unwrap_parameters(params))
+end
+
 function updateparams!(
         ::JumpSysMajParamMapper{U, V, W},
         params::Nothing

lib/ModelingToolkitBase/src/problems/odeproblem.jl

@@ -111,7 +111,7 @@ function SciMLBase.ODEFunction{iip, spec}(
     )
 
     odefn = maybe_codegen_scimlfn(Val{E}, ODEFunction{iip, spec}, args; kwargs...)
-    if !E && spec === SciMLBase.AutoSpecialize
+    if !E && spec in (SciMLBase.AutoSpecialize, SciMLBase.AutoDespecialize)
         odefn = SciMLBase.widen_bounded_type_params(odefn)
     end
     return odefn

lib/ModelingToolkitBase/src/systems/codegen_utils.jl

@@ -1041,6 +1041,11 @@ function GeneratedFunctionWrapper{P}(
 end
 
 function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg{Any, NArgs}) where {PIdx, NArgs, Split}
+    if args[PIdx] isa SciMLBase.DespecializedParameters
+        return SciMLBase.invoke_with_despecialized_parameters(
+            gfw.f_oop, args, args[PIdx], Val(PIdx)
+        )
+    end
     # non-split systems just call it as-is
     Split || return gfw.f_oop(args...)
     if args[PIdx] isa Union{Tuple, MTKParameters} && !(args[PIdx] isa Tuple{Vararg{Number}})
@@ -1059,6 +1064,11 @@ function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg
     if NArgs + 1 != N
         throw(MethodError(gfw, args))
     end
+    if args[PIdx + 1] isa SciMLBase.DespecializedParameters
+        return SciMLBase.invoke_with_despecialized_parameters(
+            gfw.f_iip, args, args[PIdx + 1], Val(PIdx + 1)
+        )
+    end
     Split || return gfw.f_iip(args...)
     if args[PIdx + 1] isa Union{Tuple, MTKParameters} && !(args[PIdx + 1] isa Tuple{Vararg{Number}})
         return gfw.f_iip(args...)

lib/ModelingToolkitBase/src/systems/nonlinear/homotopy_continuation.jl

@@ -574,7 +574,7 @@ function HomotopyContinuationProblem(
 end
 
 function HomotopyContinuationProblem{true}(sys::System, args...; kwargs...)
-    return HomotopyContinuationProblem{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+    return HomotopyContinuationProblem{true, SciMLBase.AutoDespecialize}(sys, args...; kwargs...)
 end
 
 function HomotopyContinuationProblem{false}(sys::System, args...; kwargs...)

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

@@ -699,7 +699,8 @@ function _remake_initialization_data_impl(
         circular_dependency_max_cycle_length = length(all_symbols(sys)),
     )
     kws = maybe_build_initialization_problem(
-        sys, SciMLBase.isinplace(odefn), op, t0, guesses, opts
+        sys, SciMLBase.isinplace(odefn), op, t0, guesses, opts;
+        specialize = initialization_specialization(SciMLBase.specialization(typeof(odefn)))
     )
 
     odefn = remake(odefn; kws...)
@@ -714,6 +715,11 @@ end
 function promote_type_with_nothing(::Type{T}, p::MTKParameters) where {T}
     return promote_type_with_nothing(promote_type_with_nothing(T, p.tunable), p.initials)
 end
+function promote_type_with_nothing(
+        ::Type{T}, p::SciMLBase.DespecializedParameters
+    ) where {T}
+    return promote_type_with_nothing(T, SciMLBase.unwrap_parameters(p))
+end
 
 promote_with_nothing(::Type, ::Nothing) = nothing
 promote_with_nothing(::Type, x::StaticVector{0}) = x
@@ -740,6 +746,11 @@ function promote_with_nothing(::Type{T}, p::MTKParameters) where {T}
     end
     return p
 end
+function promote_with_nothing(::Type{T}, p::SciMLBase.DespecializedParameters) where {T}
+    return SciMLBase.DespecializedParameters(
+        promote_with_nothing(T, SciMLBase.unwrap_parameters(p))
+    )
+end
 
 function promote_u0_p(u0, p, t0)
     T = Union{}
@@ -865,8 +876,9 @@ function DiffEqBase.get_updated_symbolic_problem(
 
     t0 = is_time_dependent(prob) ? current_time(prob) : nothing
 
-    if p isa MTKParameters
-        buffer = p.initials
+    unwrapped_p = _unwrap_mtk_parameters(p)
+    if unwrapped_p isa MTKParameters
+        buffer = unwrapped_p.initials
     else
         buffer = p
     end

lib/ModelingToolkitBase/src/systems/parameter_buffer.jl

@@ -40,6 +40,8 @@ struct MTKParameters{T, I, D, C, N, H}
     end
 end
 
+_unwrap_mtk_parameters(params) = SciMLBase.unwrap_parameters(params)
+
 """
     function MTKParameters(sys::AbstractSystem, p, u0 = Dict(); t0 = nothing)
 
@@ -1011,7 +1013,8 @@ end
 Base.size(::NestedGetIndex) = ()
 
 function SymbolicIndexingInterface.with_updated_parameter_timeseries_values(
-        ::AbstractSystem, ps::MTKParameters, args::Pair{<:Any, <:NestedGetIndex}...
+        ::AbstractSystem, ps::MTKParameters,
+        args::Pair{<:Any, <:NestedGetIndex}...
     )
     for (i, ngi) in args
         for (j, val) in enumerate(ngi.x)
@@ -1038,6 +1041,16 @@ function SciMLBase.create_parameter_timeseries_collection(
     return ParameterTimeseriesCollection(Tuple(buffers), copy(ps))
 end
 
+function SciMLBase.create_parameter_timeseries_collection(
+        sys::AbstractSystem, ps::SciMLBase.DespecializedParameters, tspan
+    )
+    collection = SciMLBase.create_parameter_timeseries_collection(
+        sys, SciMLBase.unwrap_parameters(ps), tspan
+    )
+    collection === nothing && return nothing
+    return ParameterTimeseriesCollection(parent(collection), copy(ps))
+end
+
 @inline __get_blocks(tsidx::Int) = ()
 @inline function __get_blocks(tsidx::Int, buffer::BlockedArray, buffers...)
     return (buffer[Block(tsidx)], __get_blocks(tsidx, buffers...)...)
@@ -1052,6 +1065,15 @@ function SciMLBase.get_saveable_values(
     return NestedGetIndex(__get_blocks(timeseries_idx, ps.discrete...))
 end
 
+
+function SciMLBase.get_saveable_values(
+        sys::AbstractSystem, ps::SciMLBase.DespecializedParameters, timeseries_idx
+    )
+    return SciMLBase.get_saveable_values(
+        sys, SciMLBase.unwrap_parameters(ps), timeseries_idx
+    )
+end
+
 function save_callback_discretes!(integ::SciMLBase.DEIntegrator, callback)
     ic = get_index_cache(indp_to_system(integ))
     ic === nothing && return