Merge pull request #82 from vpuri3/batch

Batched Diagonal Operator

Author: ChrisRackauckas

docs/src/interface.md

@@ -1,31 +1,36 @@
 # The AbstractSciMLOperator Interface
 
-## Formal Properties of DiffEqOperators
+## Formal Properties of SciMLOperators
 
 These are the formal properties that an `AbstractSciMLOperator` should obey
 for it to work in the solvers.
 
-## AbstractDiffEqOperator Interface Description
-
-1. Function call and multiplication: `L(du,u,p,t)` for inplace and `du = L(u,p,t)` for
-   out-of-place, meaning `L*u` and `mul!`.
-2. If the operator is not a constant, update it with `(u,p,t)`. A mutating form, i.e.
-   `update_coefficients!(A,u,p,t)` that changes the internal coefficients, and a
-   out-of-place form `B = update_coefficients(A,u,p,t)`.
-3. `isconstant(A)` trait for whether the operator is constant or not.
-
-## AbstractDiffEqLinearOperator Interface Description
+1. An `AbstractSciMLOperator` represents a linear or nonlinear operator with input/output
+   being `AbstractArray`s. Specifically, a SciMLOperator, `L`, of size `(M,N)` accepts
+   input argument `u` with leading length `N`, i.e. `size(u, 1) == N`, and returns an
+   `AbstractArray` of the same dimension with leading length `M`, i.e. `size(L * u, 1) == M`.
+2. SciMLOperators can be applied to an `AbstractArray` via overloaded `Base.*`, or
+   the in-place `LinearAlgebra.mul!`. Additionally, operators are allowed to be time,
+   or parameter dependent. The state of a SciMLOperator can be updated by calling
+   the mutating function `update_coefficients!(L, u, p, t)` where `p` representes
+   parameters, and `t`, time.  Calling a SciMLOperator as `L(du, u, p, t)` or out-of-place
+   `L(u, p, t)` will automatically update the state of `L` before applying it to `u`.
+   `L(u, p, t)` is the same operation as `L(u, p, t) * u`.
+3. To support the update functionality, we have lazily implemented a comprehensive operator
+   algebra. That means a user can add, subtract, scale, compose and invert SciMLOperators,
+   and the state of the resultant operator would be updated as expected upon calling
+   `L(du, u, p, t)` or `L(u, p, t)` so long as an update function is provided for the
+   component operators.
+
+## AbstractSciMLOperator Interface Description
 
 1. `AbstractSciMLLinearOperator <: AbstractSciMLOperator`
-2. Can absorb under multiplication by a scalar. In all algorithms things like
-   `dt*L` show up all the time, so the linear operator must be able to absorb
-   such constants.
-4. `isconstant(A)` trait for whether the operator is constant or not.
-5. Optional: `diagonal`, `symmetric`, etc traits from LinearMaps.jl.
-6. Optional: `exp(A)`. Required for simple exponential integration.
-7. Optional: `expv(A,u,t) = exp(t*A)*u` and `expv!(v,A::AbstractSciMLOperator,u,t)`
+2. `AbstractSciMLScalarOperator <: AbstractSciMLLinearOperator`
+3. `isconstant(A)` trait for whether the operator is constant or not.
+4. Optional: `exp(A)`. Required for simple exponential integration.
+5. Optional: `expv(A,u,t) = exp(t*A)*u` and `expv!(v,A::AbstractSciMLOperator,u,t)`
    Required for sparse-saving exponential integration.
-8. Optional: factorizations. `ldiv!`, `factorize` et. al. This is only required
+6. Optional: factorizations. `ldiv!`, `factorize` et. al. This is only required
    for algorithms which use the factorization of the operator (Crank-Nicolson),
    and only for when the default linear solve is used.
 
@@ -49,4 +54,4 @@ That same trick then can be used pretty much anywhere you would've had a linear
 the proof to affine operators, so then ``exp(A*t)*v`` operations via Krylov methods work for A being 
 affine as well, and all sorts of things. Thus affine operators have no matrix representation but they 
 are still compatible with essentially any Krylov method which would otherwise be compatible with
-matrix-free representations, hence their support in the SciMLOperators interface.
+matrix-free representations, hence their support in the SciMLOperators interface.

src/SciMLOperators.jl

@@ -41,6 +41,7 @@ include("multidim.jl")
 include("scalar.jl")
 include("basic.jl")
 include("matrix.jl")
+include("batch.jl")
 include("func.jl")
 include("tensor.jl")
 

src/batch.jl

