Merge pull request #3347 from AayushSabharwal/as/init-no-ts

feat: always build initialization problem

Author: AayushSabharwal

Project.toml

@@ -135,7 +135,7 @@ 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"

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"

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

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
 
@@ -152,6 +153,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")
@@ -258,7 +260,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
@@ -278,7 +281,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