Merge pull request #3566 from AayushSabharwal/as/type-stable-remake

fix: make `remake` type-stable

Author: ChrisRackauckas

Project.toml

@@ -139,7 +139,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
-SciMLBase = "2.75"
+SciMLBase = "2.84"
 SciMLStructures = "1.7"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
@@ -150,7 +150,7 @@ StaticArrays = "0.10, 0.11, 0.12, 1.0"
 StochasticDelayDiffEq = "1.8.1"
 StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.39"
-SymbolicUtils = "3.25.1"
+SymbolicUtils = "3.26.1"
 Symbolics = "6.37"
 URIs = "1"
 UnPack = "0.1, 1.0"

docs/src/basics/Events.md

@@ -126,8 +126,8 @@ prob = ODEProblem(ball, Pair[], tspan)
 
 sol = solve(prob, Tsit5())
 @assert 0 <= minimum(sol[x]) <= 1e-10 # the ball never went through the floor but got very close
-@assert minimum(sol[y]) > -1.5 # check wall conditions
-@assert maximum(sol[y]) < 1.5  # check wall conditions
+@assert minimum(sol[y]) >= -1.5 # check wall conditions
+@assert maximum(sol[y]) <= 1.5  # check wall conditions
 
 tv = sort([LinRange(0, 10, 200); sol.t])
 plot(sol(tv)[y], sol(tv)[x], line_z = tv)

docs/src/tutorials/disturbance_modeling.md

