Merge pull request #546 from SciML/reduction

automatic equation reduction

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -18,7 +18,9 @@ using Base: RefValue
 using RecursiveArrayTools
 
 import SymbolicUtils
-import SymbolicUtils: to_symbolic, FnType, @rule, Rewriters
+import SymbolicUtils: to_symbolic, FnType, @rule, Rewriters, Term
+
+using LinearAlgebra: LU, BlasInt
 
 import LightGraphs: SimpleDiGraph, add_edge!
 
@@ -91,6 +93,7 @@ include("function_registration.jl")
 include("simplify.jl")
 include("utils.jl")
 include("linearity.jl")
+include("solve.jl")
 include("direct.jl")
 include("domains.jl")
 
@@ -116,6 +119,8 @@ include("systems/pde/pdesystem.jl")
 include("systems/reaction/reactionsystem.jl")
 include("systems/dependency_graphs.jl")
 
+include("systems/reduction.jl")
+
 include("latexify_recipes.jl")
 include("build_function.jl")
 

src/simplify.jl

@@ -15,7 +15,7 @@ end
 SymbolicUtils.promote_symtype(f::SymbolicUtils.Sym{FnType{X,Parameter{Y}}},
                               xs...) where {X, Y} = Y
 
-SymbolicUtils.arguments(x::Operation) = x.args
+SymbolicUtils.arguments(x::Operation) = to_symbolic.(x.args)
 
 # SymbolicUtils wants raw numbers
 SymbolicUtils.to_symbolic(x::Constant) = x.value

src/solve.jl

@@ -0,0 +1,118 @@
+using SymbolicUtils: istree
+
+function nterms(t)
+    if istree(t)
+        return reduce(+, map(nterms, arguments(t)), init=0)
+    else
+        return 1
+    end
+end
+# Soft pivoted
+# Note: we call this function with a matrix of Union{SymbolicUtils.Symbolic, Any}
+# It should work as-is with Operation type too.
+function sym_lu(A)
+    m, n = size(A)
+    L = fill!(Array{Any}(undef, size(A)),0) # TODO: make sparse?
+    for i=1:min(m, n)
+        L[i,i] = 1
+    end
+    U = copy!(Array{Any}(undef, size(A)),A)
+    p = BlasInt[1:m;]
+    for k = 1:m-1
+        _, i = findmin(map(ii->iszero(U[ii, k]) ? Inf : nterms(U[ii,k]), k:n))
+        i += k - 1
+        # swap
+        U[k, k:end], U[i, k:end] = U[i, k:end], U[k, k:end]
+        L[k, 1:k-1], L[i, 1:k-1] = L[i, 1:k-1], L[k, 1:k-1]
+        p[k] = i
+
+        for j = k+1:m
+            L[j,k] = U[j, k] / U[k, k]
+            U[j,k:m] .= U[j,k:m] .- L[j,k] .* U[k,k:m]
+        end
+    end
+    for j=1:m
+        for i=j+1:n
+            U[i,j] = 0
+        end
+    end
+
+    (L, U, LinearAlgebra.ipiv2perm(p, m))
+end
+
+# Given a vector of equations and a
+# list of independent variables,
+# return the coefficient matrix `A` and a
+# vector of constants (possibly symbolic) `b` such that
+# A \ b will solve the equations for the vars
+function A_b(eqs, vars)
+    exprs = rhss(eqs) .- lhss(eqs)
+    for ex in exprs
+        @assert islinear(ex, vars)
+    end
+    A = jacobian(exprs, vars)
+    b = A * vars - exprs
+    A, b
+end
+
+"""
+    solve_for(eqs::Vector, vars::Vector)
+
+Solve the vector of equations `eqs` for a set of variables `vars`.
+
+Assumes `length(eqs) == length(vars)`
+
+Currently only works if all equations are linear.
+"""
+function solve_for(eqs, vars)
+    A, b = A_b(eqs, vars)
+    _solve(A, b)
+end
+
+function _solve(A, b)
+    A = SymbolicUtils.simplify.(to_symbolic.(A), polynorm=true)
+    b = SymbolicUtils.simplify.(to_symbolic.(b), polynorm=true)
+    map(to_mtk, SymbolicUtils.simplify.(ldiv(sym_lu(A), b)))
+end
+
+LinearAlgebra.:(\)(A::AbstractMatrix{<:Expression}, b::AbstractVector{<:Expression}) = _solve(A, b)
+LinearAlgebra.:(\)(A::AbstractMatrix{<:Expression}, b::AbstractVector) = _solve(A, b)
+LinearAlgebra.:(\)(A::AbstractMatrix, b::AbstractVector{<:Expression}) = _solve(A, b)
+
+# ldiv below
+
+_iszero(x::Number) = iszero(x)
+_isone(x::Number) = isone(x)
+_iszero(::Term) = false
+_isone(::Term) = false
+
+function simplifying_dot(x,y)
+    isempty(x) && return 0
+    muls = map(x,y) do xi,yi
+        _isone(xi) ? yi : _isone(yi) ? xi : _iszero(xi) ? 0 : _iszero(yi) ? 0 : xi * yi
+    end
+
+    reduce(muls) do acc, x
+        _iszero(acc) ? x : _iszero(x) ? acc : acc + x
+    end
+end
+
+function ldiv((L,U,p), b)
+    m, n = size(L)
+    x = Vector{Any}(undef, length(b))
+    b = b[p]
+
+    for i=n:-1:1
+        sub = simplifying_dot(x[i+1:end], U[i,i+1:end])
+        den = U[i,i]
+        x[i] = _iszero(sub) ? b[i] : b[i] - sub
+        x[i] = _isone(den) ? x[i] : _isone(-den) ? -x[i] : x[i] / den
+    end
+
+    # unit lower triangular solve first:
+    for i=1:n
+        sub = simplifying_dot(b[1:i-1], L[i, 1:i-1]) # this should be `b` not x
+        x[i] = _iszero(sub) ? x[i] : x[i] - sub
+    end
+    x
+end

