Merge pull request #210 from vpuri3/vec

allow functionoperator (batch = false) to accept/return vec'd arrays

Author: ChrisRackauckas

docs/src/tutorials/fftw.md

@@ -19,21 +19,22 @@ x  = range(start=-L/2, stop=L/2-dx, length=n) |> Array
 u  = @. sin(5x)cos(7x);
 du = @. 5cos(5x)cos(7x) - 7sin(5x)sin(7x);
 
-k  = rfftfreq(n, 2π*n/L) |> Array
-m  = length(k)
-transform = plan_rfft(x)
+k = rfftfreq(n, 2π*n/L) |> Array
+m = length(k)
+P = plan_rfft(x)
+
+F = FunctionOperator(fwd, x, im*k;
+        T=ComplexF64,
 
-T = FunctionOperator((du,u,p,t) -> mul!(du, transform, u), x, im*k;
-                     isinplace=true,
-                     T=ComplexF64,
+        op_adjoint = bwd,
+        op_inverse = bwd,
+        op_adjoint_inverse = fwd,
 
-                     op_adjoint = (du,u,p,t) -> ldiv!(du, transform, u),
-                     op_inverse = (du,u,p,t) -> ldiv!(du, transform, u),
-                     op_adjoint_inverse = (du,u,p,t) -> ldiv!(du, transform, u),
-                    )
+        islinear=true,
+        )
 
 ik = im * DiagonalOperator(k)
-Dx = T \ ik * T
+Dx = F \ ik * F
 
 Dx = cache_operator(Dx, x)
 
@@ -79,18 +80,17 @@ Now we are ready to define our wrapper for the FFT object. To `FunctionOperator`
 pass the in-place forward application of the transform,
 `(du,u,p,t) -> mul!(du, transform, u)`, its inverse application,
 `(du,u,p,t) -> ldiv!(du, transform, u)`, as well as input and output prototype vectors.
-We also set the flag `isinplace` to `true` to signal that we intend to use the operator
-in a non-allocating way, and pass in the element-type and size of the operator.
 
 ```
-T = FunctionOperator((du,u,p,t) -> mul!(du, transform, u), x, im*k;
-                     isinplace=true,
-                     T=ComplexF64,
-
-                     op_adjoint = (du,u,p,t) -> ldiv!(du, transform, u),
-                     op_inverse = (du,u,p,t) -> ldiv!(du, transform, u),
-                     op_adjoint_inverse = (du,u,p,t) -> ldiv!(du, transform, u),
-                    )
+F = FunctionOperator(fwd, x, im*k;
+        T=ComplexF64,
+
+        op_adjoint = bwd,
+        op_inverse = bwd,
+        op_adjoint_inverse = fwd,
+
+        islinear=true,
+        )
 ```
 
 After wrapping the FFT with `FunctionOperator`, we are ready to compose it with other
