BlockArrray extension

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"
 

ext/InfiniteArraysBlockArraysExt.jl

@@ -1,4 +1,45 @@
 module InfiniteArraysBlockArraysExt
+using InfiniteArrays, BlockArrays
+using InfiniteArrays.ArrayLayouts, InfiniteArrays.LazyArrays, InfiniteArrays.LinearAlgebra
+
+import Base: length, size, axes, BroadcastStyle, copy, +, -
+import Base.Broadcast: Broadcasted
+import ArrayLayouts: sub_materialize, axes_print_matrix_row
+import InfiniteArrays: OneToInf, PosInfinity, InfRanges, RealInfinity, Infinity, InfStepRange, TridiagonalToeplitzLayout
+import BlockArrays: AbstractBlockLayout, sizes_from_blocks, BlockTridiagonal, OneToCumsum, BlockSlice, AbstractBlockedUnitRange,
+                    BlockLayout
+import LazyArrays: PaddedColumns, LazyArrayStyle
+
+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}
+axes(a::OneToInfBlocks) = (a,)
+
 
 sub_materialize(_, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V
 sub_materialize(::AbstractBlockLayout, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V
@@ -7,4 +48,76 @@
     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 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

test/runtests.jl

@@ -1245,8 +1245,6 @@ end
     @test r[ind] == v
 end
 
-include("test_infconv.jl")
-include("test_block.jl")
 
 @testset "bounds-checking for StepRangeLen{<:CartesianIndex}" begin
     if VERSION >= v"1.11.0-rc3"
@@ -1259,9 +1257,11 @@ end
     a = [1+im,2+im]
     A = a * Ones{Complex{Int}}(1,∞)
     @test A[:,1:5] == a * ones(1,5)
-
+    
     @test (a*permutedims(1:∞))[:,1:5] == a*(1:5)'
     @test (a*Hcat(Zeros(1,2), permutedims(1:∞)))[1,1:5] == (a*Vcat(Hcat(Zeros(1,2), permutedims(1:∞))))[1,1:5]
 end
 
+include("test_infconv.jl")
+include("test_infblock.jl")
 include("test_infbanded.jl")

test/test_infblock.jl

@@ -0,0 +1,175 @@
+using InfiniteArrays, BlockArrays, ArrayLayouts, LazyArrays, Test
+using InfiniteArrays: OneToInf, RealInfinity
+using Base: oneto
+using LazyArrays: LazyArrayStyle
+using BlockArrays: BlockTridiagonal
+
+const InfiniteArraysBlockArraysExt = Base.get_extension(InfiniteArrays, :InfiniteArraysBlockArraysExt)
+const BlockTriPertToeplitz = InfiniteArraysBlockArraysExt.BlockTriPertToeplitz
+
+
+
+@testset "∞-block arrays" begin
+    @testset "blockedonetoinf" begin
+        b = blockedrange(OneToInf())
+        b2 = b .+ b;
+        for i in 1:10
+            @test b2[Block(i)] == b[Block(i)] + b[Block(i)]
+        end
+    end
+
+    @testset "fixed block size" begin
+        k = Base.OneTo.(oneto(∞))
+        n = Fill.(oneto(∞), oneto(∞))
+        @test broadcast(length, k) ≡ map(length, k) ≡ OneToInf()
+        @test broadcast(length, n) ≡ map(length, n) ≡ OneToInf()
+
+        b = mortar(Fill([1, 2], ∞))
+        @test blockaxes(b, 1) ≡ Block.(OneToInf())
+        @test b[Block(5)] == [1, 2]
+        @test b[Block.(2:∞)][Block.(2:10)] == b[Block.(3:11)]
+        @test exp.(b)[Block.(2:∞)][Block.(2:10)] == exp.(b[Block.(3:11)])
+
+        @test blockedrange(Vcat(2, Fill(3, ∞))) isa BlockedOneTo{<:Any,<:InfiniteArrays.InfStepRange}
+
+        c = BlockedArray(1:∞, Vcat(2, Fill(3, ∞)))
+        @test c[Block.(2:∞)][Block.(2:10)] == c[Block.(3:11)]
+
+        @test length(axes(b, 1)) ≡ ℵ₀
+        @test last(axes(b, 1)) ≡ ℵ₀
+        @test Base.BroadcastStyle(typeof(b)) isa LazyArrayStyle{1}
+    end
+
+    @testset "1:∞ blocks" begin
+        a = blockedrange(oneto(∞))
+        @test axes(a, 1) == a
+        o = Ones((a,))
+        @test Base.BroadcastStyle(typeof(a)) isa LazyArrayStyle{1}
+        b = exp.(a)
+        @test axes(b, 1) == a
+        @test o .* b isa typeof(b)
+        @test b .* o isa typeof(b)
+    end
+
+    @testset "padded" begin
+        c = BlockedArray([1; zeros(∞)], Vcat(2, Fill(3, ∞)))
+        @test c + c isa BlockedVector
+    end
+
+
+    @testset "triangle recurrences" begin
+        @testset "n and k" begin
+            n = mortar(Fill.(oneto(∞), oneto(∞)))
+            k = mortar(Base.OneTo.(oneto(∞)))
+
+            @test n[Block(5)] ≡ layout_getindex(n, Block(5)) ≡ view(n, Block(5)) ≡ Fill(5, 5)
+            @test k[Block(5)] ≡ layout_getindex(k, Block(5)) ≡ view(k, Block(5)) ≡ Base.OneTo(5)
+            @test Base.BroadcastStyle(typeof(n)) isa LazyArrays.LazyArrayStyle{1}
+            @test Base.BroadcastStyle(typeof(k)) isa LazyArrays.LazyArrayStyle{1}
+
+            N = 1000
+            v = view(n, Block.(Base.OneTo(N)))
+            @test view(v, Block(2)) ≡ Fill(2, 2)
+            @test axes(v) isa Tuple{BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,Base.OneTo{Int64}}}}
+            @test @allocated(axes(v)) ≤ 40
+
+            dest = BlockedArray{Float64}(undef, axes(v))
+            @test copyto!(dest, v) == v
+            @test @allocated(copyto!(dest, v)) ≤ 40
+
+            v = view(k, Block.(Base.OneTo(N)))
+            @test view(v, Block(2)) ≡ Base.OneTo(2)
+            @test axes(v) isa Tuple{BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,Base.OneTo{Int64}}}}
+            @test @allocated(axes(v)) ≤ 40
+            @test copyto!(dest, v) == v
+
+            @testset "stack overflow" begin
+                i = Base.to_indices(k, (Block.(2:∞),))[1].indices
+                @test last(i) == ℵ₀
+            end
+
+            v = view(k, Block.(2:∞))
+            @test Base.BroadcastStyle(typeof(v)) isa LazyArrayStyle{1}
+            @test v[Block(1)] == 1:2
+            @test v[Block(1)] ≡ k[Block(2)] ≡ Base.OneTo(2)
+
+            @test axes(n, 1) isa BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,OneToInf{Int64}}}
+        end
+    end
+
+    @testset "blockdiag" begin
+        D = Diagonal(mortar(Fill.((-(0:∞) - (0:∞) .^ 2), 1:2:∞)))
+        x = [randn(5); zeros(∞)]
+        x̃ = BlockedArray(x, (axes(D, 1),))
+        @test (D*x)[1:10] == (D*x̃)[1:10]
+    end
+
+    @testset "sortedunion" begin
+        a = cumsum(1:2:∞)
+        @test BlockArrays.sortedunion(a, a) ≡ a
+        @test BlockArrays.sortedunion([∞], a) ≡ BlockArrays.sortedunion(a, [∞]) ≡ a
+        @test BlockArrays.sortedunion([∞], [∞]) == [∞]
+
+        b = Vcat([1, 2], 3:∞)
+        c = Vcat(1, 3:∞)
+        @test BlockArrays.sortedunion(b, b) ≡ b
+        @test BlockArrays.sortedunion(c, c) ≡ c
+    end
+
+    @testset "Algebra" begin
+        @testset "Triangle OP recurrences" begin
+            k = mortar(Base.OneTo.(1:∞))
+            n = mortar(Fill.(1:∞, 1:∞))
+            @test k[Block.(2:3)] isa BlockArray
+            @test n[Block.(2:3)] isa BlockArray
+            @test k[Block.(2:3)] == [1, 2, 1, 2, 3]
+            @test n[Block.(2:3)] == [2, 2, 3, 3, 3]
+            @test blocksize(BroadcastVector(exp, k)) == (ℵ₀,)
+            @test BroadcastVector(exp, k)[Block.(2:3)] == exp.([1, 2, 1, 2, 3])
+            # BroadcastVector(+,k,n)
+        end
+        # Multivariate OPs Corollary (3)
+        # n = 5
+        # BlockTridiagonal(Zeros.(1:∞,2:∞),
+        #         (n -> Diagonal(((n+2).+(0:n)))/ (2n + 2)).(0:∞),
+        #         Zeros.(2:∞,1:∞))
+    end
+
+    @testset "findblock at +∞, HarmonicOrthogonalPolynomials#88" begin
+        @test findblock(blockedrange(1:2:∞), RealInfinity()) == Block(ℵ₀)
+    end
+
+    @testset "BlockTridiagonal Pert Toeplitz" begin
+        A = BlockTridiagonal(Vcat([fill(1.0, 2, 1), Matrix(1.0I, 2, 2), Matrix(1.0I, 2, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)),
+                                Vcat([zeros(1, 1)], Fill(zeros(2, 2), ∞)),
+                                Vcat([fill(1.0, 1, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)))
+
+        @test A isa BlockTriPertToeplitz
+        @test A[Block.(1:2), Block(1)] == A[1:3, 1:1] == reshape([0.0, 1.0, 1.0], 3, 1)
+
+        @test (A-I)[1:100, 1:100] == A[1:100, 1:100] - I
+        @test (A+I)[1:100, 1:100] == A[1:100, 1:100] + I
+        @test (I+A)[1:100, 1:100] == I + A[1:100, 1:100]
+        @test (I-A)[1:100, 1:100] == I - A[1:100, 1:100]
+
+        @test (A-im*I)[1:100, 1:100] == A[1:100, 1:100] - im * I
+        @test (A+im*I)[1:100, 1:100] == A[1:100, 1:100] + im * I
+        @test (im*I+A)[1:100, 1:100] == im * I + A[1:100, 1:100]
+        @test (im*I-A)[1:100, 1:100] == im * I - A[1:100, 1:100]
+
+        @test (A'-I)[1:100, 1:100] == A'[1:100, 1:100] - I
+        @test (A'+I)[1:100, 1:100] == A'[1:100, 1:100] + I
+        @test (I+A')[1:100, 1:100] == I + A'[1:100, 1:100]
+        @test (I-A')[1:100, 1:100] == I - A'[1:100, 1:100]
+
+        @test BlockTridiagonal(A')[Block.(1:10),Block.(1:10)] == A[Block.(1:10),Block.(1:10)]'
+    end
+
+    @testset "Bi/Diagonal mortar" begin
+        A = mortar(Diagonal(Fill([1 2; 3 4], ∞)))
+        @test A[Block(1,1)] == [1 2; 3 4]
+
+        B = mortar(Bidiagonal(Fill([1 2; 3 4], ∞), Fill([1 2; 3 4], ∞), :U))
+        @test B[Block(1,1)] == [1 2; 3 4]
+    end
+end

test/test_infblockbanded.jl

@@ -0,0 +1,22 @@
+
+@testset "BlockBanded" begin
+            a = b = c = 0.0
+            n = mortar(Fill.(oneto(∞), oneto(∞)))
+            k = mortar(Base.OneTo.(oneto(∞)))
+            Dy = BlockBandedMatrices._BandedBlockBandedMatrix((k .+ (b + c))', axes(k, 1), (-1, 1), (-1, 1))
+            N = 100
+            @test Dy[Block.(1:N), Block.(1:N)] == BlockBandedMatrices._BandedBlockBandedMatrix((k.+(b+c))[Block.(1:N)]', axes(k, 1)[Block.(1:N)], (-1, 1), (-1, 1))
+            @test colsupport(Dy, axes(Dy,2)) == 1:∞
+            @test rowsupport(Dy, axes(Dy,1)) == 2:∞
+        end
+
+
+
+        @testset "BlockTridiagonal" begin
+            A = BlockTridiagonal(Vcat([fill(1.0, 2, 1), Matrix(1.0I, 2, 2), Matrix(1.0I, 2, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)),
+                Vcat([zeros(1, 1)], Fill(zeros(2, 2), ∞)),
+                Vcat([fill(1.0, 1, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)))
+
+            @test isblockbanded(A)    
+            @test BlockBandedMatrix(A)[1:100, 1:100] == BlockBandedMatrix(A, (2, 1))[1:100, 1:100] == BlockBandedMatrix(A, (1, 1))[1:100, 1:100] == A[1:100, 1:100]
+        end