feat: add cost and coalesce to ODESystem

Author: vyudu

src/systems/abstractsystem.jl

@@ -920,7 +920,9 @@ for prop in [:eqs
              :tstops
              :index_cache
              :is_scalar_noise
-             :isscheduled]
+             :isscheduled
+             :costs
+             :coalesce]
     fname_get = Symbol(:get_, prop)
     fname_has = Symbol(:has_, prop)
     @eval begin

src/systems/diffeqs/abstractodesystem.jl

@@ -924,7 +924,7 @@ function generate_function_bc(sys::ODESystem, u0, u0_idxs, tspan; kwargs...)
     exprs = vcat(init_conds, cons)
     _p = reorder_parameters(sys, ps)
 
-    build_function_wrapper(sys, exprs, sol, _p..., t; output_type = Array, kwargs...)
+    build_function_wrapper(sys, exprs, sol, _p..., iv; output_type = Array, kwargs...)
 end
 
 """

src/systems/diffeqs/odesystem.jl

@@ -54,7 +54,7 @@ struct ODESystem <: AbstractODESystem
     """A set of expressions defining the costs of the system for optimal control."""
     costs::Vector
     """Takes the cost vector and returns a scalar for optimization."""
-    coalesce::Function
+    coalesce::Union{Nothing, Function}
     """
     Time-derivative matrix. Note: this field will not be defined until
     [`calculate_tgrad`](@ref) is called on the system.
@@ -209,7 +209,7 @@ struct ODESystem <: AbstractODESystem
     parent::Any
 
     function ODESystem(
-            tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, constraints, tgrad,
+            tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, constraints, costs, coalesce, tgrad,
             jac, ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, guesses,
             torn_matching, initializesystem, initialization_eqs, schedule,
             connector_type, preface, cevents,
@@ -233,7 +233,7 @@ struct ODESystem <: AbstractODESystem
             check_units(u, deqs)
         end
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name,
-            ctrls, observed, constraints, tgrad, jac,
+            ctrls, observed, constraints, costs, coalesce, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, description, systems, defaults, guesses, torn_matching,
             initializesystem, initialization_eqs, schedule, connector_type, preface,
             cevents, devents, parameter_dependencies, assertions, metadata,
@@ -247,6 +247,8 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         controls = Num[],
         observed = Equation[],
         constraintsystem = nothing,
+        costs = Num[],
+        coalesce = nothing,
         systems = ODESystem[],
         tspan = nothing,
         name = nothing,
@@ -327,22 +329,26 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
             cons = get_constraintsystem(sys)
             cons !== nothing && push!(conssystems, cons)
         end
-        @show conssystems
         @set! constraintsystem.systems = conssystems
     end
+    costs = wrap.(costs)
+    
+    if length(costs) > 1 && isnothing(coalesce)
+        error("Must specify a coalesce function for the costs vector.")
+    end
 
     assertions = Dict{BasicSymbolic, Any}(unwrap(k) => v for (k, v) in assertions)
 
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
-        deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, constraintsystem, tgrad, jac,
+        deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, constraintsystem, costs, coalesce, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, description, systems,
         defaults, guesses, nothing, initializesystem,
         initialization_eqs, schedule, connector_type, preface, cont_callbacks,
         disc_callbacks, parameter_dependencies, assertions,
         metadata, gui_metadata, is_dde, tstops, checks = checks)
 end
 
-function ODESystem(eqs, iv; constraints = Equation[], costs = Equation[], kwargs...)
+function ODESystem(eqs, iv; constraints = Equation[], costs = Num[], kwargs...)
     diffvars, allunknowns, ps, eqs = process_equations(eqs, iv)
 
     for eq in get(kwargs, :parameter_dependencies, Equation[])
@@ -394,9 +400,10 @@ function ODESystem(eqs, iv; constraints = Equation[], costs = Equation[], kwargs
             !in(p, new_ps) && push!(new_ps, p)
         end
     end
+    costs = wrap.(costs)
 
     return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars, consvars))),
-        collect(new_ps); constraintsystem, kwargs...)
+                     collect(new_ps); constraintsystem, costs, kwargs...)
 end
 
 # NOTE: equality does not check cached Jacobian
@@ -411,7 +418,9 @@ function Base.:(==)(sys1::ODESystem, sys2::ODESystem)
         _eq_unordered(get_ps(sys1), get_ps(sys2)) &&
         _eq_unordered(continuous_events(sys1), continuous_events(sys2)) &&
         _eq_unordered(discrete_events(sys1), discrete_events(sys2)) &&
