Expanded to include other generate functions for ODESystems & added more testing.

Author: lamorton

src/systems/diffeqs/abstractodesystem.jl

@@ -83,20 +83,23 @@ end
 function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
                         simplify = false, kwargs...)
     tgrad = calculate_tgrad(sys, simplify = simplify)
-    return build_function(tgrad, dvs, ps, get_iv(sys); kwargs...)
+    pre = get_preprocess_constants(tgrad)
+    return build_function(tgrad, dvs, ps, get_iv(sys); postprocess_fbody = pre, kwargs...)
 end
 
 function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
                            simplify = false, sparse = false, kwargs...)
     jac = calculate_jacobian(sys; simplify = simplify, sparse = sparse)
-    return build_function(jac, dvs, ps, get_iv(sys); kwargs...)
+    pre = get_preprocess_constants(jac)
+    return build_function(jac, dvs, ps, get_iv(sys); postprocess_fbody = pre, kwargs...)
 end
 
 function generate_control_jacobian(sys::AbstractODESystem, dvs = states(sys),
                                    ps = parameters(sys);
                                    simplify = false, sparse = false, kwargs...)
     jac = calculate_control_jacobian(sys; simplify = simplify, sparse = sparse)
-    return build_function(jac, dvs, ps, get_iv(sys); kwargs...)
+    pre = get_preprocess_constants(jac)
+    return build_function(jac, dvs, ps, get_iv(sys); postprocess_fbody = pre, kwargs...)
 end
 
 function generate_dae_jacobian(sys::AbstractODESystem, dvs = states(sys),
@@ -109,7 +112,8 @@ function generate_dae_jacobian(sys::AbstractODESystem, dvs = states(sys),
     dvs = states(sys)
     @variables ˍ₋gamma
     jac = ˍ₋gamma * jac_du + jac_u
-    return build_function(jac, derivatives, dvs, ps, ˍ₋gamma, get_iv(sys); kwargs...)
+    pre = get_preprocess_constants(jac)
+    return build_function(jac, derivatives, dvs, ps, ˍ₋gamma, get_iv(sys); postprocess_fbody = pre, kwargs...)
 end
 
 function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
@@ -163,8 +167,10 @@ function generate_difference_cb(sys::ODESystem, dvs = states(sys), ps = paramete
     end
 
     pre = get_postprocess_fbody(sys)
+    cpre = get_preprocess_constants(body)
+    pre2 = x -> pre(cpre(x))
     f_oop, f_iip = build_function(body, u, p, t; expression = Val{false},
-                                  postprocess_fbody = pre, kwargs...)
+                                  postprocess_fbody = pre2, kwargs...)
 
     cb_affect! = let f_oop = f_oop, f_iip = f_iip
         function cb_affect!(integ)

src/systems/diffeqs/odesystem.jl

@@ -238,14 +238,6 @@ function ODESystem(eqs, iv = nothing; kwargs...)
                      collect(Iterators.flatten((diffvars, algevars))), ps; kwargs...)
 end
 
-function collect_constants(eqs) #Does this need to be different for other system types?
-    constants = Set()
-    for eq in eqs
-        collect_constants!(constants, eq.lhs)
-        collect_constants!(constants, eq.rhs)
-    end
-    return collect(constants)
-end
 
 # NOTE: equality does not check cached Jacobian
 function Base.:(==)(sys1::ODESystem, sys2::ODESystem)

src/utils.jl

@@ -480,17 +480,6 @@ function collect_vars!(states, parameters, expr, iv)
     return nothing
 end
 
-function collect_constants!(constants, expr)
-    if expr isa Sym
-        collect_constant!(constants, expr)
-    else
-        for var in vars(expr)
-            collect_constant!(constants, var)
-        end
-    end
-    return nothing
-end
-
 function collect_vars_difference!(states, parameters, expr, iv)
     if expr isa Sym
         collect_var!(states, parameters, expr, iv)
@@ -515,10 +504,46 @@ function collect_var!(states, parameters, var, iv)
     return nothing
 end
 
+function collect_constants(eqs::Vector{Equation}) #For get_substitutions_and_solved_states
+    constants = []
+    for eq in eqs
+        collect_constants!(constants, eq.lhs)
+        collect_constants!(constants, eq.rhs)
+    end
+    return constants
+end
+
+function collect_constants(eqs::AbstractArray{T}) where T # For generate_tgrad / generate_jacobian / generate_difference_cb
+    constants = T[]
+    for eq in eqs
+        collect_constants!(constants, unwrap(eq))
+    end
+    return constants
+end
+
 function collect_constant!(constants, var)
     if isconstant(var)
-        push!(constants,var)
+        push!(constants, var)
     end
+    return nothing
+end
+
+function collect_constants!(constants, expr)
+    if expr isa Sym
+        collect_constant!(constants, expr)
+    else
+        for var in vars(expr)
+            collect_constant!(constants, var)
+        end
+    end
+    return nothing
+end
+
+function get_preprocess_constants(eqs)
+    cs = collect_constants(eqs)
+    pre = ex -> Let(Assignment[Assignment(x, getdefault(x)) for x in cs],
+                           ex, false)
+    return pre
 end
 
 function get_postprocess_fbody(sys)
@@ -561,6 +586,9 @@ end
 function get_substitutions_and_solved_states(sys; no_postprocess = false)
     #Inject substitutions for constants => values
     cs = collect_constants([sys.eqs; sys.observed]) #ctrls? what else?
+    if !empty_substitutions(sys) 
+        cs = [cs; collect_constants(sys.substitutions.subs)]
+    end
     # Swap constants for their values
     cmap = map(x -> x ~ getdefault(x), cs)
 

test/odesystem.jl

@@ -8,13 +8,14 @@ using ModelingToolkit: value
 
 # Define some variables
 @parameters t σ ρ β
+@constants κ = 1
 @variables x(t) y(t) z(t)
 D = Differential(t)
 
 # Define a differential equation
 eqs = [D(x) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y,
-    D(z) ~ x * y - β * z]
+    D(z) ~ x * y - β * z * κ]
 
 ModelingToolkit.toexpr.(eqs)[1]
 @named de = ODESystem(eqs; defaults = Dict(x => 1))
