Merge pull request #897 from SciML/myb/cleanup

Add DAEProblem overload for `ODESystem`s

Author: ChrisRackauckas

Project.toml

@@ -81,7 +81,8 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
+Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BenchmarkTools", "ForwardDiff", "GalacticOptim", "NonlinearSolve", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]
+test = ["BenchmarkTools", "ForwardDiff", "GalacticOptim", "NonlinearSolve", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]

src/ModelingToolkit.jl

@@ -137,6 +137,7 @@ include("structural_transformation/StructuralTransformations.jl")
 @reexport using .StructuralTransformations
 
 export ODESystem, ODEFunction, ODEFunctionExpr, ODEProblemExpr
+export DAEFunctionExpr, DAEProblemExpr
 export SDESystem, SDEFunction, SDEFunctionExpr, SDESystemExpr
 export SystemStructure
 export JumpSystem

src/systems/diffeqs/abstractodesystem.jl

@@ -61,23 +61,34 @@ function check_derivative_variables(eq, expr=eq.rhs)
     foreach(Base.Fix1(check_derivative_variables, eq), arguments(expr))
 end
 
-function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
+function generate_function(
+        sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys);
+        implicit_dae=false,
+        ddvs=implicit_dae ? map(Differential(independent_variable(sys)), dvs) : nothing,
+        kwargs...
+    )
     # optimization
     #obsvars = map(eq->eq.lhs, observed(sys))
     #fulldvs = [dvs; obsvars]
 
     eqs = equations(sys)
     foreach(check_derivative_variables, eqs)
     # substitute x(t) by just x
-    rhss = [deq.rhs for deq in eqs]
+    rhss = implicit_dae ? [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs] :
+                          [eq.rhs for eq in eqs]
     #obss = [makesym(value(eq.lhs)) ~ substitute(eq.rhs, sub) for eq ∈ observed(sys)]
     #rhss = Let(obss, rhss)
 
     # TODO: add an optional check on the ordering of observed equations
-    return build_function(rhss,
-                          map(x->time_varying_as_func(value(x), sys), dvs),
-                          map(x->time_varying_as_func(value(x), sys), ps),
-                          get_iv(sys); kwargs...)
+    u = map(x->time_varying_as_func(value(x), sys), dvs)
+    p = map(x->time_varying_as_func(value(x), sys), ps)
+    t = get_iv(sys)
+
+    if implicit_dae
+        build_function(rhss, ddvs, u, p, t; kwargs...)
+    else
+        build_function(rhss, u, p, t; kwargs...)
+    end
 end
 
 function time_varying_as_func(x, sys)
@@ -120,8 +131,10 @@ function isautonomous(sys::AbstractODESystem)
     all(iszero,tgrad)
 end
 
-function DiffEqBase.ODEFunction(sys::AbstractODESystem, args...; kwargs...)
-    ODEFunction{true}(sys, args...; kwargs...)
+for F in [:ODEFunction, :DAEFunction]
+    @eval function DiffEqBase.$F(sys::AbstractODESystem, args...; kwargs...)
+        $F{true}(sys, args...; kwargs...)
+    end
 end
 
 """
@@ -201,7 +214,64 @@ end
 
 """
 ```julia
-function DiffEqBase.ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
+function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing, tgrad=false,
+                                     jac = false,
+                                     sparse = false,
+                                     kwargs...) where {iip}
+```
+
+Create an `DAEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and
+`ps` are used to set the order of the dependent variable and parameter vectors,
+respectively.
+"""
+function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
+                                     ps = parameters(sys), u0 = nothing;
+                                     ddvs=map(diff2term ∘ Differential(independent_variable(sys)), dvs),
+                                     version = nothing,
+                                     #=
+                                     tgrad=false,
+                                     jac = false,
+                                     sparse = false,
+                                     =#
+                                     simplify=false,
+                                     eval_expression = true,
+                                     eval_module = @__MODULE__,
+                                     kwargs...) where {iip}
+
+    f_gen = generate_function(sys, dvs, ps; implicit_dae = true, expression=Val{eval_expression}, expression_module=eval_module, kwargs...)
+    f_oop,f_iip = eval_expression ? (@RuntimeGeneratedFunction(eval_module, ex) for ex in f_gen) : f_gen
+    f(du,u,p,t) = f_oop(du,u,p,t)
+    f(out,du,u,p,t) = f_iip(out,du,u,p,t)
+
+    # TODO: Jacobian sparsity / sparse Jacobian / dense Jacobian
+
+    #=
+    observedfun = let sys = sys, dict = Dict()
+        # TODO: We don't have enought information to reconstruct arbitrary state
+        # in general from `(u, p, t)`, e.g. `a ~ D(x)`.
+        function generated_observed(obsvar, u, p, t)
+            obs = get!(dict, value(obsvar)) do
+                build_explicit_observed_function(sys, obsvar)
+            end
+            obs(u, p, t)
+        end
+    end
+    =#
+
+    DAEFunction{iip}(
+                     f,
+                     syms = Symbol.(dvs),
+                     # missing fields in `DAEFunction`
+                     #indepsym = Symbol(independent_variable(sys)),
+                     #observed = observedfun,
+                    )
+end
+
+"""
+```julia
+function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      ps = parameters(sys);
                                      version = nothing, tgrad=false,
                                      jac = false,