-        all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
+        all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2))) &&
+        isequal(get_constraintsystem(sys1), get_constraintssystem(sys2)) &&
+        _eq_unordered(get_costs(sys1), get_costs(sys2))
 end
 
 function flatten(sys::ODESystem, noeqs = false)
@@ -768,14 +777,15 @@ end
 """
 Process the costs for the constraint system.
 """
-function process_costs(costs::Vector{Equation}, sts, ps, iv)
+function process_costs(costs::Vector, sts, ps, iv)
     coststs = OrderedSet()
     costps = OrderedSet()
     for cost in costs
         collect_vars!(coststs, costps, cost, iv)
     end
 
     validate_vars_and_find_ps!(coststs, costps, sts, iv)
+    coststs, costps
 end
 
 """
@@ -813,9 +823,34 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
     end
 end
 
-function generate_cost_function(sys::ODESystem)
+"""
+Generate a function that takes a solution object and computes the cost function obtained by coalescing the costs vector.
+"""
+function generate_cost_function(sys::ODESystem, kwargs...)
     costs = get_costs(sys)
     coalesce = get_coalesce(sys)
-    cost_fn = build_function_wrapper()
-    return (u, p, t) -> coalesce(cost_fn(u, p, t))
+    iv = get_iv(sys)
+
+    ps = parameters(sys; initial_parameters = false)
+    sts = unknowns(sys)
+    np = length(ps)
+    ns = length(sts)
+    stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
+    pidxmap = Dict([v => i for (i, v) in enumerate(ps)])
+
+    @variables sol(..)[1:ns]
+    for st in vars(costs)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+
+        costs = map(c -> Symbolics.fast_substitute(c, Dict(x(t) => sol(t)[idx])), costs)
+    end
+
+    _p = reorder_parameters(sys, ps)
+    fs = build_function_wrapper(sys, costs, sol, _p..., t; output_type = Array, kwargs...)
+    vc_oop, vc_iip = eval_or_rgf.(fs)
+
+    cost(sol, p, t) = coalesce(vc_oop(sol, p, t))
+    return cost
 end

test/bvproblem.jl

@@ -11,7 +11,7 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
 solvers = [MIRK4]
 daesolvers = [Ascher2, Ascher4, Ascher6]
 
-let
+@testset "Lotka-Volterra" begin
     @parameters α=7.5 β=4.0 γ=8.0 δ=5.0
     @variables x(t)=1.0 y(t)=2.0
 
@@ -47,7 +47,7 @@ let
 end
 
 ### Testing on pendulum
-let
+@testset "Pendulum" begin
     @parameters g=9.81 L=1.0
     @variables θ(t)=π / 2 θ_t(t)
 
@@ -86,7 +86,7 @@ end
 ##################################################################
 
 # Test generation of boundary condition function using `generate_function_bc`. Compare solutions to manually written boundary conditions
-let
+@testset "Boundary Condition Compilation" begin
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
     @variables x(..) y(..)
     eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
@@ -168,7 +168,7 @@ function test_solvers(
 end
 
 # Simple ODESystem with BVP constraints.
-let
+@testset "ODE with constraints" begin
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
     @variables x(..) y(..)
 
@@ -274,3 +274,26 @@ end
 #     bvp = SciMLBase.BVProblem{true, SciMLBase.AutoSpecialize}(pend, u0map, tspan, parammap; guesses, check_length = false)
 #     test_solvers(daesolvers, bvp, u0map, constr, get_alg_eqs(pend))
 # end
+
+@testset "Cost function compilation" begin
+    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
+    @variables x(..) y(..)
+
+    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
+        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
+
+    tspan = (0.0, 1.0)
+    u0map = [x(t) => 4.0, y(t) => 2.0]
+    parammap = [α => 7.5, β => 4, γ => 8.0, δ => 5.0]
+    costs = [x(0.6), x(0.3)^2]
+    consolidate(u) = (u[1] + 3)^2 + u[2]
+    @mtkbuild lksys = ODESystem(eqs, t; costs, coalesce = consolidate)
+    @test_throws ErrorException @mtkbuild lksys2 = ODESystem(eqs, t; costs)
+
+    prob = ODEProblem(lksys, u0map, tspan, parammap)
+    sol = solve(prob, Tsit5())
+    costfn = ModelingToolkit.generate_cost_function(lksys)
+    p = prob.p 
+    t = tspan[2]
+    @test costfn(sol, p, t) ≈ (sol(0.6)[1] + 3)^2 + sol(0.3)[1]^2
+end