Merge pull request #150 from vpuri3/func_cache

Update FunctionOperator

Author: ChrisRackauckas

src/func.jl

@@ -97,6 +97,8 @@ function FunctionOperator(op,
                           p=nothing,
                           t::Union{Number,Nothing}=nothing,
 
+                          ifcache::Bool = true,
+
                           # traits
                           islinear::Bool = false,
 
@@ -157,18 +159,20 @@ function FunctionOperator(op,
               size = sz,
              )
 
-    cache = zero.((input, output))
+    cache = nothing
 
-    FunctionOperator(
-                     op,
-                     op_adjoint,
-                     op_inverse,
-                     op_adjoint_inverse,
-                     traits,
-                     p,
-                     t,
-                     cache,
-                    )
+    L = FunctionOperator(
+                         op,
+                         op_adjoint,
+                         op_inverse,
+                         op_adjoint_inverse,
+                         traits,
+                         p,
+                         t,
+                         cache,
+                        )
+
+    ifcache ? cache_operator(L, input, output) : L
 end
 
 function update_coefficients(L::FunctionOperator, u, p, t)
@@ -200,6 +204,11 @@ function update_coefficients!(L::FunctionOperator, u, p, t)
     nothing
 end
 
+function cache_self(L::FunctionOperator, u::AbstractVecOrMat, v::AbstractVecOrMat)
+    @set! L.cache = zero.((u, v))
+    L
+end
+
 Base.size(L::FunctionOperator) = L.traits.size
 function Base.adjoint(L::FunctionOperator)
 
@@ -306,7 +315,6 @@ function getops(L::FunctionOperator)
 end
 
 #TODO - isconstant(L::FunctionOperator)
-iscached(::FunctionOperator) = true
 islinear(L::FunctionOperator) = L.traits.islinear
 has_adjoint(L::FunctionOperator) = !(L.op_adjoint isa Nothing)
 has_mul(L::FunctionOperator{iip}) where{iip} = true

src/interface.jl

@@ -40,6 +40,7 @@ function iscached(L::AbstractSciMLOperator)
     return isset & all(iscached, getops(L)) 
 end
 
+iscached(L) = true
 iscached(::Union{
                  # LinearAlgebra
                  AbstractMatrix,
@@ -57,11 +58,23 @@ Allocate caches for a SciMLOperator for fast evaluation
 
 arguments:
     L :: AbstractSciMLOperator
-    u :: AbstractVecOrMat argument to L
+    in :: AbstractVecOrMat input prototype to L
+    out :: (optional) AbstractVecOrMat output prototype to L
 """
+cache_operator
+
 cache_operator(L, u) = L
-cache_self(L::AbstractSciMLOperator, u::AbstractVecOrMat) = L
-cache_internals(L::AbstractSciMLOperator, u::AbstractVecOrMat) = L
+cache_operatro(L, u, v) = L
+cache_self(L::AbstractSciMLOperator, uv::AbstractVecOrMat...) = L
+cache_internals(L::AbstractSciMLOperator, uv::AbstractVecOrMat...) = L
+
+function cache_operator(L::AbstractSciMLOperator,
+                        u::AbstractVecOrMat,
+                        v::AbstractVecOrMat)
+    L = cache_self(L, u, v)
+    L = cache_internals(L, u, v)
+    L
+end
 
 function cache_operator(L::AbstractSciMLOperator, u::AbstractVecOrMat)
     L = cache_self(L, u)

test/func.jl

@@ -29,6 +29,8 @@ K = 12
 
                            op_inverse=f1i,
 
+                           ifcache = false,
+
                            islinear=true,
                            opnorm=true,
                            issymmetric=true,
@@ -41,6 +43,8 @@ K = 12
 
                            op_inverse=f2i,
 
+                           ifcache = false,
+
                            islinear=true,
                            opnorm=true,
                            issymmetric=true,
@@ -69,6 +73,12 @@ K = 12
     @test has_ldiv(op2)
     @test has_ldiv!(op2)
 
+    @test !iscached(op1)
+    @test !iscached(op2)
+
+    op1 = cache_operator(op1, u, A * u)
+    op2 = cache_operator(op2, u, A * u)
+
     @test iscached(op1)
     @test iscached(op2)
 
@@ -87,10 +97,35 @@ end
 
     f(du,u,p,t) = mul!(du, Diagonal(p*t), u)
 
-    op = FunctionOperator(f, u, u; p=zero(p), t=zero(t))
+    L = FunctionOperator(f, u, u; p=zero(p), t=zero(t))
 
     ans = @. u * p * t
-    @test op(u,p,t) ≈ ans
-    v=copy(u); @test op(v,u,p,t) ≈ ans
+    @test L(u,p,t) ≈ ans
+    v=copy(u); @test L(v,u,p,t) ≈ ans
+
+    # test that output isn't accidentally mutated by passing an internal cache.
+
+    A = Diagonal(p * t)
+    u1 = rand(N, K)
+    u2 = rand(N, K)
+
+    v1 = L * u1; @test v1 ≈ A * u1
+    v2 = L * u2; @test v2 ≈ A * u2; @test v1 ≈ A * u1
+    @test v1 + v2 ≈ A * (u1 + u2)
+
+    v1 .= 0.0
+    v2 .= 0.0
+
+    mul!(v1, L, u1); @test v1 ≈ A * u1
+    mul!(v2, L, u2); @test v2 ≈ A * u2; @test v1 ≈ A * u1
+    @test v1 + v2 ≈ A * (u1 + u2)
+
+    v1 = rand(N, K); w1 = copy(v1)
+    v2 = rand(N, K); w2 = copy(v2)
+    a1, a2, b1, b2 = rand(4)
+
+    mul!(v1, L, u1, a1, b1); @test v1 ≈ a1*A*u1 + b1*w1
+    mul!(v2, L, u2, a2, b2); @test v2 ≈ a2*A*u2 + b2*w2; @test v1 ≈ a1*A*u1 + b1*w1
+    @test v1 + v2 ≈ (a1*A*u1 + b1*w1) + (a2*A*u2 + b2*w2)
 end
 #