Merge pull request #43 from vpuri3/cache

ensure component operators get cached

Author: ChrisRackauckas

src/SciMLOperators.jl

@@ -11,8 +11,8 @@ import Lazy: @forward
 import Setfield: @set!
 
 # overload
-import Base: +, -, *, /, \, ∘, ==
-import Base: conj, one, iszero, inv, adjoint, transpose, size, convert, Matrix
+import Base: +, -, *, /, \, ∘, ==, one, zero
+import Base: conj, iszero, inv, adjoint, transpose, size, convert, Matrix
 import LinearAlgebra: mul!, ldiv!, lmul!, rmul!, factorize, exp, Diagonal
 import SparseArrays: sparse
 

src/basic.jl

@@ -12,7 +12,6 @@ function Base.one(A::AbstractSciMLOperator)
     IdentityOperator{N}()
 end
 
-# TODO - sparse diagonal
 Base.convert(::Type{AbstractMatrix}, ::IdentityOperator{N}) where{N} = Diagonal(ones(Bool, N))
 
 # traits
@@ -187,8 +186,10 @@ function Base.adjoint(α::ScalarOperator) # TODO - test
     ScalarOperator(val; update_func=update_func)
 end
 Base.transpose(α::ScalarOperator) = α
+Base.one(::Type{AbstractSciMLOperator}) = ScalarOperator(true)
+Base.zero(::Type{AbstractSciMLOperator}) = ScalarOperator(false)
 
-getops(α::ScalarOperator) = (α.val)
+getops(α::ScalarOperator) = (α.val,)
 islinear(L::ScalarOperator) = true
 issquare(L::ScalarOperator) = true
 isconstant(α::ScalarOperator) = α.update_func == DEFAULT_UPDATE_FUNC
