Merge pull request #96 from vpuri3/oop

allow functionoperator to specify both in-place/ out-of-place mode

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLOperators"
 uuid = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 authors = ["xtalax <[email protected]>"]
-version = "0.1.12"
+version = "0.1.13"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
@@ -11,6 +11,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [compat]
 ArrayInterfaceCore = "0.1"

src/SciMLOperators.jl

@@ -7,6 +7,7 @@ import StaticArrays
 import SparseArrays
 import ArrayInterfaceCore
 import Base: ReshapedArray
+import Tricks: static_hasmethod
 import Lazy: @forward
 import Setfield: @set!
 

src/func.jl

@@ -2,7 +2,7 @@
 """
     Matrix free operators (given by a function)
 """
-mutable struct FunctionOperator{isinplace,T,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSciMLOperator{T}
+mutable struct FunctionOperator{iip,oop,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"""
@@ -22,7 +22,8 @@ mutable struct FunctionOperator{isinplace,T,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSc
     """ Cache """
     cache::C
 
-    function FunctionOperator(op,
+    function FunctionOperator(
+                              op,
                               op_adjoint,
                               op_inverse,
                               op_adjoint_inverse,
@@ -34,11 +35,14 @@ mutable struct FunctionOperator{isinplace,T,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSc
                              )
 
         iip = traits.isinplace
+        oop = traits.outofplace
         T   = traits.T
 
         isset = cache !== nothing
 
-        new{iip,
+        new{
+            iip,
+            oop,
             T,
             typeof(op),
             typeof(op_adjoint),
@@ -62,43 +66,52 @@ mutable struct FunctionOperator{isinplace,T,F,Fa,Fi,Fai,Tr,P,Tt,C} <: AbstractSc
     end
 end
 
-function FunctionOperator(op;
+function FunctionOperator(op,
+                          input::AbstractVecOrMat,
+                          output::AbstractVecOrMat;
 
-                          # necessary
-                          isinplace=nothing,
-                          T=nothing,
-                          size=nothing,
+                          isinplace::Union{Nothing,Bool}=nothing,
+                          outofplace::Union{Nothing,Bool}=nothing,
+                          T::Union{Type{<:Number},Nothing}=nothing,
 
-                          input_prototype=nothing,
-                          output_prototype=nothing,
-
-                          # optional
                           op_adjoint=nothing,
                           op_inverse=nothing,
                           op_adjoint_inverse=nothing,
 
                           p=nothing,
-                          t=nothing,
+                          t::Union{Number,Nothing}=nothing,
 
                           # traits
                           opnorm=nothing,
-                          issymmetric=false,
-                          ishermitian=false,
-                          isposdef=false,
+                          issymmetric::Bool=false,
+                          ishermitian::Bool=false,
+                          isposdef::Bool=false,
                          )
 
-    isinplace isa Nothing  && @error "Please provide a funciton signature
-    by specifying `isinplace` as either `true`, or `false`.
-    If `isinplace = false`, the signature is `op(u, p, t)`,
-    and if `isinplace = true`, the signature is `op(du, u, p, t)`.
-    Further, it is assumed that the function call would be nonallocating
-    when called in-place"
-    T isa Nothing  && @error "Please provide a Number type for the Operator"
-    size isa Nothing  && @error "Please provide a size (m, n)"
-    if (input_prototype isa Nothing) | (output_prototype isa Nothing)
-        @error "Please provide input/out prototypes vectors/arrays."
+    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
+
+    isinplace = if isinplace isa Nothing
+        static_hasmethod(op, typeof((output, input, p, t)))
+    else
+        isinplace
+    end
+
+    outofplace = if outofplace isa Nothing
+        static_hasmethod(op, typeof((input, p, t)))
+    else
+        outofplace
     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
+        in-place application."
+    end
+
+    T isa Nothing && @error "Please provide a Number type for the Operator"
+
     isreal = T <: Real
     selfadjoint = ishermitian | (isreal & issymmetric)
     adjointable = !(op_adjoint isa Nothing) | selfadjoint
@@ -112,24 +125,20 @@ function FunctionOperator(op;
         op_adjoint_inverse = op_inverse
     end
 
-    t = t isa Nothing ? zero(T) : t
-
     traits = (;
               opnorm = opnorm,
               issymmetric = issymmetric,
               ishermitian = ishermitian,
               isposdef = isposdef,
 
               isinplace = isinplace,
+              outofplace = outofplace,
               T = T,
-              size = size,
+              size = sz,
              )
 
-    cache = (
-             zero(input_prototype),
-             zero(output_prototype),
-            )
-    isset = cache === nothing
+    cache = zero.((input, output))
+    isset = true
 
     FunctionOperator(
                      op,
@@ -260,16 +269,16 @@ has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothin
 # TODO - FunctionOperator, Base.conj, transpose
 
 # operator application
-Base.:*(L::FunctionOperator{false}, u::AbstractVecOrMat) = L.op(u, L.p, L.t)
-Base.:\(L::FunctionOperator{false}, u::AbstractVecOrMat) = L.op_inverse(u, L.p, L.t)
+Base.:*(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip} = L.op(u, L.p, L.t)
+Base.:\(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip} = L.op_inverse(u, L.p, L.t)
 
-function Base.:*(L::FunctionOperator{true}, u::AbstractVecOrMat)
+function Base.:*(L::FunctionOperator{true,false}, u::AbstractVecOrMat)
     _, co = L.cache
     du = zero(co)
     L.op(du, u, L.p, L.t)
 end
 
-function Base.:\(L::FunctionOperator{true}, u::AbstractVecOrMat)
+function Base.:\(L::FunctionOperator{true,false}, u::AbstractVecOrMat)
     ci, _ = L.cache
     du = zero(ci)
     L.op_inverse(du, u, L.p, L.t)

test/func.jl

@@ -3,7 +3,6 @@ using SciMLOperators, LinearAlgebra
 using Random
 
 using SciMLOperators: InvertibleOperator, ⊗
-using FFTW
 
 Random.seed!(0)
 N = 8
@@ -26,15 +25,7 @@ K = 12
     f2i(du, u, p, t) = ldiv!(du, F, u)
 
     # out of place
-    op1 = FunctionOperator(
-                           f1;
-
-                           isinplace=false,
-                           T=Float64,
-                           size=(N,N),
-
-                           input_prototype=u,
-                           output_prototype=A*u,
+    op1 = FunctionOperator(f1, u, A*u;
 
                            op_inverse=f1i,
 
@@ -45,15 +36,7 @@ K = 12
                           )
 
     # in place
-    op2 = FunctionOperator(
-                           f2;
-
-                           isinplace=true,
-                           T=Float64,
-                           size=(N,N),
-
-                           input_prototype=u,
-                           output_prototype=A*u,
+    op2 = FunctionOperator(f2, u, A*u;
 
                            op_inverse=f2i,
 
@@ -96,19 +79,7 @@ end
 
     f(du,u,p,t) = mul!(du, Diagonal(p*t), u)
 
-    op = FunctionOperator(
-                          f;
-
-                          isinplace=true,
-                          T=Float64,
-                          size=(N,N),
-
-                          input_prototype=u,
-                          output_prototype=u,
-
-                          p=p*0.0,
-                          t=0.0,
-                         )
+    op = FunctionOperator(f, u, u; p=zero(p), t=zero(t))
 
     ans = @. u * p * t
     @test op(u,p,t) ≈ ans

test/matrix.jl

@@ -78,6 +78,23 @@ end
     v=copy(u); @test D(v,u,p,t) ≈ ans
 end
 
+@testset "Batched Diagonal Operator" begin
+    u = rand(N,K)
+    d = rand(N,K)
+    α = rand()
+    β = rand()
+
+    L = DiagonalOperator(d)
+
+    @test L * u ≈ d .* u
+    v=rand(N,K); @test mul!(v, L, u) ≈ d .* u
+    v=rand(N,K); w=copy(v); @test mul!(v, L, u, α, β) ≈ α*(d .* u) + β*w
+
+    @test L \ u ≈ d .\ u
+    v=rand(N,K); @test ldiv!(v, L, u) ≈ d .\ u
+    v=copy(u); @test ldiv!(L, u) ≈ d .\ v
+end
+
 @testset "AffineOperator" begin
     u = rand(N,K)
     A = rand(N,N)