add method cache_internals(::FunctionOp, ::AbstractArray)

src/func.jl

@@ -384,6 +384,18 @@ function _cache_self(L::FunctionOperator, u::AbstractArray)
     @set! L.cache = (_u, _v)
 end
 
+# fix method amg bw AbstractArray, AbstractVecOrMat
+cache_internals(L::FunctionOperator, u::AbstractArray) = _cache_internals(L, u)
+cache_internals(L::FunctionOperator, u::AbstractVecOrMat) = _cache_internals(L, u)
+
+function _cache_internals(L::FunctionOperator, u::AbstractArray)
+
+    @set! L.op = cache_operator(L.op, u)
+    @set! L.op_adjoint = cache_operator(L.op_adjoint, u)
+    @set! L.op_inverse = cache_operator(L.op_inverse, u)
+    @set! L.op_adjoint_inverse = cache_operator(L.op_adjoint_inverse, u)
+end
+
 function Base.show(io::IO, L::FunctionOperator)
     M, N = size(L)
     print(io, "FunctionOperator($M × $N)")
@@ -542,7 +554,7 @@ function _sizecheck(L::FunctionOperator, u, v)
         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(u))."""
+                    Recievd array of size $(size(v))."""
                 DimensionMismatch(msg) |> throw
             end
         end
@@ -558,7 +570,7 @@ function _sizecheck(L::FunctionOperator, u, v)
         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(u))."""
+                    Recievd array of size $(size(v))."""
                 DimensionMismatch(msg) |> throw
             end
         end
@@ -618,8 +630,7 @@ function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{true, oop, fal
 
     copy!(co, v)
     mul!(v, L, u)
-    lmul!(α, v)
-    axpy!(β, co, v)
+    axpby!(β, co, α, v)
 end
 
 function LinearAlgebra.mul!(v::AbstractArray, L::FunctionOperator{true, oop, true}, u::AbstractArray, α, β) where{oop}
@@ -639,8 +650,6 @@ end
 function LinearAlgebra.ldiv!(L::FunctionOperator{true}, u::AbstractArray)
     ci, _ = L.cache
 
-    _sizecheck(L, nothing, u)
-
     copy!(ci, u)
     ldiv!(u, L, ci)
 end

test/func.jl

@@ -9,6 +9,54 @@ N = 8
 K = 12
 NK = N * K
 
+@testset "(Unbatched) FunctionOperator ND array" begin
+    N1, N2, N3 = 3, 4, 5
+    M1, M2, M3 = 4, 5, 6
+
+    p = nothing
+    t = 0.0
+    α = rand()
+    β = rand()
+
+    for (sz_in, sz_out) in (
+        ((N1, N2, N3), (N1, N2, N3)), # equal size
+        ((N1, N2, N3), (M1, M2, M3)), # different size
+    )
+        N = prod(sz_in)
+        M = prod(sz_out)
+
+        A = rand(M, N)
+        u = rand(sz_in... )
+        v = rand(sz_out...)
+
+        _mul(A, u) = reshape(A * vec(u), sz_out)
+        f(u, p, t) = _mul(A, u)
+        f(du, u, p, t)  = (mul!( vec(du), A, vec(u)); du)
+
+        kw = (;) # FunctionOp kwargs
+
+        if sz_in == sz_out
+            F = lu(A)
+            _div(A, v) = reshape(A \ vec(v), sz_in)
+            fi(u, p, t) = _div(A, u)
+            fi(du, u, p, t) = (ldiv!(vec(du), F, vec(u)); du)
+
+            kw = (; op_inverse = fi)
+        end
+
+        L = FunctionOperator(f, u, v; kw...)
+        L = cache_operator(L, u)
+
+        @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
+
+end
+
 @testset "(Unbatched) FunctionOperator" begin
     u = rand(N, K)
     p = nothing