Merge pull request #199 from vpuri3/docs

Overload Base.kron, clean up tensor prod benchmark

Author: ChrisRackauckas

benchmarks/tensor.jl

@@ -1,43 +1,48 @@
 using SciMLOperators, LinearAlgebra, BenchmarkTools
 using SciMLOperators: IdentityOperator, ⊗
 
-Id = IdentityOperator{12}()
-A = rand(12,12)
-B = rand(12,12)
-C = rand(12,12)
+N = 12
+K = 100
+Id = IdentityOperator(N)
+A = rand(N,N)
+B = rand(N,N)
+C = rand(N,N)
 
 println("#===============================#")
 println("2D Tensor Products")
 println("#===============================#")
 
 println("⊗(A, B)")
 
-u = rand(12^2, 100)
-v = rand(12^2, 100)
+u = rand(N^2, K)
+v = rand(N^2, K)
 
 T = ⊗(A, B)
 T = cache_operator(T, u)
 
+@btime *($T, $u)
 @btime mul!($v, $T, $u)
 
 println("⊗(I, B)")
 
-u = rand(12^2, 100)
-v = rand(12^2, 100)
+u = rand(N^2, K)
+v = rand(N^2, K)
 
 T = ⊗(Id, B)
 T = cache_operator(T, u)
 
+@btime *($T, $u)
 @btime mul!($v, $T, $u)
 
 println("⊗(A, I)")
 
-u = rand(12^2, 100)
-v = rand(12^2, 100)
+u = rand(N^2, K)
+v = rand(N^2, K)
 
 T = ⊗(A, Id)
 T = cache_operator(T, u)
 
+@btime *($T, $u)
 @btime mul!($v, $T, $u)
 
 println("#===============================#")
@@ -46,25 +51,25 @@ println("#===============================#")
 
 println("⊗(⊗(A, B), C)")
 
-u = rand(12^3, 100)
-v = rand(12^3, 100)
+u = rand(N^3, K)
+v = rand(N^3, K)
 
 T = ⊗(⊗(A, B), C)
 T = cache_operator(T, u)
 
-mul!(v, T, u) # dunny
+@btime *($T, $u)
 @btime mul!($v, $T, $u); #
 
 println("⊗(A, ⊗(B, C))")
 
-u = rand(12^3, 100)
-v = rand(12^3, 100)
+u = rand(N^3, K)
+v = rand(N^3, K)
 
 T = ⊗(A, ⊗(B, C))
 T = cache_operator(T, u)
 
-mul!(v, T, u) # dunny
+@btime *($T, $u)
 @btime mul!($v, $T, $u); #
 
 println("#===============================#")
-nothing
+#

src/func.jl

