Merge pull request #142 from JuliaDiffEq/myb/lu

Add static version of Wfact and negate the W operator

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -126,22 +126,32 @@ function generate_factorized_W(sys::ODESystem, simplify=true; version::FunctionV
 
     gam = Variable(:gam; known = true)()
 
-    W = LinearAlgebra.I - gam*jac
+    W = - LinearAlgebra.I + gam*jac
     Wfact = lu(W, Val(false), check=false).factors
 
     if simplify
         Wfact = simplify_constants.(Wfact)
     end
 
-    W_t = LinearAlgebra.I/gam - jac
+    W_t = - LinearAlgebra.I/gam + jac
     Wfact_t = lu(W_t, Val(false), check=false).factors
     if simplify
         Wfact_t = simplify_constants.(Wfact_t)
     end
 
+    if version === SArrayFunction
+        siz = size(Wfact)
+        constructor = :(x -> begin
+                            A = SMatrix{$siz...}(x)
+                            StaticArrays.LU(LowerTriangular( SMatrix{$siz...}(UnitLowerTriangular(A)) ), UpperTriangular(A), SVector(ntuple(n->n, max($siz...))))
+                        end)
+    else
+        constructor = nothing
+    end
+
     vs, ps = sys.dvs, sys.ps
-    Wfact_func   = build_function(Wfact  , vs, ps, (:gam,:t), ODEToExpr(sys); version = version)
-    Wfact_t_func = build_function(Wfact_t, vs, ps, (:gam,:t), ODEToExpr(sys); version = version)
+    Wfact_func   = build_function(Wfact  , vs, ps, (:gam,:t), ODEToExpr(sys); version = version, constructor=constructor)
+    Wfact_t_func = build_function(Wfact_t, vs, ps, (:gam,:t), ODEToExpr(sys); version = version, constructor=constructor)
 
     return (Wfact_func, Wfact_t_func)
 end

src/utils.jl

@@ -31,7 +31,7 @@ function flatten_expr!(x)
     x
 end
 
-function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); version::FunctionVersion)
+function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); version::FunctionVersion, constructor=nothing)
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
     param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
     (ls, rs) = zip(var_pairs..., param_pairs...)
@@ -48,7 +48,7 @@ function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs
         let_expr = Expr(:let, var_eqs, sys_expr)
         :((u,p,$(args...)) -> begin
             X = $let_expr
-            T = StaticArrays.similar_type(typeof(u), eltype(X))
+            T = $(constructor === nothing ? :(StaticArrays.similar_type(typeof(u), eltype(X))) : constructor)
             T(X)
         end)
     end

test/system_construction.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit
+using ModelingToolkit, StaticArrays, LinearAlgebra
 using Test
 
 # Define some variables
@@ -35,16 +35,40 @@ generate_function(de, [x,y,z], [σ,ρ,β])
 generate_function(de, [x,y,z], [σ,ρ,β]; version=ModelingToolkit.SArrayFunction)
 jac_expr = generate_jacobian(de)
 jac = calculate_jacobian(de)
+jacfun = eval(jac_expr)
+# iip
 f = ODEFunction(de, [x,y,z], [σ,ρ,β])
 fw, fwt = map(eval, ModelingToolkit.generate_factorized_W(de))
 du = zeros(3)
 u  = collect(1:3)
 p  = collect(4:6)
 f(du, u, p, 0.1)
 @test du == [4, 0, -16]
+J = zeros(3, 3)
+jacfun(J, u, p, t)
 FW = zeros(3, 3)
+FWt = zeros(3, 3)
 fw(FW, u, p, 0.2, 0.1)
-fwt(FW, u, p, 0.2, 0.1)
+fwt(FWt, u, p, 0.2, 0.1)
+# oop
+f = ODEFunction(de, [x,y,z], [σ,ρ,β]; version=ModelingToolkit.SArrayFunction)
+fw, fwt = map(eval, ModelingToolkit.generate_factorized_W(de; version=ModelingToolkit.SArrayFunction))
+du = @SArray zeros(3)
+u  = SVector(1:3...)
+p  = SVector(4:6...)
+@test f(u, p, 0.1) === @SArray [4, 0, -16]
+Sfw = fw(u, p, 0.2, 0.1)
+@test Sfw.L ≈ UnitLowerTriangular(FW)
+@test Sfw.U ≈ UpperTriangular(FW)
+sol = Sfw \ @SArray ones(3)
+@test sol isa SArray
+@test sol ≈ -(I - 0.2*J)\ones(3)
+Sfw_t = fwt(u, p, 0.2, 0.1)
+@test Sfw_t.L ≈ UnitLowerTriangular(FWt)
+@test Sfw_t.U ≈ UpperTriangular(FWt)
+sol = Sfw_t \ @SArray ones(3)
+@test sol isa SArray
+@test sol ≈ -(I/0.2 - J)\ones(3)
 
 @testset "time-varying parameters" begin
     @parameters σ′(t-1)