Merge pull request #157 from vpuri3/5arg

Support `L(v, u, p, t, a, b)` corresponding to 5 arg `mul!`

Author: xtalax

src/func.jl

@@ -2,7 +2,7 @@
 """
     Matrix free operators (given by a function)
 """
-mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSciMLOperator{T}
+mutable struct FunctionOperator{iip,oop,mul5,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSciMLOperator{T}
     """ Function with signature op(u, p, t) and (if isinplace) op(du, u, p, t) """
     op::F
     """ Adjoint operator"""
@@ -33,11 +33,13 @@ mutable struct FunctionOperator{iip,oop,T<:Number,F,Fa,Fi,Fai,Tr,P,Tt,C} <: Abst
 
         iip = traits.isinplace
         oop = traits.outofplace
+        mul5 = traits.has_mul5
         T   = traits.T
 
         new{
             iip,
             oop,
+            mul5,
             T,
             typeof(op),
             typeof(op_adjoint),
@@ -88,6 +90,8 @@ function FunctionOperator(op,
 
                           isinplace::Union{Nothing,Bool}=nothing,
                           outofplace::Union{Nothing,Bool}=nothing,
+                          has_mul5::Union{Nothing,Bool}=nothing,
+                          cache::Union{Nothing, NTuple{2}}=nothing,
                           T::Union{Type{<:Number},Nothing}=nothing,
 
                           op_adjoint=nothing,
@@ -109,21 +113,34 @@ function FunctionOperator(op,
                          )
 
     sz = (size(output, 1), size(input, 1))
-    T  = T isa Nothing ? promote_type(eltype.((input, output))...) : T
-    t  = t isa Nothing ? zero(real(T)) : t
+    T  = isnothing(T) ? promote_type(eltype.((input, output))...) : T
+    t  = isnothing(t) ? zero(real(T)) : t
 
-    isinplace = if isinplace isa Nothing
+    isinplace = if isnothing(isinplace)
         static_hasmethod(op, typeof((output, input, p, t)))
     else
         isinplace
     end
 
-    outofplace = if outofplace isa Nothing
+    outofplace = if isnothing(outofplace)
         static_hasmethod(op, typeof((input, p, t)))
     else
         outofplace
     end
 
+    has_mul5 = if isnothing(has_mul5)
+        has_mul5 = true
+        for f in (
+                  op, op_adjoint, op_inverse, op_adjoint_inverse,
+                 )
+            if !isnothing(f)
+                has_mul5 *= static_hasmethod(f, typeof((output, input, p, t, t, t)))
+            end
+        end
+
+        has_mul5
+    end
+
     if !isinplace & !outofplace
         @error "Please provide a funciton with signatures `op(u, p, t)` for applying
         the operator out-of-place, and/or the signature is `op(du, u, p, t)` for
@@ -155,12 +172,12 @@ function FunctionOperator(op,
 
               isinplace = isinplace,
               outofplace = outofplace,
+              has_mul5 = has_mul5,
+              ifcache = ifcache,
               T = T,
               size = sz,
              )
 
-    cache = nothing
-
     L = FunctionOperator(
                          op,
                          op_adjoint,
@@ -172,7 +189,11 @@ function FunctionOperator(op,
                          cache,
                         )
 
-    ifcache ? cache_operator(L, input, output) : L
+    if ifcache & isnothing(L.cache)
+        L = cache_operator(L, input, output)
+    end
+
+    L
 end
 
 function update_coefficients(L::FunctionOperator, u, p, t)
@@ -204,7 +225,13 @@ function update_coefficients!(L::FunctionOperator, u, p, t)
     nothing
 end
 
+function iscached(L::FunctionOperator)
+    L.traits.ifcache ? !isnothing(L.cache) : !L.traits.ifcache
+    !isnothing(L.cache)
+end
+
 function cache_self(L::FunctionOperator, u::AbstractVecOrMat, v::AbstractVecOrMat)
+    L.traits.ifcache && @warn "you are allocating cache for a FunctionOperator for which ifcache = false."
     @set! L.cache = zero.((u, v))
     L
 end
@@ -365,7 +392,7 @@ function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{false}, u::
     @error "LinearAlgebra.mul! not defined for out-of-place FunctionOperators"
 end
 
-function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::AbstractVecOrMat, α, β)
+function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true, oop, false}, u::AbstractVecOrMat, α, β) where{oop}
     _, co = L.cache
 
     copy!(co, v)
@@ -374,6 +401,10 @@ function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::A
     axpy!(β, co, v)
 end
 
+function LinearAlgebra.mul!(v::AbstractVecOrMat, L::FunctionOperator{true, oop, true}, u::AbstractVecOrMat, α, β) where{oop}
+    L.op(v, u, L.p, L.t, α, β)
+end
+
 function LinearAlgebra.ldiv!(v::AbstractVecOrMat, L::FunctionOperator{true}, u::AbstractVecOrMat)
     L.op_inverse(v, u, L.p, L.t)
 end

src/interface.jl

@@ -28,19 +28,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))
+(L::AbstractSciMLOperator)(du, u, p, t, α, β) = (update_coefficients!(L, u, p, t); mul!(du, L, u, α, β))
 
 ###
 # caching interface
 ###
 