src/systems/abstractsystem.jl

@@ -193,7 +193,7 @@ pins(sys::AbstractSystem) = isempty(sys.systems) ? sys.pins : [sys.pins;reduce(v
 function observed(sys::AbstractSystem)
     [sys.observed;
      reduce(vcat,
-            (namespace_equation.(s.observed, s.name, s.iv.name) for s in sys.systems),
+            (namespace_equation.(observed(s), s.name, s.iv.name) for s in sys.systems),
             init=Equation[])]
 end
 
@@ -215,7 +215,7 @@ end
 lhss(xs) = map(x->x.lhs, xs)
 rhss(xs) = map(x->x.rhs, xs)
 
-function equations(sys::ModelingToolkit.AbstractSystem; remove_aliases = true)
+function equations(sys::ModelingToolkit.AbstractSystem)
     if isempty(sys.systems)
         return sys.eqs
     else
@@ -224,14 +224,7 @@ function equations(sys::ModelingToolkit.AbstractSystem; remove_aliases = true)
                       namespace_equations.(sys.systems);
                       init=Equation[])]
 
-        if !remove_aliases
-            return eqs
-        end
-        aliases = observed(sys)
-        dict = Dict(lhss(aliases) .=> rhss(aliases))
-
-        # Substitute aliases
-        return Equation.(lhss(eqs), Rewriters.Fixpoint(x->substitute(x, dict)).(rhss(eqs)))
+        return eqs
     end
 end
 

src/systems/reduction.jl

@@ -0,0 +1,65 @@
+export alias_elimination
+
+function flatten(sys::ODESystem)
+    ODESystem(equations(sys),
+              independent_variable(sys),
+              observed=observed(sys))
+end
+
+
+using SymbolicUtils: Rewriters
+
+function fixpoint_sub(x, dict)
+    y = substitute(x, dict)
+    while !isequal(x, y)
+        y = x
+        x = substitute(y, dict)
+    end
+
+    return x
+end
+
+function substitute_aliases(diffeqs, dict)
+    lhss(diffeqs) .~ fixpoint_sub.(rhss(diffeqs), (dict,))
+end
+
+function make_lhs_0(eq)
+    if eq.lhs isa Constant && iszero(eq.lhs)
+        return eq
+    else
+        0 ~ eq.lhs - eq.rhs
+    end
+end
+
+function alias_elimination(sys::ODESystem)
+    eqs = vcat(equations(sys),
+               make_lhs_0.(observed(sys)))
+
+    new_stateops = map(eqs) do eq
+        if eq.lhs isa Operation && eq.lhs.op isa Differential
+            get_variables(eq.lhs)
+        else
+            []
+        end
+    end |> Iterators.flatten |> collect |> unique
+
+    all_vars = map(eqs) do eq
+        filter(x->!isparameter(x.op), get_variables(eq.rhs))
+    end |> Iterators.flatten |> collect |> unique
+
+    newstates = convert.(Variable, new_stateops)
+
+    alg_idxs = findall(x->x.lhs isa Constant && iszero(x.lhs), eqs)
+
+    eliminate = setdiff(convert.(Variable, all_vars), newstates)
+
+    vars = map(x->x(sys.iv()), eliminate)
+    outputs = solve_for(eqs[alg_idxs], vars)
+
+    diffeqs = eqs[setdiff(1:length(eqs), alg_idxs)]
+
+    diffeqs′ = substitute_aliases(diffeqs, Dict(vars .=> outputs))
+
+    ODESystem(diffeqs′, sys.iv(), new_stateops, parameters(sys), observed=vars .~ outputs)
+end
+

src/utils.jl

@@ -131,7 +131,18 @@ substitute(expr::Operation, s::Union{Vector, Dict}) = substituter(s)(expr)
 
 function substituter(pairs)
     dict = Dict(to_symbolic(k) => to_symbolic(v)  for (k, v) in pairs)
-    expr -> to_mtk(SymbolicUtils.simplify(SymbolicUtils.substitute(expr, dict)))
+    expr -> to_mtk(SymbolicUtils.substitute(expr, dict))
+end
+
+macro showarr(x)
+    n = string(x)
+    quote
+        y = $(esc(x))
+        println($n, " = ", summary(y))
+        Base.print_array(stdout, y)
+        println()
+        y
+    end
 end
 
 @deprecate substitute_expr!(expr,s) substitute(expr,s)