@@ -71,7 +72,7 @@ end
 
 eqs = [D(x) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y * t,
-    D(z) ~ x * y - β * z]
+    D(z) ~ x * y - β * z * κ]
 @named de = ODESystem(eqs)
 ModelingToolkit.calculate_tgrad(de)
 
@@ -87,7 +88,7 @@ tgrad_iip(du, u, p, t)
 @parameters σ′(t - 1)
 eqs = [D(x) ~ σ′ * (y - x),
     D(y) ~ x * (ρ - z) - y,
-    D(z) ~ x * y - β * z]
+    D(z) ~ x * y - β * z * κ]
 @named de = ODESystem(eqs)
 test_diffeq_inference("global iv-varying", de, t, (x, y, z), (σ′, ρ, β))
 
@@ -99,7 +100,7 @@ f(du, [1.0, 2.0, 3.0], [x -> x + 7, 2, 3], 5.0)
 @parameters σ(..)
 eqs = [D(x) ~ σ(t - 1) * (y - x),
     D(y) ~ x * (ρ - z) - y,
-    D(z) ~ x * y - β * z]
+    D(z) ~ x * y - β * z * κ]
 @named de = ODESystem(eqs)
 test_diffeq_inference("single internal iv-varying", de, t, (x, y, z), (σ(t - 1), ρ, β))
 f = eval(generate_function(de, [x, y, z], [σ, ρ, β])[2])
@@ -146,7 +147,7 @@ ODEFunction(de1, [uˍtt, xˍt, uˍt, u, x], [])(du, ones(5), nothing, 0.1)
 a = y - x
 eqs = [D(x) ~ σ * a,
     D(y) ~ x * (ρ - z) - y,
-    D(z) ~ x * y - β * z]
+    D(z) ~ x * y - β * z * κ]
 @named de = ODESystem(eqs)
 generate_function(de, [x, y, z], [σ, ρ, β])
 jac = calculate_jacobian(de)
@@ -201,7 +202,7 @@ D = Differential(t)
 # reorder the system just to be a little spicier
 eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃,
     0 ~ y₁ + y₂ + y₃ - 1,
-    D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃]
+    D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃ * κ]
 @named sys = ODESystem(eqs, defaults = [k₁ => 100, k₂ => 3e7, y₁ => 1.0])
 u0 = Pair[]
 push!(u0, y₂ => 0.0)
@@ -222,7 +223,7 @@ for p in [prob1, prob14]
     @test Set(Num.(states(sys)) .=> p.u0) == Set([y₁ => 1, y₂ => 0, y₃ => 0])
 end
 prob2 = ODEProblem(sys, u0, tspan, p, jac = true)
-prob3 = ODEProblem(sys, u0, tspan, p, jac = true, sparse = true)
+prob3 = ODEProblem(sys, u0, tspan, p, jac = true, sparse = true) #SparseMatrixCSC need to handle
 @test prob3.f.jac_prototype isa SparseMatrixCSC
 prob3 = ODEProblem(sys, u0, tspan, p, jac = true, sparsity = true)
 @test prob3.f.sparsity isa SparseMatrixCSC