+getops(L) = ()
+
 function iscached(L::AbstractSciMLOperator)
+
     has_cache = hasfield(typeof(L), :cache) # TODO - confirm this is static
-    isset = has_cache ? L.cache !== nothing : true
+    isset = has_cache ? !isnothing(L.cache) : true
 
     return isset & all(iscached, getops(L)) 
 end
 
 iscached(L) = true
+
 iscached(::Union{
                  # LinearAlgebra
                  AbstractMatrix,
@@ -61,22 +66,22 @@ arguments:
     in :: AbstractVecOrMat input prototype to L
     out :: (optional) AbstractVecOrMat output prototype to L
 """
-cache_operator
+function cache_operator end
 
 cache_operator(L, u) = L
-cache_operatro(L, u, v) = L
-cache_self(L::AbstractSciMLOperator, uv::AbstractVecOrMat...) = L
-cache_internals(L::AbstractSciMLOperator, uv::AbstractVecOrMat...) = L
+cache_operator(L, u, v) = L
+cache_self(L::AbstractSciMLOperator, ::AbstractVecOrMat...) = L
+cache_internals(L::AbstractSciMLOperator, ::AbstractVecOrMat...) = L
+
+function cache_operator(L::AbstractSciMLOperator, u::AbstractVecOrMat, v::AbstractVecOrMat)
 
-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)
     L = cache_internals(L, u)
     L

test/func.jl

@@ -17,12 +17,15 @@ K = 12
 
     A = rand(N,N) |> Symmetric
     F = lu(A)
+    Ai = inv(A)
 
     f1(u, p, t)  = A * u
     f1i(u, p, t) = A \ u
 
     f2(du, u, p, t)  = mul!(du, A, u)
+    f2(du, u, p, t, α, β)  = mul!(du, A, u, α, β)
     f2i(du, u, p, t) = ldiv!(du, F, u)
+    f2i(du, u, p, t, α, β) = mul!(du, Ai, u, α, β)
 
     # out of place
     op1 = FunctionOperator(f1, u, A*u;
@@ -51,6 +54,7 @@ K = 12
                            ishermitian=true,
                            isposdef=true,
                           )
+
     @test issquare(op1)
     @test issquare(op2)
 
@@ -76,6 +80,12 @@ K = 12
     @test !iscached(op1)
     @test !iscached(op2)
 
+    @test !op1.traits.has_mul5
+    @test op2.traits.has_mul5
+
+    # 5-arg mul! (w/o cache)
+    v = rand(N,K); w=copy(v); @test α*(A*u)+ β*w ≈ mul!(v, op2, u, α, β)
+
     op1 = cache_operator(op1, u, A * u)
     op2 = cache_operator(op2, u, A * u)
 

test/matrix.jl

@@ -65,22 +65,27 @@ end
     p = rand(N)
     t = rand()
 
+    α = rand()
+    β = rand()
+
     L = MatrixOperator(zeros(N,N);
                        update_func = (A,u,p,t) -> (A .= p*p'; nothing)
                       )
 
     @test !isconstant(L)
 
-    A = p*p'
-    ans = A * u
-    @test L(u,p,t) ≈ ans
-    v=copy(u); @test L(v,u,p,t) ≈ ans
+    A = p * p'
+    @test L(u, p, t) ≈ A * u
+    v=copy(u); @test L(v, u, p, t) ≈ A * u
+    v=rand(N,K); w=copy(v); @test L(v, u, p, t, α, β) ≈ α*A*u + β*w
 end
 
 @testset "DiagonalOperator update test" begin
     u = rand(N,K)
     p = rand(N)
     t = rand()
+    α = rand()
+    β = rand()
 
     D = DiagonalOperator(zeros(N);
                          update_func = (diag,u,p,t) -> (diag .= p*t; nothing)
@@ -93,6 +98,7 @@ end
     ans = Diagonal(p*t) * u
     @test D(u,p,t) ≈ ans
     v=copy(u); @test D(v,u,p,t) ≈ ans
+    v=rand(N,K); w=copy(v); @test D(v, u, p, t, α, β) ≈ α*ans + β*w
 end
 
 @testset "Batched Diagonal Operator" begin
@@ -173,6 +179,8 @@ end
     u = rand(N,K)
     p = rand(N)
     t = rand()
+    α = rand()
+    β = rand()
 
     L = AffineOperator(A, B, b;
                        update_func = (b,u,p,t) -> (b .= Diagonal(p*t)*b; nothing)
@@ -186,6 +194,9 @@ end
     b = Diagonal(p*t)*b
     ans = A * u + B * b
     v=copy(u); @test L(v,u,p,t) ≈ ans
+    b = Diagonal(p*t)*b
+    ans = A * u + B * b
+    v=rand(N,K); w=copy(v); @test L(v, u, p, t, α, β) ≈ α*ans + β*w
 end
 
 @testset "TensorProductOperator" begin

test/scalar.jl

@@ -61,26 +61,44 @@ K = 12
 end
 
 @testset "ScalarOperator update test" begin
-    u = ones(N,K)
-    v = zeros(N,K)
+    u = rand(N,K)
+    v = rand(N,K)
     p = rand()
     t = rand()
+    a = rand()
+    b = rand()
 
     α = ScalarOperator(0.0; update_func=(a,u,p,t) -> p)
     β = ScalarOperator(0.0; update_func=(a,u,p,t) -> t)
 
     @test !isconstant(α)
     @test !isconstant(β)
 
-    @test α(u,p,t)   ≈ p * u
-    @test α(v,u,p,t) ≈ p * u
+    @test convert(Number, α) ≈ 0.0
+    @test convert(Number, β) ≈ 0.0
+
+    update_coefficients!(α, u, p, t)
+    update_coefficients!(β, u, p, t)
+
+    @test convert(Number, α) ≈ p
+    @test convert(Number, β) ≈ t
+
+    @test α(u, p, t) ≈ p * u
+    v=rand(N,K); @test α(v, u, p, t) ≈ p * u
+    v=rand(N,K); w=copy(v); @test α(v, u, p, t, a, b) ≈ a*p*u + b*w
+
+    @test β(u, p, t) ≈ t * u
+    v=rand(N,K); @test β(v, u, p, t) ≈ t * u
+    v=rand(N,K); w=copy(v); @test β(v, u, p, t, a, b) ≈ a*t*u + b*w
 
     num = α + 2 / β * 3 - 4
     val = p + 2 / t * 3 - 4
 
-    @test num(u,p,t)   ≈ val * u
-    @test num(v,u,p,t) ≈ val * u
-
     @test convert(Number, num) ≈ val
+
+    @test num(u, p, t) ≈ val * u
+    v=rand(N,K); @test num(v, u, p, t) ≈ val * u
+    v=rand(N,K); w=copy(v); @test num(v, u, p, t, a, b) ≈ a*val*u + b*w
+
 end
 #