@@ -86,7 +86,6 @@ function FunctionOperator(op,
     FunctionOperator(op, input, output; kwargs...)
 end
 
-# TODO: document constructor and revisit design as needed (e.g. for "accepted_kwargs")
 """
 $(SIGNATURES)
 
@@ -127,7 +126,6 @@ uniform across `op`, `op_adjoint`, `op_inverse`, `op_adjoint_inverse`.
 * `p` - Prototype of parameter struct passed to the operator during evaluation, i.e. `L(u, p, t)`. `p` is set to `nothing` if no value is provided.
 * `t` - Protype of scalar time variable passed to the operator during evaluation. `t` is set to `zero(T)` if no value is provided.
 * `accepted_kwargs` - `Tuple` of `Symbol`s corresponding to the keyword arguments accepted by `op*`, and `update_coefficients[!]`. For example, if `op` accepts kwarg `scale`, as in `op(u, p, t; scale)`, then `accepted_kwargs = (:scale,)`.
-
 * `T` - `eltype` of the operator. If no value is provided, the constructor inferrs the value from types of `input`, and `output`
 * `isinplace` - `true` if the operator can be used is a mutating way with in-place allocations. This trait is inferred if no value is provided.
 * `outofplace` - `true` if the operator can be used is a non-mutating way with in-place allocations. This trait is inferred if no value is provided.
@@ -440,8 +438,6 @@ has_mul!(::FunctionOperator{iip}) where{iip} = iip
 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)
 
-# TODO - FunctionOperator, Base.conj, transpose
-
 # operator application
 function Base.:*(L::FunctionOperator{iip,true}, u::AbstractVecOrMat) where{iip}
     L.op(u, L.p, L.t; L.traits.kwargs...)

src/tensor.jl

@@ -1,7 +1,5 @@
 #
-"""
-$SIGNATURES
-
+TENSOR_PROD_DOC = """
 Computes the lazy pairwise Kronecker product, or tensor product,
 operator of `AbstractMatrix`, and `AbstractSciMLOperator` subtypes.
 Calling `⊗(ops...)` is equivalent to `Base.kron(ops...)`. Fast
@@ -10,11 +8,18 @@ product operator.
 
 ```
 TensorProductOperator(A, B) = A ⊗ B
+TensorProductOperator(A, B, C) = A ⊗ B ⊗ C
 
 (A ⊗ B)(u) = vec(B * reshape(u, M, N) * transpose(A))
 ```
 where `M = size(B, 2)`, and `N = size(A, 2)`
 """
+
+"""
+$SIGNATURES
+
+$TENSOR_PROD_DOC
+"""
 struct TensorProductOperator{T,O,C} <: AbstractSciMLOperator{T}
     ops::O
     cache::C
@@ -54,26 +59,24 @@ TensorProductOperator(ii1::IdentityOperator, ii2::IdentityOperator) = IdentityOp
 """
 $SIGNATURES
 
-Computes the lazy pairwise Kronecker product, or tensor product,
-operator of `AbstractMatrix`, and `AbstractSciMLOperator` subtypes.
-Calling `⊗(ops...)` is equivalent to `Base.kron(ops...)`. Fast
-operator evaluation is performed without forming the full tensor
-product operator.
-
-```
-TensorProductOperator(A, B) = A ⊗ B
-
-(A ⊗ B)(u) = vec(B * reshape(u, M, N) * transpose(A))
-```
-where `M = size(B, 2)`, and `N = size(A, 2)`
+$TENSOR_PROD_DOC
 """
 ⊗(ops::Union{AbstractMatrix,AbstractSciMLOperator}...) = TensorProductOperator(ops...)
 
-# TODO - overload Base.kron for tensor product operators
-#Base.kron(ops::Union{AbstractMatrix,AbstractSciMLOperator}...) = TensorProductOperator(ops...)
+"""
+$SIGNATURES
+
+Construct a lazy representation of the Kronecker product `A ⊗ B`. One of the
+two factors can be an `AbstractMatrix`, which is then promoted to a
+`MatrixOperator` automatically. To avoid fallback to the generic
+[`Base.kron`](@ref), at least one of `A` and `B` must be an
+`AbstractSciMLOperator`.
+"""
+Base.kron(A::AbstractSciMLOperator, B::AbstractSciMLOperator) = TensorProductOperator(A, B)
+Base.kron(A::AbstractMatrix, B::AbstractSciMLOperator) = TensorProductOperator(A, B)
+Base.kron(A::AbstractSciMLOperator, B::AbstractMatrix) = TensorProductOperator(A, B)
 
-# convert to matrix
-Base.kron(ops::AbstractSciMLOperator...) = kron(convert.(AbstractMatrix, ops)...)
+Base.kron(ops::AbstractSciMLOperator...) = TensorProductOperator(ops...)
 
 function Base.convert(::Type{AbstractMatrix}, L::TensorProductOperator)
     kron(convert.(AbstractMatrix, L.ops)...)

test/matrix.jl

@@ -1,7 +1,7 @@
 using SciMLOperators, LinearAlgebra
 using Random
 
-using SciMLOperators: InvertibleOperator, InvertedOperator, ⊗
+using SciMLOperators: InvertibleOperator, InvertedOperator, ⊗, AbstractSciMLOperator
 using FFTW
 
 Random.seed!(0)
@@ -263,6 +263,18 @@ for square in [false, true] #for K in [1, K]
     AB  = kron(A, B)
     ABC = kron(A, B, C)
 
+    # test Base.kron overload
+    # ensure kron(mat, mat) is not a TensorProductOperator
+    @test !isa(AB, AbstractSciMLOperator)
+    @test !isa(ABC, AbstractSciMLOperator)
+
+    # test Base.kron overload
+    _A = rand(N, N)
+    @test kron(_A, MatrixOperator(_A)) isa TensorProductOperator
+    @test kron(MatrixOperator(_A), _A) isa TensorProductOperator
+
+    @test kron(MatrixOperator(_A), MatrixOperator(_A)) isa TensorProductOperator
+
     # Inputs
     u2 = rand(n1*n2, K)
     u3 = rand(n1*n2*n3, K)