@@ -224,7 +224,7 @@ To see full examples that perform state estimation with ModelingToolkit models,
 Pages = ["disturbance_modeling.md"]
 ```
 
-```@autodocs
+```@autodocs; canonical = false
 Modules = [ModelingToolkit]
 Pages   = ["systems/analysis_points.jl"]
 Order   = [:function, :type]

src/systems/abstractsystem.jl

@@ -915,6 +915,7 @@ for prop in [:eqs
              :substitutions
              :metadata
              :gui_metadata
+             :is_initializesystem
              :discrete_subsystems
              :parameter_dependencies
              :assertions

src/systems/diffeqs/abstractodesystem.jl

@@ -1457,12 +1457,12 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
         isys = generate_initializesystem(
             sys; initialization_eqs, check_units, pmap = parammap,
-            guesses, extra_metadata = (; use_scc), algebraic_only)
+            guesses, algebraic_only)
         simplify_system = true
     else
         isys = generate_initializesystem(
             sys; u0map, initialization_eqs, check_units,
-            pmap = parammap, guesses, extra_metadata = (; use_scc), algebraic_only)
+            pmap = parammap, guesses, algebraic_only)
         simplify_system = true
     end
 
@@ -1477,12 +1477,6 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         isys = structural_simplify(isys; fully_determined)
     end
 
-    meta = get_metadata(isys)
-    if meta isa InitializationSystemMetadata
-        @set! isys.metadata.oop_reconstruct_u0_p = ReconstructInitializeprob(
-            sys, isys)
-    end
-
     ts = get_tearing_state(isys)
     unassigned_vars = StructuralTransformations.singular_check(ts)
     if warn_initialize_determined && !isempty(unassigned_vars)

src/systems/index_cache.jl

@@ -502,7 +502,7 @@ function reorder_parameters(
     end
 end
 
-function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
+function reorder_parameters(ic::IndexCache, ps; drop_missing = false, flatten = true)
     isempty(ps) && return ()
     param_buf = if ic.tunable_buffer_size.length == 0
         ()
@@ -555,20 +555,37 @@ function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
         end
     end
 
-    result = broadcast.(
-        unwrap, (
-            param_buf..., initials_buf..., disc_buf..., const_buf..., nonnumeric_buf...))
+    param_buf = broadcast.(unwrap, param_buf)
+    initials_buf = broadcast.(unwrap, initials_buf)
+    disc_buf = broadcast.(unwrap, disc_buf)
+    const_buf = broadcast.(unwrap, const_buf)
+    nonnumeric_buf = broadcast.(unwrap, nonnumeric_buf)
+
     if drop_missing
-        result = map(result) do buf
-            filter(buf) do sym
-                return !isequal(sym, unwrap(variable(:DEF)))
-            end
+        filterer = !isequal(unwrap(variable(:DEF)))
+        param_buf = filter.(filterer, param_buf)
+        initials_buf = filter.(filterer, initials_buf)
+        disc_buf = filter.(filterer, disc_buf)
+        const_buf = filter.(filterer, const_buf)
+        nonnumeric_buf = filter.(filterer, nonnumeric_buf)
+    end
+
+    if flatten
+        result = (
+            param_buf..., initials_buf..., disc_buf..., const_buf..., nonnumeric_buf...)
+        if all(isempty, result)
+            return ()
         end
+        return result
+    else
+        if isempty(param_buf)
+            param_buf = ((),)
+        end
+        if isempty(initials_buf)
+            initials_buf = ((),)
+        end
+        return (param_buf..., initials_buf..., disc_buf, const_buf, nonnumeric_buf)
     end
-    if all(isempty, result)
-        return ()
-    end
-    return result
 end
 
 # Given a parameter index, find the index of the buffer it is in when

src/systems/nonlinear/initializesystem.jl

@@ -11,7 +11,7 @@ function generate_initializesystem(sys::AbstractTimeDependentSystem;
         default_dd_guess = Bool(0),
         algebraic_only = false,
         check_units = true, check_defguess = false,
-        name = nameof(sys), extra_metadata = (;), kwargs...)
+        name = nameof(sys), kwargs...)
     eqs = equations(sys)
     if !(eqs isa Vector{Equation})
         eqs = Equation[x for x in eqs if x isa Equation]
@@ -143,17 +143,15 @@ function generate_initializesystem(sys::AbstractTimeDependentSystem;
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
     end
-    meta = InitializationSystemMetadata(
-        anydict(u0map), anydict(pmap), additional_guesses,
-        additional_initialization_eqs, extra_metadata, nothing)
+
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
         defaults = defs,
         checks = check_units,
         parameter_dependencies = new_parameter_deps,
         name,
-        metadata = meta,
+        is_initializesystem = true,
         kwargs...)
 end
 
@@ -169,7 +167,7 @@ function generate_initializesystem(sys::AbstractTimeIndependentSystem;
         guesses = Dict(),
         algebraic_only = false,
         check_units = true, check_defguess = false,
-        name = nameof(sys), extra_metadata = (;), kwargs...)
+        name = nameof(sys), kwargs...)
     eqs = equations(sys)
     trueobs, eqs = unhack_observed(observed(sys), eqs)
     vars = unique([unknowns(sys); getfield.(trueobs, :lhs)])
@@ -244,17 +242,15 @@ function generate_initializesystem(sys::AbstractTimeIndependentSystem;
     for k in keys(defs)
         defs[k] = substitute(defs[k], paramsubs)
     end
-    meta = InitializationSystemMetadata(
-        anydict(u0map), anydict(pmap), additional_guesses,
-        additional_initialization_eqs, extra_metadata, nothing)
+
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
         defaults = defs,
         checks = check_units,
         parameter_dependencies = new_parameter_deps,
         name,
-        metadata = meta,
+        is_initializesystem = true,
         kwargs...)
 end
 
@@ -436,64 +432,6 @@ function _has_delays(sys::AbstractSystem, ex, banned)
     return any(x -> _has_delays(sys, x, banned), args)
 end
 
-struct ReconstructInitializeprob
-    getter::Any
-    setter::Any
-end
-
-function ReconstructInitializeprob(
-        srcsys::AbstractSystem, dstsys::AbstractSystem)
-    syms = reduce(
-        vcat, reorder_parameters(dstsys, parameters(dstsys));
-        init = [])
-    getter = getu(srcsys, syms)
-    setter = setp_oop(dstsys, syms)
-    return ReconstructInitializeprob(getter, setter)
-end
-
-function (rip::ReconstructInitializeprob)(srcvalp, dstvalp)
-    newp = rip.setter(dstvalp, rip.getter(srcvalp))
-    if state_values(dstvalp) === nothing
-        return nothing, newp
-    end
-    srcu0 = state_values(srcvalp)
-    T = srcu0 === nothing || isempty(srcu0) ? Union{} : eltype(srcu0)
-    if parameter_values(dstvalp) isa MTKParameters
-        if !isempty(newp.tunable)
-            T = promote_type(eltype(newp.tunable), T)
-        end
-    elseif !isempty(newp)
-        T = promote_type(eltype(newp), T)
-    end
-    if T == eltype(state_values(dstvalp))
-        u0 = state_values(dstvalp)
-    elseif T != Union{}
-        u0 = T.(state_values(dstvalp))
-    end
-    buf, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Tunable(), newp)
-    if eltype(buf) != T
-        newbuf = similar(buf, T)
-        copyto!(newbuf, buf)
-        newp = repack(newbuf)
-    end
-    buf, repack, alias = SciMLStructures.canonicalize(SciMLStructures.Initials(), newp)
-    if eltype(buf) != T
-        newbuf = similar(buf, T)
-        copyto!(newbuf, buf)
-        newp = repack(newbuf)
-    end
-    return u0, newp
-end
-
-struct InitializationSystemMetadata
-    u0map::Dict{Any, Any}
-    pmap::Dict{Any, Any}
-    additional_guesses::Dict{Any, Any}
-    additional_initialization_eqs::Vector{Equation}
-    extra_metadata::NamedTuple
-    oop_reconstruct_u0_p::Union{Nothing, ReconstructInitializeprob}
-end
-
 function get_possibly_array_fallback_singletons(varmap, p)
     if haskey(varmap, p)
         return varmap[p]
@@ -543,22 +481,19 @@ function SciMLBase.remake_initialization_data(
     if u0 === missing && p === missing
         return odefn.initialization_data
     end
+
+    oldinitdata = odefn.initialization_data
+
     if !(eltype(u0) <: Pair) && !(eltype(p) <: Pair)
-        oldinitdata = odefn.initialization_data
         oldinitdata === nothing && return nothing
 
         oldinitprob = oldinitdata.initializeprob
         oldinitprob === nothing && return nothing
-        if !SciMLBase.has_sys(oldinitprob.f) || !(oldinitprob.f.sys isa NonlinearSystem)
-            return oldinitdata
-        end
-        oldinitsys = oldinitprob.f.sys
-        meta = get_metadata(oldinitsys)
-        if meta isa InitializationSystemMetadata && meta.oop_reconstruct_u0_p !== nothing
-            reconstruct_fn = meta.oop_reconstruct_u0_p
-        else
-            reconstruct_fn = ReconstructInitializeprob(sys, oldinitsys)
-        end
+
+        meta = oldinitdata.metadata
+        meta isa InitializationMetadata || return oldinitdata
+
+        reconstruct_fn = meta.oop_reconstruct_u0_p
         # the history function doesn't matter because `reconstruct_fn` is only going to
         # update the values of parameters, which aren't time dependent. The reason it
         # is called is because `Initial` parameters are calculated from the corresponding
@@ -569,16 +504,15 @@ function SciMLBase.remake_initialization_data(
         if oldinitprob.f.resid_prototype === nothing
             newf = oldinitprob.f
         else
-            newf = NonlinearFunction{
-                SciMLBase.isinplace(oldinitprob.f), SciMLBase.specialization(oldinitprob.f)}(
-                oldinitprob.f;
+            newf = remake(oldinitprob.f;
                 resid_prototype = calculate_resid_prototype(
                     length(oldinitprob.f.resid_prototype), new_initu0, new_initp))
         end
         initprob = remake(oldinitprob; f = newf, u0 = new_initu0, p = new_initp)
         return SciMLBase.OverrideInitData(initprob, oldinitdata.update_initializeprob!,
-            oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap)
+            oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap; metadata = oldinitdata.metadata)
     end
+
     dvs = unknowns(sys)
     ps = parameters(sys)
     u0map = to_varmap(u0, dvs)
@@ -592,16 +526,13 @@ function SciMLBase.remake_initialization_data(
     use_scc = true
     initialization_eqs = Equation[]
 
-    if SciMLBase.has_initializeprob(odefn)
-        oldsys = odefn.initialization_data.initializeprob.f.sys
-        meta = get_metadata(oldsys)
-        if meta isa InitializationSystemMetadata
-            u0map = merge(meta.u0map, u0map)
-            pmap = merge(meta.pmap, pmap)
-            merge!(guesses, meta.additional_guesses)
-            use_scc = get(meta.extra_metadata, :use_scc, true)
-            initialization_eqs = meta.additional_initialization_eqs
-        end
+    if oldinitdata !== nothing && oldinitdata.metadata isa InitializationMetadata
+        meta = oldinitdata.metadata
+        u0map = merge(meta.u0map, u0map)
+        pmap = merge(meta.pmap, pmap)
+        merge!(guesses, meta.guesses)
+        use_scc = meta.use_scc
+        initialization_eqs = meta.additional_initialization_eqs
     else
         # there is no initializeprob, so the original problem construction
         # had no solvable parameters and had the differential variables
@@ -662,19 +593,22 @@ function SciMLBase.late_binding_update_u0_p(
     if !(eltype(u0) <: Pair)
         # if `p` is not provided or is symbolic
         p === missing || eltype(p) <: Pair || return newu0, newp
-        newu0 === nothing && return newu0, newp
-        all(is_parameter(sys, Initial(x)) for x in unknowns(sys)) || return newu0, newp
+        (newu0 === nothing || isempty(newu0)) && return newu0, newp
+        initdata = prob.f.initialization_data
+        initdata === nothing && return newu0, newp
+        meta = initdata.metadata
+        meta isa InitializationMetadata || return newu0, newp
         newp = p === missing ? copy(newp) : newp
         initials, repack, alias = SciMLStructures.canonicalize(
             SciMLStructures.Initials(), newp)
         if eltype(initials) != eltype(newu0)
             initials = DiffEqBase.promote_u0(initials, newu0, t0)
             newp = repack(initials)
         end
-        if length(newu0) != length(unknowns(sys))
-            throw(ArgumentError("Expected `newu0` to be of same length as unknowns ($(length(unknowns(sys)))). Got $(typeof(newu0)) of length $(length(newu0))"))
+        if length(newu0) != length(prob.u0)
+            throw(ArgumentError("Expected `newu0` to be of same length as unknowns ($(length(prob.u0))). Got $(typeof(newu0)) of length $(length(newu0))"))
         end
-        setp(sys, Initial.(unknowns(sys)))(newp, newu0)
+        meta.set_initial_unknowns!(newp, newu0)
         return newu0, newp
     end
 
@@ -714,7 +648,7 @@ end
 Check if the given system is an initialization system.
 """
 function is_initializesystem(sys::AbstractSystem)
-    sys isa NonlinearSystem && get_metadata(sys) isa InitializationSystemMetadata
+    has_is_initializesystem(sys) && get_is_initializesystem(sys)
 end
 
 """