@@ -0,0 +1,112 @@
+#
+"""
+    BatchedDiagonalOperator(diag, [; update_func])
+
+Represents a time-dependent elementwise scaling (diagonal-scaling) operation.
+Acts on `AbstractArray`s of the same size as `diag`. The update function is called
+by `update_coefficients!` and is assumed to have the following signature:
+
+    update_func(diag::AbstractVector,u,p,t) -> [modifies diag]
+"""
+struct BatchedDiagonalOperator{T,D,F} <: AbstractSciMLLinearOperator{T}
+    diag::D
+    update_func::F
+
+    function BatchedDiagonalOperator(
+                                     diag::AbstractArray;
+                                     update_func=DEFAULT_UPDATE_FUNC
+                                    )
+        new{
+            eltype(diag),
+            typeof(diag),
+            typeof(update_func)
+           }(
+             diag, update_func,
+            )
+    end
+end
+
+function DiagonalOperator(u::AbstractArray; update_func=DEFAULT_UPDATE_FUNC)
+    BatchedDiagonalOperator(u; update_func=update_func)
+end
+
+# traits
+Base.size(L::BatchedDiagonalOperator) = (N = size(L.diag, 1); (N, N))
+Base.iszero(L::BatchedDiagonalOperator) = iszero(L.diag)
+Base.transpose(L::BatchedDiagonalOperator) = L
+Base.adjoint(L::BatchedDiagonalOperator) = conj(L)
+function Base.conj(L::BatchedDiagonalOperator) # TODO - test this thoroughly
+    diag = conj(L.diag)
+    update_func = if isreal(L)
+        L.update_func
+    else
+        (L,u,p,t) -> conj(L.update_func(conj(L.diag),u,p,t))
+    end
+    BatchedDiagonalOperator(diag; update_func=update_func)
+end
+
+LinearAlgebra.issymmetric(L::BatchedDiagonalOperator) = true
+function LinearAlgebra.ishermitian(L::BatchedDiagonalOperator)
+    if isreal(L)
+        true
+    else
+        d = _vec(L.diag)
+        D = Diagonal(d)
+        ishermitian(d)
+    end
+end
+LinearAlgebra.isposdef(L::BatchedDiagonalOperator) = isposdef(Diagonal(_vec(L.diag)))
+
+isconstant(L::BatchedDiagonalOperator) = L.update_func == DEFAULT_UPDATE_FUNC
+issquare(L::BatchedDiagonalOperator) = true
+has_adjoint(L::BatchedDiagonalOperator) = true
+has_ldiv(L::BatchedDiagonalOperator) = all(x -> !iszero(x), L.diag)
+has_ldiv!(L::BatchedDiagonalOperator) = has_ldiv(L)
+
+getops(L::BatchedDiagonalOperator) = (L.diag,)
+
+update_coefficients!(L::BatchedDiagonalOperator,u,p,t) = (L.update_func(L.diag,u,p,t); nothing)
+
+# operator application
+Base.:*(L::BatchedDiagonalOperator, u::AbstractVecOrMat) = L.diag .* u
+Base.:\(L::BatchedDiagonalOperator, u::AbstractVecOrMat) = L.diag .\ u
+
+function LinearAlgebra.mul!(v::AbstractVecOrMat, L::BatchedDiagonalOperator, u::AbstractVecOrMat)
+    V = _vec(v)
+    U = _vec(u)
+    d = _vec(L.diag)
+    D = Diagonal(d)
+    mul!(V, D, U)
+
+    v
+end
+
+function LinearAlgebra.mul!(v::AbstractVecOrMat, L::BatchedDiagonalOperator, u::AbstractVecOrMat, α, β)
+    V = _vec(v)
+    U = _vec(u)
+    d = _vec(L.diag)
+    D = Diagonal(d)
+    mul!(V, D, U, α, β)
+
+    v
+end
+
+function LinearAlgebra.ldiv!(v::AbstractVecOrMat, L::BatchedDiagonalOperator, u::AbstractVecOrMat)
+    V = _vec(v)
+    U = _vec(u)
+    d = _vec(L.diag)
+    D = Diagonal(d)
+    ldiv!(V, D, U)
+
+    v
+end
+
+function LinearAlgebra.ldiv!(L::BatchedDiagonalOperator, u::AbstractVecOrMat)
+    U = _vec(u)
+    d = _vec(L.diag)
+    D = Diagonal(d)
+    ldiv!(D, U)
+
+    u
+end
+#

src/interface.jl

@@ -89,24 +89,8 @@ has_mul!(L::AbstractSciMLOperator) = false # mul!(du, L, u)
 has_ldiv(L::AbstractSciMLOperator) = false # du = L\u
 has_ldiv!(L::AbstractSciMLOperator) = false # ldiv!(du, L, u)
 
