Merge remote-tracking branch 'origin' into implicit_discrete_system

merge test fixes

Author: vyudu

.github/workflows/Downstream.yml

@@ -37,6 +37,7 @@ jobs:
           - {user: SciML, repo: MethodOfLines.jl, group: Interface}
           - {user: SciML, repo: MethodOfLines.jl, group: 2D_Diffusion}
           - {user: SciML, repo: MethodOfLines.jl, group: DAE}
+          - {user: SciML, repo: ModelingToolkitNeuralNets.jl, group: All}
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Yingbo Ma <[email protected]>", "Chris Rackauckas <[email protected]> and contributors"]
-version = "9.62.0"
+version = "9.64.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -104,11 +104,11 @@ DynamicQuantities = "^0.11.2, 0.12, 0.13, 1"
 EnumX = "1.0.4"
 ExprTools = "0.1.10"
 Expronicon = "0.8"
+FMI = "0.14"
 FindFirstFunctions = "1"
 ForwardDiff = "0.10.3"
 FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
-FMI = "0.14"
 Graphs = "1.5.2"
 HomotopyContinuation = "2.11"
 InfiniteOpt = "0.5"
@@ -119,6 +119,7 @@ LabelledArrays = "1.3"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16"
 Libdl = "1"
 LinearAlgebra = "1"
+Logging = "1"
 MLStyle = "0.4.17"
 ModelingToolkitStandardLibrary = "2.19"
 NaNMath = "0.3, 1"
@@ -128,26 +129,26 @@ OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
 OrdinaryDiffEqCore = "1.15.0"
 OrdinaryDiffEqDefault = "1.2"
-OrdinaryDiffEqNonlinearSolve = "1.3.0"
+OrdinaryDiffEqNonlinearSolve = "1.5.0"
 PrecompileTools = "1"
 REPL = "1"
 RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
-SciMLBase = "2.72.1"
+SciMLBase = "2.73"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
-StochasticDiffEq = "6.72.1"
 StochasticDelayDiffEq = "1.8.1"
+StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.37"
 SymbolicUtils = "3.14"
-Symbolics = "6.29"
+Symbolics = "6.29.1"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
@@ -164,6 +165,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -187,4 +189,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve", "Logging"]

docs/Project.toml

@@ -15,6 +15,7 @@ Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
@@ -38,6 +39,7 @@ Optimization = "3.9, 4"
 OptimizationOptimJL = "0.1, 0.4"
 OrdinaryDiffEq = "6.31"
 Plots = "1.36"
+PreallocationTools = "0.4"
 SciMLStructures = "1.1"
 Setfield = "1"
 StochasticDiffEq = "6"

docs/src/basics/Variable_metadata.md

@@ -201,6 +201,7 @@ state_priority(important_dof)
 Units for variables can be designated using symbolic metadata. For more information, please see the [model validation and units](@ref units) section of the docs. Note that `getunit` is not equivalent to `get_unit` - the former is a metadata getter for individual variables (and is provided so the same interface function for `unit` exists like other metadata), while the latter is used to handle more general symbolic expressions.
 
 ```@example metadata
+using DynamicQuantities
 @variables speed [unit = u"m/s"]
 hasunit(speed)
 ```

docs/src/examples/remake.md

@@ -45,12 +45,22 @@ parameters to optimize.
 
 ```@example Remake
 using SymbolicIndexingInterface: parameter_values, state_values
-using SciMLStructures: Tunable, replace, replace!
+using SciMLStructures: Tunable, canonicalize, replace, replace!
+using PreallocationTools
 
 function loss(x, p)
     odeprob = p[1] # ODEProblem stored as parameters to avoid using global variables
     ps = parameter_values(odeprob) # obtain the parameter object from the problem
-    ps = replace(Tunable(), ps, x) # create a copy with the values passed to the loss function
+    diffcache = p[5]
+    # get an appropriately typed preallocated buffer to store the `x` values in
+    buffer = get_tmp(diffcache, x)
+    # copy the current values to this buffer
+    copyto!(buffer, canonicalize(Tunable(), ps)[1])
+    # create a copy of the parameter object with the buffer
+    ps = replace(Tunable(), ps, buffer)
+    # set the updated values in the parameter object
+    setter = p[4]
+    setter(ps, x)
     # remake the problem, passing in our new parameter object
     newprob = remake(odeprob; p = ps)
     timesteps = p[2]
@@ -81,49 +91,22 @@ We can perform the optimization as below:
 ```@example Remake
 using Optimization
 using OptimizationOptimJL
+using SymbolicIndexingInterface
 
 # manually create an OptimizationFunction to ensure usage of `ForwardDiff`, which will
 # require changing the types of parameters from `Float64` to `ForwardDiff.Dual`
 optfn = OptimizationFunction(loss, Optimization.AutoForwardDiff())