@@ -100,7 +100,7 @@ both in-place, and out-of-place by comparing its output to the analytical deriva
 
 ```
 ik = im * DiagonalOperator(k)
-Dx = T \ ik * T
+Dx = F \ ik * F
 
 @show ≈(Dx * u, du; atol=1e-8)
 @show ≈(mul!(copy(u), Dx, u), du; atol=1e-8)

src/func.jl

@@ -173,7 +173,7 @@ function FunctionOperator(op,
             msg = """`FunctionOperator` constructed with `batch = true` only
                 accepts `AbstractVecOrMat` types with
                 `size(L, 2) == size(u, 1)`."""
-            ArgumentError(msg) |> throw
+            throw(ArgumentError(msg))
         end
 
         if input isa AbstractMatrix
@@ -184,7 +184,7 @@ function FunctionOperator(op,
                     array, $(typeof(input)), has size $(size(input)), whereas
                     output array, $(typeof(output)), has size
                     $(size(output))."""
-                ArgumentError(msg) |> throw
+                throw(ArgumentError(msg))
             end
         end
     end
@@ -340,14 +340,14 @@ function _cache_operator(L::FunctionOperator, u::AbstractArray)
         if !isa(u, AbstractVecOrMat)
             msg = """$L constructed with `batch = true` only accepts
                 `AbstractVecOrMat` types with `size(L, 2) == size(u, 1)`."""
-            ArgumentError(msg) |> throw
+            throw(ArgumentError(msg))
         end
 
         if size(L, 2) != size(u, 1)
             msg = """Second dimension of $L of size $(size(L))
                 is not consistent with first dimension of input array `u`
                 of size $(size(u))."""
-            DimensionMismatch(msg) |> throw
+            throw(DimensionMismatch(msg))
         end
 
         M = size(L, 1)
@@ -486,7 +486,7 @@ function Base.resize!(L::FunctionOperator, n::Integer)
     if length(L.traits.sizes[1]) != 1
         msg = """`Base.resize!` is only supported by $L whose input/output
             arrays are `AbstractVector`s."""
-        MethodError(msg) |> throw
+        throw(MethodError(msg))
     end
 
     for op in getops(L)
@@ -534,131 +534,187 @@ has_ldiv(L::FunctionOperator{iip}) where{iip} = !(L.op_inverse isa Nothing)
 has_ldiv!(L::FunctionOperator{iip}) where{iip} = iip & !(L.op_inverse isa Nothing)
 
 function _sizecheck(L::FunctionOperator, u, v)
-
+    sizes = L.traits.sizes
     if L.traits.batch
         if !isnothing(u)
             if !isa(u, AbstractVecOrMat)
                 msg = """$L constructed with `batch = true` only
                     accept input arrays that are `AbstractVecOrMat`s with
-                    `size(L, 2) == size(u, 1)`."""
-                ArgumentError(msg) |> throw
+                    `size(L, 2) == size(u, 1)`. Recieved $(typeof(u))."""
+                throw(ArgumentError(msg))
             end
 
-            if size(u) != L.traits.sizes[1]
-                msg = """$L expects input arrays of size $(L.traits.sizes[1]).
-                    Recievd array of size $(size(u))."""
-                DimensionMismatch(msg) |> throw
+            if size(L, 2) != size(u, 1)
+                msg = """$L accepts input `AbstractVecOrMat`s of size
+                    ($(size(L, 2)), K). Recievd array of size $(size(u))."""
+                throw(DimensionMismatch(msg))
             end
-        end
+        end # u
 
         if !isnothing(v)
-            if size(v) != L.traits.sizes[2]
-                msg = """$L expects input arrays of size $(L.traits.sizes[1]).
-                    Recievd array of size $(size(v))."""
-                DimensionMismatch(msg) |> throw
+            if !isa(v, AbstractVecOrMat)
+                msg = """$L constructed with `batch = true` only
+                    returns output arrays that are `AbstractVecOrMat`s with
+                    `size(L, 1) == size(v, 1)`. Recieved $(typeof(v))."""
+                throw(ArgumentError(msg))
             end
-        end
+
+            if size(L, 1) != size(v, 1)
+                msg = """$L accepts output `AbstractVecOrMat`s of size
+                    ($(size(L, 1)), K). Recievd array of size $(size(v))."""
+                throw(DimensionMismatch(msg))
+            end
+        end # v
+
+        if !isnothing(u) & !isnothing(v)
+            if size(u, 2) != size(v, 2)
+                msg = """input array $u, and output array, $v, must have the
+                    same batch size (i.e. length of second dimension). Got
+                    $(size(u)), $(size(v)). If you encounter this error during
+                    an in-place evaluation (`LinearAlgebra.mul!`, `ldiv!`),
+                    ensure that the operator $L has been cached with an input
+                    array of the correct size. Do so by calling
+                    `L = cache_operator(L, u)`."""
+                throw(DimensionMismatch(msg))
+            end
+        end # u, v
+
     else # !batch
+
         if !isnothing(u)
-            if size(u) != L.traits.sizes[1]
-                msg = """$L expects input arrays of size $(L.traits.sizes[1]).
-                    Recievd array of size $(size(u))."""
-                DimensionMismatch(msg) |> throw
+            if size(u) ∉ (sizes[1], tuple(size(L, 2)),)
+                msg = """$L recievd input array of size $(size(u)), but only
+                    accepts input arrays of size $(sizes[1]), or vectors like
+                    `vec(u)` of size $(tuple(prod(sizes[1])))."""
+                throw(DimensionMismatch(msg))
             end
-        end
+        end # u
 
         if !isnothing(v)
-            if size(v) != L.traits.sizes[2]
-                msg = """$L expects input arrays of size $(L.traits.sizes[1]).
-                    Recievd array of size $(size(v))."""
-                DimensionMismatch(msg) |> throw
+            if size(v) ∉ (sizes[2], tuple(size(L, 1)),)
+                msg = """$L recievd output array of size $(size(v)), but only
+                    accepts output arrays of size $(sizes[2]), or vectors like
+                    `vec(u)` of size $(tuple(prod(sizes[2])))"""
+                throw(DimensionMismatch(msg))
             end
-        end
+        end # v
     end # batch
 
     return
 end
 
-# operator application
-function Base.:*(L::FunctionOperator{iip,true}, u::AbstractArray) where{iip}
-    _sizecheck(L, u, nothing)
+function _unvec(L::FunctionOperator, u, v)
+    if L.traits.batch
+        return u, v, false
+    else
+        sizes = L.traits.sizes
 
-    L.op(u, L.p, L.t; L.traits.kwargs...)
-end
+        # no need to vec since expected input/output are AbstractVectors
+        if length(sizes[1]) == 1
+            return u, v, false
+        end
 
-function Base.:\(L::FunctionOperator{iip,true}, u::AbstractArray) where{iip}
-    _sizecheck(L, nothing, u)
+        vec_u = isnothing(u) ? false : size(u) != sizes[1]
+        vec_v = isnothing(v) ? false : size(v) != sizes[2]
 
-    L.op_inverse(u, L.p, L.t; L.traits.kwargs...)
-end
+        if !isnothing(u) & !isnothing(v)
+            if (vec_u & !vec_v) | (!vec_u & vec_v)
+                msg = """Input / output to $L can either be of sizes
+                    $(sizes[1]) / $(sizes[2]), or
+                    $(tuple(prod(sizes[1]))) / $(tuple(prod(sizes[2]))). Got
+                    $(size(u)), $(size(v))."""
+                throw(DimensionMismatch(msg))
+            end
+        end
 
-# fallback *, \ for FunctionOperator with no OOP method
+        U = vec_u ? reshape(u, sizes[1]) : u
+        V = vec_v ? reshape(v, sizes[2]) : v
+        vec_output = vec_u | vec_v
 
-function Base.:*(L::FunctionOperator{true,false}, u::AbstractArray)
-    _, co = L.cache
-    v = zero(co)
+        return U, V, vec_output
+    end
+end
 
-    _sizecheck(L, u, v)
+# operator application
+function Base.:*(L::FunctionOperator{iip,true}, u::AbstractArray) where{iip}
+    _sizecheck(L, u, nothing)
+    U, _, vec_output = _unvec(L, u, nothing)
 
-    L.op(v, u, L.p, L.t; L.traits.kwargs...)
+    V = L.op(U, L.p, L.t; L.traits.kwargs...)
+
+    vec_output ? vec(V) : V
 end
 
-function Base.:\(L::FunctionOperator{true,false}, u::AbstractArray)
-    ci, _ = L.cache
-    v = zero(ci)
+function Base.:\(L::FunctionOperator{iip,true}, v::AbstractArray) where{iip}
+    _sizecheck(L, nothing, v)
+    _, V, vec_output = _unvec(L, nothing, v)
 
-    _sizecheck(L, v, u)
+    U = L.op_inverse(V, L.p, L.t; L.traits.kwargs...)
 
-    L.op_inverse(v, u, L.p, L.t; L.traits.kwargs...)
+    vec_output ? vec(U) : U
 end
 
 function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{true}, u::AbstractArray)
-
     _sizecheck(L, u, v)
+    U, V, vec_output = _unvec(L, u, v)
+
+    L.op(V, U, L.p, L.t; L.traits.kwargs...)
 
-    L.op(v, u, L.p, L.t; L.traits.kwargs...)
+    vec_output ? vec(V) : V
 end
 
-function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{false}, u::AbstractArray, args...)
-    @error "LinearAlgebra.mul! not defined for out-of-place FunctionOperators"
+function LinearAlgebra.mul!(::AbstractArray, L::FunctionOperator{false}, ::AbstractArray, args...)
+    @error "LinearAlgebra.mul! not defined for out-of-place operator $L"
 end
 
 function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{true, oop, false}, u::AbstractArray, α, β) where{oop}
-    _, co = L.cache
+    _, Co = L.cache
 
     _sizecheck(L, u, v)
+    U, V, _ = _unvec(L, u, v)
 
-    copy!(co, v)
-    mul!(v, L, u)
-    axpby!(β, co, α, v)
+    copy!(Co, V)
+    L.op(V, U, L.p, L.t; L.traits.kwargs...) # mul!(V, L, U)
+    axpby!(β, Co, α, V)
+
+    v
 end
 
 function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{true, oop, true}, u::AbstractArray, α, β) where{oop}
-
     _sizecheck(L, u, v)
+    U, V, _ = _unvec(L, u, v)
 
-    L.op(v, u, L.p, L.t, α, β; L.traits.kwargs...)
+    L.op(V, U, L.p, L.t, α, β; L.traits.kwargs...)
+
+    v
 end
 
-function LinearAlgebra.ldiv!(v::AbstractArray, L::FunctionOperator{true}, u::AbstractArray)
+function LinearAlgebra.ldiv!(u::AbstractArray, L::FunctionOperator{true}, v::AbstractArray)
+    _sizecheck(L, u, v)
+    U, V, _ = _unvec(L, u, v)
 
-    _sizecheck(L, v, u)
+    L.op_inverse(U, V, L.p, L.t; L.traits.kwargs...)
 
-    L.op_inverse(v, u, L.p, L.t; L.traits.kwargs...)
+    u
 end
 
 function LinearAlgebra.ldiv!(L::FunctionOperator{true}, u::AbstractArray)
-    ci, _ = L.cache
+    V, _ = L.cache
+
+    _sizecheck(L, u, V)
+    U, _, _ = _unvec(L, u, nothing)
+
+    copy!(V, U)
+    L.op_inverse(U, V, L.p, L.t; L.traits.kwargs...) # ldiv!(U, L, V)
 
-    copy!(ci, u)
-    ldiv!(u, L, ci)
+    u
 end
 
 function LinearAlgebra.ldiv!(v::AbstractArray, L::FunctionOperator{false}, u::AbstractArray)
-    @error "LinearAlgebra.ldiv! not defined for out-of-place FunctionOperators"
+    @error "LinearAlgebra.ldiv! not defined for out-of-place $L"
 end
 
 function LinearAlgebra.ldiv!(L::FunctionOperator{false}, u::AbstractArray)
-    @error "LinearAlgebra.ldiv! not defined for out-of-place FunctionOperators"
+    @error "LinearAlgebra.ldiv! not defined for out-of-place $L"
 end
 #

test/func.jl

@@ -47,13 +47,25 @@ NK = N * K
         L = FunctionOperator(f, u, v; kw...)
         L = cache_operator(L, u)
 
+        # test with ND-arrays
         @test _mul(A, u) ≈ L(u, p, t) ≈ L * u ≈ mul!(zero(v), L, u)
         @test α * _mul(A, u)+ β * v ≈ mul!(copy(v), L, u, α, β)
 
         if sz_in == sz_out
             @test _div(A, v) ≈ L \ v ≈ ldiv!(zero(u), L, v) ≈ ldiv!(L, copy(v))
         end
-    end
+
+        # test with vec(Array)
+        @test vec(_mul(A, u)) ≈ L(vec(u), p, t) ≈ L * vec(u) ≈ mul!(vec(zero(v)), L, vec(u))
+        @test vec(α * _mul(A, u)+ β * v) ≈ mul!(vec(copy(v)), L, vec(u), α, β)
+
+        if sz_in == sz_out
+            @test vec(_div(A, v)) ≈ L \ vec(v) ≈ ldiv!(vec(zero(u)), L, vec(v)) ≈ ldiv!(L, vec(copy(v)))
+        end
+
+        @test_throws DimensionMismatch mul!(vec(v), L, u)
+        @test_throws DimensionMismatch mul!(v, L, vec(u))
+    end # for
 
 end