mass matrix ODE support

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -107,7 +107,8 @@ export Equation, ConstrainedEquation
 export simplify_constants
 
 export Operation, Expression
-export calculate_jacobian, generate_jacobian, generate_function, generate_diffusion_function
+export calculate_jacobian, generate_jacobian, generate_function
+export calculate_massmatrix, generate_diffusion_function
 export independent_variables, dependent_variables, parameters
 export simplified_expr, eval_function
 export @register, @I

src/systems/diffeqs/abstractodesystem.jl

@@ -89,6 +89,23 @@ function generate_factorized_W(sys::AbstractODESystem, vs = sys.dvs, ps = sys.ps
     return (Wfact_func, Wfact_t_func)
 end
 
+function calculate_massmatrix(sys::AbstractODESystem, simplify=true)
+    eqs = sys.eqs
+    M = zeros(length(eqs),length(eqs))
+    for (i,eq) in enumerate(eqs)
+        if eq.lhs isa Constant
+            @assert eq.lhs.value == 0
+        elseif eq.lhs.op isa Differential
+            j = findfirst(x->isequal(x.name,var_from_nested_derivative(eq.lhs)[1].name),sys.dvs)
+            M[i,j] = 1
+        else
+            error("Only semi-explicit mass matrices are currently supported")
+        end
+    end
+    M = simplify ? simplify_constants.(M) : M
+    M == I ? I : M
+end
+
 """
 $(SIGNATURES)
 
@@ -132,10 +149,13 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = sys.dvs, ps =
         _Wfact,_Wfact_t = nothing,nothing
     end
 
+    M = calculate_massmatrix(sys)
+
     ODEFunction{iip}(f,jac=_jac,
                       tgrad = _tgrad,
                       Wfact = _Wfact,
                       Wfact_t = _Wfact_t,
+                      mass_matrix = M,
                       syms = Symbol.(sys.dvs))
 end
 

src/systems/diffeqs/sdesystem.jl

@@ -98,10 +98,13 @@ function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.dvs, ps = sys.ps;
         _Wfact,_Wfact_t = nothing,nothing
     end
 
+    M = calculate_massmatrix(sys)
+
     SDEFunction{iip}(f,g,jac=_jac,
                       tgrad = _tgrad,
                       Wfact = _Wfact,
                       Wfact_t = _Wfact_t,
+                      mass_matrix = M,
                       syms = Symbol.(sys.dvs))
 end
 

test/mass_matrix.jl

@@ -0,0 +1,33 @@
+using OrdinaryDiffEq, ModelingToolkit, Test
+@parameters t
+@variables y[1:3](t)
+@parameters k[1:3]
+@derivatives D'~t
+
+eqs = [D(y[1]) ~ -k[1]*y[1] + k[3]*y[2]*y[3],
+       D(y[2]) ~  k[1]*y[1] - k[3]*y[2]*y[3] - k[2]*y[2]^2,
+       0 ~  y[1] + y[2] + y[3] - 1]
+
+sys = ODESystem(eqs,t,y,k)
+M = calculate_massmatrix(sys)
+M == [1 0 0
+      0 1 0
+      0 0 0]
+
+f = ODEFunction(sys)
+prob_mm = ODEProblem(f,[1.0,0.0,0.0],(0.0,1e5),(0.04,3e7,1e4))
+sol = solve(prob_mm,Rodas5(),reltol=1e-8,abstol=1e-8)
+
+function rober(du,u,p,t)
+  y₁,y₂,y₃ = u
+  k₁,k₂,k₃ = p
+  du[1] = -k₁*y₁ + k₃*y₂*y₃
+  du[2] =  k₁*y₁ - k₃*y₂*y₃ - k₂*y₂^2
+  du[3] =  y₁ + y₂ + y₃ - 1
+  nothing
+end
+f = ODEFunction(rober,mass_matrix=M)
+prob_mm2 = ODEProblem(f,[1.0,0.0,0.0],(0.0,1e5),(0.04,3e7,1e4))
+sol2 = solve(prob_mm2,Rodas5(),reltol=1e-8,abstol=1e-8)
+
+@test Array(sol) == Array(sol2)

test/odesystem.jl

@@ -7,6 +7,16 @@ using Test
 @variables x(t) y(t) z(t)
 @derivatives D'~t
 
+# Define a differential equation
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
+
+ModelingToolkit.simplified_expr.(eqs)[1]
+:(derivative(x(t), t) = σ * (y(t) - x(t))).args
+de = ODESystem(eqs)
+
+
 function _clean(O::Operation)
     @assert isa(O.op, Variable)
     return O.op
@@ -26,14 +36,6 @@ function test_nlsys_inference(name, sys, vs, ps)
     end
 end
 
-# Define a differential equation
-eqs = [D(x) ~ σ*(y-x),
-       D(y) ~ x*(ρ-z)-y,
-       D(z) ~ x*y - β*z]
-
-ModelingToolkit.simplified_expr.(eqs)[1]
-:(derivative(x(t), t) = σ * (y(t) - x(t))).args
-de = ODESystem(eqs)
 test_diffeq_inference("standard", de, t, (x, y, z), (σ, ρ, β))
 generate_function(de, [x,y,z], [σ,ρ,β])
 jac_expr = generate_jacobian(de)

test/runtests.jl

@@ -6,6 +6,7 @@ using ModelingToolkit, Test
 @testset "Operation Overloads Test" begin include("operation_overloads.jl") end
 @testset "Direct Usage Test" begin include("direct.jl") end
 @testset "ODESystem Test" begin include("odesystem.jl") end
+@testset "Mass Matrix Test" begin include("odesystem.jl") end
 @testset "SDESystem Test" begin include("sdesystem.jl") end
 @testset "Build Targets Test" begin include("build_targets.jl") end
 @testset "Domain Test" begin include("domains.jl") end