-### AbstractSciMLLinearOperator Interface
-
-#=
-1. AbstractSciMLLinearOperator <: AbstractSciMLOperator
-2. Can absorb under multiplication by a scalar. In all algorithms things like
-   dt*L show up all the time, so the linear operator must be able to absorb
-   such constants.
-4. isconstant(A) trait for whether the operator is constant or not.
-5. Optional: diagonal, symmetric, etc traits from LinearMaps.jl.
-6. Optional: exp(A). Required for simple exponential integration.
-7. Optional: expmv(A,u,p,t) = exp(t*A)*u and expmv!(v,A::SciMLOperator,u,p,t)
-   Required for sparse-saving exponential integration.
-8. Optional: factorizations. A_ldiv_B, factorize et. al. This is only required
-   for algorithms which use the factorization of the operator (Crank-Nicholson),
-   and only for when the default linear solve is used.
-=#
-
-# Extra standard assumptions
+### Extra standard assumptions
+
 isconstant(L) = true
 isconstant(L::AbstractSciMLOperator) = all(isconstant, getops(L))
 #isconstant(L::AbstractSciMLOperator) = L.update_func = DEFAULT_UPDATE_FUNC

src/matrix.jl

@@ -1,3 +1,4 @@
+#
 """
     MatrixOperator(A[; update_func])
 
@@ -32,7 +33,7 @@ for op in (
            :adjoint,
            :transpose,
           )
-    @eval function Base.$op(L::MatrixOperator)
+    @eval function Base.$op(L::MatrixOperator) # TODO - test this thoroughly
         MatrixOperator(
                        $op(L.A);
                        update_func= (A,u,p,t) -> $op(L.update_func($op(L.A),u,p,t)) # TODO - test
@@ -82,13 +83,30 @@ LinearAlgebra.mul!(v::AbstractVecOrMat, L::MatrixOperator, u::AbstractVecOrMat,
 LinearAlgebra.ldiv!(v::AbstractVecOrMat, L::MatrixOperator, u::AbstractVecOrMat) = ldiv!(v, L.A, u)
 LinearAlgebra.ldiv!(L::MatrixOperator, u::AbstractVecOrMat) = ldiv!(L.A, u)
 
-""" Diagonal Operator """
-function DiagonalOperator(u::AbstractVector; update_func=DEFAULT_UPDATE_FUNC)
+"""
+    DiagonalOperator(diag, [; update_func])
+
+Represents a time-dependent elementwise scaling (diagonal-scaling) operation.
+The update function is called by `update_coefficients!` and is assumed to have
+the following signature:
+
+    update_func(diag::AbstractVector,u,p,t) -> [modifies diag]
+
+When `diag` is an `AbstractVector` of length N, `L=DiagonalOpeator(diag, ...)`
+can be applied to `AbstractArray`s with `size(u, 1) == N`. Each column of the `u`
+will be scaled by `diag`, as in `LinearAlgebra.Diagonal(diag) * u`.
+
+When `diag` is a multidimensional array, `L = DiagonalOperator(diag, ...)` forms
+an operator of size `(N, N)` where `N = size(diag, 1)` is the leading length of `diag`.
+`L` then is the elementwise-scaling operation on arrays of `length(u) = length(diag)`
+with leading length `size(u, 1) = N`.
+"""
+function DiagonalOperator(diag::AbstractVector; update_func=DEFAULT_UPDATE_FUNC)
     function diag_update_func(A, u, p, t)
         update_func(A.diag, u, p, t)
         A
     end
-    MatrixOperator(Diagonal(u); update_func=diag_update_func)
+    MatrixOperator(Diagonal(diag); update_func=diag_update_func)
 end
 LinearAlgebra.Diagonal(L::MatrixOperator) = MatrixOperator(Diagonal(L.A))
 
@@ -205,7 +223,7 @@ end
 
 function AffineOperator(A::Union{AbstractMatrix,AbstractSciMLOperator},
                         B::Union{AbstractMatrix,AbstractSciMLOperator},
-                        b::AbstractVecOrMat;
+                        b::AbstractArray;
                         update_func=DEFAULT_UPDATE_FUNC,
                        )
     @assert size(A, 1) == size(B, 1) "Dimension mismatch: A, B don't output vectors

src/utils.jl

@@ -21,4 +21,8 @@ function _mat_sizes(L::AbstractSciMLOperator, u::AbstractArray)
 
     size_in, size_out
 end
+
+dims(A) = length(size(A))
+dims(::AbstractArray{<:Any,N}) where{N} = N
+dims(::AbstractSciMLOperator) = 2
 #