From fb98c9dc47a9c5b0e84de9f8884ecae23f2e5ddb Mon Sep 17 00:00:00 2001 From: Sheehan Olver Date: Fri, 29 Nov 2024 23:25:06 +0000 Subject: [PATCH] BlockArrray extension --- Project.toml | 2 + ext/InfiniteArraysBlockArraysExt.jl | 113 ++++++++++++++++++++++++++++ 2 files changed, 115 insertions(+) diff --git a/Project.toml b/Project.toml index e2050fb..c6d3638 100644 --- a/Project.toml +++ b/Project.toml @@ -12,11 +12,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [weakdeps] BandedMatrices = "aae01518-5342-5314-be14-df237901396f" +BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [extensions] InfiniteArraysBandedMatricesExt = "BandedMatrices" +InfiniteArraysBlockArraysExt = "BlockArrays" InfiniteArraysDSPExt = "DSP" InfiniteArraysStatisticsExt = "Statistics" diff --git a/ext/InfiniteArraysBlockArraysExt.jl b/ext/InfiniteArraysBlockArraysExt.jl index 5c2388f..109fc69 100644 --- a/ext/InfiniteArraysBlockArraysExt.jl +++ b/ext/InfiniteArraysBlockArraysExt.jl @@ -1,4 +1,41 @@ module InfiniteArraysBlockArraysExt +using InfiniteArrays, BlockArrays +using InfiniteArrays.ArrayLayouts, InfiniteArrays.LazyArrays +import ArrayLayouts: sub_materialize + +const OneToInfCumsum = RangeCumsum{Int,OneToInf{Int}} + +BlockArrays.sortedunion(::AbstractVector{<:PosInfinity}, ::AbstractVector{<:PosInfinity}) = [∞] +function BlockArrays.sortedunion(::AbstractVector{<:PosInfinity}, b) + @assert isinf(length(b)) + b +end + +function BlockArrays.sortedunion(b, ::AbstractVector{<:PosInfinity}) + @assert isinf(length(b)) + b +end +BlockArrays.sortedunion(a::OneToInfCumsum, ::OneToInfCumsum) = a + +BlockArrays.blocklasts(a::InfRanges) = Fill(length(a),1) + +BlockArrays.findblock(::BlockedOneTo, ::RealInfinity) = Block(ℵ₀) + +function BlockArrays.sortedunion(a::Vcat{Int,1,<:Tuple{Union{Int,AbstractVector{Int}},<:AbstractRange}}, + b::Vcat{Int,1,<:Tuple{Union{Int,AbstractVector{Int}},<:AbstractRange}}) + @assert a == b # TODO: generailse? Not sure how to do so in a type stable fashion + a +end + +sizes_from_blocks(A::AbstractVector, ::Tuple{OneToInf{Int}}) = (map(length,A),) +length(::BlockedOneTo{Int,<:InfRanges}) = ℵ₀ + +const OneToInfBlocks = BlockedOneTo{Int,OneToInfCumsum} +const OneToBlocks = BlockedOneTo{Int,OneToCumsum} + +axes(a::OneToInfBlocks) = (a,) +axes(a::OneToBlocks) = (a,) + sub_materialize(_, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V sub_materialize(::AbstractBlockLayout, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V @@ -7,4 +44,80 @@ function sub_materialize(::PaddedColumns, v::AbstractVector{T}, ax::Tuple{Blocke BlockedVector(Vcat(sub_materialize(dat), Zeros{T}(∞)), ax) end +BlockArrays.dimlength(start, ::Infinity) = ℵ₀ + +function copy(bc::Broadcasted{<:BroadcastStyle,<:Any,typeof(*),<:Tuple{Ones{T,1,Tuple{OneToInfBlocks}},AbstractArray{V,N}}}) where {N,T,V} + a,b = bc.args + @assert bc.axes == axes(b) + convert(AbstractArray{promote_type(T,V),N}, b) +end + +function copy(bc::Broadcasted{<:BroadcastStyle,<:Any,typeof(*),<:Tuple{AbstractArray{T,N},Ones{V,1,Tuple{OneToInfBlocks}}}}) where {N,T,V} + a,b = bc.args + @assert bc.axes == axes(a) + convert(AbstractArray{promote_type(T,V),N}, a) +end + +_block_interlace_axes(::Int, ax::Tuple{BlockedOneTo{Int,OneToInf{Int}}}...) = (blockedrange(Fill(length(ax), ∞)),) + +_block_interlace_axes(nbc::Int, ax::NTuple{2,BlockedOneTo{Int,OneToInf{Int}}}...) = + (blockedrange(Fill(length(ax) ÷ nbc, ∞)),blockedrange(Fill(mod1(length(ax),nbc), ∞))) + +####### +# block broadcasted +###### + + +BroadcastStyle(::Type{<:SubArray{T,N,Arr,<:NTuple{N,BlockSlice{BlockRange{1,Tuple{II}}}},false}}) where {T,N,Arr<:BlockArray,II<:InfRanges} = + LazyArrayStyle{N}() + +# TODO: generalise following +BroadcastStyle(::Type{<:BlockArray{T,N,<:AbstractArray{<:AbstractArray{T,N},N},<:NTuple{N,BlockedOneTo{Int,<:InfRanges}}}}) where {T,N} = LazyArrayStyle{N}() +# BroadcastStyle(::Type{<:BlockedArray{T,N,<:AbstractArray{T,N},<:NTuple{N,BlockedOneTo{Int,<:InfRanges}}}}) where {T,N} = LazyArrayStyle{N}() +BroadcastStyle(::Type{<:BlockArray{T,N,<:AbstractArray{<:AbstractArray{T,N},N},<:NTuple{N,BlockedOneTo{Int,<:RangeCumsum{Int,<:InfRanges}}}}}) where {T,N} = LazyArrayStyle{N}() +# BroadcastStyle(::Type{<:BlockedArray{T,N,<:AbstractArray{T,N},<:NTuple{N,BlockedOneTo{Int,<:RangeCumsum{Int,<:InfRanges}}}}}) where {T,N} = LazyArrayStyle{N}() + + +# Block banded support + +sizes_from_blocks(A::Diagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.diag, 1), size.(A.diag,2) +sizes_from_blocks(A::Tridiagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.d, 1), size.(A.d,2) +sizes_from_blocks(A::Bidiagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.dv, 1), size.(A.dv,2) + + +axes_print_matrix_row(::NTuple{2,AbstractBlockedUnitRange}, io, X, A, i, ::AbstractVector{<:PosInfinity}, sep) = nothing + + +const BlockTriPertToeplitz{T} = BlockMatrix{T,Tridiagonal{Matrix{T},Vcat{Matrix{T},1,Tuple{Vector{Matrix{T}},Fill{Matrix{T},1,Tuple{OneToInf{Int}}}}}}, + NTuple{2,BlockedOneTo{Int,Vcat{Int,1,Tuple{Vector{Int},InfStepRange{Int,Int}}}}}} + +const BlockTridiagonalToeplitzLayout = BlockLayout{TridiagonalToeplitzLayout,DenseColumnMajor} + +function BlockTridiagonal(adjA::Adjoint{T,BlockTriPertToeplitz{T}}) where T + A = parent(adjA) + BlockTridiagonal(Matrix.(adjoint.(A.blocks.du)), + Matrix.(adjoint.(A.blocks.d)), + Matrix.(adjoint.(A.blocks.dl))) +end + +for op in (:-, :+) + @eval begin + function $op(A::BlockTriPertToeplitz{T}, λ::UniformScaling) where T + TV = promote_type(T,eltype(λ)) + BlockTridiagonal(Vcat(convert.(AbstractVector{Matrix{TV}}, A.blocks.dl.args)...), + Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.d, Ref(λ)).args)...), + Vcat(convert.(AbstractVector{Matrix{TV}}, A.blocks.du.args)...)) + end + function $op(λ::UniformScaling, A::BlockTriPertToeplitz{V}) where V + TV = promote_type(eltype(λ),V) + BlockTridiagonal(Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.dl.args))...), + Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, Ref(λ), A.blocks.d).args)...), + Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.du.args))...)) + end + $op(adjA::Adjoint{T,BlockTriPertToeplitz{T}}, λ::UniformScaling) where T = $op(BlockTridiagonal(adjA), λ) + $op(λ::UniformScaling, adjA::Adjoint{T,BlockTriPertToeplitz{T}}) where T = $op(λ, BlockTridiagonal(adjA)) + end +end + + end # module \ No newline at end of file