Merge pull request #290 from SciML/problem

add a direct problem building interface for AbstractSystems

Author: ChrisRackauckas

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
 version = "2.0.0"
 
 [deps]
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -21,7 +22,7 @@ TreeViews = "a2a6695c-b41b-5b7d-aed9-dbfdeacea5d7"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]
-DiffEqBase = "6"
+DiffEqBase = "6.28"
 DiffRules = "0.1, 1.0"
 DocStringExtensions = "0.7, 0.8"
 GeneralizedGenerated = "0.1.4, 0.2"

src/ModelingToolkit.jl

@@ -2,7 +2,7 @@ module ModelingToolkit
 
 using DiffEqBase, Distributed
 using StaticArrays, LinearAlgebra, SparseArrays
-using Latexify, Unitful
+using Latexify, Unitful, ArrayInterface
 using MacroTools
 
 using Base.Threads
@@ -122,6 +122,7 @@ export simplify_constants
 export Operation, Expression, Variable
 export calculate_jacobian, generate_jacobian, generate_function
 export calculate_tgrad, generate_tgrad
+export calculate_gradient, generate_gradient
 export calculate_hessian, generate_hessian
 export calculate_massmatrix, generate_diffusion_function
 export independent_variable, states, parameters, equations

src/systems/diffeqs/abstractodesystem.jl

@@ -100,7 +100,7 @@ function calculate_massmatrix(sys::AbstractODESystem, simplify=true)
             j = findfirst(x->isequal(x.name,var_from_nested_derivative(eq.lhs)[1].name),dvs)
             M[i,j] = 1
         else
-            error("Only semi-explicit mass matrices are currently supported")
+            error("Only semi-explicit constant mass matrices are currently supported")
         end
     end
     M = simplify ? simplify_constants.(M) : M
@@ -121,30 +121,30 @@ respectively.
 function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      ps = parameters(sys);
                                      version = nothing, tgrad=false,
-                                     jac = false, Wfact = false) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
+                                     jac = false, Wfact = false, kwargs...) where {iip}
+    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false}; kwargs...)
 
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
 
     if tgrad
-        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false}; kwargs...)
         _tgrad(u,p,t) = tgrad_oop(u,p,t)
         _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
     else
         _tgrad = nothing
     end
 
     if jac
-        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false}; kwargs...)
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
         _jac = nothing
     end
 
     if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false})
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false}; kwargs...)
         Wfact_oop, Wfact_iip = tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
         _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
@@ -162,5 +162,23 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                       Wfact = _Wfact,
                       Wfact_t = _Wfact_t,
                       mass_matrix = M,
-                      syms = Symbol.(sys.states))
+                      syms = Symbol.(states(sys)))
+end
+
+function DiffEqBase.ODEProblem(sys::AbstractODESystem, args...; kwargs...)
+    ODEProblem{true}(sys, args...; kwargs...)
+end
+
+function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
+                                    parammap=DiffEqBase.NullParameters();
+                                    version = nothing, tgrad=false,
+                                    jac = false, Wfact = false,
+                                    checkbounds = false,
+                                    linenumbers = true, multithread=false,
+                                    kwargs...) where iip
+    f = ODEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
+                        linenumbers=linenumbers,multithread=multithread)
+    u0 = varmap_to_vars(u0map,states(sys))
+    p = varmap_to_vars(parammap,parameters(sys))
+    ODEProblem(f,u0,tspan,p;kwargs...)
 end

src/systems/diffeqs/sdesystem.jl

@@ -121,3 +121,17 @@ end
 function rename(sys::SDESystem,name)
     ODESystem(sys.eqs, sys.noiseeqs, sys.iv, sys.states, sys.ps, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
 end
+
+function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,p=parammap;
+                                    version = nothing, tgrad=false,
+                                    jac = false, Wfact = false,
+                                    checkbounds = false,
+                                    linenumbers = true, multithread=false,
+                                    kwargs...) where iip
+
+    f = SDEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
+                        linenumbers=linenumbers,multithread=multithread)
+    u0 = varmap_to_vars(u0map,states(sys))
+    p = varmap_to_vars(parammap,parameters(sys))
+    SDEProblem(f,f.g,u0,tspan,p;kwargs...)
+end

src/systems/nonlinear/nonlinearsystem.jl

@@ -71,3 +71,19 @@ function generate_function(sys::NonlinearSystem, vs = states(sys), ps = paramete
     ps′ = convert.(Variable,ps)
     return build_function(rhss, vs′, ps′, (), NLSysToExpr(sys), expression; kwargs...)
 end
+
+function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
+                                          parammap=DiffEqBase.NullParameters();
+                                          jac = false,
+                                          checkbounds = false,
+                                          linenumbers = true, multithread=false,
+                                          kwargs...) where iip
+    dvs = states(sys)
+    ps = parameters(sys)
+
+    f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
+                              multithread=multithread)
+    u0 = varmap_to_vars(u0map,dvs)
+    p = varmap_to_vars(parammap,ps)
+    NonlinearProblem(f,u0,tspan,p;kwargs...)
+end

src/systems/optimization/optimizationsystem.jl