test/operation_overloads.jl

@@ -70,3 +70,8 @@ M \ b
 M \ reshape(b,2,1)
 M = [1 1; 0 2]
 M \ reshape(b,2,1)
+
+
+M = [1 a; 0 2]
+M \ b
+M \ [1, 2]

test/reduction.jl

@@ -0,0 +1,97 @@
+using ModelingToolkit, OrdinaryDiffEq, Test
+
+@parameters t σ ρ β
+@variables x(t) y(t) z(t) a(t) u(t) F(t)
+@derivatives D'~t
+
+test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynorm=true))
+
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ a*y - β*z,
+       0 ~ x - a]
+
+lorenz1 = ODESystem(eqs,t,[x,y,z,a],[σ,ρ,β],name=:lorenz1)
+
+lorenz1_aliased = alias_elimination(lorenz1)
+@test length(equations(lorenz1_aliased)) == 3
+@test length(states(lorenz1_aliased)) == 3
+
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
+
+@test lorenz1_aliased == ODESystem(eqs,t,[x,y,z],[σ,ρ,β],observed=[a ~ x],name=:lorenz1)
+
+# Multi-System Reduction
+
+eqs1 = [D(x) ~ σ*(y-x) + F,
+       D(y) ~ x*(ρ-z)-u,
+       D(z) ~ x*y - β*z]
+
+aliases = [u ~ x + y - z]
+
+lorenz1 = ODESystem(eqs1,pins=[F],observed=aliases,name=:lorenz1)
+
+eqs2 = [D(x) ~ F,
+       D(y) ~ x*(ρ-z)-x,
+       D(z) ~ x*y - β*z]
+
+aliases2 = [u ~ x - y - z]
+
+lorenz2 = ODESystem(eqs2,pins=[F],observed=aliases2,name=:lorenz2)
+
+connections = [lorenz1.F ~ lorenz2.u,
+               lorenz2.F ~ lorenz1.u]
+
+connected = ODESystem([0 ~ a + lorenz1.x - lorenz2.y],t,[a],[],observed=connections,systems=[lorenz1,lorenz2])
+
+# Reduced Unflattened System
+#=
+
+connections2 = [lorenz1.F ~ lorenz2.u,
+                lorenz2.F ~ lorenz1.u,
+                a ~ -lorenz1.x + lorenz2.y]
+connected = ODESystem(Equation[],t,[],[],observed=connections2,systems=[lorenz1,lorenz2])
+=#
+
+# Reduced Flattened System
+
+flattened_system = ModelingToolkit.flatten(connected)
+
+aliased_flattened_system = alias_elimination(flattened_system)
+
+@test states(aliased_flattened_system) == convert.(Variable, [
+        lorenz1.x
+        lorenz1.y
+        lorenz1.z
+        lorenz2.x
+        lorenz2.y
+        lorenz2.z
+       ])
+
+@test setdiff(parameters(aliased_flattened_system), convert.(Variable, [
+        lorenz1.σ
+        lorenz1.ρ
+        lorenz1.β
+        lorenz1.F
+        lorenz2.F
+        lorenz2.ρ
+        lorenz2.β
+       ])) |> isempty
+
+test_equal.(equations(aliased_flattened_system), [
+       D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - lorenz2.y - lorenz2.z,
+       D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-(lorenz1.x + lorenz1.y - lorenz1.z),
+       D(lorenz1.z) ~ lorenz1.x*lorenz1.y - lorenz1.β*lorenz1.z,
+       D(lorenz2.x) ~ lorenz1.x + lorenz1.y - lorenz1.z,
+       D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ-lorenz2.z)-lorenz2.x,
+       D(lorenz2.z) ~ lorenz2.x*lorenz2.y - lorenz2.β*lorenz2.z])
+
+test_equal.(observed(aliased_flattened_system), [
+       lorenz1.F ~ lorenz2.x + -1 * (lorenz2.y + lorenz2.z),
+       lorenz1.u ~ lorenz1.x + lorenz1.y + -1 * lorenz1.z,
+       lorenz2.F ~ lorenz1.x + lorenz1.y + -1 * lorenz1.z,
+       a ~ lorenz2.y + -1 * lorenz1.x,
+       lorenz2.u ~ lorenz2.x + -1 * (lorenz2.y + lorenz2.z),
+])

test/runtests.jl

@@ -20,6 +20,7 @@ using SafeTestsets, Test
 @safetestset "Build Targets Test" begin include("build_targets.jl") end
 @safetestset "Domain Test" begin include("domains.jl") end
 @safetestset "Constraints Test" begin include("constraints.jl") end
+@safetestset "Reduction Test" begin include("reduction.jl") end
 @safetestset "PDE Construction Test" begin include("pde.jl") end
 @safetestset "Lowering Integration Test" begin include("lowering_solving.jl") end
 @safetestset "Test Big System Usage" begin include("bigsystem.jl") end