Merge pull request #2382 from MarcBerliner/initial_conditions

Custom initial condition equations for an `ODESystem`

Author: ChrisRackauckas

docs/src/basics/Variable_metadata.md

@@ -83,6 +83,19 @@ hasbounds(u)
 getbounds(u)
 ```
 
+## Guess
+
+Specify an initial guess for custom initial conditions of an `ODESystem`.
+
+```@example metadata
+@variables u [guess = 1]
+hasguess(u)
+```
+
+```@example metadata
+getguess(u)
+```
+
 ## Mark input as a disturbance
 
 Indicate that an input is not available for control, i.e., it's a disturbance input.

src/ModelingToolkit.jl

@@ -142,6 +142,7 @@ include("systems/jumps/jumpsystem.jl")
 
 include("systems/nonlinear/nonlinearsystem.jl")
 include("systems/nonlinear/modelingtoolkitize.jl")
+include("systems/nonlinear/initializesystem.jl")
 
 include("systems/optimization/constraints_system.jl")
 include("systems/optimization/optimizationsystem.jl")
@@ -201,7 +202,8 @@ export NonlinearSystem, OptimizationSystem, ConstraintsSystem
 export alias_elimination, flatten
 export connect, domain_connect, @connector, Connection, Flow, Stream, instream
 export @component, @mtkmodel, @mtkbuild
-export isinput, isoutput, getbounds, hasbounds, isdisturbance, istunable, getdist, hasdist,
+export isinput, isoutput, getbounds, hasbounds, getguess, hasguess, isdisturbance,
+    istunable, getdist, hasdist,
     tunable_parameters, isirreducible, getdescription, hasdescription, isbinaryvar,
     isintegervar
 export ode_order_lowering, dae_order_lowering, liouville_transform
@@ -236,6 +238,7 @@ export toexpr, get_variables
 export simplify, substitute
 export build_function
 export modelingtoolkitize
+export initializesystem
 
 export @variables, @parameters, @constants, @brownian
 export @named, @nonamespace, @namespace, extend, compose, complete

src/systems/nonlinear/initializesystem.jl

@@ -0,0 +1,55 @@
+"""
+$(TYPEDSIGNATURES)
+
+Generate `NonlinearSystem` which initializes an ODE problem from specified initial conditions of an `ODESystem`.
+"""
+function initializesystem(sys::ODESystem; name = nameof(sys), kwargs...)
+    if has_parent(sys) && (parent = get_parent(sys); parent !== nothing)
+        sys = parent
+    end
+    sts, eqs = states(sys), equations(sys)
+
+    idxs_diff = isdiffeq.(eqs)
+    idxs_alge = .!idxs_diff
+
+    # Algebraic equations and initial guesses are unchanged
+    eqs_ics = similar(eqs)
+    u0 = Vector{Any}(undef, length(sts))
+
+    eqs_ics[idxs_alge] .= eqs[idxs_alge]
+    u0[idxs_alge] .= getmetadata.(unwrap.(sts[idxs_alge]),
+        Symbolics.VariableDefaultValue,
+        nothing)
+
+    for idx in findall(idxs_diff)
+        st = sts[idx]
+        if !hasdefault(st)
+            error("Invalid setup: unknown $(st) has no default value or equation.")
+        end
+
+        def = getdefault(st)
+        if def isa Equation
+            if !hasguess(st)
+                error("Invalid setup: unknown $(st) has an initial condition equation with no guess.")
+            end
+            guess = getguess(st)
+            eqs_ics[idx] = def
+
+            u0[idx] = guess
+        else
+            eqs_ics[idx] = st ~ def
+
+            u0[idx] = def
+        end
+    end
+
+    pars = parameters(sys)
+    sys_nl = NonlinearSystem(eqs_ics,
+        sts,
+        pars;
+        defaults = Dict(sts .=> u0),
+        name,
+        kwargs...)
+
+    return sys_nl
+end

src/variables.jl

@@ -425,3 +425,42 @@ macro brownian(xs...)
         xs,
         tobrownian) |> esc
 end
+
+## Guess ======================================================================
+struct VariableGuess end
+Symbolics.option_to_metadata_type(::Val{:guess}) = VariableGuess
+getguess(x::Num) = getguess(Symbolics.unwrap(x))
+
+"""
+    getguess(x)
+
+Get the guess for the initial value associated with symbolic variable `x`.
+Create variables with a guess like this
+
+```
+@variables x [guess=1]
+```
+"""
+function getguess(x)
+    p = Symbolics.getparent(x, nothing)
+    p === nothing || (x = p)
+    Symbolics.getmetadata(x, VariableGuess, nothing)
+end
+
+"""
+    hasguess(x)
+
+Determine whether symbolic variable `x` has a guess associated with it.
+See also [`getguess`](@ref).
+"""
+function hasguess(x)
+    getguess(x) !== nothing
+end
+
+function get_default_or_guess(x)
+    if hasdefault(x) && !((def = getdefault(x)) isa Equation)
+        return def
+    else
+        return getguess(x)
+    end
+end

test/nonlinearsystem.jl

@@ -4,6 +4,7 @@ using DiffEqBase, SparseArrays
 using Test
 using NonlinearSolve
 using ModelingToolkit: value
+using ModelingToolkit: get_default_or_guess
 
 canonequal(a, b) = isequal(simplify(a), simplify(b))
 
@@ -232,3 +233,45 @@ testdict = Dict([:test => 1])
     @test prob_.u0 == [1.0, 2.0, 1.0]
     @test prob_.p == [2.0, 1.0, 1.0]
 end
+
+@testset "Initialization System" begin
+    # Define the Lotka Volterra system which begins at steady state
+    @parameters t
+    pars = @parameters a=1.5 b=1.0 c=3.0 d=1.0 dx_ss = 1e-5
+
+    vars = @variables begin
+        dx(t),
+        dy(t),
+        (x(t) = dx ~ dx_ss), [guess = 0.5]
+        (y(t) = dy ~ 0), [guess = -0.5]
+    end
+
+    D = Differential(t)
+
+    eqs = [dx ~ a * x - b * x * y
+        dy ~ -c * y + d * x * y
+        D(x) ~ dx
+        D(y) ~ dy]
+
+    @named sys = ODESystem(eqs, t, vars, pars)
+
+    sys_simple = structural_simplify(sys)
+
+    # Set up the initialization system
+    sys_init = initializesystem(sys_simple)
+
+    sys_init_simple = structural_simplify(sys_init)
+
+    prob = NonlinearProblem(sys_init_simple, get_default_or_guess.(states(sys_init_simple)))
+
+    @test prob.u0 == [0.5, -0.5]
+
+    sol = solve(prob)
+    @test sol.retcode == SciMLBase.ReturnCode.Success
+
+    # Confirm for all the states of the non-simplified system
+    @test all(.≈(sol[states(sys)], [1e-5, 0, 1e-5 / 1.5, 0]; atol = 1e-8))
+
+    # Confirm for all the states of the simplified system
+    @test all(.≈(sol[states(sys_simple)], [1e-5 / 1.5, 0]; atol = 1e-8))
+end

test/test_variable_metadata.jl

@@ -8,6 +8,14 @@ using ModelingToolkit
 @variables y
 @test !hasbounds(y)
 
+# Guess
+@variables y [guess = 0]
+@test getguess(y) === 0
+@test hasguess(y) === true
+
+@variables y
+@test hasguess(y) === false
+
 # Disturbance
 @variables u [disturbance = true]
 @test isdisturbance(u)