LU-based equation solve

Author: shashi

src/ModelingToolkit.jl

@@ -17,7 +17,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!
 
@@ -90,6 +92,7 @@ include("function_registration.jl")
 include("simplify.jl")
 include("utils.jl")
 include("linearity.jl")
+include("solve.jl")
 include("direct.jl")
 include("domains.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,99 @@
+# 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 = Array{Any}(undef, size(A)) # TODO: make sparse?
+    for i=1:min(m, n)
+        L[i,i] = 1
+    end
+    U = copy(A)
+    p = BlasInt[1:m;]
+    for k = 1:m-1
+        i = k
+        for ii=k:n
+            if !iszero(U[ii, k])
+                i = ii
+                break
+            end
+        end
+        # 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] .- simplifying_dot(L[j,k],  U[k,k:m])
+        end
+    end
+    factors = copy(U)
+    for j=1:m
+        for i=j+1:n
+            factors[i,j] = L[i,j]
+        end
+    end
+
+    LU(factors, p, BlasInt(0))
+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
+
+function solve_for(eqs, vars)
+    A, b = A_b(eqs, vars)
+    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
+
+# 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(A::LU, b)
+    # unit lower triangular solve first:
+    L = A.L
+    U = A.U
+    m, n = size(L)
+    x = Vector{Any}(undef, length(b))
+    b = b[A.p]
+    for i=1:n
+        sub = simplifying_dot(b[1:i-1], L[i, 1:i-1])
+        x[i] = _iszero(sub) ? b[i] : b[i] - sub
+    end
+
+    for i=n:-1:1
+        sub = simplifying_dot(b[i+1:end], U[i,i+1:end])
+        den = U[i,i]
+        x[i] = _iszero(sub) ? x[i] : x[i] - sub
+        x[i] = _isone(U[i,i]) ? x[i] : x[i] / U[i,i]
+    end
+    x
+end

src/utils.jl

@@ -135,38 +135,3 @@ function substituter(pairs)
 end
 
 @deprecate substitute_expr!(expr,s) substitute(expr,s)
-
-# Really bad solve for vars
-function solve_for(eqs, vars)
-    @assert length(eqs) >= length(vars)
-    @assert all(iszero(eq.lhs) for eq in eqs)
-    neweqs = []
-    for (i, eq) in enumerate(eqs)
-        rhs = eq.rhs
-        if rhs.op == (-)
-            if any(isequal(rhs.args[1]), vars) && any(isequal(rhs.args[2]), vars)
-                push!(neweqs, rhs.args[1] ~ rhs.args[2]) # pick one?
-                @warn("todo")
-            elseif any(isequal(rhs.args[1]), vars)
-                push!(neweqs, rhs.args[1] ~ rhs.args[2])
-            elseif any(isequal(rhs.args[2]), vars)
-                push!(neweqs, rhs.args[2] ~ rhs.args[1])
-            else
-                @warn("may require unimplemented solve")
-                #error("todo 2")
-                push!(neweqs, eq)
-            end
-        elseif rhs.op == (+)
-            eqs[i] = 0 ~ rhs.args[1] - (-rhs.args[2])
-        else
-            error("todo")
-        end
-    end
-    if length(neweqs) >= length(vars)
-        return neweqs
-    else
-        # substitute
-        eqs′ = Equation.(0, substitute.(rhss(eqs), (Dict(lhss(neweqs) .=> rhss(neweqs)),)))
-        solve_for(eqs′, vars)
-    end
-end