@@ -39,13 +39,22 @@ function OptimizationSystem(op, states, ps;
     OptimizationSystem(op, convert.(Variable,states), convert.(Variable,ps), name, systems)
 end
 
+function calculate_gradient(sys::OptimizationSystem)
+    expand_derivatives.(gradient(equations(sys), [dv() for dv in states(sys)]))
+end
+
+function generate_gradient(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+    grad = calculate_gradient(sys)
+    return build_function(grad, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x); kwargs...)
+end
+
 function calculate_hessian(sys::OptimizationSystem)
     expand_derivatives.(hessian(equations(sys), [dv() for dv in states(sys)]))
 end
 
 function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     hes = calculate_hessian(sys)
-    return build_function(hes, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x))
+    return build_function(hes, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x); kwargs...)
 end
 
 function generate_function(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
@@ -56,3 +65,22 @@ end
 
 equations(sys::OptimizationSystem) = isempty(sys.systems) ? sys.op : sys.op + reduce(+,namespace_operation.(sys.systems))
 namespace_operation(sys::OptimizationSystem) = namespace_operation(sys.op,sys.name,nothing)
+
+function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
+                                          parammap=DiffEqBase.NullParameters();
+                                          u0=nothing, lb=nothing, ub=nothing,
+                                          hes = false,
+                                          checkbounds = false,
+                                          linenumbers = true, multithread=false,
+                                          kwargs...) where iip
+    dvs = states(sys)
+    ps = parameters(sys)
+
+    f = generate_function(sys,checkbounds=checkbounds,linenumbers=linenumbers,
+                              multithread=multithread)
+    u0 = varmap_to_vars(u0,dvs)
+    p = varmap_to_vars(parammap,ps)
+    lb = varmap_to_vars(lb,dvs)
+    ub = varmap_to_vars(ub,dvs)
+    OptimizationProblem(f,p;u0=u0,lb=lb,ub=ub,kwargs...)
+end

src/variables.jl

@@ -171,3 +171,25 @@ function TreeViews.treelabel(io::IO,x::Variable,
                              mime::MIME"text/plain" = MIME"text/plain"())
   show(io,mime,Text(x.name))
 end
+
+"""
+varmap_to_vars(varmap,varlist)
+
+Takes a list of pairs of variables=>values and an ordered list of variables and
+creates the array of values in the correct order
+"""
+function varmap_to_vars(varmap,varlist)
+    out = similar(varmap,typeof(last(first(varmap))))
+    for i in 1:length(varmap)
+        ivar = convert(Variable,varmap[i][1])
+        j = findfirst(x->ivar.name == convert(Variable,x).name,varlist)
+        out[j] = varmap[i][2]
+    end
+
+    # Make output match varmap in type and shape
+    # Does things like MArray->SArray
+    ArrayInterface.restructure(varmap,out)
+end
+
+varmap_to_vars(varmap::DiffEqBase.NullParameters,varlist) = varmap
+varmap_to_vars(varmap::Nothing,varlist) = varmap

test/components.jl

@@ -20,6 +20,30 @@ connnectedeqs = [D(a) ~ a*lorenz1.x]
 
 connected1 = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=:connected1)
 
+u0map = [lorenz1.x => 1.0,
+         lorenz2.x => 2.0,
+         lorenz1.y => 3.0,
+         lorenz2.y => 4.0,
+         lorenz1.z => 5.0,
+         lorenz2.z => 6.0,
+         a => 7.0]
+
+parammap = SA[lorenz1.σ => 1.0,
+              lorenz1.ρ => 2.0,
+              lorenz1.β => 3.0,
+              lorenz2.σ => 4.0,
+              lorenz2.ρ => 5.0,
+              lorenz2.β => 6.0,
+              α => 7.0]
+
+prob = ODEProblem(connected1,u0map,(0.0,100.0),parammap,jac=true)
+
+@test prob.u0 == [7.0,1.0,3.0,5.0,2.0,4.0,6.0]
+@test prob.p == [7.0,1.0,2.0,3.0,4.0,5.0,6.0]
+@test prob.u0 isa Vector
+@test prob.p isa SVector
+@test prob.f.syms == [:a,:lorenz1′x,:lorenz1′y,:lorenz1′z,:lorenz2′x,:lorenz2′y,:lorenz2′z]
+
 eqs_flat = [D(a) ~ a*lorenz1.x,
             D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x),
             D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-lorenz1.y,

test/optimizationsystem.jl

@@ -8,12 +8,14 @@ sys2 = OptimizationSystem(loss,[x,y],[a,b],name=:sys2)
 @variables z
 @parameters β
 loss2 = sys1.x - sys2.y + z*β
-combinedsys = OptimizationSystem(loss,[z],[β],systems=[sys1,sys2],name=:combinedsys)
+combinedsys = OptimizationSystem(loss2,[z],[β],systems=[sys1,sys2],name=:combinedsys)
 
 equations(combinedsys)
 states(combinedsys)
 parameters(combinedsys)
 
+calculate_gradient(combinedsys)
 calculate_hessian(combinedsys)
 generate_function(combinedsys)
+generate_gradient(combinedsys)
 generate_hessian(combinedsys)