@@ -249,9 +250,11 @@ struct ScaledOperator{T,
     L::LType
     cache::T
 
-    function ScaledOperator(λ::ScalarOperator, L::AbstractSciMLOperator)
-        T = promote_type(eltype.((λ, L))...)
-        cache = zero(T)
+    function ScaledOperator(λ::ScalarOperator{Tλ},
+                            L::AbstractSciMLOperator{TL},
+                            cache = zero(promote_type(Tλ,TL))
+                           ) where{Tλ,TL}
+        T = promote_type(Tλ, TL)
         new{T,typeof(λ),typeof(L)}(λ, L, cache)
     end
 end
@@ -297,7 +300,7 @@ for op in (
 end
 LinearAlgebra.opnorm(L::ScaledOperator, p::Real=2) = abs(L.λ) * opnorm(L.L, p)
 
-getops(L::ScaledOperator) = (L.λ, L.L)
+getops(L::ScaledOperator) = (L.λ, L.L,)
 islinear(L::ScaledOperator) = all(islinear, L.ops)
 isconstant(L::ScaledOperator) = isconstant(L.L) & isconstant(L.λ)
 Base.iszero(L::ScaledOperator) = iszero(L.L) | iszero(L.λ)
@@ -306,6 +309,12 @@ has_mul!(L::ScaledOperator) = has_mul!(L.L)
 has_ldiv(L::ScaledOperator) = has_ldiv(L.L) & !iszero(L.λ)
 has_ldiv!(L::ScaledOperator) = has_ldiv!(L.L) & !iszero(L.λ)
 
+function cache_internals(L::ScaledOperator, u::AbstractVector)
+    @set! L.L = cache_operator(L.L, u)
+    @set! L.λ = cache_operator(L.λ, u)
+    L
+end
+
 # getindex
 Base.getindex(L::ScaledOperator, i::Int) = L.coeff * L.op[i]
 Base.getindex(L::ScaledOperator, I::Vararg{Int, N}) where {N} = L.λ * L.L[I...]
@@ -348,7 +357,7 @@ function LinearAlgebra.ldiv!(L::ScaledOperator, u::AbstractVector)
 end
 
 """
-Lazy operator addition (A + B)
+Lazy operator addition
 
     (A1 + A2 + A3...)u = A1*u + A2*u + A3*u ....
 """
@@ -357,17 +366,22 @@ struct AddedOperator{T,
                     } <: AbstractSciMLOperator{T}
     ops::O
 
-    function AddedOperator(ops...)
-        sz = size(first(ops))
-        for op in ops[2:end]
-            @assert size(op) == sz "Size mismatich in operators $ops"
-        end
-
+    function AddedOperator(ops)
         T = promote_type(eltype.(ops)...)
         new{T,typeof(ops)}(ops)
     end
 end
 
+function AddedOperator(ops::AbstractSciMLOperator...)
+    sz = size(first(ops))
+    for op in ops[2:end]
+        @assert size(op) == sz "Size mismatich in operators $ops"
+    end
+    AddedOperator(ops)
+end
+
+AddedOperator(L::AbstractSciMLOperator) = L
+
 # constructors
 Base.:+(ops::AbstractSciMLOperator...) = AddedOperator(ops...)
 Base.:-(A::AbstractSciMLOperator, B::AbstractSciMLOperator) = AddedOperator(A, -B)
@@ -423,6 +437,13 @@ getops(L::AddedOperator) = L.ops
 Base.iszero(L::AddedOperator) = all(iszero, getops(L))
 has_adjoint(L::AddedOperator) = all(has_adjoint, L.ops)
 
+function cache_internals(L::AddedOperator, u::AbstractVector)
+    for i=1:length(L.ops)
+        @set! L.ops[i] = cache_operator(L.ops[i], u)
+    end
+    L
+end
+
 getindex(L::AddedOperator, i::Int) = sum(op -> op[i], L.ops)
 getindex(L::AddedOperator, I::Vararg{Int, N}) where {N} = sum(op -> op[I...], L.ops)
 
@@ -532,30 +553,35 @@ end
 Base.:*(L::ComposedOperator, u::AbstractVector) = foldl((acc, op) -> op * acc, reverse(L.ops); init=u)
 Base.:\(L::ComposedOperator, u::AbstractVector) = foldl((acc, op) -> op \ acc, L.ops; init=u)
 
-function cache_operator(L::ComposedOperator, u::AbstractVector)
-    # for 3 arg mul!
-    # Tuple of N-1 cache vectors. cache[N-1] = op[N] * u and so on
-    vec = u
-    c3 = ()
+function cache_self(L::ComposedOperator, u::AbstractVector)
+    vec = similar(u)
+    cache = (vec,)
     for i in reverse(2:length(L.ops))
-        vec = L.ops[i] * vec
-        c3 = (c3..., vec)
+        vec   = L.ops[i] * vec
+        cache = (vec, cache...)
     end
 
-    # for 5 arg mul!
-    c5 = similar(u)
+    @set! L.cache = cache
+    L
+end
 
-    cache = (;c3=c3, c5=c5)
+function cache_internals(L::ComposedOperator, u::AbstractVector)
+    if !(L.isset)
+        L = cache_self(L, u)
+    end
+
+    vecs = L.cache
+    for i in reverse(1:length(L.ops))
+        @set! L.ops[i] = cache_operator(L.ops[i], vecs[i])
+    end
 
-    @set! L.cache = cache
     L
 end
 
 function LinearAlgebra.mul!(v::AbstractVector, L::ComposedOperator, u::AbstractVector)
     @assert L.isset "cache needs to be set up to use LinearAlgebra.mul!"
 
-    cache = L.cache.c3
-    vecs = (v, cache..., u)
+    vecs = (v, L.cache[1:end-1]..., u)
     for i in reverse(1:length(L.ops))
         mul!(vecs[i], L.ops[i], vecs[i+1])
     end
@@ -565,7 +591,7 @@ end
 function LinearAlgebra.mul!(v::AbstractVector, L::ComposedOperator, u::AbstractVector, α, β)
     @assert L.isset "cache needs to be set up to use LinearAlgebra.mul!"
 
-    cache = L.cache.c5
+    cache = L.cache[end]
     copy!(cache, v)
 
     mul!(v, L, u)
@@ -576,8 +602,7 @@ end
 function LinearAlgebra.ldiv!(v::AbstractVector, L::ComposedOperator, u::AbstractVector)
     @assert L.isset "cache needs to be set up to use 3 arg LinearAlgebra.ldiv!"
 
-    cache = L.cache.c3
-    vecs = (u, reverse(cache)..., v)
+    vecs = (u, reverse(L.cache[1:end-1])..., v)
     for i in 1:length(L.ops)
         ldiv!(vecs[i+1], L.ops[i], vecs[i])
     end
@@ -645,6 +670,11 @@ for (op, LType, VType) in (
                              has_ldiv!,
                             )
 
+    @eval function cache_internals(L::$LType, u::AbstractVector)
+        @set! L.L = cache_operator(L.L, _reshape(u, size(L,1)))
+        L
+    end
+
     # oeprator application
     @eval Base.:*(u::$VType, L::$LType) = $op(L.L * u.parent)
     @eval Base.:/(u::$VType, L::$LType) = $op(L.L \ u.parent)
@@ -724,12 +754,17 @@ has_ldiv!(L::InvertedOperator) = has_mul!(L.L)
 Base.:*(L::InvertedOperator, u::AbstractVector) = L.L \ u
 Base.:\(L::InvertedOperator, u::AbstractVector) = L.L * u
 
-function cache_operator(L::InvertedOperator, u::AbstractVector)
+function cache_self(L::InvertedOperator, u::AbstractVector)
     cache = similar(u)
     @set! L.cache = cache
     L
 end
 
+function cache_internals(L::InvertedOperator, u::AbstractVector)
+    @set! L.L = cache_operator(L.L, u)
+    L
+end
+
 function LinearAlgebra.mul!(v::AbstractVector, L::InvertedOperator, u::AbstractVector)
     ldiv!(v, L.L, u)
 end

src/interface.jl

@@ -29,14 +29,24 @@ end
 (L::AbstractSciMLOperator)(u, p, t) = (update_coefficients!(L, u, p, t); L * u)
 (L::AbstractSciMLOperator)(du, u, p, t) = (update_coefficients!(L, u, p, t); mul!(du, L, u))
 
+###
+# caching interface
+###
+
 """
 Allocate caches for a SciMLOperator for fast evaluation
 
 arguments:
     L :: AbstractSciMLOperator
     u :: AbstractVector argument to L
 """
+cache_operator(L, u) = L
+cache_self(L, u) = L
+cache_internals(L, u) = L
+
 function cache_operator(L::AbstractSciMLOperator, u::AbstractVector)
+    L = cache_self(L, u)
+    L = cache_internals(L, u)
     L
 end
 

src/sciml.jl

@@ -457,6 +457,7 @@ LinearAlgebra.issymmetric(L::FunctionOperator) = L.traits.issymmetric
 LinearAlgebra.ishermitian(L::FunctionOperator) = L.traits.ishermitian
 LinearAlgebra.isposdef(L::FunctionOperator) = L.traits.isposdef
 
+getops(::FunctionOperator) = ()
 has_adjoint(L::FunctionOperator) = !(L.op_adjoint isa Nothing)
 has_mul(L::FunctionOperator{iip}) where{iip} = !iip
 has_mul!(L::FunctionOperator{iip}) where{iip} = iip
@@ -467,7 +468,7 @@ has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothin
 Base.:*(L::FunctionOperator, u::AbstractVector) = L.op(u, L.p, L.t)
 Base.:\(L::FunctionOperator, u::AbstractVector) = L.op_inverse(u, L.p, L.t)
 
-function cache_operator(L::FunctionOperator, u::AbstractVector)
+function cache_self(L::FunctionOperator, u::AbstractVector)
     @set! L.cache = similar(u)
     L
 end
@@ -584,16 +585,28 @@ function Base.:\(L::TensorProductOperator, u::AbstractVector)
     _vec(V)
 end
 
-function cache_operator(L::TensorProductOperator, u::AbstractVector)
+function cache_self(L::TensorProductOperator, u::AbstractVector)
     sz = (size(L.inner, 2), size(L.outer, 2))
     U  = _reshape(u, sz)
     cache = L.inner * U
 
     @set! L.cache = cache
+    L
+end
+
+function cache_internals(L::TensorProductOperator, u::AbstractVector)
+    if !(L.isset)
+        L = cache_self(L, u)
+    end
+
+    sz = (size(L.inner, 2), size(L.outer, 2))
+    U  = _reshape(u, sz)
 
-    L.inner isa AbstractSciMLOperator && @set! L.inner = cache_operator(L.inner)
-    L.outer isa AbstractSciMLOperator && @set! L.outer = cache_operator(L.outer)
+    uinner = U
+    uouter = transpose(L.cache)
 
+    @set! L.inner = cache_operator(L.inner, uinner)
+    @set! L.outer = cache_operator(L.outer, uouter)
     L
 end
 

test/sciml.jl

@@ -1,7 +1,7 @@
 using SciMLOperators, LinearAlgebra
 using Random
 
-using SciMLOperators: InvertibleOperator, ⊗
+using SciMLOperators: AbstractSciMLOperator, InvertibleOperator, ⊗
 
 Random.seed!(0)
 N = 8
@@ -197,6 +197,30 @@ end
 end
 
 @testset "Operator Algebra" begin
-    # try out array arithmatic
+    N2 = N*N
+    A = rand(N,N)
+    B = rand(N,N)
+    C = rand(N,N)
+    D = rand(N,N)
+
+    u = rand(N2)
+    α = rand()
+    β = rand()
+
+    T1 = ⊗(A, B)
+    T2 = ⊗(C, D)
+
+    D1  = DiagonalOperator(rand(N2))
+    D2  = DiagonalOperator(rand(N2))
+
+    TT = AbstractSciMLOperator[T1, T2]
+    DD = Diagonal(AbstractSciMLOperator[D1, D2])
+
+    op = TT' * DD * TT
+
+    op = cache_operator(op, u)
+
+    v=rand(N2); @test mul!(v, op, u) ≈ op * u
+    v=rand(N2); w=copy(v); @test mul!(v, op, u, α, β) ≈ α*(op * u) + β * w
 end
 #