+# function to set the parameters we are optimizing
+setter = setp(odeprob, [α, β, γ, δ])
+# `DiffCache` to avoid allocations
+diffcache = DiffCache(canonicalize(Tunable(), parameter_values(odeprob))[1])
 # parameter object is a tuple, to store differently typed objects together
 optprob = OptimizationProblem(
-    optfn, rand(4), (odeprob, timesteps, data), lb = 0.1zeros(4), ub = 3ones(4))
+    optfn, rand(4), (odeprob, timesteps, data, setter, diffcache),
+    lb = 0.1zeros(4), ub = 3ones(4))
 sol = solve(optprob, BFGS())
 ```
 
-To identify which values correspond to which parameters, we can `replace!` them into the
-`ODEProblem`:
-
-```@example Remake
-replace!(Tunable(), parameter_values(odeprob), sol.u)
-odeprob.ps[[α, β, γ, δ]]
-```
-
-`replace!` operates in-place, so the values being replaced must be of the same type as those
-stored in the parameter object, or convertible to that type. For demonstration purposes, we
-can construct a loss function that uses `replace!`, and calculate gradients using
-`AutoFiniteDiff` rather than `AutoForwardDiff`.
-
-```@example Remake
-function loss2(x, p)
-    odeprob = p[1] # ODEProblem stored as parameters to avoid using global variables
-    newprob = remake(odeprob) # copy the problem with `remake`
-    # update the parameter values in-place
-    replace!(Tunable(), parameter_values(newprob), x)
-    timesteps = p[2]
-    sol = solve(newprob, AutoTsit5(Rosenbrock23()); saveat = timesteps)
-    truth = p[3]
-    data = Array(sol)
-    return sum((truth .- data) .^ 2) / length(truth)
-end
-
-# use finite-differencing to calculate derivatives
-optfn2 = OptimizationFunction(loss2, Optimization.AutoFiniteDiff())
-optprob2 = OptimizationProblem(
-    optfn2, rand(4), (odeprob, timesteps, data), lb = 0.1zeros(4), ub = 3ones(4))
-sol = solve(optprob2, BFGS())
-```
-
 # Re-creating the problem
 
 There are multiple ways to re-create a problem with new state/parameter values. We will go

docs/src/tutorials/initialization.md

@@ -203,6 +203,73 @@ long enough you will see that `λ = 0` is required for this equation, but since
     problem constructor. Additionally, any warning about not being fully determined can
     be suppressed via passing `warn_initialize_determined = false`.
 
+## Constant constraints in initialization
+
+Consider the pendulum system again:
+
+```@repl init
+equations(pend)
+observed(pend)
+```
+
+Suppose we want to solve the same system with multiple different initial
+y-velocities from a given position.
+
+```@example init
+prob = ODEProblem(
+    pend, [x => 1, D(y) => 0], (0.0, 1.5), [g => 1], guesses = [λ => 0, y => 1, x => 1])
+sol1 = solve(prob, Rodas5P())
+```
+
+```@example init
+sol1[D(y), 1]
+```
+
+Repeatedly re-creating the `ODEProblem` with different values of `D(y)` and `x` or
+repeatedly calling `remake` is slow. Instead, for any `variable => constant` constraint
+in the `ODEProblem` initialization (whether provided to the `ODEProblem` constructor or
+a default value) we can update the `constant` value. ModelingToolkit refers to these
+values using the `Initial` operator. For example:
+
+```@example init
+prob.ps[[Initial(x), Initial(D(y))]]
+```
+
+To solve with a different starting y-velocity, we can simply do
+
+```@example init
+prob.ps[Initial(D(y))] = -0.1
+sol2 = solve(prob, Rodas5P())
+```
+
+```@example init
+sol2[D(y), 1]
+```
+
+Note that this _only_ applies for constant constraints for the current ODEProblem.
+For example, `D(x)` does not have a constant constraint - it is solved for by
+initialization. Thus, mutating `Initial(D(x))` does not have any effect:
+
+```@repl init
+sol2[D(x), 1]
+prob.ps[Initial(D(x))] = 1.0
+sol3 = solve(prob, Rodas5P())
+sol3[D(x), 1]
+```
+
+To enforce this constraint, we would have to `remake` the problem (or construct a new one).
+
+```@repl init
+prob2 = remake(prob; u0 = [y => 0.0, D(x) => 0.0, x => nothing, D(y) => nothing]);
+sol4 = solve(prob2, Rodas5P())
+sol4[D(x), 1]
+```
+
+Note the need to provide `x => nothing, D(y) => nothing` to override the previously
+provided initial conditions. Since `remake` is a partial update, the constraints provided
+to it are merged with the ones already present in the problem. Existing constraints can be
+removed by providing a value of `nothing`.
+
 ## Initialization of parameters
 
 Parameters may also be treated as unknowns in the initialization system. Doing so works
@@ -231,6 +298,9 @@ constraints provided to it. The new values will be combined with the original
 variable-value mapping provided to `ODEProblem` and used to construct the initialization
 problem.
 
+The variable on the left hand side of all parameter dependencies also has an `Initial`
+variant, which is used if a constant constraint is provided for the variable.
+
 ### Parameter initialization by example
 
 Consider the following system, where the sum of two unknowns is a constant parameter

ext/MTKBifurcationKitExt.jl

@@ -97,7 +97,10 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
     @set! nsys.index_cache = nothing # force usage of a parameter vector instead of `MTKParameters`
     # Creates F and J functions.
     ofun = NonlinearFunction(nsys; jac = jac)
-    F = ofun.f
+    F = let f = ofun.f
+        _f(resid, u, p) = (f(resid, u, p); resid)
+        _f(u, p) = f(u, p)
+    end
     J = jac ? ofun.jac : nothing
 
     # Converts the input state guess.
@@ -136,6 +139,7 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
         args...;
         record_from_solution = record_from_solution,
         J = J,
+        inplace = true,
         kwargs...)
 end
 

ext/MTKChainRulesCoreExt.jl

@@ -16,7 +16,12 @@ end
 
 function subset_idxs(idxs, portion, template)
     ntuple(Val(length(template))) do subi
-        [Base.tail(idx.idx) for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+        result = [Base.tail(idx.idx)
+                  for idx in idxs if idx.portion == portion && idx.idx[1] == subi]
+        if isempty(result)
+            result = []
+        end
+        result
     end
 end
 

src/ModelingToolkit.jl

@@ -45,7 +45,7 @@ using Compat
 using AbstractTrees
 using DiffEqBase, SciMLBase, ForwardDiff
 using SciMLBase: StandardODEProblem, StandardNonlinearProblem, handle_varmap, TimeDomain,
-                 PeriodicClock, Clock, SolverStepClock, Continuous
+                 PeriodicClock, Clock, SolverStepClock, Continuous, OverrideInit, NoInit
 using Distributed
 import JuliaFormatter
 using MLStyle
@@ -56,6 +56,7 @@ using RecursiveArrayTools
 import Graphs: SimpleDiGraph, add_edge!, incidence_matrix
 import BlockArrays: BlockArray, BlockedArray, Block, blocksize, blocksizes, blockpush!,
                     undef_blocks, blocks
+using OffsetArrays: Origin
 import CommonSolve
 import EnumX
 
@@ -153,6 +154,7 @@ include("systems/imperative_affect.jl")
 include("systems/callbacks.jl")
 include("systems/codegen_utils.jl")
 include("systems/problem_utils.jl")
+include("linearization.jl")
 
 include("systems/nonlinear/nonlinearsystem.jl")
 include("systems/nonlinear/homotopy_continuation.jl")
@@ -261,7 +263,8 @@ export Term, Sym
 export SymScope, LocalScope, ParentScope, DelayParentScope, GlobalScope
 export independent_variable, equations, controls, observed, full_equations
 export initialization_equations, guesses, defaults, parameter_dependencies, hierarchy
-export structural_simplify, expand_connections, linearize, linearization_function
+export structural_simplify, expand_connections, linearize, linearization_function,
+       LinearizationProblem
 
 export calculate_jacobian, generate_jacobian, generate_function, generate_custom_function
 export calculate_control_jacobian, generate_control_jacobian
@@ -281,7 +284,7 @@ export toexpr, get_variables
 export simplify, substitute
 export build_function
 export modelingtoolkitize
-export generate_initializesystem
+export generate_initializesystem, Initial
 
 export alg_equations, diff_equations, has_alg_equations, has_diff_equations
 export get_alg_eqs, get_diff_eqs, has_alg_eqs, has_diff_eqs

src/inputoutput.jl

@@ -211,7 +211,7 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     sys, _ = io_preprocessing(sys, inputs, []; simplify, kwargs...)
 
     dvs = unknowns(sys)
-    ps = parameters(sys)
+    ps = parameters(sys; initial_parameters = true)
     ps = setdiff(ps, inputs)
     if disturbance_inputs !== nothing
         # remove from inputs since we do not want them as actual inputs to the dynamics
@@ -234,16 +234,14 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
            [eq.rhs for eq in eqs]
 
     # TODO: add an optional check on the ordering of observed equations
-    u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = map(x -> time_varying_as_func(value(x), sys), ps)
-    p = reorder_parameters(sys, p)
+    p = reorder_parameters(sys, ps)
     t = get_iv(sys)
 
     # pre = has_difference ? (ex -> ex) : get_postprocess_fbody(sys)
     if disturbance_argument
-        args = (u, inputs, p..., t, disturbance_inputs)
+        args = (dvs, inputs, p..., t, disturbance_inputs)
     else
-        args = (u, inputs, p..., t)
+        args = (dvs, inputs, p..., t)
     end
     if implicit_dae
         ddvs = map(Differential(get_iv(sys)), dvs)
@@ -252,6 +250,7 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
         p_end = length(p) + 2 + implicit_dae)
     f = eval_or_rgf.(f; eval_expression, eval_module)
+    ps = setdiff(parameters(sys), inputs, disturbance_inputs)
     (; f, dvs, ps, io_sys = sys)
 end