SciML / SciMLOperators.jl

SciMLOperators.jl: Matrix-Free Operators for the SciML Scientific Machine Learning Common Interface in Julia
https://docs.sciml.ai/SciMLOperators/stable
MIT License
42 stars 9 forks source link

Lazy concatenation / block diagonal forms of operators #161

Open gaurav-arya opened 1 year ago

gaurav-arya commented 1 year ago

Does there exist functionality analogous to those listed in https://julialinearalgebra.github.io/LinearMaps.jl/stable/types/#BlockMap-and-BlockDiagonalMap?

In particular, lazy hcat/vcat/hvcat/diagonal op formation, etc.

We should also revisit this line of code: https://github.com/SciML/SciMLOperators.jl/blob/31275d4e2b2f3671fb4c2811656dd555c4e1f24a/src/matrix.jl#L116

vpuri3 commented 1 year ago

No, we don't have that functionality yet. It is on my to-do list. Would make writing vector-calculus operators really convenient. And will be very useful for solving coupled PDE systems.

vpuri3 commented 1 year ago

Potential API:

abstract type Abstract BlockOperator{T} <: AbstractSciMLOperator{T} end

###
# Block Operator
###

struct BlockOperator{T, Ti, Tb}
  Mblocks:Ti
  Nblocks::Ti
  blocks::Tb
  size_info

  function BlockOperator(Mblocks::Integer, Mblocks::Integer, blocks)
    @assert Mblocks * Nblocks == length(blocks)

    size_consistent = true
    for 
    # loop over rows, cols, to ensure size consistency.
    end

    @assert size_consistent
    T = promote_type(eltype.(blocks)...)

    BlockOperator{T, typeof(M), typeof(blocks)}(Mblocks, Nblocks, blocks)
  end
end

function BlockOperator(A::AbstractVecOrMat{<:AbstractSciMLOperator})
  BlockOperator(size(A)..., vec(A))
end

function Base.:*(L::BlockOperator, u::AbstractVecOrMat)
end

L = BlockOperator([A1, A2])

###
# Block Diagonal
###

struct BlockDiagonal{T, Tb}
  blocks::Tb

  function BlockDiagonal(blocks... :: AbstractSciMLOperator)
    T = promote_type(eltype.(blocks)...)
    BlockDiagonal{T, typeof(blocks)}(blocks)
  end
end

function Base.:*(L::BlockOperator, u::AbstractVecOrMat)
  vcat(Tuple(B * u for B in L.blocks)...)
end
gaurav-arya commented 1 year ago

if we support sparse block patterns, then block diagonal becomes a special case

vpuri3 commented 1 year ago

Yeah for sparse patterns, we can have a constructor that accepts pairs (L, (m, n)). And a special constructor for BlockDiagonal.

vpuri3 commented 1 year ago

It might make sense to implement BlockDiagonal as its own class so we can define has_ldiv, has_ldiv! on it when each block is invertible.

gaurav-arya commented 1 year ago

Can always do const BlockDiagonal = BlockOperator{AA} where {AA<:Diagonal}, with AA imagined to be the block pattern

gaurav-arya commented 1 year ago

Also want to be mindful of code dup with kronecker products. $I \otimes A$ is a block diagonal matrix, and $A \otimes B = (A \otimes I)(I \otimes B)$, although that might not be the best perspective for efficiency

vpuri3 commented 1 year ago

I remember putting in short-circuits in TensorProductOperator for when the outer is IdentityOperator. But we can always just write an overload if our lazy block version is faster.

vpuri3 commented 1 year ago

BlockOperators also need to support lazy matrix algebra. So we can multiply blocks and compose their components. Example

G = BlockOperator([Dx, Dy])
T = DiagonalBlock([Tx, Ty])
Lapl = G' * T * G # assume sizes match
Lapl == Dx' * Tx * Dx + Dy' * Ty * Dy # true

so we need to overload

Base.:*(A::AbstractBlockOperator, B::AbstractBlockOperator)
vpuri3 commented 1 year ago

linking old block matrix issue here for completeness. https://github.com/SciML/SciMLOperators.jl/issues/45

vpuri3 commented 1 year ago

some inspiration: https://github.com/JuliaLinearAlgebra/LinearMaps.jl/blob/master/src/blockmap.jl