Batched Operators

[vpuri3]

In solving nonlinear PDEs, you often have to form the product, u * u where u is the solution vector. This operation can be cast as a a SciMLOperator as follows:

function quadratic_update_func(L, u, p, t)
    copy!(L.A.diag, u)
end

D = MatrixOperator(Diagonal(zero(u)); update_func=quadratic_update_func)

For a single trajectory, the size of D would be (N,N) where N is the problem size. However, when you are evolving K trajectories at the same time the size of solution vector is (N,K). When you form D by the above approach, Diagonal(zero(u)) returns a Diagonal matrix of size (N,N) or (K,K) whichever's smaller, not respecting the intention of intention of our operation.

julia> Diagonal(rand(2,3))
2×2 Diagonal{Float64, Vector{Float64}}:
 0.166451   ⋅ 
  ⋅        0.0533492

julia> Diagonal(rand(3,2))
2×2 Diagonal{Float64, Vector{Float64}}:
 0.780478   ⋅ 
  ⋅        0.239675

For this reason, I have restricted DiagonalOperator to accept only AbstractVector arguments. You could indeed pass a vec(u) to Diagonal, but then the size of D would be (N*K, N*K) which would be incompatible with your other operations of size (N,N).

As diagonal scaling is a fundamental operation, I think we should address this problem. What we need is an elementwise-scaling operator that has size (N,N) but can act on objects of size (N,K) (SciMLOperators.jl has convention that the input to an (M,N) operator must be an AbstractArray of size (N,...)). I propose creating an object called BatchedDiagonalOperator.

struct BatchedDiagonalOperator{T,D,F} <: AbstractSciMLOperator
    diag::D
    update_func::F
end

Base.size(L::BatchedDiagonalOperator) = (size(L.diag, 1), size(L.diag, 1))

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

This should solve the problem of performing nonlinear operations with intra-ODE parallelism

[vpuri3]

We could have just added a couple of vecs to the definition of DiagonalOperator but that would change Base.size of the operator itself and then we'd have to store metadata somewhere. The only option is to make a new operator type.

[vpuri3]

AffineOperator, the only other operator that stores some solution vector type, doesn't need to be changed as it can have a matrix type as L.b because it doesn't affect its Base.size.