@@ -277,6 +347,7 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
 end
 
 function process_DEProblem(constructor, sys::AbstractODESystem,u0map,parammap;
+                           implicit_dae = false, du0map = nothing,
                            version = nothing, tgrad=false,
                            jac = false,
                            checkbounds = false, sparse = false,
@@ -287,27 +358,80 @@ function process_DEProblem(constructor, sys::AbstractODESystem,u0map,parammap;
     dvs = states(sys)
     ps = parameters(sys)
     defs = defaults(sys)
+    iv = independent_variable(sys)
 
     u0 = varmap_to_vars(u0map,dvs; defaults=defs)
+    if implicit_dae && du0map !== nothing
+        ddvs = map(Differential(iv), dvs)
+        du0 = varmap_to_vars(du0map, ddvs; defaults=defaults, toterm=identity)
+    else
+        du0 = nothing
+        ddvs = nothing
+    end
     p = varmap_to_vars(parammap,ps; defaults=defs)
 
     if u0 !== nothing
         length(dvs) == length(u0) || throw(ArgumentError("States ($(length(dvs))) and initial conditions ($(length(u0))) are of different lengths."))
     end
 
-    f = constructor(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,checkbounds=checkbounds,
+    f = constructor(sys,dvs,ps,u0;ddvs=ddvs,tgrad=tgrad,jac=jac,checkbounds=checkbounds,
                     linenumbers=linenumbers,parallel=parallel,simplify=simplify,
                     sparse=sparse,eval_expression=eval_expression,kwargs...)
-    return f, u0, p
+    implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end
 
 function ODEFunctionExpr(sys::AbstractODESystem, args...; kwargs...)
     ODEFunctionExpr{true}(sys, args...; kwargs...)
 end
 
+"""
+```julia
+function DAEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing, tgrad=false,
+                                     jac = false,
+                                     sparse = false,
+                                     kwargs...) where {iip}
+```
+
+Create a Julia expression for an `ODEFunction` from the [`ODESystem`](@ref).
+The arguments `dvs` and `ps` are used to set the order of the dependent
+variable and parameter vectors, respectively.
+"""
+struct DAEFunctionExpr{iip} end
 
-function DiffEqBase.ODEProblem(sys::AbstractODESystem, args...; kwargs...)
-    ODEProblem{true}(sys, args...; kwargs...)
+struct DAEFunctionClosure{O, I} <: Function
+    f_oop::O
+    f_iip::I
+end
+(f::DAEFunctionClosure)(du, u, p, t) = f.f_oop(du, u, p, t)
+(f::DAEFunctionClosure)(out, du, u, p, t) = f.f_iip(out, du, u, p, t)
+
+function DAEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
+                                     ps = parameters(sys), u0 = nothing;
+                                     version = nothing, tgrad=false,
+                                     jac = false,
+                                     linenumbers = false,
+                                     sparse = false, simplify=false,
+                                     kwargs...) where {iip}
+    f_oop, f_iip = generate_function(sys, dvs, ps; expression=Val{true}, implicit_dae = true, kwargs...)
+    fsym = gensym(:f)
+    _f = :($fsym = $DAEFunctionClosure($f_oop, $f_iip))
+    ex = quote
+        $_f
+        ODEFunction{$iip}($fsym,)
+    end
+    !linenumbers ? striplines(ex) : ex
+end
+
+function DAEFunctionExpr(sys::AbstractODESystem, args...; kwargs...)
+    DAEFunctionExpr{true}(sys, args...; kwargs...)
+end
+
+for P in [:ODEProblem, :DAEProblem]
+    @eval function DiffEqBase.$P(sys::AbstractODESystem, args...; kwargs...)
+        $P{true}(sys, args...; kwargs...)
+    end
 end
 
 """
@@ -333,7 +457,34 @@ end
 
 """
 ```julia
-function DiffEqBase.ODEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
+function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
+                                    parammap=DiffEqBase.NullParameters();
+                                    version = nothing, tgrad=false,
+                                    jac = false,
+                                    checkbounds = false, sparse = false,
+                                    simplify=false,
+                                    linenumbers = true, parallel=SerialForm(),
+                                    kwargs...) where iip
+```
+
+Generates an DAEProblem from an ODESystem and allows for automatically
+symbolically calculating numerical enhancements.
+"""
+function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem,du0map,u0map,tspan,
+                                    parammap=DiffEqBase.NullParameters();kwargs...) where iip
+    f, du0, u0, p = process_DEProblem(
+        DAEFunction{iip}, sys, u0map, parammap;
+        implicit_dae=true, du0map=du0map, kwargs...
+    )
+    diffvars = collect_differential_variables(sys)
+    sts = states(sys)
+    differential_vars = map(Base.Fix2(in, diffvars), sts)
+    DAEProblem{iip}(f,du0,u0,tspan,p;differential_vars=differential_vars,kwargs...)
+end
+
+"""
+```julia
+function ODEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
                                     parammap=DiffEqBase.NullParameters();
                                     version = nothing, tgrad=false,
                                     jac = false,
@@ -371,6 +522,53 @@ function ODEProblemExpr(sys::AbstractODESystem, args...; kwargs...)
     ODEProblemExpr{true}(sys, args...; kwargs...)
 end
 
+"""
+```julia
+function DAEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
+                                    parammap=DiffEqBase.NullParameters();
+                                    version = nothing, tgrad=false,
+                                    jac = false,
+                                    checkbounds = false, sparse = false,
+                                    linenumbers = true, parallel=SerialForm(),
+                                    skipzeros=true, fillzeros=true,
+                                    simplify=false,
+                                    kwargs...) where iip
+```
+
+Generates a Julia expression for constructing an ODEProblem from an
+ODESystem and allows for automatically symbolically calculating
+numerical enhancements.
+"""
+struct DAEProblemExpr{iip} end
+
+function DAEProblemExpr{iip}(sys::AbstractODESystem,du0map,u0map,tspan,
+                             parammap=DiffEqBase.NullParameters();
+                             kwargs...) where iip
+    f, du0, u0, p = process_DEProblem(
+        DAEFunctionExpr{iip}, sys, u0map, parammap;
+        implicit_dae=true, du0map=du0map, kwargs...
+    )
+    linenumbers = get(kwargs, :linenumbers, true)
+    diffvars = collect_differential_variables(sys)
+    sts = states(sys)
+    differential_vars = map(Base.Fix2(in, diffvars), sts)
+
+    ex = quote
+        f = $f
+        u0 = $u0
+        du0 = $du0
+        tspan = $tspan
+        p = $p
+        differential_vars = $differential_vars
+        DAEProblem{$iip}(f,du0,u0,tspan,p;differential_vars=differential_vars,$(kwargs...))
+    end
+    !linenumbers ? striplines(ex) : ex
+end
+
+function DAEProblemExpr(sys::AbstractODESystem, args...; kwargs...)
+    DAEProblemExpr{true}(sys, args...; kwargs...)
+end
+
 
 ### Enables Steady State Problems ###
 function DiffEqBase.SteadyStateProblem(sys::AbstractODESystem, args...; kwargs...)

src/systems/diffeqs/odesystem.jl

@@ -280,3 +280,18 @@ function _eq_unordered(a, b)
     end
     return true
 end
+
+function collect_differential_variables(sys::ODESystem)
+    eqs = equations(sys)
+    vars = Set()
+    diffvars = Set()
+    for eq in eqs
+        vars!(vars, eq)
+        for v in vars
+            isdifferential(v) || continue
+            push!(diffvars, arguments(v)[1])
+        end
+        empty!(vars)
+    end
+    return diffvars
+end

src/variables.jl

@@ -10,7 +10,7 @@ Takes a list of pairs of `variables=>values` and an ordered list of variables
 and creates the array of values in the correct order with default values when
 applicable.
 """
-function varmap_to_vars(varmap, varlist; defaults=Dict(), check=true)
+function varmap_to_vars(varmap, varlist; defaults=Dict(), check=true, toterm=Symbolics.diff2term)
     # Edge cases where one of the arguments is effectively empty.
     is_incomplete_initialization = varmap isa DiffEqBase.NullParameters || varmap === nothing
     if is_incomplete_initialization || isempty(varmap)
@@ -31,7 +31,7 @@ function varmap_to_vars(varmap, varlist; defaults=Dict(), check=true)
     if eltype(varmap) <: Pair # `varmap` is a dict or an array of pairs
         varmap = todict(varmap)
         rules = Dict(varmap)
-        vals = _varmap_to_vars(varmap, varlist; defaults=defaults, check=check)
+        vals = _varmap_to_vars(varmap, varlist; defaults=defaults, check=check, toterm=toterm)
     else # plain array-like initialization
         vals = varmap
     end
@@ -45,9 +45,9 @@ function varmap_to_vars(varmap, varlist; defaults=Dict(), check=true)
     end
 end
 
-function _varmap_to_vars(varmap::Dict, varlist; defaults=Dict(), check=false)
+function _varmap_to_vars(varmap::Dict, varlist; defaults=Dict(), check=false, toterm=Symbolics.diff2term)
     varmap = merge(defaults, varmap) # prefers the `varmap`
-    varmap = Dict(Symbolics.diff2term(value(k))=>value(varmap[k]) for k in keys(varmap))
+    varmap = Dict(toterm(value(k))=>value(varmap[k]) for k in keys(varmap))
     # resolve symbolic parameter expressions
     for (p, v) in pairs(varmap)
         varmap[p] = fixpoint_sub(v, varmap)