From 9b91a5a24f76a0c12095a643b641c65af71422e4 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 9 Aug 2026 19:05:01 -0400 Subject: [PATCH 01/12] Reduce MTK parameter specialization Co-Authored-By: Chris Rackauckas --- benchmark/benchmarks.jl | 4 +- docs/src/API/problems.md | 11 +- docs/src/basics/FAQ.md | 39 +++--- .../src/ModelingToolkitBase.jl | 3 +- .../src/modelingtoolkitize/common.jl | 8 +- .../src/problems/jumpproblem.jl | 7 ++ .../src/problems/odeproblem.jl | 4 + .../src/systems/codegen_utils.jl | 22 ++++ .../src/systems/nonlinear/initializesystem.jl | 11 +- .../src/systems/parameter_buffer.jl | 114 +++++++++++++++++- .../src/systems/problem_utils.jl | 25 ++-- .../extensions/opaque_mtkparameters_ad.jl | 28 +++++ lib/ModelingToolkitBase/test/index_cache.jl | 11 +- .../test/jacobiansparsity.jl | 7 +- lib/ModelingToolkitBase/test/mtkparameters.jl | 6 +- lib/ModelingToolkitBase/test/odesystem.jl | 11 +- .../test/opaque_mtkparameters.jl | 86 +++++++++++++ lib/ModelingToolkitBase/test/runtests.jl | 2 + .../test/split_parameters.jl | 28 ++++- .../test/symbolic_indexing_interface.jl | 4 +- src/linearization.jl | 8 +- src/problems/sccnonlinearproblem.jl | 2 + test/qa/qa.jl | 1 + 23 files changed, 381 insertions(+), 61 deletions(-) create mode 100644 lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl create mode 100644 lib/ModelingToolkitBase/test/opaque_mtkparameters.jl diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 5685330295..c2b54af875 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -90,5 +90,5 @@ prob = ODEProblem(model, u0, tspan; jac, sparse) out = similar(prob.f.jac_prototype) sparse_analytical_jacobian["ODEProblem"] = @benchmarkable ODEProblem($model, $u0, $tspan; jac, sparse) -sparse_analytical_jacobian["f_oop"] = @benchmarkable $(prob.f.jac.f_oop)($(prob.u0), $(prob.p), $(first(tspan))) -sparse_analytical_jacobian["f_iip"] = @benchmarkable $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), $(first(tspan))) +sparse_analytical_jacobian["f_oop"] = @benchmarkable $(prob.f.jac)($(prob.u0), $(prob.p), $(first(tspan))) +sparse_analytical_jacobian["f_iip"] = @benchmarkable $(prob.f.jac)($out, $(prob.u0), $(prob.p), $(first(tspan))) diff --git a/docs/src/API/problems.md b/docs/src/API/problems.md index c3827582fa..99f33dd4c2 100644 --- a/docs/src/API/problems.md +++ b/docs/src/API/problems.md @@ -89,13 +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`. `AutoSpecialize` ODE problems wrap it in `OpaqueMTKParameters` so solvers +can reuse compiled code across parameter-buffer layouts. `FullSpecialize` problems keep the +concrete `MTKParameters` object and avoid the wrapper's dynamic function barrier. These +objects should only be interacted with through their defined public API. +SymbolicIndexingInterface.jl contains functionality useful for this purpose. ```@docs MTKParameters +OpaqueMTKParameters ModelingToolkit.get_p ``` diff --git a/docs/src/basics/FAQ.md b/docs/src/basics/FAQ.md index 7b712395b3..56250bec8c 100644 --- a/docs/src/basics/FAQ.md +++ b/docs/src/basics/FAQ.md @@ -3,9 +3,11 @@ ## 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. `AutoSpecialize` ODE problems expose an `OpaqueMTKParameters` +wrapper around that object so solver compilation can be reused across parameter layouts. +The internals of both types 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 +18,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 `OpaqueMTKParameters`: - 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. + - The wrapped `MTKParameters` object is available as the `params` field of an + `OpaqueMTKParameters` 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 +71,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 diff --git a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl index b8840cfb0b..e61a3e5537 100644 --- a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl +++ b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl @@ -395,7 +395,8 @@ export Shift, ShiftIndex export Sample, Hold, SampleTime export Clock, SolverStepClock, TimeDomain -export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables +export MTKParameters, OpaqueMTKParameters, reorder_dimension_by_tunables!, + reorder_dimension_by_tunables export HomotopyContinuationProblem diff --git a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl index ac63bfe59b..f2b183aebe 100644 --- a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl +++ b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl @@ -219,6 +219,9 @@ function define_params(p::MTKParameters, t, names = nothing) end end +define_params(p::OpaqueMTKParameters, t, names = nothing) = + define_params(p.params, t, names) + """ $(TYPEDSIGNATURES) @@ -233,6 +236,8 @@ function to_paramvec(p::MTKParameters) return reduce(vcat, collect(p); init = []) end +to_paramvec(p::OpaqueMTKParameters) = to_paramvec(p.params) + """ $(TYPEDSIGNATURES) @@ -390,7 +395,8 @@ function defaults_from_u0_p(prob, vars, paramobj, paramvec) if !(p isa Union{SciMLBase.NullParameters, Nothing}) if p isa Union{NamedTuple, AbstractDict} merge!(defaults, Dict(v => p[k] for (k, v) in pairs(paramobj))) - elseif p isa MTKParameters + elseif p isa Union{MTKParameters, OpaqueMTKParameters} + p = _unwrap_mtk_parameters(p) pvals = [ p.tunable; reduce(vcat, p.discrete; init = []); reduce(vcat, p.constant; init = []); diff --git a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl index 6b9b0a9be6..dd778840b1 100644 --- a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl @@ -188,6 +188,13 @@ function updateparams!( return nothing end +function updateparams!( + ratemap::JumpSysMajParamMapper{U, V, W}, + params::OpaqueMTKParameters + ) where {U <: AbstractArray, V <: AbstractArray, W} + return updateparams!(ratemap, params.params) +end + function updateparams!( ::JumpSysMajParamMapper{U, V, W}, params::Nothing diff --git a/lib/ModelingToolkitBase/src/problems/odeproblem.jl b/lib/ModelingToolkitBase/src/problems/odeproblem.jl index 171ba0d998..26921c2125 100644 --- a/lib/ModelingToolkitBase/src/problems/odeproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/odeproblem.jl @@ -145,6 +145,10 @@ Base.@nospecializeinfer function _ode_problem( sys; expression, callback, eval_expression, eval_module, op, _skip_events, tspan, kwargs... ) + if spec === SciMLBase.AutoSpecialize && p isa MTKParameters + p = OpaqueMTKParameters(p) + end + ptype = getmetadata(sys, ProblemTypeCtx, StandardODEProblem()) args = (; f, u0, tspan, p, ptype) return maybe_codegen_scimlproblem(expression, ODEProblem{_iip}, args; kwargs...) diff --git a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl index b7fcc2f498..12c937e5f3 100644 --- a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl @@ -1026,6 +1026,18 @@ end # function internals, which improves `juliac` trimmability. Base.deepcopy_internal(gfw::GeneratedFunctionWrapper, ::IdDict) = gfw +Base.@noinline _invoke_with_unwrapped_parameters(f, a, b) = f(a, b) +Base.@noinline _invoke_with_unwrapped_parameters(f, a, b, c) = f(a, b, c) +Base.@noinline _invoke_with_unwrapped_parameters(f, a, b, c, d) = f(a, b, c, d) +Base.@noinline _invoke_with_unwrapped_parameters(f, args...) = f(args...) + +Base.@inline @generated function _call_with_unwrapped_parameters( + f, args::Tuple{Vararg{Any, N}}, p::OpaqueMTKParameters, ::Val{PIdx} + ) where {N, PIdx} + call_args = [i == PIdx ? :(p.params) : :(args[$i]) for i in 1:N] + return :(_invoke_with_unwrapped_parameters(f, $(call_args...))) +end + function GeneratedFunctionWrapper{P}(::Type{Val{true}}, foop, fiip; kwargs...) where {P} return :($(GeneratedFunctionWrapper{_gfw_params_type(P)})($foop, $fiip)) end @@ -1043,6 +1055,11 @@ end function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg{Any, NArgs}) where {PIdx, NArgs, Split} # non-split systems just call it as-is Split || return gfw.f_oop(args...) + if args[PIdx] isa OpaqueMTKParameters + return _call_with_unwrapped_parameters( + gfw.f_oop, args, args[PIdx], Val(PIdx) + ) + end if args[PIdx] isa Union{Tuple, MTKParameters} && !(args[PIdx] isa Tuple{Vararg{Number}}) # for split systems, call it as-is if the parameter object is a tuple or MTKParameters # but not if it is a tuple of numbers @@ -1060,6 +1077,11 @@ function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg throw(MethodError(gfw, args)) end Split || return gfw.f_iip(args...) + if args[PIdx + 1] isa OpaqueMTKParameters + return _call_with_unwrapped_parameters( + gfw.f_iip, args, args[PIdx + 1], Val(PIdx + 1) + ) + end if args[PIdx + 1] isa Union{Tuple, MTKParameters} && !(args[PIdx + 1] isa Tuple{Vararg{Number}}) return gfw.f_iip(args...) end diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index 6765d19bcc..aa3c3799b1 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -714,6 +714,9 @@ 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::OpaqueMTKParameters) where {T} + return promote_type_with_nothing(T, p.params) +end promote_with_nothing(::Type, ::Nothing) = nothing promote_with_nothing(::Type, x::StaticVector{0}) = x @@ -740,6 +743,9 @@ function promote_with_nothing(::Type{T}, p::MTKParameters) where {T} end return p end +function promote_with_nothing(::Type{T}, p::OpaqueMTKParameters) where {T} + return OpaqueMTKParameters(promote_with_nothing(T, p.params)) +end function promote_u0_p(u0, p, t0) T = Union{} @@ -865,8 +871,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 diff --git a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl index 2273ce3c47..1fcc56aada 100644 --- a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl +++ b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl @@ -40,6 +40,36 @@ struct MTKParameters{T, I, D, C, N, H} end end +""" + OpaqueMTKParameters(params) + +Wrap an [`MTKParameters`](@ref) object in a stable outer type. ModelingToolkit uses this +for `AutoSpecialize` ODE problems so that solver compilation does not specialize on the +concrete parameter-buffer layout. Symbolic indexing and SciMLStructures operations are +forwarded to the wrapped parameter object. + +The wrapped object is available through the `params` field. Construct the problem with +`FullSpecialize` to keep the concrete `MTKParameters` type and avoid the dynamic function +barrier when runtime performance takes priority over compilation reuse. +""" +struct OpaqueMTKParameters + params::Any +end + +OpaqueMTKParameters(params::OpaqueMTKParameters) = params + +_unwrap_mtk_parameters(params) = params +_unwrap_mtk_parameters(params::OpaqueMTKParameters) = params.params + +function Base.getproperty(params::OpaqueMTKParameters, name::Symbol) + name === :params && return getfield(params, :params) + return getproperty(getfield(params, :params), name) +end + +function Base.propertynames(params::OpaqueMTKParameters, private::Bool = false) + return (:params, propertynames(getfield(params, :params), private)...) +end + """ function MTKParameters(sys::AbstractSystem, p, u0 = Dict(); t0 = nothing) @@ -367,8 +397,34 @@ end SciMLStructures.isscimlstructure(::MTKParameters) = true +SciMLStructures.isscimlstructure(::OpaqueMTKParameters) = true + SciMLStructures.ismutablescimlstructure(::MTKParameters) = true +SciMLStructures.ismutablescimlstructure(::OpaqueMTKParameters) = true + +function SciMLStructures.canonicalize( + portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters + ) + values, repack, aliases = SciMLStructures.canonicalize(portion, p.params) + if repack === nothing + return values, nothing, aliases + end + return values, new_values -> OpaqueMTKParameters(repack(new_values)), aliases +end + +function SciMLStructures.replace( + portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters, new_values + ) + return OpaqueMTKParameters(SciMLStructures.replace(portion, p.params, new_values)) +end + +function SciMLStructures.replace!( + portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters, new_values + ) + return SciMLStructures.replace!(portion, p.params, new_values) +end + function SciMLStructures.canonicalize(::SciMLStructures.Tunable, p::MTKParameters) arr = p.tunable repack = let p = p @@ -453,6 +509,10 @@ function Base.copy(p::MTKParameters) ) end +Base.copy(p::OpaqueMTKParameters) = OpaqueMTKParameters(copy(p.params)) + +ArrayInterface.ismutable(::Type{OpaqueMTKParameters}) = true + function ArrayInterface.ismutable( ::Type{ MTKParameters{ @@ -469,6 +529,11 @@ end function SymbolicIndexingInterface.parameter_values(p::MTKParameters, pind::ParameterIndex) return _ducktyped_parameter_values(p, pind) end +function SymbolicIndexingInterface.parameter_values( + p::OpaqueMTKParameters, pind::ParameterIndex + ) + return parameter_values(p.params, pind) +end function _ducktyped_parameter_values(p, pind::ParameterIndex) @unpack portion, idx = pind if portion isa SciMLStructures.Tunable @@ -540,6 +605,12 @@ function SymbolicIndexingInterface.set_parameter!( return nothing end +function SymbolicIndexingInterface.set_parameter!( + p::OpaqueMTKParameters, val, pidx::ParameterIndex + ) + return set_parameter!(p.params, val, pidx) +end + function narrow_buffer_type_and_fallback_undefs( oldbuf::AbstractVector, newbuf::AbstractVector ) @@ -659,6 +730,12 @@ function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, id return _remake_buffer(indp, oldbuf, idxs, vals) end +function SymbolicIndexingInterface.remake_buffer( + indp, oldbuf::OpaqueMTKParameters, idxs, vals + ) + return OpaqueMTKParameters(remake_buffer(indp, oldbuf.params, idxs, vals)) +end + function _remake_buffer(indp, oldbuf::MTKParameters, idxs, vals; validate = true) return __remake_buffer(indp, oldbuf, idxs, vals; validate) end @@ -1011,22 +1088,25 @@ end Base.size(::NestedGetIndex) = () function SymbolicIndexingInterface.with_updated_parameter_timeseries_values( - ::AbstractSystem, ps::MTKParameters, args::Pair{<:Any, <:NestedGetIndex}... + ::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, + args::Pair{<:Any, <:NestedGetIndex}... ) + unwrapped_ps = _unwrap_mtk_parameters(ps) for (i, ngi) in args for (j, val) in enumerate(ngi.x) - copyto!(view(ps.discrete[j], Block(i)), val) + copyto!(view(unwrapped_ps.discrete[j], Block(i)), val) end end return ps end function SciMLBase.create_parameter_timeseries_collection( - sys::AbstractSystem, ps::MTKParameters, tspan + sys::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, tspan ) ic = get_index_cache(sys) # this exists because the parameters are `MTKParameters` - isempty(ps.discrete) && return nothing - num_discretes = only(blocksize(ps.discrete[1])) + unwrapped_ps = _unwrap_mtk_parameters(ps) + isempty(unwrapped_ps.discrete) && return nothing + num_discretes = only(blocksize(unwrapped_ps.discrete[1])) buffers = [] partition_type = typeof(SciMLBase.get_saveable_values(sys, ps, 1)) for i in 1:num_discretes @@ -1047,8 +1127,9 @@ end end function SciMLBase.get_saveable_values( - sys::AbstractSystem, ps::MTKParameters, timeseries_idx + sys::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, timeseries_idx ) + ps = _unwrap_mtk_parameters(ps) return NestedGetIndex(__get_blocks(timeseries_idx, ps.discrete...)) end @@ -1069,6 +1150,11 @@ function DiffEqBase.anyeltypedual( ) where {counter} return DiffEqBase.anyeltypedual(p.tunable) end +function DiffEqBase.anyeltypedual( + p::OpaqueMTKParameters, ::Type{Val{counter}} = Val{0} + ) where {counter} + return DiffEqBase.anyeltypedual(p.params, Val{counter}) +end function DiffEqBase.anyeltypedual( p::Type{<:MTKParameters{T}}, ::Type{Val{counter}} = Val{0} @@ -1130,10 +1216,24 @@ Base.size(ps::MTKParameters) = (length(ps),) Base.IndexStyle(::Type{T}) where {T <: MTKParameters} = IndexLinear() +Base.IndexStyle(::Type{OpaqueMTKParameters}) = IndexLinear() + +Base.eltype(::Type{OpaqueMTKParameters}) = Any + Base.getindex(p::MTKParameters, pind::ParameterIndex) = parameter_values(p, pind) +Base.getindex(p::OpaqueMTKParameters, idx) = getindex(p.params, idx) + Base.setindex!(p::MTKParameters, val, pind::ParameterIndex) = set_parameter!(p, val, pind) +Base.setindex!(p::OpaqueMTKParameters, val, idx) = setindex!(p.params, val, idx) + +Base.length(p::OpaqueMTKParameters) = length(p.params) + +Base.size(p::OpaqueMTKParameters) = size(p.params) + +Base.iterate(p::OpaqueMTKParameters, state = 1) = iterate(p.params, state) + function Base.iterate(buf::MTKParameters, state = 1) total_len = length(buf) if state <= total_len @@ -1153,6 +1253,8 @@ function Base.:(==)(a::MTKParameters, b::MTKParameters) ) end +Base.:(==)(a::OpaqueMTKParameters, b::OpaqueMTKParameters) = a.params == b.params + const MISSING_PARAMETERS_MESSAGE = """ Some parameters are missing from the variable map. Please provide a value or default for the following variables: diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index ff6e35ab18..92a24f557c 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -817,7 +817,7 @@ struct FallbackSlice end function __apply_copy_template(valp, template) - p = parameter_values(valp) + p = _unwrap_mtk_parameters(parameter_values(valp)) u = state_values(valp) if template isa ParameterIndex{SciMLStructures.Tunable, UnitRange{Int}} if p isa MTKParameters @@ -1113,7 +1113,8 @@ end function (recon::MTKParametersReconstructor)(src, dst) src_ps = parameter_values(src) dst_ps = parameter_values(dst) - oldcache = dst_ps.caches + unwrapped_dst_ps = _unwrap_mtk_parameters(dst_ps) + oldcache = unwrapped_dst_ps.caches # I don't know why but this makes it infer properly if recon.tunables_fn isa ComposedFunction tunablevals = recon.tunables_fn.outer(recon.tunables_fn.inner(src)) @@ -1121,7 +1122,7 @@ function (recon::MTKParametersReconstructor)(src, dst) tunablevals = recon.tunables_fn(src) end initialvals = recon.initials_fn(src) - nonnumerics = recon.nonnumerics_fn(src)::typeof(dst_ps.nonnumeric) + nonnumerics = recon.nonnumerics_fn(src)::typeof(unwrapped_dst_ps.nonnumeric) (; diffcache_buffer_idx) = recon if !iszero(diffcache_buffer_idx) @set! nonnumerics[diffcache_buffer_idx] = DiffCacheAllocatorAPIWrapper{ForwardDiff.valtype(eltype(initialvals))}.(nonnumerics[diffcache_buffer_idx]) @@ -1129,11 +1130,12 @@ function (recon::MTKParametersReconstructor)(src, dst) # This `convert` exists because a `Real` discrete might get its value from an # integer function of integer parameters/discretes. This ends up creating a # `BlockedArray{Int, ...}` instead of a `BlockedArray{Float64, ...}`. - return MTKParameters( + new_ps = MTKParameters( tunablevals, initialvals, - convert(typeof(parameter_values(dst).discrete), recon.discretes_fn(src)), + convert(typeof(unwrapped_dst_ps.discrete), recon.discretes_fn(src)), recon.consts_fn(src), nonnumerics, oldcache isa Tuple{} ? () : copy.(oldcache) ) + return dst_ps isa OpaqueMTKParameters ? OpaqueMTKParameters(new_ps) : new_ps end """ @@ -1310,9 +1312,10 @@ function (rip::ReconstructInitializeprob)(srcvalp, dstvalp) srcu0 = state_values(srcvalp) T = srcu0 === nothing ? Union{} : eltype(srcu0) # promote with the tunable eltype - if parameter_values(dstvalp) isa MTKParameters - if !isempty(newp.tunable) - T = promote_type(eltype(newp.tunable), T) + if _unwrap_mtk_parameters(parameter_values(dstvalp)) isa MTKParameters + unwrapped_newp = _unwrap_mtk_parameters(newp) + if !isempty(unwrapped_newp.tunable) + T = promote_type(eltype(unwrapped_newp.tunable), T) end elseif !isempty(newp) T = promote_type(eltype(newp), T) @@ -1330,7 +1333,7 @@ function (rip::ReconstructInitializeprob)(srcvalp, dstvalp) copyto!(newbuf, buf) newp = repack(newbuf) end - if newp isa MTKParameters + if _unwrap_mtk_parameters(newp) isa MTKParameters # and initials portion buf, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Initials(), newp) if eltype(buf) != T && !(buf isa SVector{0}) @@ -1533,6 +1536,10 @@ function (siu::SetInitialUnknowns)(p::MTKParameters, u0) return p end +function (siu::SetInitialUnknowns)(p::OpaqueMTKParameters, u0) + return OpaqueMTKParameters(siu(p.params, u0)) +end + function (siu::SetInitialUnknowns)(p::AbstractVector, u0) if ArrayInterface.ismutable(p) siu.setter!(p, u0) diff --git a/lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl b/lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl new file mode 100644 index 0000000000..b1280bbae1 --- /dev/null +++ b/lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl @@ -0,0 +1,28 @@ +using ForwardDiff +using ModelingToolkitBase +using ModelingToolkitBase: t_nounits as t, D_nounits as D +using OrdinaryDiffEqTsit5 +using SciMLSensitivity +using SciMLStructures +using Test +using Zygote + +@testset "opaque parameter sensitivities" begin + @parameters a = 2.0 + @variables x(t) = 1.0 + sys = mtkcompile(System([D(x) ~ -a * x], t; name = :opaque_parameter_ad)) + prob = ODEProblem(sys, [], (0.0, 1.0)) + + function terminal_value(a_value) + p = SciMLStructures.replace(SciMLStructures.Tunable(), prob.p, [a_value]) + remade = remake(prob; p) + sol = solve( + remade, Tsit5(); saveat = [1.0], abstol = 1.0e-10, reltol = 1.0e-10 + ) + return sol[x][end] + end + + expected = -exp(-2) + @test ForwardDiff.derivative(terminal_value, 2.0) ≈ expected rtol = 1.0e-6 + @test only(Zygote.gradient(terminal_value, 2.0)) ≈ expected rtol = 1.0e-6 +end diff --git a/lib/ModelingToolkitBase/test/index_cache.jl b/lib/ModelingToolkitBase/test/index_cache.jl index 24e886f48b..6d5d9a5194 100644 --- a/lib/ModelingToolkitBase/test/index_cache.jl +++ b/lib/ModelingToolkitBase/test/index_cache.jl @@ -168,7 +168,16 @@ end @test SciMLBase.successful_retcode(solve(prob, Tsit5())) idata = prob.f.initialization_data - @test_nowarn @inferred idata.metadata.oop_reconstruct_u0_p.pgetter(prob, idata.initializeprob) + reconstructed_p = idata.metadata.oop_reconstruct_u0_p.pgetter( + prob, idata.initializeprob + ) + @test reconstructed_p isa MTKParameters + + full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [x => 1.0], (0.0, 1.0)) + idata = full_prob.f.initialization_data + @test_nowarn @inferred idata.metadata.oop_reconstruct_u0_p.pgetter( + full_prob, idata.initializeprob + ) end function costfn(theta, ps) diff --git a/lib/ModelingToolkitBase/test/jacobiansparsity.jl b/lib/ModelingToolkitBase/test/jacobiansparsity.jl index 3282dd6457..a5a00e8e0b 100644 --- a/lib/ModelingToolkitBase/test/jacobiansparsity.jl +++ b/lib/ModelingToolkitBase/test/jacobiansparsity.jl @@ -74,8 +74,11 @@ prob = ODEProblem(sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = true, jac #@test_nowarn solve(prob, Rosenbrock23()) @test findnz(calculate_jacobian(sys, sparse = true))[1:2] == findnz(prob.f.jac_prototype)[1:2] -out = similar(prob.f.jac_prototype) -@test (@ballocated $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), 0.0)) == 0 # should not allocate +full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( + sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = true, jac = true +) +out = similar(full_prob.f.jac_prototype) +@test (@ballocated $(full_prob.f.jac)($out, $(full_prob.u0), $(full_prob.p), 0.0)) == 0 # test when not sparse prob = ODEProblem(sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = false, jac = true) diff --git a/lib/ModelingToolkitBase/test/mtkparameters.jl b/lib/ModelingToolkitBase/test/mtkparameters.jl index 1f1c3fffcd..7ff56948a1 100644 --- a/lib/ModelingToolkitBase/test/mtkparameters.jl +++ b/lib/ModelingToolkitBase/test/mtkparameters.jl @@ -125,7 +125,7 @@ end @test getp(sys, g)(newps) isa Vector{Float32} @testset "Type-stability of `remake_buffer`" begin - prob = ODEProblem(sys, ivs, (0.0, 1.0)) + prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, ivs, (0.0, 1.0)) idxs = (a, c, d, e, f, g, h) vals = (1.0, 2.0, 3, ones(3), ones(Int, 3, 3), ones(2), "a") @@ -482,7 +482,9 @@ end @parameters p @named sys = System(D(x) ~ x * p, t) sys = complete(sys) - prob = ODEProblem(sys, SA[x => 1.0, p => 1.0], (0.0, 1.0)) + prob = ODEProblem{true, SciMLBase.FullSpecialize}( + sys, SA[x => 1.0, p => 1.0], (0.0, 1.0) + ) @test isbits(prob.p) @test isbits(prob.f.initialization_data.initializeprob.p) end diff --git a/lib/ModelingToolkitBase/test/odesystem.jl b/lib/ModelingToolkitBase/test/odesystem.jl index c1902fa361..20de35f498 100644 --- a/lib/ModelingToolkitBase/test/odesystem.jl +++ b/lib/ModelingToolkitBase/test/odesystem.jl @@ -1,5 +1,6 @@ using ModelingToolkitBase, StaticArrays, LinearAlgebra -using ModelingToolkitBase: get_metadata, MTKParameters, SymbolicDiscreteCallback, +using ModelingToolkitBase: get_metadata, MTKParameters, OpaqueMTKParameters, + SymbolicDiscreteCallback, SymbolicContinuousCallback using SymbolicIndexingInterface using OrdinaryDiffEq, Sundials @@ -269,7 +270,7 @@ prob12 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tsp prob13 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tspan) prob14 = ODEProblem(sys, [u0; p2], tspan) for p in [prob1, prob14] - @test p.p isa MTKParameters + @test p.p isa OpaqueMTKParameters p.ps[k₁] ≈ 0.04 p.ps[k₂] ≈ 3.0e7 p.ps[k₃] ≈ 1.0e-4 @@ -285,7 +286,7 @@ u01 = [y₁ => 1, y₂ => 1, y₃ => 1] prob_pmap = remake(prob14; p = p3, u0 = u01) prob_dpmap = remake(prob14; p = Dict(p3), u0 = Dict(u01)) for p in [prob_pmap, prob_dpmap] - @test p.p isa MTKParameters + @test p.p isa OpaqueMTKParameters p.ps[k₁] ≈ 0.05 p.ps[k₂] ≈ 2.0e7 p.ps[k₃] ≈ 1.1e-4 @@ -320,7 +321,7 @@ sol_dpmap = solve(prob_dpmap, Rodas5()) prob, p = Dict(sys1.a => 3.0, b => 4.0), u0 = Dict(sys1.x => 1.0) ) - @test prob_new.p isa MTKParameters + @test prob_new.p isa OpaqueMTKParameters @test prob_new.ps[b] ≈ 4.0 @test prob_new.ps[sys1.a] ≈ 3.0 @test prob_new.ps[sys.sys2.a] ≈ 1.0 @@ -665,7 +666,7 @@ let ivmap = [A => 1.0, k1 => 1.0, k2 => 1.0] tspan = (0.0, 1.0) prob = ODEProblem(sys, ivmap, tspan; tofloat = false) - @test prob.p isa MTKParameters + @test prob.p isa OpaqueMTKParameters @test prob.ps[k1] ≈ 1.0 @test prob.ps[k2] == 1 && prob.ps[k2] isa Int end diff --git a/lib/ModelingToolkitBase/test/opaque_mtkparameters.jl b/lib/ModelingToolkitBase/test/opaque_mtkparameters.jl new file mode 100644 index 0000000000..36d026271a --- /dev/null +++ b/lib/ModelingToolkitBase/test/opaque_mtkparameters.jl @@ -0,0 +1,86 @@ +using ModelingToolkitBase +using ModelingToolkitBase: t_nounits as t, D_nounits as D, MTKParameters +using OrdinaryDiffEq +using Random +using SciMLStructures +using SymbolicIndexingInterface +using Test + +@testset "opaque AutoSpecialize parameters" begin + @parameters a = 2.0 + @variables x(t) = 1.0 z(t) + sys = mtkcompile( + System( + [D(x) ~ -a * x], t; + observed = [z ~ a * x], name = :opaque_parameters + ) + ) + + auto_prob = ODEProblem(sys, [], (0.0, 1.0); jac = true) + full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0.0, 1.0)) + expression_prob = eval(ODEProblem(sys, [], (0.0, 1.0); expression = Val{true})) + + @test auto_prob.p isa ModelingToolkitBase.OpaqueMTKParameters + @test auto_prob.p.params isa MTKParameters + @test auto_prob.p.tunable === auto_prob.p.params.tunable + @test eltype(typeof(auto_prob.p)) === Any + @test SciMLStructures.isscimlstructure(auto_prob.p) + @test SciMLStructures.ismutablescimlstructure(auto_prob.p) + @test full_prob.p isa MTKParameters + @test expression_prob.p isa ModelingToolkitBase.OpaqueMTKParameters + + du = similar(auto_prob.u0) + auto_prob.f(du, auto_prob.u0, auto_prob.p, first(auto_prob.tspan)) + @test du == [-2.0] + jac = zeros(1, 1) + auto_prob.f.jac(jac, auto_prob.u0, auto_prob.p, first(auto_prob.tspan)) + @test jac == [-2.0;;] + + get_a = getp(sys, a) + set_a! = setp(sys, a) + @test get_a(auto_prob.p) == 2.0 + set_a!(auto_prob.p, 3.0) + @test get_a(auto_prob.p) == 3.0 + + tunables, repack, aliases = SciMLStructures.canonicalize( + SciMLStructures.Tunable(), auto_prob.p + ) + @test aliases + @test tunables == [3.0] + new_p = repack([4.0]) + @test new_p isa ModelingToolkitBase.OpaqueMTKParameters + @test get_a(new_p) == 4.0 + @test copy(new_p) == new_p + @test SciMLStructures.replace!(SciMLStructures.Tunable(), new_p, [5.0]) === nothing + @test get_a(new_p) == 5.0 + replaced_p = SciMLStructures.replace(SciMLStructures.Tunable(), new_p, [4.0]) + @test replaced_p isa ModelingToolkitBase.OpaqueMTKParameters + @test get_a(replaced_p) == 4.0 + + other_inner = SciMLStructures.replace( + SciMLStructures.Tunable(), new_p.params, Float32[5] + ) + @test typeof(new_p) === typeof(ModelingToolkitBase.OpaqueMTKParameters(other_inner)) + @test typeof(new_p.params) !== typeof(other_inner) + other_prob = remake(auto_prob; p = ModelingToolkitBase.OpaqueMTKParameters(other_inner)) + @test typeof(other_prob) === typeof(auto_prob) + + remade = remake(auto_prob; p = [a => 4.0]) + @test remade.p isa ModelingToolkitBase.OpaqueMTKParameters + sol = solve(remade, Tsit5(); abstol = 1.0e-10, reltol = 1.0e-10) + @test sol[x][end] ≈ exp(-4) + @test sol[z][end] ≈ 4exp(-4) + + @parameters drift = 0.0 rate = 1.0 + @variables population(t) = 10.0 + jump = SymbolicMassActionJump(rate, [population => 1], [population => -1]) + hybrid_sys = mtkcompile( + System( + [D(population) ~ drift], t, [population], [drift, rate]; jumps = [jump], + name = :opaque_jump_parameters + ) + ) + hybrid_prob = JumpProblem(hybrid_sys, [], (0.0, 0.1); rng = Xoshiro(1)) + @test hybrid_prob.prob.p isa ModelingToolkitBase.OpaqueMTKParameters + @test SciMLBase.successful_retcode(solve(hybrid_prob, Tsit5())) +end diff --git a/lib/ModelingToolkitBase/test/runtests.jl b/lib/ModelingToolkitBase/test/runtests.jl index 9349e53a6b..bbd9f1e36c 100644 --- a/lib/ModelingToolkitBase/test/runtests.jl +++ b/lib/ModelingToolkitBase/test/runtests.jl @@ -115,6 +115,7 @@ end @safetestset "SymbolicIndexingInterface test" include("symbolic_indexing_interface.jl") @safetestset "SciML Problem Input Test" include("sciml_problem_inputs.jl") @safetestset "MTKParameters Test" include("mtkparameters.jl") + @safetestset "Opaque MTKParameters Test" include("opaque_mtkparameters.jl") end if GROUP == "All" || GROUP == "Extended" @@ -135,6 +136,7 @@ end @safetestset "HomotopyContinuation Extension Test" include("extensions/homotopy_continuation.jl") @safetestset "LabelledArrays Test" include("extensions/labelledarrays.jl") @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl") + @safetestset "Opaque MTKParameters AD Test" include("extensions/opaque_mtkparameters_ad.jl") # @safetestset "Auto Differentiation Test" include("extensions/ad.jl") end diff --git a/lib/ModelingToolkitBase/test/split_parameters.jl b/lib/ModelingToolkitBase/test/split_parameters.jl index 1c4f44ea2a..767d8b2d07 100644 --- a/lib/ModelingToolkitBase/test/split_parameters.jl +++ b/lib/ModelingToolkitBase/test/split_parameters.jl @@ -254,15 +254,27 @@ end getter = getp(sys, fn) prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0)) - @inferred getter(prob) + @test getter(prob)(1.0) == 2.0 + full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( + sys, [x => 1.0], (0.0, 1.0) + ) + @inferred getter(full_prob) # cannot be inferred better since `FunctionWrapper` is only known to return `Real` - @inferred Vector{<:Real} prob.f(prob.u0, prob.p, prob.tspan[1]) + @inferred Vector{<:Real} full_prob.f( + full_prob.u0, full_prob.p, full_prob.tspan[1] + ) sol = solve(prob, Tsit5(); abstol = 1.0e-10, reltol = 1.0e-10) @test sol.u[end][] ≈ 2.0 prob = ODEProblem(sys, [x => 1.0, fn => Foo()], (0.0, 1.0)) - @inferred getter(prob) - @inferred Vector{<:Real} prob.f(prob.u0, prob.p, prob.tspan[1]) + @test getter(prob)(1.0) == 3.0 + full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( + sys, [x => 1.0, fn => Foo()], (0.0, 1.0) + ) + @inferred getter(full_prob) + @inferred Vector{<:Real} full_prob.f( + full_prob.u0, full_prob.p, full_prob.tspan[1] + ) sol = solve(prob; abstol = 1.0e-10, reltol = 1.0e-10) @test sol.u[end][] ≈ 2.5 end @@ -278,8 +290,12 @@ end @test is_parameter(sys, fn) getter = getp(sys, fn) prob = ODEProblem(sys, [x => 1.0, fn => interp], (0.0, 1.0)) - @inferred getter(prob) - @inferred prob.f(prob.u0, prob.p, prob.tspan[1]) + @test getter(prob) === interp + full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( + sys, [x => 1.0, fn => interp], (0.0, 1.0) + ) + @inferred getter(full_prob) + @inferred full_prob.f(full_prob.u0, full_prob.p, full_prob.tspan[1]) @test_nowarn sol = solve(prob, Tsit5()) @test_nowarn prob.ps[fn] = LinearInterpolation( ts .^ 3, ts; extrapolation = ExtrapolationType.Extension diff --git a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl index 6b3197cd0a..d3623d4d1d 100644 --- a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl +++ b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl @@ -45,7 +45,9 @@ using Test @test default_values(odesys)[y] == 2.0 @test isequal(default_values(odesys)[xy], x + y) - prob = ODEProblem(odesys, [a => 1.0, b => 2.0], (0.0, 1.0)) + prob = ODEProblem{true, SciMLBase.FullSpecialize}( + odesys, [a => 1.0, b => 2.0], (0.0, 1.0) + ) getter = getu(odesys, (x + 1, x + 2)) @test getter(prob) isa Tuple @test_nowarn @inferred getter(prob) diff --git a/src/linearization.jl b/src/linearization.jl index 7519dd0f7d..08b3ef6c59 100644 --- a/src/linearization.jl +++ b/src/linearization.jl @@ -258,7 +258,13 @@ function linearization_function( if u0 === nothing uf_jac = h_jac = pf_jac = nothing - Tp = promote_type(p isa MTKParameters ? eltype(p.tunable) : eltype(p), typeof(t0)) + parameter_eltype = if p isa Union{MTKParameters, OpaqueMTKParameters} + tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p) + eltype(tunables) + else + eltype(p) + end + Tp = promote_type(parameter_eltype, typeof(t0)) hp_jac = PreparedJacobian{true}( hp_fun, zeros(Tp, size(outputs)), autodiff, inputvals, cu0T, cp, DI.Constant(t0) diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl index 9769ef6cab..82e45b63c7 100644 --- a/src/problems/sccnonlinearproblem.jl +++ b/src/problems/sccnonlinearproblem.jl @@ -6,6 +6,8 @@ function (cw::CacheWriter)(p::MTKParameters, sols) return cw.fn(p.caches, sols, p) end +(cw::CacheWriter)(p::OpaqueMTKParameters, sols) = cw(p.params, sols) + const SCCCacheVarsExprsElT = Dict{TypeT, Vector{SymbolicT}} const SCC_EXPLICITFUN_CACHE_OUT = unwrap(only(@parameters __outₘₜₖ::Vector{Vector{Any}})) diff --git a/test/qa/qa.jl b/test/qa/qa.jl index 3d0ba2a7a3..ed825737a9 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -134,6 +134,7 @@ const REEXPORTED_API = ( :JumpSystem, :linear_fractional_to_ordinary, :liouville_transform, :LocalScope, :maybe_zeros, :MissingGuessValue, :ModelingToolkitBase, :modelingtoolkitize, :modified_unknowns!, :mtkcompile, :MTKParameters, + :OpaqueMTKParameters, :MTKVariableTypeCtx, :namespace_equations, :noise_to_brownians, :NonlinearSystem, :observables, :observed, :ODESystem, :open_loop, :OptimizationSystem, :outputs, :parameters, :parameters_toplevel, :ParentScope, :PDESystem, :Pre, :ProblemTypeCtx, From b6545c9d97447b8a164e58be629bd72fe152ae10 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 10 Aug 2026 04:23:20 -0400 Subject: [PATCH 02/12] Preserve zero-allocation specialization test Co-Authored-By: Chris Rackauckas --- test/structural_transformation/tearing.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/structural_transformation/tearing.jl b/test/structural_transformation/tearing.jl index 3960dda5dd..a0650138a3 100644 --- a/test/structural_transformation/tearing.jl +++ b/test/structural_transformation/tearing.jl @@ -181,7 +181,9 @@ newdaesys = mtkcompile(daesys) equations(tearing_substitution(newdaesys)), [D(x) ~ h * z; 0 ~ x + sin(z) - p * t] ) @test issetequal(unknowns(newdaesys), [x, z]) -prob = ODEProblem(newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0)) +prob = ODEProblem{true, SciMLBase.FullSpecialize}( + newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0) +) du = [0.0, 0.0]; u = [1.0, -0.5π]; pr = prob.p; From 42755d23f19f8ea374e2779f9f9d8b60f619d2d0 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 10 Aug 2026 06:16:20 -0400 Subject: [PATCH 03/12] Use SciMLBase parameter despecialization Co-Authored-By: Chris Rackauckas --- Project.toml | 2 +- docs/src/API/problems.md | 10 +- docs/src/basics/FAQ.md | 17 +-- .../src/ModelingToolkitBase.jl | 2 +- .../src/modelingtoolkitize/common.jl | 6 +- .../src/problems/jumpproblem.jl | 2 +- .../src/problems/odeproblem.jl | 4 - .../src/systems/codegen_utils.jl | 20 +--- .../src/systems/nonlinear/initializesystem.jl | 8 +- .../src/systems/parameter_buffer.jl | 106 +----------------- .../src/systems/problem_utils.jl | 17 ++- ...ameters.jl => despecialized_parameters.jl} | 35 ++++-- ...s_ad.jl => despecialized_parameters_ad.jl} | 4 +- lib/ModelingToolkitBase/test/odesystem.jl | 10 +- lib/ModelingToolkitBase/test/runtests.jl | 4 +- src/linearization.jl | 2 +- src/problems/sccnonlinearproblem.jl | 2 +- test/qa/qa.jl | 1 - 18 files changed, 84 insertions(+), 168 deletions(-) rename lib/ModelingToolkitBase/test/{opaque_mtkparameters.jl => despecialized_parameters.jl} (67%) rename lib/ModelingToolkitBase/test/extensions/{opaque_mtkparameters_ad.jl => despecialized_parameters_ad.jl} (84%) diff --git a/Project.toml b/Project.toml index 6cff8918ba..a07e02a9ce 100644 --- a/Project.toml +++ b/Project.toml @@ -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" diff --git a/docs/src/API/problems.md b/docs/src/API/problems.md index 99f33dd4c2..19417cac50 100644 --- a/docs/src/API/problems.md +++ b/docs/src/API/problems.md @@ -90,15 +90,15 @@ ModelingToolkit.varmap_to_vars ``` The parameters of a split system are stored in a custom data structure called -`MTKParameters`. `AutoSpecialize` ODE problems wrap it in `OpaqueMTKParameters` so solvers -can reuse compiled code across parameter-buffer layouts. `FullSpecialize` problems keep the -concrete `MTKParameters` object and avoid the wrapper's dynamic function barrier. These -objects should only be interacted with through their defined public API. +`MTKParameters`. `AutoSpecialize` problems wrap it in +[`SciMLBase.DespecializedParameters`](@ref) so solvers can reuse compiled code across +parameter-buffer layouts. `FullSpecialize` problems keep the concrete `MTKParameters` +object and avoid the wrapper's dynamic function barrier. These objects should only be +interacted with through their defined public API. SymbolicIndexingInterface.jl contains functionality useful for this purpose. ```@docs MTKParameters -OpaqueMTKParameters ModelingToolkit.get_p ``` diff --git a/docs/src/basics/FAQ.md b/docs/src/basics/FAQ.md index 56250bec8c..a03d08df73 100644 --- a/docs/src/basics/FAQ.md +++ b/docs/src/basics/FAQ.md @@ -3,11 +3,12 @@ ## Why are my parameters some obscure object? In ModelingToolkit.jl version 9, the parameter vector was replaced with a custom -`MTKParameters` object. `AutoSpecialize` ODE problems expose an `OpaqueMTKParameters` -wrapper around that object so solver compilation can be reused across parameter layouts. -The internals of both types 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 +`MTKParameters` object. `AutoSpecialize` problems expose a +[`SciMLBase.DespecializedParameters`](@ref) wrapper around that object so solver compilation +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: @@ -19,8 +20,8 @@ getβ(prob) ``` 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 `OpaqueMTKParameters`: +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 @@ -30,7 +31,7 @@ similarly undocumented. The following behaviors can be relied on for both - `setindex!` with a `ParameterIndex` can be used to set the value of a parameter with the given index. - The wrapped `MTKParameters` object is available as the `params` field of an - `OpaqueMTKParameters` object. + `SciMLBase.DespecializedParameters` object. - `parameter_index(sys, sym)` will return a `ParameterIndex` object if `sys` has been `complete`d (through `mtkcompile`, `complete` or `@mtkcompile`). - `copy` duplicates the parameter object, including the memory used by the underlying diff --git a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl index e61a3e5537..f2f05d18d8 100644 --- a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl +++ b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl @@ -395,7 +395,7 @@ export Shift, ShiftIndex export Sample, Hold, SampleTime export Clock, SolverStepClock, TimeDomain -export MTKParameters, OpaqueMTKParameters, reorder_dimension_by_tunables!, +export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables export HomotopyContinuationProblem diff --git a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl index f2b183aebe..83c3e44913 100644 --- a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl +++ b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl @@ -219,7 +219,7 @@ function define_params(p::MTKParameters, t, names = nothing) end end -define_params(p::OpaqueMTKParameters, t, names = nothing) = +define_params(p::SciMLBase.DespecializedParameters, t, names = nothing) = define_params(p.params, t, names) """ @@ -236,7 +236,7 @@ function to_paramvec(p::MTKParameters) return reduce(vcat, collect(p); init = []) end -to_paramvec(p::OpaqueMTKParameters) = to_paramvec(p.params) +to_paramvec(p::SciMLBase.DespecializedParameters) = to_paramvec(p.params) """ $(TYPEDSIGNATURES) @@ -395,7 +395,7 @@ function defaults_from_u0_p(prob, vars, paramobj, paramvec) if !(p isa Union{SciMLBase.NullParameters, Nothing}) if p isa Union{NamedTuple, AbstractDict} merge!(defaults, Dict(v => p[k] for (k, v) in pairs(paramobj))) - elseif p isa Union{MTKParameters, OpaqueMTKParameters} + elseif p isa Union{MTKParameters, SciMLBase.DespecializedParameters} p = _unwrap_mtk_parameters(p) pvals = [ p.tunable; reduce(vcat, p.discrete; init = []); diff --git a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl index dd778840b1..e9ff39443e 100644 --- a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl @@ -190,7 +190,7 @@ end function updateparams!( ratemap::JumpSysMajParamMapper{U, V, W}, - params::OpaqueMTKParameters + params::SciMLBase.DespecializedParameters ) where {U <: AbstractArray, V <: AbstractArray, W} return updateparams!(ratemap, params.params) end diff --git a/lib/ModelingToolkitBase/src/problems/odeproblem.jl b/lib/ModelingToolkitBase/src/problems/odeproblem.jl index 26921c2125..171ba0d998 100644 --- a/lib/ModelingToolkitBase/src/problems/odeproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/odeproblem.jl @@ -145,10 +145,6 @@ Base.@nospecializeinfer function _ode_problem( sys; expression, callback, eval_expression, eval_module, op, _skip_events, tspan, kwargs... ) - if spec === SciMLBase.AutoSpecialize && p isa MTKParameters - p = OpaqueMTKParameters(p) - end - ptype = getmetadata(sys, ProblemTypeCtx, StandardODEProblem()) args = (; f, u0, tspan, p, ptype) return maybe_codegen_scimlproblem(expression, ODEProblem{_iip}, args; kwargs...) diff --git a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl index 12c937e5f3..c1265c9ec0 100644 --- a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl @@ -1026,18 +1026,6 @@ end # function internals, which improves `juliac` trimmability. Base.deepcopy_internal(gfw::GeneratedFunctionWrapper, ::IdDict) = gfw -Base.@noinline _invoke_with_unwrapped_parameters(f, a, b) = f(a, b) -Base.@noinline _invoke_with_unwrapped_parameters(f, a, b, c) = f(a, b, c) -Base.@noinline _invoke_with_unwrapped_parameters(f, a, b, c, d) = f(a, b, c, d) -Base.@noinline _invoke_with_unwrapped_parameters(f, args...) = f(args...) - -Base.@inline @generated function _call_with_unwrapped_parameters( - f, args::Tuple{Vararg{Any, N}}, p::OpaqueMTKParameters, ::Val{PIdx} - ) where {N, PIdx} - call_args = [i == PIdx ? :(p.params) : :(args[$i]) for i in 1:N] - return :(_invoke_with_unwrapped_parameters(f, $(call_args...))) -end - function GeneratedFunctionWrapper{P}(::Type{Val{true}}, foop, fiip; kwargs...) where {P} return :($(GeneratedFunctionWrapper{_gfw_params_type(P)})($foop, $fiip)) end @@ -1055,8 +1043,8 @@ end function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg{Any, NArgs}) where {PIdx, NArgs, Split} # non-split systems just call it as-is Split || return gfw.f_oop(args...) - if args[PIdx] isa OpaqueMTKParameters - return _call_with_unwrapped_parameters( + if args[PIdx] isa SciMLBase.DespecializedParameters + return SciMLBase.invoke_with_despecialized_parameters( gfw.f_oop, args, args[PIdx], Val(PIdx) ) end @@ -1077,8 +1065,8 @@ function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg throw(MethodError(gfw, args)) end Split || return gfw.f_iip(args...) - if args[PIdx + 1] isa OpaqueMTKParameters - return _call_with_unwrapped_parameters( + if args[PIdx + 1] isa SciMLBase.DespecializedParameters + return SciMLBase.invoke_with_despecialized_parameters( gfw.f_iip, args, args[PIdx + 1], Val(PIdx + 1) ) end diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index aa3c3799b1..a95cdf884c 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -714,7 +714,9 @@ 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::OpaqueMTKParameters) where {T} +function promote_type_with_nothing( + ::Type{T}, p::SciMLBase.DespecializedParameters + ) where {T} return promote_type_with_nothing(T, p.params) end @@ -743,8 +745,8 @@ function promote_with_nothing(::Type{T}, p::MTKParameters) where {T} end return p end -function promote_with_nothing(::Type{T}, p::OpaqueMTKParameters) where {T} - return OpaqueMTKParameters(promote_with_nothing(T, p.params)) +function promote_with_nothing(::Type{T}, p::SciMLBase.DespecializedParameters) where {T} + return SciMLBase.DespecializedParameters(promote_with_nothing(T, p.params)) end function promote_u0_p(u0, p, t0) diff --git a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl index 1fcc56aada..c6056e1a13 100644 --- a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl +++ b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl @@ -40,35 +40,7 @@ struct MTKParameters{T, I, D, C, N, H} end end -""" - OpaqueMTKParameters(params) - -Wrap an [`MTKParameters`](@ref) object in a stable outer type. ModelingToolkit uses this -for `AutoSpecialize` ODE problems so that solver compilation does not specialize on the -concrete parameter-buffer layout. Symbolic indexing and SciMLStructures operations are -forwarded to the wrapped parameter object. - -The wrapped object is available through the `params` field. Construct the problem with -`FullSpecialize` to keep the concrete `MTKParameters` type and avoid the dynamic function -barrier when runtime performance takes priority over compilation reuse. -""" -struct OpaqueMTKParameters - params::Any -end - -OpaqueMTKParameters(params::OpaqueMTKParameters) = params - -_unwrap_mtk_parameters(params) = params -_unwrap_mtk_parameters(params::OpaqueMTKParameters) = params.params - -function Base.getproperty(params::OpaqueMTKParameters, name::Symbol) - name === :params && return getfield(params, :params) - return getproperty(getfield(params, :params), name) -end - -function Base.propertynames(params::OpaqueMTKParameters, private::Bool = false) - return (:params, propertynames(getfield(params, :params), private)...) -end +_unwrap_mtk_parameters(params) = SciMLBase.unwrap_parameters(params) """ function MTKParameters(sys::AbstractSystem, p, u0 = Dict(); t0 = nothing) @@ -397,34 +369,8 @@ end SciMLStructures.isscimlstructure(::MTKParameters) = true -SciMLStructures.isscimlstructure(::OpaqueMTKParameters) = true - SciMLStructures.ismutablescimlstructure(::MTKParameters) = true -SciMLStructures.ismutablescimlstructure(::OpaqueMTKParameters) = true - -function SciMLStructures.canonicalize( - portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters - ) - values, repack, aliases = SciMLStructures.canonicalize(portion, p.params) - if repack === nothing - return values, nothing, aliases - end - return values, new_values -> OpaqueMTKParameters(repack(new_values)), aliases -end - -function SciMLStructures.replace( - portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters, new_values - ) - return OpaqueMTKParameters(SciMLStructures.replace(portion, p.params, new_values)) -end - -function SciMLStructures.replace!( - portion::SciMLStructures.AbstractPortion, p::OpaqueMTKParameters, new_values - ) - return SciMLStructures.replace!(portion, p.params, new_values) -end - function SciMLStructures.canonicalize(::SciMLStructures.Tunable, p::MTKParameters) arr = p.tunable repack = let p = p @@ -509,10 +455,6 @@ function Base.copy(p::MTKParameters) ) end -Base.copy(p::OpaqueMTKParameters) = OpaqueMTKParameters(copy(p.params)) - -ArrayInterface.ismutable(::Type{OpaqueMTKParameters}) = true - function ArrayInterface.ismutable( ::Type{ MTKParameters{ @@ -529,11 +471,6 @@ end function SymbolicIndexingInterface.parameter_values(p::MTKParameters, pind::ParameterIndex) return _ducktyped_parameter_values(p, pind) end -function SymbolicIndexingInterface.parameter_values( - p::OpaqueMTKParameters, pind::ParameterIndex - ) - return parameter_values(p.params, pind) -end function _ducktyped_parameter_values(p, pind::ParameterIndex) @unpack portion, idx = pind if portion isa SciMLStructures.Tunable @@ -605,12 +542,6 @@ function SymbolicIndexingInterface.set_parameter!( return nothing end -function SymbolicIndexingInterface.set_parameter!( - p::OpaqueMTKParameters, val, pidx::ParameterIndex - ) - return set_parameter!(p.params, val, pidx) -end - function narrow_buffer_type_and_fallback_undefs( oldbuf::AbstractVector, newbuf::AbstractVector ) @@ -730,12 +661,6 @@ function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, id return _remake_buffer(indp, oldbuf, idxs, vals) end -function SymbolicIndexingInterface.remake_buffer( - indp, oldbuf::OpaqueMTKParameters, idxs, vals - ) - return OpaqueMTKParameters(remake_buffer(indp, oldbuf.params, idxs, vals)) -end - function _remake_buffer(indp, oldbuf::MTKParameters, idxs, vals; validate = true) return __remake_buffer(indp, oldbuf, idxs, vals; validate) end @@ -1088,7 +1013,7 @@ end Base.size(::NestedGetIndex) = () function SymbolicIndexingInterface.with_updated_parameter_timeseries_values( - ::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, + ::AbstractSystem, ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, args::Pair{<:Any, <:NestedGetIndex}... ) unwrapped_ps = _unwrap_mtk_parameters(ps) @@ -1101,7 +1026,8 @@ function SymbolicIndexingInterface.with_updated_parameter_timeseries_values( end function SciMLBase.create_parameter_timeseries_collection( - sys::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, tspan + sys::AbstractSystem, + ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, tspan ) ic = get_index_cache(sys) # this exists because the parameters are `MTKParameters` unwrapped_ps = _unwrap_mtk_parameters(ps) @@ -1127,7 +1053,8 @@ end end function SciMLBase.get_saveable_values( - sys::AbstractSystem, ps::Union{MTKParameters, OpaqueMTKParameters}, timeseries_idx + sys::AbstractSystem, + ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, timeseries_idx ) ps = _unwrap_mtk_parameters(ps) return NestedGetIndex(__get_blocks(timeseries_idx, ps.discrete...)) @@ -1150,11 +1077,6 @@ function DiffEqBase.anyeltypedual( ) where {counter} return DiffEqBase.anyeltypedual(p.tunable) end -function DiffEqBase.anyeltypedual( - p::OpaqueMTKParameters, ::Type{Val{counter}} = Val{0} - ) where {counter} - return DiffEqBase.anyeltypedual(p.params, Val{counter}) -end function DiffEqBase.anyeltypedual( p::Type{<:MTKParameters{T}}, ::Type{Val{counter}} = Val{0} @@ -1216,24 +1138,10 @@ Base.size(ps::MTKParameters) = (length(ps),) Base.IndexStyle(::Type{T}) where {T <: MTKParameters} = IndexLinear() -Base.IndexStyle(::Type{OpaqueMTKParameters}) = IndexLinear() - -Base.eltype(::Type{OpaqueMTKParameters}) = Any - Base.getindex(p::MTKParameters, pind::ParameterIndex) = parameter_values(p, pind) -Base.getindex(p::OpaqueMTKParameters, idx) = getindex(p.params, idx) - Base.setindex!(p::MTKParameters, val, pind::ParameterIndex) = set_parameter!(p, val, pind) -Base.setindex!(p::OpaqueMTKParameters, val, idx) = setindex!(p.params, val, idx) - -Base.length(p::OpaqueMTKParameters) = length(p.params) - -Base.size(p::OpaqueMTKParameters) = size(p.params) - -Base.iterate(p::OpaqueMTKParameters, state = 1) = iterate(p.params, state) - function Base.iterate(buf::MTKParameters, state = 1) total_len = length(buf) if state <= total_len @@ -1253,8 +1161,6 @@ function Base.:(==)(a::MTKParameters, b::MTKParameters) ) end -Base.:(==)(a::OpaqueMTKParameters, b::OpaqueMTKParameters) = a.params == b.params - const MISSING_PARAMETERS_MESSAGE = """ Some parameters are missing from the variable map. Please provide a value or default for the following variables: diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index 92a24f557c..b18de54437 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -1135,7 +1135,8 @@ function (recon::MTKParametersReconstructor)(src, dst) convert(typeof(unwrapped_dst_ps.discrete), recon.discretes_fn(src)), recon.consts_fn(src), nonnumerics, oldcache isa Tuple{} ? () : copy.(oldcache) ) - return dst_ps isa OpaqueMTKParameters ? OpaqueMTKParameters(new_ps) : new_ps + return dst_ps isa SciMLBase.DespecializedParameters ? + SciMLBase.DespecializedParameters(new_ps) : new_ps end """ @@ -1536,8 +1537,8 @@ function (siu::SetInitialUnknowns)(p::MTKParameters, u0) return p end -function (siu::SetInitialUnknowns)(p::OpaqueMTKParameters, u0) - return OpaqueMTKParameters(siu(p.params, u0)) +function (siu::SetInitialUnknowns)(p::SciMLBase.DespecializedParameters, u0) + return SciMLBase.DespecializedParameters(siu(p.params, u0)) end function (siu::SetInitialUnknowns)(p::AbstractVector, u0) @@ -2140,6 +2141,15 @@ Base.@nospecializeinfer function process_SciMLProblem( __process_SciMLProblem(constructor, sys, op, opts; kwargs...) end +function maybe_despecialize_parameters(constructor, p) + return if p isa MTKParameters && constructor <: SciMLBase.AbstractSciMLFunction && + SciMLBase.specialization(constructor) === SciMLBase.AutoSpecialize + SciMLBase.DespecializedParameters(p) + else + p + end +end + function __process_SciMLProblem( @nospecialize(constructor), sys::AbstractSystem, op::AnyDict, opts::SciMLProblemOptions; kwargs... @@ -2296,6 +2306,7 @@ function __process_SciMLProblem( compiler_options, kwargs... ) + p = maybe_despecialize_parameters(constructor, p) if return_operating_point return implicit_dae ? (f, du0, u0, p, op) : (f, u0, p, op) else diff --git a/lib/ModelingToolkitBase/test/opaque_mtkparameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl similarity index 67% rename from lib/ModelingToolkitBase/test/opaque_mtkparameters.jl rename to lib/ModelingToolkitBase/test/despecialized_parameters.jl index 36d026271a..358fc307ea 100644 --- a/lib/ModelingToolkitBase/test/opaque_mtkparameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -6,13 +6,13 @@ using SciMLStructures using SymbolicIndexingInterface using Test -@testset "opaque AutoSpecialize parameters" begin +@testset "despecialized AutoSpecialize parameters" begin @parameters a = 2.0 @variables x(t) = 1.0 z(t) sys = mtkcompile( System( [D(x) ~ -a * x], t; - observed = [z ~ a * x], name = :opaque_parameters + observed = [z ~ a * x], name = :despecialized_parameters ) ) @@ -20,14 +20,27 @@ using Test full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0.0, 1.0)) expression_prob = eval(ODEProblem(sys, [], (0.0, 1.0); expression = Val{true})) - @test auto_prob.p isa ModelingToolkitBase.OpaqueMTKParameters + @test auto_prob.p isa SciMLBase.DespecializedParameters @test auto_prob.p.params isa MTKParameters @test auto_prob.p.tunable === auto_prob.p.params.tunable @test eltype(typeof(auto_prob.p)) === Any @test SciMLStructures.isscimlstructure(auto_prob.p) @test SciMLStructures.ismutablescimlstructure(auto_prob.p) @test full_prob.p isa MTKParameters - @test expression_prob.p isa ModelingToolkitBase.OpaqueMTKParameters + @test expression_prob.p isa SciMLBase.DespecializedParameters + + @parameters b = 2.0 + @variables y = 0.0 + nonlinear_sys = mtkcompile( + System([0 ~ y - b]; name = :despecialized_nonlinear_parameters) + ) + nonlinear_prob = NonlinearProblem(nonlinear_sys, [y => 1.0]) + full_nonlinear_prob = NonlinearProblem{true, SciMLBase.FullSpecialize}( + nonlinear_sys, [y => 1.0] + ) + @test nonlinear_prob.p isa SciMLBase.DespecializedParameters + @test full_nonlinear_prob.p isa MTKParameters + @test nonlinear_prob.f(nonlinear_prob.u0, nonlinear_prob.p) == [-1.0] du = similar(auto_prob.u0) auto_prob.f(du, auto_prob.u0, auto_prob.p, first(auto_prob.tspan)) @@ -48,25 +61,25 @@ using Test @test aliases @test tunables == [3.0] new_p = repack([4.0]) - @test new_p isa ModelingToolkitBase.OpaqueMTKParameters + @test new_p isa SciMLBase.DespecializedParameters @test get_a(new_p) == 4.0 @test copy(new_p) == new_p @test SciMLStructures.replace!(SciMLStructures.Tunable(), new_p, [5.0]) === nothing @test get_a(new_p) == 5.0 replaced_p = SciMLStructures.replace(SciMLStructures.Tunable(), new_p, [4.0]) - @test replaced_p isa ModelingToolkitBase.OpaqueMTKParameters + @test replaced_p isa SciMLBase.DespecializedParameters @test get_a(replaced_p) == 4.0 other_inner = SciMLStructures.replace( SciMLStructures.Tunable(), new_p.params, Float32[5] ) - @test typeof(new_p) === typeof(ModelingToolkitBase.OpaqueMTKParameters(other_inner)) + @test typeof(new_p) === typeof(SciMLBase.DespecializedParameters(other_inner)) @test typeof(new_p.params) !== typeof(other_inner) - other_prob = remake(auto_prob; p = ModelingToolkitBase.OpaqueMTKParameters(other_inner)) + other_prob = remake(auto_prob; p = SciMLBase.DespecializedParameters(other_inner)) @test typeof(other_prob) === typeof(auto_prob) remade = remake(auto_prob; p = [a => 4.0]) - @test remade.p isa ModelingToolkitBase.OpaqueMTKParameters + @test remade.p isa SciMLBase.DespecializedParameters sol = solve(remade, Tsit5(); abstol = 1.0e-10, reltol = 1.0e-10) @test sol[x][end] ≈ exp(-4) @test sol[z][end] ≈ 4exp(-4) @@ -77,10 +90,10 @@ using Test hybrid_sys = mtkcompile( System( [D(population) ~ drift], t, [population], [drift, rate]; jumps = [jump], - name = :opaque_jump_parameters + name = :despecialized_jump_parameters ) ) hybrid_prob = JumpProblem(hybrid_sys, [], (0.0, 0.1); rng = Xoshiro(1)) - @test hybrid_prob.prob.p isa ModelingToolkitBase.OpaqueMTKParameters + @test hybrid_prob.prob.p isa SciMLBase.DespecializedParameters @test SciMLBase.successful_retcode(solve(hybrid_prob, Tsit5())) end diff --git a/lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl b/lib/ModelingToolkitBase/test/extensions/despecialized_parameters_ad.jl similarity index 84% rename from lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl rename to lib/ModelingToolkitBase/test/extensions/despecialized_parameters_ad.jl index b1280bbae1..358bc7c669 100644 --- a/lib/ModelingToolkitBase/test/extensions/opaque_mtkparameters_ad.jl +++ b/lib/ModelingToolkitBase/test/extensions/despecialized_parameters_ad.jl @@ -7,10 +7,10 @@ using SciMLStructures using Test using Zygote -@testset "opaque parameter sensitivities" begin +@testset "despecialized parameter sensitivities" begin @parameters a = 2.0 @variables x(t) = 1.0 - sys = mtkcompile(System([D(x) ~ -a * x], t; name = :opaque_parameter_ad)) + sys = mtkcompile(System([D(x) ~ -a * x], t; name = :despecialized_parameter_ad)) prob = ODEProblem(sys, [], (0.0, 1.0)) function terminal_value(a_value) diff --git a/lib/ModelingToolkitBase/test/odesystem.jl b/lib/ModelingToolkitBase/test/odesystem.jl index 20de35f498..b19def5d7c 100644 --- a/lib/ModelingToolkitBase/test/odesystem.jl +++ b/lib/ModelingToolkitBase/test/odesystem.jl @@ -1,5 +1,5 @@ using ModelingToolkitBase, StaticArrays, LinearAlgebra -using ModelingToolkitBase: get_metadata, MTKParameters, OpaqueMTKParameters, +using ModelingToolkitBase: get_metadata, MTKParameters, SymbolicDiscreteCallback, SymbolicContinuousCallback using SymbolicIndexingInterface @@ -270,7 +270,7 @@ prob12 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tsp prob13 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tspan) prob14 = ODEProblem(sys, [u0; p2], tspan) for p in [prob1, prob14] - @test p.p isa OpaqueMTKParameters + @test p.p isa SciMLBase.DespecializedParameters p.ps[k₁] ≈ 0.04 p.ps[k₂] ≈ 3.0e7 p.ps[k₃] ≈ 1.0e-4 @@ -286,7 +286,7 @@ u01 = [y₁ => 1, y₂ => 1, y₃ => 1] prob_pmap = remake(prob14; p = p3, u0 = u01) prob_dpmap = remake(prob14; p = Dict(p3), u0 = Dict(u01)) for p in [prob_pmap, prob_dpmap] - @test p.p isa OpaqueMTKParameters + @test p.p isa SciMLBase.DespecializedParameters p.ps[k₁] ≈ 0.05 p.ps[k₂] ≈ 2.0e7 p.ps[k₃] ≈ 1.1e-4 @@ -321,7 +321,7 @@ sol_dpmap = solve(prob_dpmap, Rodas5()) prob, p = Dict(sys1.a => 3.0, b => 4.0), u0 = Dict(sys1.x => 1.0) ) - @test prob_new.p isa OpaqueMTKParameters + @test prob_new.p isa SciMLBase.DespecializedParameters @test prob_new.ps[b] ≈ 4.0 @test prob_new.ps[sys1.a] ≈ 3.0 @test prob_new.ps[sys.sys2.a] ≈ 1.0 @@ -666,7 +666,7 @@ let ivmap = [A => 1.0, k1 => 1.0, k2 => 1.0] tspan = (0.0, 1.0) prob = ODEProblem(sys, ivmap, tspan; tofloat = false) - @test prob.p isa OpaqueMTKParameters + @test prob.p isa SciMLBase.DespecializedParameters @test prob.ps[k1] ≈ 1.0 @test prob.ps[k2] == 1 && prob.ps[k2] isa Int end diff --git a/lib/ModelingToolkitBase/test/runtests.jl b/lib/ModelingToolkitBase/test/runtests.jl index bbd9f1e36c..22b4d44148 100644 --- a/lib/ModelingToolkitBase/test/runtests.jl +++ b/lib/ModelingToolkitBase/test/runtests.jl @@ -115,7 +115,7 @@ end @safetestset "SymbolicIndexingInterface test" include("symbolic_indexing_interface.jl") @safetestset "SciML Problem Input Test" include("sciml_problem_inputs.jl") @safetestset "MTKParameters Test" include("mtkparameters.jl") - @safetestset "Opaque MTKParameters Test" include("opaque_mtkparameters.jl") + @safetestset "Despecialized MTKParameters Test" include("despecialized_parameters.jl") end if GROUP == "All" || GROUP == "Extended" @@ -136,7 +136,7 @@ end @safetestset "HomotopyContinuation Extension Test" include("extensions/homotopy_continuation.jl") @safetestset "LabelledArrays Test" include("extensions/labelledarrays.jl") @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl") - @safetestset "Opaque MTKParameters AD Test" include("extensions/opaque_mtkparameters_ad.jl") + @safetestset "Despecialized MTKParameters AD Test" include("extensions/despecialized_parameters_ad.jl") # @safetestset "Auto Differentiation Test" include("extensions/ad.jl") end diff --git a/src/linearization.jl b/src/linearization.jl index 08b3ef6c59..747a51b62b 100644 --- a/src/linearization.jl +++ b/src/linearization.jl @@ -258,7 +258,7 @@ function linearization_function( if u0 === nothing uf_jac = h_jac = pf_jac = nothing - parameter_eltype = if p isa Union{MTKParameters, OpaqueMTKParameters} + parameter_eltype = if p isa Union{MTKParameters, SciMLBase.DespecializedParameters} tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p) eltype(tunables) else diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl index 82e45b63c7..1c03e60d54 100644 --- a/src/problems/sccnonlinearproblem.jl +++ b/src/problems/sccnonlinearproblem.jl @@ -6,7 +6,7 @@ function (cw::CacheWriter)(p::MTKParameters, sols) return cw.fn(p.caches, sols, p) end -(cw::CacheWriter)(p::OpaqueMTKParameters, sols) = cw(p.params, sols) +(cw::CacheWriter)(p::SciMLBase.DespecializedParameters, sols) = cw(p.params, sols) const SCCCacheVarsExprsElT = Dict{TypeT, Vector{SymbolicT}} diff --git a/test/qa/qa.jl b/test/qa/qa.jl index ed825737a9..3d0ba2a7a3 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -134,7 +134,6 @@ const REEXPORTED_API = ( :JumpSystem, :linear_fractional_to_ordinary, :liouville_transform, :LocalScope, :maybe_zeros, :MissingGuessValue, :ModelingToolkitBase, :modelingtoolkitize, :modified_unknowns!, :mtkcompile, :MTKParameters, - :OpaqueMTKParameters, :MTKVariableTypeCtx, :namespace_equations, :noise_to_brownians, :NonlinearSystem, :observables, :observed, :ODESystem, :open_loop, :OptimizationSystem, :outputs, :parameters, :parameters_toplevel, :ParentScope, :PDESystem, :Pre, :ProblemTypeCtx, From 4bf089d11c78862da67f2c005e09abedcd2fe8ab Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 01:54:18 -0400 Subject: [PATCH 04/12] Preserve parameter specialization through MTK initialization Co-Authored-By: Chris Rackauckas --- Project.toml | 4 +-- lib/ModelingToolkitBase/Project.toml | 2 +- .../src/problems/initializationproblem.jl | 7 +++- .../src/systems/codegen_utils.jl | 6 ++-- .../src/systems/nonlinear/initializesystem.jl | 3 +- .../src/systems/problem_utils.jl | 34 +++++++++++++++---- .../test/despecialized_parameters.jl | 19 +++++++++++ .../test/symbolic_indexing_interface.jl | 2 +- 8 files changed, 62 insertions(+), 15 deletions(-) diff --git a/Project.toml b/Project.toml index a07e02a9ce..69c71ce9c2 100644 --- a/Project.toml +++ b/Project.toml @@ -112,7 +112,7 @@ REPL = "1" Reexport = "1" RuntimeGeneratedFunctions = "0.5.12" SCCNonlinearSolve = "1.13" -SciMLBase = "3.46" +SciMLBase = "3.45" 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" diff --git a/lib/ModelingToolkitBase/Project.toml b/lib/ModelingToolkitBase/Project.toml index d3bd18fd22..4100a2e57d 100644 --- a/lib/ModelingToolkitBase/Project.toml +++ b/lib/ModelingToolkitBase/Project.toml @@ -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" diff --git a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl index 35d1aa2439..d2249afd05 100644 --- a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl @@ -168,7 +168,12 @@ 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{_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, diff --git a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl index c1265c9ec0..7254fad361 100644 --- a/lib/ModelingToolkitBase/src/systems/codegen_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/codegen_utils.jl @@ -1041,13 +1041,13 @@ function GeneratedFunctionWrapper{P}( end function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg{Any, NArgs}) where {PIdx, NArgs, Split} - # non-split systems just call it as-is - Split || return gfw.f_oop(args...) 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}}) # for split systems, call it as-is if the parameter object is a tuple or MTKParameters # but not if it is a tuple of numbers @@ -1064,12 +1064,12 @@ function (gfw::GeneratedFunctionWrapper{Tuple{PIdx, NArgs, Split}})(args::Vararg if NArgs + 1 != N throw(MethodError(gfw, args)) end - Split || return gfw.f_iip(args...) 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...) end diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index a95cdf884c..34e57d0586 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/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 = SciMLBase.specialization(typeof(odefn)) ) odefn = remake(odefn; kws...) diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index b18de54437..62115143fa 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -1606,6 +1606,20 @@ this function in extensions. """ __iip_u0_ad_wrapper(x) = x +struct InitializationMap{IIP, U, F} + u0_constructor::U + map::F +end + +function InitializationMap{IIP}(u0_constructor::U, map::F) where {IIP, U, F} + return InitializationMap{IIP, U, F}(u0_constructor, map) +end + +(map::InitializationMap{false})(x) = map.u0_constructor(map.map(x)) +function (map::InitializationMap{true})(x) + return __iip_u0_ad_wrapper(map.u0_constructor(map.map(x))) +end + """ $(TYPEDSIGNATURES) @@ -1871,6 +1885,7 @@ constructed is in implicit DAE form (`DAEProblem`). `opts.check_initialization_u function maybe_build_initialization_problem( sys::AbstractSystem, iip::Bool, op::SymmapT, t, guesses, opts::SciMLProblemOptions; + specialize = SciMLBase.AutoSpecialize, # Intercept `expression` because we don't support it here yet expression = Val{false}, kwargs... ) @@ -1891,7 +1906,7 @@ function maybe_build_initialization_problem( end orig_op = copy(op) - initializeprob = ModelingToolkitBase.InitializationProblem{iip}( + initializeprob = ModelingToolkitBase.InitializationProblem{iip, specialize}( sys, t, op, opts; guesses, fast_path = true, kwargs... ) initsys = initializeprob.f.sys::System @@ -1963,10 +1978,16 @@ function maybe_build_initialization_problem( if isempty(solved_unknowns) initializeprobmap = nothing else - initializeprobmap = u0_constructor ∘ PromoteToTunableEltype(CopyParamsByTemplate(initializeprob.f.sys, solved_unknowns; eval_expression, eval_module, kwargs...), floatT) - if iip - initializeprobmap = __iip_u0_ad_wrapper ∘ initializeprobmap - end + initializeprobmap = InitializationMap{iip}( + u0_constructor, + PromoteToTunableEltype( + CopyParamsByTemplate( + initializeprob.f.sys, solved_unknowns; + eval_expression, eval_module, kwargs... + ), + floatT + ) + ) end else initializeprobmap = nothing @@ -2196,7 +2217,8 @@ function __process_SciMLProblem( if build_initializeprob kws = maybe_build_initialization_problem( sys, constructor <: SciMLBase.AbstractSciMLFunction{true}, - op, t, guesses, opts; kwargs... + op, t, guesses, opts; + specialize = SciMLBase.specialization(constructor), kwargs... ) kwargs = merge(kwargs, kws) diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index 358fc307ea..a1b2e75a74 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -17,6 +17,9 @@ using Test ) auto_prob = ODEProblem(sys, [], (0.0, 1.0); jac = true) + respecialized_prob = ODEProblem{true, SciMLBase.AutoRespecialize}( + sys, [], (0.0, 1.0) + ) full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0.0, 1.0)) expression_prob = eval(ODEProblem(sys, [], (0.0, 1.0); expression = Val{true})) @@ -26,9 +29,25 @@ using Test @test eltype(typeof(auto_prob.p)) === Any @test SciMLStructures.isscimlstructure(auto_prob.p) @test SciMLStructures.ismutablescimlstructure(auto_prob.p) + @test respecialized_prob.p isa MTKParameters @test full_prob.p isa MTKParameters + @test auto_prob.f.initialization_data.initializeprob.p isa + SciMLBase.DespecializedParameters + @test respecialized_prob.f.initialization_data.initializeprob.p isa MTKParameters + @test full_prob.f.initialization_data.initializeprob.p isa MTKParameters @test expression_prob.p isa SciMLBase.DespecializedParameters + seen_parameter_type = Ref{DataType}() + gfw = ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, false}}( + (u, p, t) -> (seen_parameter_type[] = typeof(p); u), + (du, u, p, t) -> (seen_parameter_type[] = typeof(p); copyto!(du, u)) + ) + @test gfw(auto_prob.u0, auto_prob.p, 0.0) === auto_prob.u0 + @test seen_parameter_type[] === typeof(auto_prob.p.params) + du = similar(auto_prob.u0) + gfw(du, auto_prob.u0, auto_prob.p, 0.0) + @test seen_parameter_type[] === typeof(auto_prob.p.params) + @parameters b = 2.0 @variables y = 0.0 nonlinear_sys = mtkcompile( diff --git a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl index d3623d4d1d..4967de01a8 100644 --- a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl +++ b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl @@ -136,7 +136,7 @@ end @test isempty(get_all_timeseries_indexes(ns, [σ + ρ, ρ + β])) @test pobs(ps) == [3.0, 5.0] - prob = NonlinearProblem( + prob = NonlinearProblem{true, SciMLBase.FullSpecialize}( ns, [x => 1.0, y => 2.0, z => 3.0, σ => 1.0, ρ => 2.0, β => 3.0] ) getter = getu(ns, (x + 1, x + 2)) From 9d4f1a7503b827d0ecbc5c971e117db690b09874 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 04:19:39 -0400 Subject: [PATCH 05/12] Move parameter despecialization to solve time Co-Authored-By: Chris Rackauckas --- Project.toml | 6 +- benchmark/benchmarks.jl | 4 +- docs/src/API/problems.md | 9 +- docs/src/basics/FAQ.md | 14 ++-- lib/ModelingToolkitBase/Project.toml | 2 +- .../src/ModelingToolkitBase.jl | 3 +- .../src/modelingtoolkitize/common.jl | 5 +- lib/ModelingToolkitBase/src/precompile.jl | 22 ++--- .../src/problems/initializationproblem.jl | 2 +- .../src/problems/jumpproblem.jl | 2 +- .../src/problems/odeproblem.jl | 2 +- .../nonlinear/homotopy_continuation.jl | 2 +- .../src/systems/nonlinear/initializesystem.jl | 6 +- .../src/systems/parameter_buffer.jl | 36 +++++--- .../src/systems/problem_utils.jl | 45 +++++----- .../test/despecialized_parameters.jl | 83 ++++++++++++------- lib/ModelingToolkitBase/test/index_cache.jl | 11 +-- .../test/jacobiansparsity.jl | 7 +- lib/ModelingToolkitBase/test/mtkparameters.jl | 6 +- lib/ModelingToolkitBase/test/odesystem.jl | 11 ++- .../test/split_parameters.jl | 28 ++----- .../test/symbolic_indexing_interface.jl | 6 +- src/problems/sccnonlinearproblem.jl | 3 +- test/structural_transformation/tearing.jl | 4 +- 24 files changed, 158 insertions(+), 161 deletions(-) diff --git a/Project.toml b/Project.toml index 69c71ce9c2..922df2b73e 100644 --- a/Project.toml +++ b/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.45" +SciMLBase = "3.46" SciMLPublic = "1.0.0" Serialization = "1" Setfield = "1" diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index c2b54af875..5685330295 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -90,5 +90,5 @@ prob = ODEProblem(model, u0, tspan; jac, sparse) out = similar(prob.f.jac_prototype) sparse_analytical_jacobian["ODEProblem"] = @benchmarkable ODEProblem($model, $u0, $tspan; jac, sparse) -sparse_analytical_jacobian["f_oop"] = @benchmarkable $(prob.f.jac)($(prob.u0), $(prob.p), $(first(tspan))) -sparse_analytical_jacobian["f_iip"] = @benchmarkable $(prob.f.jac)($out, $(prob.u0), $(prob.p), $(first(tspan))) +sparse_analytical_jacobian["f_oop"] = @benchmarkable $(prob.f.jac.f_oop)($(prob.u0), $(prob.p), $(first(tspan))) +sparse_analytical_jacobian["f_iip"] = @benchmarkable $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), $(first(tspan))) diff --git a/docs/src/API/problems.md b/docs/src/API/problems.md index 19417cac50..8f394780a0 100644 --- a/docs/src/API/problems.md +++ b/docs/src/API/problems.md @@ -90,10 +90,11 @@ ModelingToolkit.varmap_to_vars ``` The parameters of a split system are stored in a custom data structure called -`MTKParameters`. `AutoSpecialize` problems wrap it in -[`SciMLBase.DespecializedParameters`](@ref) so solvers can reuse compiled code across -parameter-buffer layouts. `FullSpecialize` problems keep the concrete `MTKParameters` -object and avoid the wrapper's dynamic function barrier. These objects should only be +`MTKParameters`. ModelingToolkit problem constructors use +[`SciMLBase.AutoDespecialize`](@ref) by default. Solvers that support this policy wrap the +parameters in [`SciMLBase.DespecializedParameters`](@ref) 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. diff --git a/docs/src/basics/FAQ.md b/docs/src/basics/FAQ.md index a03d08df73..0aac1371dc 100644 --- a/docs/src/basics/FAQ.md +++ b/docs/src/basics/FAQ.md @@ -3,12 +3,12 @@ ## Why are my parameters some obscure object? In ModelingToolkit.jl version 9, the parameter vector was replaced with a custom -`MTKParameters` object. `AutoSpecialize` problems expose a -[`SciMLBase.DespecializedParameters`](@ref) wrapper around that object so solver compilation -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 +`MTKParameters` object. ModelingToolkit problems use [`SciMLBase.AutoDespecialize`](@ref) +by default. Supporting solvers expose a [`SciMLBase.DespecializedParameters`](@ref) 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: @@ -30,7 +30,7 @@ and `SciMLBase.DespecializedParameters`: the given index. - `setindex!` with a `ParameterIndex` can be used to set the value of a parameter with the given index. - - The wrapped `MTKParameters` object is available as the `params` field of an + - `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`). diff --git a/lib/ModelingToolkitBase/Project.toml b/lib/ModelingToolkitBase/Project.toml index 4100a2e57d..44f0e649aa 100644 --- a/lib/ModelingToolkitBase/Project.toml +++ b/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" diff --git a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl index f2f05d18d8..b8840cfb0b 100644 --- a/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl +++ b/lib/ModelingToolkitBase/src/ModelingToolkitBase.jl @@ -395,8 +395,7 @@ export Shift, ShiftIndex export Sample, Hold, SampleTime export Clock, SolverStepClock, TimeDomain -export MTKParameters, reorder_dimension_by_tunables!, - reorder_dimension_by_tunables +export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables export HomotopyContinuationProblem diff --git a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl index 83c3e44913..c08016213d 100644 --- a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl +++ b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl @@ -220,7 +220,7 @@ function define_params(p::MTKParameters, t, names = nothing) end define_params(p::SciMLBase.DespecializedParameters, t, names = nothing) = - define_params(p.params, t, names) + define_params(SciMLBase.unwrap_parameters(p), t, names) """ $(TYPEDSIGNATURES) @@ -236,7 +236,8 @@ function to_paramvec(p::MTKParameters) return reduce(vcat, collect(p); init = []) end -to_paramvec(p::SciMLBase.DespecializedParameters) = to_paramvec(p.params) +to_paramvec(p::SciMLBase.DespecializedParameters) = + to_paramvec(SciMLBase.unwrap_parameters(p)) """ $(TYPEDSIGNATURES) diff --git a/lib/ModelingToolkitBase/src/precompile.jl b/lib/ModelingToolkitBase/src/precompile.jl index 0e4b2e9ed7..1066af85c4 100644 --- a/lib/ModelingToolkitBase/src/precompile.jl +++ b/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}) diff --git a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl index d2249afd05..30efd3b9a1 100644 --- a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl +++ b/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}( diff --git a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl index e9ff39443e..1b6a603a12 100644 --- a/lib/ModelingToolkitBase/src/problems/jumpproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/jumpproblem.jl @@ -192,7 +192,7 @@ function updateparams!( ratemap::JumpSysMajParamMapper{U, V, W}, params::SciMLBase.DespecializedParameters ) where {U <: AbstractArray, V <: AbstractArray, W} - return updateparams!(ratemap, params.params) + return updateparams!(ratemap, SciMLBase.unwrap_parameters(params)) end function updateparams!( diff --git a/lib/ModelingToolkitBase/src/problems/odeproblem.jl b/lib/ModelingToolkitBase/src/problems/odeproblem.jl index 171ba0d998..a020be55b9 100644 --- a/lib/ModelingToolkitBase/src/problems/odeproblem.jl +++ b/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 diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/homotopy_continuation.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/homotopy_continuation.jl index 8c908f2ddc..adc824c295 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/homotopy_continuation.jl +++ b/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...) diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index 34e57d0586..ff5334bf1b 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -718,7 +718,7 @@ end function promote_type_with_nothing( ::Type{T}, p::SciMLBase.DespecializedParameters ) where {T} - return promote_type_with_nothing(T, p.params) + return promote_type_with_nothing(T, SciMLBase.unwrap_parameters(p)) end promote_with_nothing(::Type, ::Nothing) = nothing @@ -747,7 +747,9 @@ function promote_with_nothing(::Type{T}, p::MTKParameters) where {T} return p end function promote_with_nothing(::Type{T}, p::SciMLBase.DespecializedParameters) where {T} - return SciMLBase.DespecializedParameters(promote_with_nothing(T, p.params)) + return SciMLBase.DespecializedParameters( + promote_with_nothing(T, SciMLBase.unwrap_parameters(p)) + ) end function promote_u0_p(u0, p, t0) diff --git a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl index c6056e1a13..9b84d91a9e 100644 --- a/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl +++ b/lib/ModelingToolkitBase/src/systems/parameter_buffer.jl @@ -1013,26 +1013,23 @@ end Base.size(::NestedGetIndex) = () function SymbolicIndexingInterface.with_updated_parameter_timeseries_values( - ::AbstractSystem, ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, + ::AbstractSystem, ps::MTKParameters, args::Pair{<:Any, <:NestedGetIndex}... ) - unwrapped_ps = _unwrap_mtk_parameters(ps) for (i, ngi) in args for (j, val) in enumerate(ngi.x) - copyto!(view(unwrapped_ps.discrete[j], Block(i)), val) + copyto!(view(ps.discrete[j], Block(i)), val) end end return ps end function SciMLBase.create_parameter_timeseries_collection( - sys::AbstractSystem, - ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, tspan + sys::AbstractSystem, ps::MTKParameters, tspan ) ic = get_index_cache(sys) # this exists because the parameters are `MTKParameters` - unwrapped_ps = _unwrap_mtk_parameters(ps) - isempty(unwrapped_ps.discrete) && return nothing - num_discretes = only(blocksize(unwrapped_ps.discrete[1])) + isempty(ps.discrete) && return nothing + num_discretes = only(blocksize(ps.discrete[1])) buffers = [] partition_type = typeof(SciMLBase.get_saveable_values(sys, ps, 1)) for i in 1:num_discretes @@ -1044,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...)...) @@ -1053,13 +1060,20 @@ end end function SciMLBase.get_saveable_values( - sys::AbstractSystem, - ps::Union{MTKParameters, SciMLBase.DespecializedParameters}, timeseries_idx + sys::AbstractSystem, ps::MTKParameters, timeseries_idx ) - ps = _unwrap_mtk_parameters(ps) 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 diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index 62115143fa..098e56c81c 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -1111,10 +1111,19 @@ end # TODO: make this infer when the nonnumerics are non-trivial function (recon::MTKParametersReconstructor)(src, dst) - src_ps = parameter_values(src) - dst_ps = parameter_values(dst) - unwrapped_dst_ps = _unwrap_mtk_parameters(dst_ps) - oldcache = unwrapped_dst_ps.caches + return recon(src, parameter_values(dst)) +end + +function (recon::MTKParametersReconstructor)( + src, dst_ps::SciMLBase.DespecializedParameters + ) + return SciMLBase.DespecializedParameters( + recon(src, SciMLBase.unwrap_parameters(dst_ps)) + ) +end + +function (recon::MTKParametersReconstructor)(src, dst_ps::MTKParameters) + oldcache = dst_ps.caches # I don't know why but this makes it infer properly if recon.tunables_fn isa ComposedFunction tunablevals = recon.tunables_fn.outer(recon.tunables_fn.inner(src)) @@ -1122,7 +1131,7 @@ function (recon::MTKParametersReconstructor)(src, dst) tunablevals = recon.tunables_fn(src) end initialvals = recon.initials_fn(src) - nonnumerics = recon.nonnumerics_fn(src)::typeof(unwrapped_dst_ps.nonnumeric) + nonnumerics = recon.nonnumerics_fn(src)::typeof(dst_ps.nonnumeric) (; diffcache_buffer_idx) = recon if !iszero(diffcache_buffer_idx) @set! nonnumerics[diffcache_buffer_idx] = DiffCacheAllocatorAPIWrapper{ForwardDiff.valtype(eltype(initialvals))}.(nonnumerics[diffcache_buffer_idx]) @@ -1130,13 +1139,11 @@ function (recon::MTKParametersReconstructor)(src, dst) # This `convert` exists because a `Real` discrete might get its value from an # integer function of integer parameters/discretes. This ends up creating a # `BlockedArray{Int, ...}` instead of a `BlockedArray{Float64, ...}`. - new_ps = MTKParameters( + return MTKParameters( tunablevals, initialvals, - convert(typeof(unwrapped_dst_ps.discrete), recon.discretes_fn(src)), + convert(typeof(dst_ps.discrete), recon.discretes_fn(src)), recon.consts_fn(src), nonnumerics, oldcache isa Tuple{} ? () : copy.(oldcache) ) - return dst_ps isa SciMLBase.DespecializedParameters ? - SciMLBase.DespecializedParameters(new_ps) : new_ps end """ @@ -1538,7 +1545,7 @@ function (siu::SetInitialUnknowns)(p::MTKParameters, u0) end function (siu::SetInitialUnknowns)(p::SciMLBase.DespecializedParameters, u0) - return SciMLBase.DespecializedParameters(siu(p.params, u0)) + return SciMLBase.DespecializedParameters(siu(SciMLBase.unwrap_parameters(p), u0)) end function (siu::SetInitialUnknowns)(p::AbstractVector, u0) @@ -1885,7 +1892,7 @@ constructed is in implicit DAE form (`DAEProblem`). `opts.check_initialization_u function maybe_build_initialization_problem( sys::AbstractSystem, iip::Bool, op::SymmapT, t, guesses, opts::SciMLProblemOptions; - specialize = SciMLBase.AutoSpecialize, + specialize = SciMLBase.AutoDespecialize, # Intercept `expression` because we don't support it here yet expression = Val{false}, kwargs... ) @@ -2162,15 +2169,6 @@ Base.@nospecializeinfer function process_SciMLProblem( __process_SciMLProblem(constructor, sys, op, opts; kwargs...) end -function maybe_despecialize_parameters(constructor, p) - return if p isa MTKParameters && constructor <: SciMLBase.AbstractSciMLFunction && - SciMLBase.specialization(constructor) === SciMLBase.AutoSpecialize - SciMLBase.DespecializedParameters(p) - else - p - end -end - function __process_SciMLProblem( @nospecialize(constructor), sys::AbstractSystem, op::AnyDict, opts::SciMLProblemOptions; kwargs... @@ -2328,7 +2326,6 @@ function __process_SciMLProblem( compiler_options, kwargs... ) - p = maybe_despecialize_parameters(constructor, p) if return_operating_point return implicit_dae ? (f, du0, u0, p, op) : (f, u0, p, op) else @@ -2498,7 +2495,7 @@ resolve_iip(::Type{Both}, @nospecialize(op)) = !(op isa StaticArray) Macro for writing problem/function constructors. Expects a function definition with type parameters for `iip` and `specialize`. Generates fallbacks with -`specialize = SciMLBase.AutoSpecialize` and `iip = Both` (resolved at construction time). +`specialize = SciMLBase.AutoDespecialize` and `iip = Both` (resolved at construction time). """ # Unwrap `@nospecialize(arg)` to get the underlying argument expression. # Returns the argument unchanged if not wrapped in @nospecialize. @@ -2571,9 +2568,9 @@ macro fallback_iip_specialize(ex) fnname_name, curly_args... = fnname_curly.args @assert curly_args == where_args - # callexpr_iip is `ODEProblem{iip, AutoSpecialize}(call_args...)` + # callexpr_iip is `ODEProblem{iip, AutoDespecialize}(call_args...)` callexpr_iip = Expr( - :call, Expr(:curly, fnname_name, curly_args[1], SciMLBase.AutoSpecialize), call_args... + :call, Expr(:curly, fnname_name, curly_args[1], SciMLBase.AutoDespecialize), call_args... ) # `ODEProblem{iip}` fnname_iip = Expr(:curly, fnname_name, curly_args[1]) diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index a1b2e75a74..6ba086ffd8 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -1,12 +1,13 @@ using ModelingToolkitBase using ModelingToolkitBase: t_nounits as t, D_nounits as D, MTKParameters +using DiffEqBase using OrdinaryDiffEq using Random using SciMLStructures using SymbolicIndexingInterface using Test -@testset "despecialized AutoSpecialize parameters" begin +@testset "AutoDespecialize parameters" begin @parameters a = 2.0 @variables x(t) = 1.0 z(t) sys = mtkcompile( @@ -16,37 +17,53 @@ using Test ) ) - auto_prob = ODEProblem(sys, [], (0.0, 1.0); jac = true) + despecialized_prob = ODEProblem(sys, [], (0.0, 1.0); jac = true) + auto_prob = ODEProblem{true, SciMLBase.AutoSpecialize}(sys, [], (0.0, 1.0)) respecialized_prob = ODEProblem{true, SciMLBase.AutoRespecialize}( sys, [], (0.0, 1.0) ) full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0.0, 1.0)) expression_prob = eval(ODEProblem(sys, [], (0.0, 1.0); expression = Val{true})) - @test auto_prob.p isa SciMLBase.DespecializedParameters - @test auto_prob.p.params isa MTKParameters - @test auto_prob.p.tunable === auto_prob.p.params.tunable - @test eltype(typeof(auto_prob.p)) === Any - @test SciMLStructures.isscimlstructure(auto_prob.p) - @test SciMLStructures.ismutablescimlstructure(auto_prob.p) + @test SciMLBase.specialization(despecialized_prob.f) === SciMLBase.AutoDespecialize + @test SciMLBase.specialization(auto_prob.f) === SciMLBase.AutoSpecialize + @test despecialized_prob.p isa MTKParameters + @test auto_prob.p isa MTKParameters @test respecialized_prob.p isa MTKParameters @test full_prob.p isa MTKParameters - @test auto_prob.f.initialization_data.initializeprob.p isa - SciMLBase.DespecializedParameters + @test despecialized_prob.f.initialization_data.initializeprob.p isa MTKParameters @test respecialized_prob.f.initialization_data.initializeprob.p isa MTKParameters @test full_prob.f.initialization_data.initializeprob.p isa MTKParameters - @test expression_prob.p isa SciMLBase.DespecializedParameters + @test SciMLBase.specialization( + despecialized_prob.f.initialization_data.initializeprob.f + ) === SciMLBase.AutoDespecialize + @test SciMLBase.specialization(auto_prob.f.initialization_data.initializeprob.f) === + SciMLBase.AutoSpecialize + @test SciMLBase.specialization(full_prob.f.initialization_data.initializeprob.f) === + SciMLBase.FullSpecialize + @test expression_prob.p isa MTKParameters + + concrete_prob = DiffEqBase.get_concrete_problem( + despecialized_prob, true; alg = Tsit5() + ) + concrete_parameters = SciMLBase.unwrap_parameters(concrete_prob.p) + @test concrete_prob.p isa SciMLBase.DespecializedParameters + @test concrete_parameters isa MTKParameters + @test concrete_prob.p.tunable === concrete_parameters.tunable + @test eltype(typeof(concrete_prob.p)) === Any + @test SciMLStructures.isscimlstructure(concrete_prob.p) + @test SciMLStructures.ismutablescimlstructure(concrete_prob.p) seen_parameter_type = Ref{DataType}() gfw = ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, false}}( (u, p, t) -> (seen_parameter_type[] = typeof(p); u), (du, u, p, t) -> (seen_parameter_type[] = typeof(p); copyto!(du, u)) ) - @test gfw(auto_prob.u0, auto_prob.p, 0.0) === auto_prob.u0 - @test seen_parameter_type[] === typeof(auto_prob.p.params) - du = similar(auto_prob.u0) - gfw(du, auto_prob.u0, auto_prob.p, 0.0) - @test seen_parameter_type[] === typeof(auto_prob.p.params) + @test gfw(concrete_prob.u0, concrete_prob.p, 0.0) === concrete_prob.u0 + @test seen_parameter_type[] === typeof(concrete_parameters) + du = similar(concrete_prob.u0) + gfw(du, concrete_prob.u0, concrete_prob.p, 0.0) + @test seen_parameter_type[] === typeof(concrete_parameters) @parameters b = 2.0 @variables y = 0.0 @@ -57,25 +74,28 @@ using Test full_nonlinear_prob = NonlinearProblem{true, SciMLBase.FullSpecialize}( nonlinear_sys, [y => 1.0] ) - @test nonlinear_prob.p isa SciMLBase.DespecializedParameters + @test SciMLBase.specialization(nonlinear_prob.f) === SciMLBase.AutoDespecialize + @test nonlinear_prob.p isa MTKParameters @test full_nonlinear_prob.p isa MTKParameters @test nonlinear_prob.f(nonlinear_prob.u0, nonlinear_prob.p) == [-1.0] - du = similar(auto_prob.u0) - auto_prob.f(du, auto_prob.u0, auto_prob.p, first(auto_prob.tspan)) + du = similar(concrete_prob.u0) + concrete_prob.f(du, concrete_prob.u0, concrete_prob.p, first(concrete_prob.tspan)) @test du == [-2.0] jac = zeros(1, 1) - auto_prob.f.jac(jac, auto_prob.u0, auto_prob.p, first(auto_prob.tspan)) + concrete_prob.f.jac( + jac, concrete_prob.u0, concrete_prob.p, first(concrete_prob.tspan) + ) @test jac == [-2.0;;] get_a = getp(sys, a) set_a! = setp(sys, a) - @test get_a(auto_prob.p) == 2.0 - set_a!(auto_prob.p, 3.0) - @test get_a(auto_prob.p) == 3.0 + @test get_a(concrete_prob.p) == 2.0 + set_a!(concrete_prob.p, 3.0) + @test get_a(concrete_prob.p) == 3.0 tunables, repack, aliases = SciMLStructures.canonicalize( - SciMLStructures.Tunable(), auto_prob.p + SciMLStructures.Tunable(), concrete_prob.p ) @test aliases @test tunables == [3.0] @@ -90,16 +110,17 @@ using Test @test get_a(replaced_p) == 4.0 other_inner = SciMLStructures.replace( - SciMLStructures.Tunable(), new_p.params, Float32[5] + SciMLStructures.Tunable(), SciMLBase.unwrap_parameters(new_p), Float32[5] ) @test typeof(new_p) === typeof(SciMLBase.DespecializedParameters(other_inner)) - @test typeof(new_p.params) !== typeof(other_inner) - other_prob = remake(auto_prob; p = SciMLBase.DespecializedParameters(other_inner)) - @test typeof(other_prob) === typeof(auto_prob) + @test typeof(SciMLBase.unwrap_parameters(new_p)) !== typeof(other_inner) + other_prob = remake(concrete_prob; p = SciMLBase.DespecializedParameters(other_inner)) + @test typeof(other_prob) === typeof(concrete_prob) - remade = remake(auto_prob; p = [a => 4.0]) - @test remade.p isa SciMLBase.DespecializedParameters + remade = remake(despecialized_prob; p = [a => 4.0]) + @test remade.p isa MTKParameters sol = solve(remade, Tsit5(); abstol = 1.0e-10, reltol = 1.0e-10) + @test sol.prob.p isa SciMLBase.DespecializedParameters @test sol[x][end] ≈ exp(-4) @test sol[z][end] ≈ 4exp(-4) @@ -113,6 +134,6 @@ using Test ) ) hybrid_prob = JumpProblem(hybrid_sys, [], (0.0, 0.1); rng = Xoshiro(1)) - @test hybrid_prob.prob.p isa SciMLBase.DespecializedParameters + @test hybrid_prob.prob.p isa MTKParameters @test SciMLBase.successful_retcode(solve(hybrid_prob, Tsit5())) end diff --git a/lib/ModelingToolkitBase/test/index_cache.jl b/lib/ModelingToolkitBase/test/index_cache.jl index 6d5d9a5194..24e886f48b 100644 --- a/lib/ModelingToolkitBase/test/index_cache.jl +++ b/lib/ModelingToolkitBase/test/index_cache.jl @@ -168,16 +168,7 @@ end @test SciMLBase.successful_retcode(solve(prob, Tsit5())) idata = prob.f.initialization_data - reconstructed_p = idata.metadata.oop_reconstruct_u0_p.pgetter( - prob, idata.initializeprob - ) - @test reconstructed_p isa MTKParameters - - full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [x => 1.0], (0.0, 1.0)) - idata = full_prob.f.initialization_data - @test_nowarn @inferred idata.metadata.oop_reconstruct_u0_p.pgetter( - full_prob, idata.initializeprob - ) + @test_nowarn @inferred idata.metadata.oop_reconstruct_u0_p.pgetter(prob, idata.initializeprob) end function costfn(theta, ps) diff --git a/lib/ModelingToolkitBase/test/jacobiansparsity.jl b/lib/ModelingToolkitBase/test/jacobiansparsity.jl index a5a00e8e0b..3282dd6457 100644 --- a/lib/ModelingToolkitBase/test/jacobiansparsity.jl +++ b/lib/ModelingToolkitBase/test/jacobiansparsity.jl @@ -74,11 +74,8 @@ prob = ODEProblem(sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = true, jac #@test_nowarn solve(prob, Rosenbrock23()) @test findnz(calculate_jacobian(sys, sparse = true))[1:2] == findnz(prob.f.jac_prototype)[1:2] -full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( - sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = true, jac = true -) -out = similar(full_prob.f.jac_prototype) -@test (@ballocated $(full_prob.f.jac)($out, $(full_prob.u0), $(full_prob.p), 0.0)) == 0 +out = similar(prob.f.jac_prototype) +@test (@ballocated $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), 0.0)) == 0 # should not allocate # test when not sparse prob = ODEProblem(sys, unknowns(sys) .=> vec(u0), (0, 11.5), sparse = false, jac = true) diff --git a/lib/ModelingToolkitBase/test/mtkparameters.jl b/lib/ModelingToolkitBase/test/mtkparameters.jl index 7ff56948a1..1f1c3fffcd 100644 --- a/lib/ModelingToolkitBase/test/mtkparameters.jl +++ b/lib/ModelingToolkitBase/test/mtkparameters.jl @@ -125,7 +125,7 @@ end @test getp(sys, g)(newps) isa Vector{Float32} @testset "Type-stability of `remake_buffer`" begin - prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, ivs, (0.0, 1.0)) + prob = ODEProblem(sys, ivs, (0.0, 1.0)) idxs = (a, c, d, e, f, g, h) vals = (1.0, 2.0, 3, ones(3), ones(Int, 3, 3), ones(2), "a") @@ -482,9 +482,7 @@ end @parameters p @named sys = System(D(x) ~ x * p, t) sys = complete(sys) - prob = ODEProblem{true, SciMLBase.FullSpecialize}( - sys, SA[x => 1.0, p => 1.0], (0.0, 1.0) - ) + prob = ODEProblem(sys, SA[x => 1.0, p => 1.0], (0.0, 1.0)) @test isbits(prob.p) @test isbits(prob.f.initialization_data.initializeprob.p) end diff --git a/lib/ModelingToolkitBase/test/odesystem.jl b/lib/ModelingToolkitBase/test/odesystem.jl index b19def5d7c..c1902fa361 100644 --- a/lib/ModelingToolkitBase/test/odesystem.jl +++ b/lib/ModelingToolkitBase/test/odesystem.jl @@ -1,6 +1,5 @@ using ModelingToolkitBase, StaticArrays, LinearAlgebra -using ModelingToolkitBase: get_metadata, MTKParameters, - SymbolicDiscreteCallback, +using ModelingToolkitBase: get_metadata, MTKParameters, SymbolicDiscreteCallback, SymbolicContinuousCallback using SymbolicIndexingInterface using OrdinaryDiffEq, Sundials @@ -270,7 +269,7 @@ prob12 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tsp prob13 = ODEProblem(sys, [u0; [k₁ => 0.04, k₂ => 3.0e7, k₃ => 1.0e4]], tspan) prob14 = ODEProblem(sys, [u0; p2], tspan) for p in [prob1, prob14] - @test p.p isa SciMLBase.DespecializedParameters + @test p.p isa MTKParameters p.ps[k₁] ≈ 0.04 p.ps[k₂] ≈ 3.0e7 p.ps[k₃] ≈ 1.0e-4 @@ -286,7 +285,7 @@ u01 = [y₁ => 1, y₂ => 1, y₃ => 1] prob_pmap = remake(prob14; p = p3, u0 = u01) prob_dpmap = remake(prob14; p = Dict(p3), u0 = Dict(u01)) for p in [prob_pmap, prob_dpmap] - @test p.p isa SciMLBase.DespecializedParameters + @test p.p isa MTKParameters p.ps[k₁] ≈ 0.05 p.ps[k₂] ≈ 2.0e7 p.ps[k₃] ≈ 1.1e-4 @@ -321,7 +320,7 @@ sol_dpmap = solve(prob_dpmap, Rodas5()) prob, p = Dict(sys1.a => 3.0, b => 4.0), u0 = Dict(sys1.x => 1.0) ) - @test prob_new.p isa SciMLBase.DespecializedParameters + @test prob_new.p isa MTKParameters @test prob_new.ps[b] ≈ 4.0 @test prob_new.ps[sys1.a] ≈ 3.0 @test prob_new.ps[sys.sys2.a] ≈ 1.0 @@ -666,7 +665,7 @@ let ivmap = [A => 1.0, k1 => 1.0, k2 => 1.0] tspan = (0.0, 1.0) prob = ODEProblem(sys, ivmap, tspan; tofloat = false) - @test prob.p isa SciMLBase.DespecializedParameters + @test prob.p isa MTKParameters @test prob.ps[k1] ≈ 1.0 @test prob.ps[k2] == 1 && prob.ps[k2] isa Int end diff --git a/lib/ModelingToolkitBase/test/split_parameters.jl b/lib/ModelingToolkitBase/test/split_parameters.jl index 767d8b2d07..1c4f44ea2a 100644 --- a/lib/ModelingToolkitBase/test/split_parameters.jl +++ b/lib/ModelingToolkitBase/test/split_parameters.jl @@ -254,27 +254,15 @@ end getter = getp(sys, fn) prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0)) - @test getter(prob)(1.0) == 2.0 - full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( - sys, [x => 1.0], (0.0, 1.0) - ) - @inferred getter(full_prob) + @inferred getter(prob) # cannot be inferred better since `FunctionWrapper` is only known to return `Real` - @inferred Vector{<:Real} full_prob.f( - full_prob.u0, full_prob.p, full_prob.tspan[1] - ) + @inferred Vector{<:Real} prob.f(prob.u0, prob.p, prob.tspan[1]) sol = solve(prob, Tsit5(); abstol = 1.0e-10, reltol = 1.0e-10) @test sol.u[end][] ≈ 2.0 prob = ODEProblem(sys, [x => 1.0, fn => Foo()], (0.0, 1.0)) - @test getter(prob)(1.0) == 3.0 - full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( - sys, [x => 1.0, fn => Foo()], (0.0, 1.0) - ) - @inferred getter(full_prob) - @inferred Vector{<:Real} full_prob.f( - full_prob.u0, full_prob.p, full_prob.tspan[1] - ) + @inferred getter(prob) + @inferred Vector{<:Real} prob.f(prob.u0, prob.p, prob.tspan[1]) sol = solve(prob; abstol = 1.0e-10, reltol = 1.0e-10) @test sol.u[end][] ≈ 2.5 end @@ -290,12 +278,8 @@ end @test is_parameter(sys, fn) getter = getp(sys, fn) prob = ODEProblem(sys, [x => 1.0, fn => interp], (0.0, 1.0)) - @test getter(prob) === interp - full_prob = ODEProblem{true, SciMLBase.FullSpecialize}( - sys, [x => 1.0, fn => interp], (0.0, 1.0) - ) - @inferred getter(full_prob) - @inferred full_prob.f(full_prob.u0, full_prob.p, full_prob.tspan[1]) + @inferred getter(prob) + @inferred prob.f(prob.u0, prob.p, prob.tspan[1]) @test_nowarn sol = solve(prob, Tsit5()) @test_nowarn prob.ps[fn] = LinearInterpolation( ts .^ 3, ts; extrapolation = ExtrapolationType.Extension diff --git a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl index 4967de01a8..6b3197cd0a 100644 --- a/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl +++ b/lib/ModelingToolkitBase/test/symbolic_indexing_interface.jl @@ -45,9 +45,7 @@ using Test @test default_values(odesys)[y] == 2.0 @test isequal(default_values(odesys)[xy], x + y) - prob = ODEProblem{true, SciMLBase.FullSpecialize}( - odesys, [a => 1.0, b => 2.0], (0.0, 1.0) - ) + prob = ODEProblem(odesys, [a => 1.0, b => 2.0], (0.0, 1.0)) getter = getu(odesys, (x + 1, x + 2)) @test getter(prob) isa Tuple @test_nowarn @inferred getter(prob) @@ -136,7 +134,7 @@ end @test isempty(get_all_timeseries_indexes(ns, [σ + ρ, ρ + β])) @test pobs(ps) == [3.0, 5.0] - prob = NonlinearProblem{true, SciMLBase.FullSpecialize}( + prob = NonlinearProblem( ns, [x => 1.0, y => 2.0, z => 3.0, σ => 1.0, ρ => 2.0, β => 3.0] ) getter = getu(ns, (x + 1, x + 2)) diff --git a/src/problems/sccnonlinearproblem.jl b/src/problems/sccnonlinearproblem.jl index 1c03e60d54..632de8d45d 100644 --- a/src/problems/sccnonlinearproblem.jl +++ b/src/problems/sccnonlinearproblem.jl @@ -6,7 +6,8 @@ function (cw::CacheWriter)(p::MTKParameters, sols) return cw.fn(p.caches, sols, p) end -(cw::CacheWriter)(p::SciMLBase.DespecializedParameters, sols) = cw(p.params, sols) +(cw::CacheWriter)(p::SciMLBase.DespecializedParameters, sols) = + cw(SciMLBase.unwrap_parameters(p), sols) const SCCCacheVarsExprsElT = Dict{TypeT, Vector{SymbolicT}} diff --git a/test/structural_transformation/tearing.jl b/test/structural_transformation/tearing.jl index a0650138a3..3960dda5dd 100644 --- a/test/structural_transformation/tearing.jl +++ b/test/structural_transformation/tearing.jl @@ -181,9 +181,7 @@ newdaesys = mtkcompile(daesys) equations(tearing_substitution(newdaesys)), [D(x) ~ h * z; 0 ~ x + sin(z) - p * t] ) @test issetequal(unknowns(newdaesys), [x, z]) -prob = ODEProblem{true, SciMLBase.FullSpecialize}( - newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0) -) +prob = ODEProblem(newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0)) du = [0.0, 0.0]; u = [1.0, -0.5π]; pr = prob.p; From fa75ed00a2be0d965f717c9ab79894a4ee7b501e Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 05:44:23 -0400 Subject: [PATCH 06/12] Preserve legacy initialization specialization Co-Authored-By: Chris Rackauckas --- .../src/systems/nonlinear/initializesystem.jl | 2 +- lib/ModelingToolkitBase/src/systems/problem_utils.jl | 8 +++++++- .../test/despecialized_parameters.jl | 11 ++++++++++- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index ff5334bf1b..07def4c1d4 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -700,7 +700,7 @@ function _remake_initialization_data_impl( ) kws = maybe_build_initialization_problem( sys, SciMLBase.isinplace(odefn), op, t0, guesses, opts; - specialize = SciMLBase.specialization(typeof(odefn)) + specialize = initialization_specialization(SciMLBase.specialization(typeof(odefn))) ) odefn = remake(odefn; kws...) diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index 098e56c81c..cd323386ff 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -2080,6 +2080,10 @@ function maybe_build_initialization_problem( ) end +initialization_specialization(::Type{SciMLBase.AutoDespecialize}) = + SciMLBase.AutoDespecialize +initialization_specialization(::Type) = SciMLBase.AutoSpecialize + """ $(TYPEDSIGNATURES) @@ -2216,7 +2220,9 @@ function __process_SciMLProblem( kws = maybe_build_initialization_problem( sys, constructor <: SciMLBase.AbstractSciMLFunction{true}, op, t, guesses, opts; - specialize = SciMLBase.specialization(constructor), kwargs... + specialize = initialization_specialization( + SciMLBase.specialization(constructor) + ), kwargs... ) kwargs = merge(kwargs, kws) diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index 6ba086ffd8..5bc3df4dc1 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -23,6 +23,9 @@ using Test sys, [], (0.0, 1.0) ) full_prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys, [], (0.0, 1.0)) + function_wrapper_prob = ODEProblem{true, SciMLBase.FunctionWrapperSpecialize}( + sys, [], (0.0, 1.0) + ) expression_prob = eval(ODEProblem(sys, [], (0.0, 1.0); expression = Val{true})) @test SciMLBase.specialization(despecialized_prob.f) === SciMLBase.AutoDespecialize @@ -39,8 +42,14 @@ using Test ) === SciMLBase.AutoDespecialize @test SciMLBase.specialization(auto_prob.f.initialization_data.initializeprob.f) === SciMLBase.AutoSpecialize + @test SciMLBase.specialization( + respecialized_prob.f.initialization_data.initializeprob.f + ) === SciMLBase.AutoSpecialize @test SciMLBase.specialization(full_prob.f.initialization_data.initializeprob.f) === - SciMLBase.FullSpecialize + SciMLBase.AutoSpecialize + @test SciMLBase.specialization( + function_wrapper_prob.f.initialization_data.initializeprob.f + ) === SciMLBase.AutoSpecialize @test expression_prob.p isa MTKParameters concrete_prob = DiffEqBase.get_concrete_problem( From 1be7aaa270cb3e7f55963510a1d6c20cabf39e25 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 13:10:21 -0400 Subject: [PATCH 07/12] Preserve AutoDespecialize SCC initialization Co-Authored-By: Chris Rackauckas --- docs/src/API/problems.md | 6 ++++-- docs/src/basics/FAQ.md | 6 ++++-- .../src/problems/initializationproblem.jl | 3 ++- test/scc_nonlinear_problem.jl | 1 + 4 files changed, 11 insertions(+), 5 deletions(-) diff --git a/docs/src/API/problems.md b/docs/src/API/problems.md index 8f394780a0..a93c4376fe 100644 --- a/docs/src/API/problems.md +++ b/docs/src/API/problems.md @@ -91,8 +91,10 @@ ModelingToolkit.varmap_to_vars The parameters of a split system are stored in a custom data structure called `MTKParameters`. ModelingToolkit problem constructors use -[`SciMLBase.AutoDespecialize`](@ref) by default. Solvers that support this policy wrap the -parameters in [`SciMLBase.DespecializedParameters`](@ref) at solve time so compiled code +[`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. diff --git a/docs/src/basics/FAQ.md b/docs/src/basics/FAQ.md index 0aac1371dc..f73fe400c8 100644 --- a/docs/src/basics/FAQ.md +++ b/docs/src/basics/FAQ.md @@ -3,8 +3,10 @@ ## Why are my parameters some obscure object? In ModelingToolkit.jl version 9, the parameter vector was replaced with a custom -`MTKParameters` object. ModelingToolkit problems use [`SciMLBase.AutoDespecialize`](@ref) -by default. Supporting solvers expose a [`SciMLBase.DespecializedParameters`](@ref) wrapper +`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 diff --git a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl index 30efd3b9a1..496d67e8c5 100644 --- a/lib/ModelingToolkitBase/src/problems/initializationproblem.jl +++ b/lib/ModelingToolkitBase/src/problems/initializationproblem.jl @@ -168,7 +168,8 @@ 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) - problem_constructor = if TProb === LinearInitializationProblem + problem_constructor = if TProb === LinearInitializationProblem || + TProb === SCCNonlinearProblem TProb{_iip} else TProb{_iip, specialize} diff --git a/test/scc_nonlinear_problem.jl b/test/scc_nonlinear_problem.jl index 3a6a3e9fd8..69eab71de8 100644 --- a/test/scc_nonlinear_problem.jl +++ b/test/scc_nonlinear_problem.jl @@ -413,6 +413,7 @@ end [D(a) ~ b, 0 ~ b^3 + b + a - 2, 0 ~ c^3 + c - b, 0 ~ d - c * b], t ) prob = ODEProblem(sys, [a => 0.5], (0.0, 1.0)) + @test SciMLBase.specialization(prob.f) === SciMLBase.AutoDespecialize @test prob.f.initialization_data.initializeprob isa SCCNonlinearProblem sol = solve(prob, Rodas5P()) @test SciMLBase.successful_retcode(sol) From 61bb71a6a45377ce907b1017b0066f3e6af6eb13 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 13:42:27 -0400 Subject: [PATCH 08/12] Expose semilinear specialization policy Co-Authored-By: Chris Rackauckas --- src/problems/docs.jl | 1 + test/semilinearodeproblem.jl | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/src/problems/docs.jl b/src/problems/docs.jl index 5e9739bddd..bc188ad4b9 100644 --- a/src/problems/docs.jl +++ b/src/problems/docs.jl @@ -1,4 +1,5 @@ struct SemilinearODEFunction{iip, spec} end +SciMLBase.specialization(::Type{<:SemilinearODEFunction{iip, spec}}) where {iip, spec} = spec struct SemilinearODEProblem{iip, spec} end const SEMILINEAR_EXTRA_BODY = """ diff --git a/test/semilinearodeproblem.jl b/test/semilinearodeproblem.jl index 15863e01e0..72a0e3f36e 100644 --- a/test/semilinearodeproblem.jl +++ b/test/semilinearodeproblem.jl @@ -6,6 +6,10 @@ using LinearAlgebra using Test using ModelingToolkit: t_nounits as t, D_nounits as D +@test SciMLBase.specialization( + SemilinearODEFunction{true, SciMLBase.AutoDespecialize} +) === SciMLBase.AutoDespecialize + # from https://docs.sciml.ai/SciMLBenchmarksOutput/dev/AstroChem/nelson/ @testset "Astrochem model" begin function Nelson!(du, u, p, t) From b1d23506595a952ca0727fd2346dc17efe557da7 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 15:54:52 -0400 Subject: [PATCH 09/12] Preserve despecialized event parameters after initialization Co-Authored-By: Chris Rackauckas --- .../src/systems/problem_utils.jl | 6 +++++- .../test/despecialized_parameters.jl | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index cd323386ff..c31b5fcfac 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -1380,7 +1380,11 @@ function construct_initializeprobpmap( ), p_constructor = p_constructor function initprobpmap_nosplit(prob, initsol) - return p_constructor(getter(initsol)) + p = p_constructor(getter(initsol)) + if parameter_values(prob) isa SciMLBase.DespecializedParameters + p = SciMLBase.DespecializedParameters(p) + end + return p end end end diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index 5bc3df4dc1..93264d2d9d 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -133,6 +133,20 @@ using Test @test sol[x][end] ≈ exp(-4) @test sol[z][end] ≈ 4exp(-4) + @variables event_x(t) = 0.0 + event_sys = complete( + System( + [D(event_x) ~ 1], t; continuous_events = [event_x ~ 1], + name = :despecialized_nonsplit_event + ); split = false + ) + event_prob = ODEProblem(event_sys, [], (0.0, 2.0)) + concrete_event_prob = DiffEqBase.get_concrete_problem(event_prob, true; alg = Tsit5()) + @test concrete_event_prob.p isa SciMLBase.DespecializedParameters + event_sol = solve(event_prob, Tsit5()) + @test event_sol.prob.p isa SciMLBase.DespecializedParameters + @test minimum(t -> abs(t - 1), event_sol.t) < 1.0e-10 + @parameters drift = 0.0 rate = 1.0 @variables population(t) = 10.0 jump = SymbolicMassActionJump(rate, [population => 1], [population => -1]) From 272470ea4f0a361301e4f7fd6ec4bc9bdf7d69b8 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 11 Aug 2026 15:54:52 -0400 Subject: [PATCH 10/12] Test initialization inference on Julia LTS Co-Authored-By: Chris Rackauckas --- lib/ModelingToolkitBase/test/initializationsystem.jl | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/lib/ModelingToolkitBase/test/initializationsystem.jl b/lib/ModelingToolkitBase/test/initializationsystem.jl index 5254886678..5d51886628 100644 --- a/lib/ModelingToolkitBase/test/initializationsystem.jl +++ b/lib/ModelingToolkitBase/test/initializationsystem.jl @@ -1823,19 +1823,9 @@ end # AutoSpecialize uses Union types for compilation sharing, so @inferred # is only expected to pass with FullSpecialize. prob = ODEProblem{true, SciMLBase.FullSpecialize}(complete(sys), [], (0.0, 1)) - if v"1.13-" <= VERSION - # Don't run the inference test on pre-release versions - if VERSION >= v"1.13.0" - @inferred remake(prob; u0 = 2 .* prob.u0, p = prob.p) - @inferred solve(prob) - end - elseif VERSION >= v"1.11" + if VERSION < v"1.13-" || VERSION >= v"1.13.0" @inferred remake(prob; u0 = 2 .* prob.u0, p = prob.p) @inferred solve(prob) - else - # This passes locally but fails in CI for some reason - @test_broken @inferred remake(prob; u0 = 2 .* prob.u0, p = prob.p) - @test_broken @inferred solve(prob) end end From 14e90b67558133106d565e4abd9dedafa982ebf6 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 15 Aug 2026 08:12:08 -0400 Subject: [PATCH 11/12] Recurse on concrete initialization parameters Co-Authored-By: Chris Rackauckas --- .../src/systems/nonlinear/initializesystem.jl | 11 ++++++++-- .../test/despecialized_parameters.jl | 20 +++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index 07def4c1d4..6abe1a9086 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -877,8 +877,15 @@ function DiffEqBase.get_updated_symbolic_problem( t0 = is_time_dependent(prob) ? current_time(prob) : nothing unwrapped_p = _unwrap_mtk_parameters(p) - if unwrapped_p isa MTKParameters - buffer = unwrapped_p.initials + if unwrapped_p !== p + updated_prob = DiffEqBase.get_updated_symbolic_problem( + sys, prob; u0, p = unwrapped_p, kw... + ) + return @set updated_prob.p = p + end + + if p isa MTKParameters + buffer = p.initials else buffer = p end diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index 93264d2d9d..1d156fccb6 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -1,12 +1,18 @@ using ModelingToolkitBase using ModelingToolkitBase: t_nounits as t, D_nounits as D, MTKParameters using DiffEqBase +using ForwardDiff using OrdinaryDiffEq using Random +using SciMLBase using SciMLStructures using SymbolicIndexingInterface using Test +struct _OpaqueParameterWrapper end +const _opaque_parameters = Ref{Any}() +SciMLBase.unwrap_parameters(::_OpaqueParameterWrapper) = _opaque_parameters[] + @testset "AutoDespecialize parameters" begin @parameters a = 2.0 @variables x(t) = 1.0 z(t) @@ -63,6 +69,20 @@ using Test @test SciMLStructures.isscimlstructure(concrete_prob.p) @test SciMLStructures.ismutablescimlstructure(concrete_prob.p) + dual_initials = map( + value -> ForwardDiff.Dual{Nothing}(value, one(value)), concrete_parameters.initials + ) + dual_parameters = SciMLStructures.replace( + SciMLStructures.Initials(), concrete_parameters, dual_initials + ) + _opaque_parameters[] = SciMLBase.DespecializedParameters(dual_parameters) + opaque_parameters = _OpaqueParameterWrapper() + updated_prob = DiffEqBase.get_updated_symbolic_problem( + sys, concrete_prob; u0 = concrete_prob.u0, p = opaque_parameters + ) + @test eltype(updated_prob.u0) <: ForwardDiff.Dual + @test updated_prob.p === opaque_parameters + seen_parameter_type = Ref{DataType}() gfw = ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, false}}( (u, p, t) -> (seen_parameter_type[] = typeof(p); u), From 0678b4c15249cc50b820d61ab48d039207e648c1 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 17 Aug 2026 05:03:53 -0400 Subject: [PATCH 12/12] Recurse on concrete modelingtoolkitize parameters Co-Authored-By: Chris Rackauckas --- .../src/modelingtoolkitize/common.jl | 16 ++++++++++++--- .../src/systems/nonlinear/initializesystem.jl | 11 ++-------- .../test/despecialized_parameters.jl | 20 ------------------- .../test/modelingtoolkitize.jl | 13 +++++++++++- 4 files changed, 27 insertions(+), 33 deletions(-) diff --git a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl index c08016213d..741bbff8bd 100644 --- a/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl +++ b/lib/ModelingToolkitBase/src/modelingtoolkitize/common.jl @@ -390,14 +390,24 @@ 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} merge!(defaults, Dict(v => p[k] for (k, v) in pairs(paramobj))) - elseif p isa Union{MTKParameters, SciMLBase.DespecializedParameters} - p = _unwrap_mtk_parameters(p) + elseif p isa MTKParameters pvals = [ p.tunable; reduce(vcat, p.discrete; init = []); reduce(vcat, p.constant; init = []); diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index 6abe1a9086..07def4c1d4 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -877,15 +877,8 @@ function DiffEqBase.get_updated_symbolic_problem( t0 = is_time_dependent(prob) ? current_time(prob) : nothing unwrapped_p = _unwrap_mtk_parameters(p) - if unwrapped_p !== p - updated_prob = DiffEqBase.get_updated_symbolic_problem( - sys, prob; u0, p = unwrapped_p, kw... - ) - return @set updated_prob.p = p - end - - if p isa MTKParameters - buffer = p.initials + if unwrapped_p isa MTKParameters + buffer = unwrapped_p.initials else buffer = p end diff --git a/lib/ModelingToolkitBase/test/despecialized_parameters.jl b/lib/ModelingToolkitBase/test/despecialized_parameters.jl index 1d156fccb6..93264d2d9d 100644 --- a/lib/ModelingToolkitBase/test/despecialized_parameters.jl +++ b/lib/ModelingToolkitBase/test/despecialized_parameters.jl @@ -1,18 +1,12 @@ using ModelingToolkitBase using ModelingToolkitBase: t_nounits as t, D_nounits as D, MTKParameters using DiffEqBase -using ForwardDiff using OrdinaryDiffEq using Random -using SciMLBase using SciMLStructures using SymbolicIndexingInterface using Test -struct _OpaqueParameterWrapper end -const _opaque_parameters = Ref{Any}() -SciMLBase.unwrap_parameters(::_OpaqueParameterWrapper) = _opaque_parameters[] - @testset "AutoDespecialize parameters" begin @parameters a = 2.0 @variables x(t) = 1.0 z(t) @@ -69,20 +63,6 @@ SciMLBase.unwrap_parameters(::_OpaqueParameterWrapper) = _opaque_parameters[] @test SciMLStructures.isscimlstructure(concrete_prob.p) @test SciMLStructures.ismutablescimlstructure(concrete_prob.p) - dual_initials = map( - value -> ForwardDiff.Dual{Nothing}(value, one(value)), concrete_parameters.initials - ) - dual_parameters = SciMLStructures.replace( - SciMLStructures.Initials(), concrete_parameters, dual_initials - ) - _opaque_parameters[] = SciMLBase.DespecializedParameters(dual_parameters) - opaque_parameters = _OpaqueParameterWrapper() - updated_prob = DiffEqBase.get_updated_symbolic_problem( - sys, concrete_prob; u0 = concrete_prob.u0, p = opaque_parameters - ) - @test eltype(updated_prob.u0) <: ForwardDiff.Dual - @test updated_prob.p === opaque_parameters - seen_parameter_type = Ref{DataType}() gfw = ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, false}}( (u, p, t) -> (seen_parameter_type[] = typeof(p); u), diff --git a/lib/ModelingToolkitBase/test/modelingtoolkitize.jl b/lib/ModelingToolkitBase/test/modelingtoolkitize.jl index d923bed7f0..3c2e88c491 100644 --- a/lib/ModelingToolkitBase/test/modelingtoolkitize.jl +++ b/lib/ModelingToolkitBase/test/modelingtoolkitize.jl @@ -4,7 +4,7 @@ using Optimization, RecursiveArrayTools, OptimizationOptimJL using SymbolicIndexingInterface using ModelingToolkitBase: t_nounits as t, D_nounits as D using Symbolics: value -using SciMLBase: parameterless_type, successful_retcode +using SciMLBase: DespecializedParameters, parameterless_type, successful_retcode N = 32 const xyd_brusselator = range(0, stop = 1, length = N) @@ -280,6 +280,17 @@ sys = modelingtoolkitize(prob) @test [value(ModelingToolkitBase.initial_conditions(sys)[s]) for s in unknowns(sys)] == u0 @test [value(ModelingToolkitBase.initial_conditions(sys)[s]) for s in parameters(sys)] == [10, 20] +function ode_prob_namedtuple(du, u, p, t) + du[1] = u[1] + p.a + du[2] = u[2] + p.b + return nothing +end +params = DespecializedParameters((a = 10, b = 20)) +prob = ODEProblem(ode_prob_namedtuple, u0, (0.0, 1.0), params) +sys = modelingtoolkitize(prob) +@test [value(ModelingToolkitBase.initial_conditions(sys)[s]) for s in unknowns(sys)] == u0 +@test [value(ModelingToolkitBase.initial_conditions(sys)[s]) for s in parameters(sys)] == [10, 20] + @parameters sig = 10 rho = 28.0 beta = 8 / 3 @variables x(t) = 100 y(t) = 1.0 z(t) = 1