Skip to content

Commit a189f2f

Browse files
dlfivefiftydlfivefifty
authored
BlockArrray extension (#200)
* BlockArrray extension * Update InfiniteArraysBlockArraysExt.jl * Update InfiniteArraysBlockArraysExt.jl * add tests * Update test_infblock.jl * increase coverage * fix test * increase cov
1 parent 1a4f215 commit a189f2f

File tree

6 files changed

+315
-13
lines changed

6 files changed

+315
-13
lines changed

‎Project.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
1212

1313
[weakdeps]
1414
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
15+
BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
1516
DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
1617
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
1718

1819
[extensions]
1920
InfiniteArraysBandedMatricesExt = "BandedMatrices"
21+
InfiniteArraysBlockArraysExt = "BlockArrays"
2022
InfiniteArraysDSPExt = "DSP"
2123
InfiniteArraysStatisticsExt = "Statistics"
2224

‎ext/InfiniteArraysBlockArraysExt.jl‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,45 @@
11
module InfiniteArraysBlockArraysExt
2+
using InfiniteArrays, BlockArrays
3+
using InfiniteArrays.ArrayLayouts, InfiniteArrays.LazyArrays, InfiniteArrays.LinearAlgebra
4+
5+
import Base: length, size, axes, BroadcastStyle, copy, +, -
6+
import Base.Broadcast: Broadcasted
7+
import ArrayLayouts: sub_materialize, axes_print_matrix_row
8+
import InfiniteArrays: OneToInf, PosInfinity, InfRanges, RealInfinity, Infinity, InfStepRange, TridiagonalToeplitzLayout
9+
import BlockArrays: AbstractBlockLayout, sizes_from_blocks, BlockTridiagonal, OneToCumsum, BlockSlice, AbstractBlockedUnitRange,
10+
BlockLayout
11+
import LazyArrays: PaddedColumns, LazyArrayStyle
12+
13+
const OneToInfCumsum = RangeCumsum{Int,OneToInf{Int}}
14+
15+
BlockArrays.sortedunion(::AbstractVector{<:PosInfinity}, ::AbstractVector{<:PosInfinity}) = [∞]
16+
function BlockArrays.sortedunion(::AbstractVector{<:PosInfinity}, b)
17+
@assert isinf(length(b))
18+
b
19+
end
20+
21+
function BlockArrays.sortedunion(b, ::AbstractVector{<:PosInfinity})
22+
@assert isinf(length(b))
23+
b
24+
end
25+
BlockArrays.sortedunion(a::OneToInfCumsum, ::OneToInfCumsum) = a
26+
27+
BlockArrays.blocklasts(a::InfRanges) = Fill(length(a),1)
28+
29+
BlockArrays.findblock(::BlockedOneTo, ::RealInfinity) = Block(ℵ₀)
30+
31+
function BlockArrays.sortedunion(a::Vcat{Int,1,<:Tuple{Union{Int,AbstractVector{Int}},<:AbstractRange}},
32+
b::Vcat{Int,1,<:Tuple{Union{Int,AbstractVector{Int}},<:AbstractRange}})
33+
@assert a == b # TODO: generailse? Not sure how to do so in a type stable fashion
34+
a
35+
end
36+
37+
sizes_from_blocks(A::AbstractVector, ::Tuple{OneToInf{Int}}) = (map(length,A),)
38+
length(::BlockedOneTo{Int,<:InfRanges}) = ℵ₀
39+
40+
const OneToInfBlocks = BlockedOneTo{Int,OneToInfCumsum}
41+
axes(a::OneToInfBlocks) = (a,)
42+
243

344
sub_materialize(_, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V
445
sub_materialize(::AbstractBlockLayout, V, ::Tuple{BlockedOneTo{Int,<:InfRanges}}) = V
@@ -7,4 +48,76 @@ function sub_materialize(::PaddedColumns, v::AbstractVector{T}, ax::Tuple{Blocke
748
BlockedVector(Vcat(sub_materialize(dat), Zeros{T}(∞)), ax)
849
end
950

51+
# BlockArrays.dimlength(start, ::Infinity) = ℵ₀
52+
53+
function copy(bc::Broadcasted{<:BroadcastStyle,<:Any,typeof(*),<:Tuple{Ones{T,1,Tuple{OneToInfBlocks}},AbstractArray{V,N}}}) where {N,T,V}
54+
a,b = bc.args
55+
@assert bc.axes == axes(b)
56+
convert(AbstractArray{promote_type(T,V),N}, b)
57+
end
58+
59+
function copy(bc::Broadcasted{<:BroadcastStyle,<:Any,typeof(*),<:Tuple{AbstractArray{T,N},Ones{V,1,Tuple{OneToInfBlocks}}}}) where {N,T,V}
60+
a,b = bc.args
61+
@assert bc.axes == axes(a)
62+
convert(AbstractArray{promote_type(T,V),N}, a)
63+
end
64+
65+
66+
#######
67+
# block broadcasted
68+
######
69+
70+
71+
BroadcastStyle(::Type{<:SubArray{T,N,Arr,<:NTuple{N,BlockSlice{BlockRange{1,Tuple{II}}}},false}}) where {T,N,Arr<:BlockArray,II<:InfRanges} =
72+
LazyArrayStyle{N}()
73+
74+
# TODO: generalise following
75+
BroadcastStyle(::Type{<:BlockArray{T,N,<:AbstractArray{<:AbstractArray{T,N},N},<:NTuple{N,BlockedOneTo{Int,<:InfRanges}}}}) where {T,N} = LazyArrayStyle{N}()
76+
# BroadcastStyle(::Type{<:BlockedArray{T,N,<:AbstractArray{T,N},<:NTuple{N,BlockedOneTo{Int,<:InfRanges}}}}) where {T,N} = LazyArrayStyle{N}()
77+
BroadcastStyle(::Type{<:BlockArray{T,N,<:AbstractArray{<:AbstractArray{T,N},N},<:NTuple{N,BlockedOneTo{Int,<:RangeCumsum{Int,<:InfRanges}}}}}) where {T,N} = LazyArrayStyle{N}()
78+
# BroadcastStyle(::Type{<:BlockedArray{T,N,<:AbstractArray{T,N},<:NTuple{N,BlockedOneTo{Int,<:RangeCumsum{Int,<:InfRanges}}}}}) where {T,N} = LazyArrayStyle{N}()
79+
80+
81+
# Block banded support
82+
83+
sizes_from_blocks(A::Diagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.diag, 1), size.(A.diag,2)
84+
sizes_from_blocks(A::Tridiagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.d, 1), size.(A.d,2)
85+
sizes_from_blocks(A::Bidiagonal, ::NTuple{2,OneToInf{Int}}) = size.(A.dv, 1), size.(A.dv,2)
86+
87+
88+
axes_print_matrix_row(::NTuple{2,AbstractBlockedUnitRange}, io, X, A, i, ::AbstractVector{<:PosInfinity}, sep) = nothing
89+
90+
91+
const BlockTriPertToeplitz{T} = BlockMatrix{T,Tridiagonal{Matrix{T},Vcat{Matrix{T},1,Tuple{Vector{Matrix{T}},Fill{Matrix{T},1,Tuple{OneToInf{Int}}}}}},
92+
NTuple{2,BlockedOneTo{Int,Vcat{Int,1,Tuple{Vector{Int},InfStepRange{Int,Int}}}}}}
93+
94+
const BlockTridiagonalToeplitzLayout = BlockLayout{TridiagonalToeplitzLayout,DenseColumnMajor}
95+
96+
function BlockTridiagonal(adjA::Adjoint{T,BlockTriPertToeplitz{T}}) where T
97+
A = parent(adjA)
98+
BlockTridiagonal(Matrix.(adjoint.(A.blocks.du)),
99+
Matrix.(adjoint.(A.blocks.d)),
100+
Matrix.(adjoint.(A.blocks.dl)))
101+
end
102+
103+
for op in (:-, :+)
104+
@eval begin
105+
function $op(A::BlockTriPertToeplitz{T}, λ::UniformScaling) where T
106+
TV = promote_type(T,eltype(λ))
107+
BlockTridiagonal(Vcat(convert.(AbstractVector{Matrix{TV}}, A.blocks.dl.args)...),
108+
Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.d, Ref(λ)).args)...),
109+
Vcat(convert.(AbstractVector{Matrix{TV}}, A.blocks.du.args)...))
110+
end
111+
function $op(λ::UniformScaling, A::BlockTriPertToeplitz{V}) where V
112+
TV = promote_type(eltype(λ),V)
113+
BlockTridiagonal(Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.dl.args))...),
114+
Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, Ref(λ), A.blocks.d).args)...),
115+
Vcat(convert.(AbstractVector{Matrix{TV}}, broadcast($op, A.blocks.du.args))...))
116+
end
117+
$op(adjA::Adjoint{T,BlockTriPertToeplitz{T}}, λ::UniformScaling) where T = $op(BlockTridiagonal(adjA), λ)
118+
$op(λ::UniformScaling, adjA::Adjoint{T,BlockTriPertToeplitz{T}}) where T = $op(λ, BlockTridiagonal(adjA))
119+
end
120+
end
121+
122+
10123
end # module

‎test/runtests.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1245,8 +1245,6 @@ end
12451245
@test r[ind] == v
12461246
end
12471247

1248-
include("test_infconv.jl")
1249-
include("test_block.jl")
12501248

12511249
@testset "bounds-checking for StepRangeLen{<:CartesianIndex}" begin
12521250
if VERSION >= v"1.11.0-rc3"
@@ -1259,9 +1257,11 @@ end
12591257
a = [1+im,2+im]
12601258
A = a * Ones{Complex{Int}}(1,∞)
12611259
@test A[:,1:5] == a * ones(1,5)
1262-
1260+
12631261
@test (a*permutedims(1:∞))[:,1:5] == a*(1:5)'
12641262
@test (a*Hcat(Zeros(1,2), permutedims(1:∞)))[1,1:5] == (a*Vcat(Hcat(Zeros(1,2), permutedims(1:∞))))[1,1:5]
12651263
end
12661264

1265+
include("test_infconv.jl")
1266+
include("test_infblock.jl")
12671267
include("test_infbanded.jl")

‎test/test_block.jl‎

Lines changed: 0 additions & 10 deletions
This file was deleted.

‎test/test_infblock.jl‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
using InfiniteArrays, BlockArrays, ArrayLayouts, LazyArrays, Test
2+
using InfiniteArrays: OneToInf, RealInfinity
3+
using Base: oneto
4+
using LazyArrays: LazyArrayStyle
5+
using BlockArrays: BlockTridiagonal
6+
7+
const InfiniteArraysBlockArraysExt = Base.get_extension(InfiniteArrays, :InfiniteArraysBlockArraysExt)
8+
const BlockTriPertToeplitz = InfiniteArraysBlockArraysExt.BlockTriPertToeplitz
9+
10+
11+
12+
@testset "∞-block arrays" begin
13+
@testset "blockedonetoinf" begin
14+
b = blockedrange(OneToInf())
15+
b2 = b .+ b;
16+
for i in 1:10
17+
@test b2[Block(i)] == b[Block(i)] + b[Block(i)]
18+
end
19+
end
20+
21+
@testset "fixed block size" begin
22+
k = Base.OneTo.(oneto(∞))
23+
n = Fill.(oneto(∞), oneto(∞))
24+
@test broadcast(length, k) ≡ map(length, k) ≡ OneToInf()
25+
@test broadcast(length, n) ≡ map(length, n) ≡ OneToInf()
26+
27+
b = mortar(Fill([1, 2], ∞))
28+
@test blockaxes(b, 1) ≡ Block.(OneToInf())
29+
@test b[Block(5)] == [1, 2]
30+
@test b[Block.(2:∞)][Block.(2:10)] == b[Block.(3:11)]
31+
@test exp.(b)[Block.(2:∞)][Block.(2:10)] == exp.(b[Block.(3:11)])
32+
33+
@test blockedrange(Vcat(2, Fill(3, ∞))) isa BlockedOneTo{<:Any,<:InfiniteArrays.InfStepRange}
34+
35+
c = BlockedArray(1:∞, Vcat(2, Fill(3, ∞)))
36+
@test c[Block.(2:∞)][Block.(2:10)] == c[Block.(3:11)]
37+
38+
@test length(axes(b, 1)) ≡ ℵ₀
39+
@test last(axes(b, 1)) ≡ ℵ₀
40+
@test Base.BroadcastStyle(typeof(b)) isa LazyArrayStyle{1}
41+
end
42+
43+
@testset "1:∞ blocks" begin
44+
a = blockedrange(oneto(∞))
45+
@test axes(a, 1) == a
46+
o = Ones((a,))
47+
@test Base.BroadcastStyle(typeof(a)) isa LazyArrayStyle{1}
48+
b = exp.(a)
49+
@test axes(b, 1) == a
50+
@test o .* b isa typeof(b)
51+
@test b .* o isa typeof(b)
52+
end
53+
54+
@testset "padded" begin
55+
c = BlockedArray([1; zeros(∞)], Vcat(2, Fill(3, ∞)))
56+
@test c + c isa BlockedVector
57+
end
58+
59+
60+
@testset "triangle recurrences" begin
61+
@testset "n and k" begin
62+
n = mortar(Fill.(oneto(∞), oneto(∞)))
63+
k = mortar(Base.OneTo.(oneto(∞)))
64+
65+
@test n[Block(5)] ≡ layout_getindex(n, Block(5)) ≡ view(n, Block(5)) ≡ Fill(5, 5)
66+
@test k[Block(5)] ≡ layout_getindex(k, Block(5)) ≡ view(k, Block(5)) ≡ Base.OneTo(5)
67+
@test Base.BroadcastStyle(typeof(n)) isa LazyArrays.LazyArrayStyle{1}
68+
@test Base.BroadcastStyle(typeof(k)) isa LazyArrays.LazyArrayStyle{1}
69+
70+
N = 1000
71+
v = view(n, Block.(Base.OneTo(N)))
72+
@test view(v, Block(2)) ≡ Fill(2, 2)
73+
@test axes(v) isa Tuple{BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,Base.OneTo{Int64}}}}
74+
@test @allocated(axes(v)) ≤ 40
75+
76+
dest = BlockedArray{Float64}(undef, axes(v))
77+
@test copyto!(dest, v) == v
78+
@test @allocated(copyto!(dest, v)) ≤ 40
79+
80+
v = view(k, Block.(Base.OneTo(N)))
81+
@test view(v, Block(2)) ≡ Base.OneTo(2)
82+
@test axes(v) isa Tuple{BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,Base.OneTo{Int64}}}}
83+
@test @allocated(axes(v)) ≤ 40
84+
@test copyto!(dest, v) == v
85+
86+
@testset "stack overflow" begin
87+
i = Base.to_indices(k, (Block.(2:∞),))[1].indices
88+
@test last(i) == ℵ₀
89+
end
90+
91+
v = view(k, Block.(2:∞))
92+
@test Base.BroadcastStyle(typeof(v)) isa LazyArrayStyle{1}
93+
@test v[Block(1)] == 1:2
94+
@test v[Block(1)] ≡ k[Block(2)] ≡ Base.OneTo(2)
95+
96+
@test axes(n, 1) isa BlockedOneTo{Int,ArrayLayouts.RangeCumsum{Int64,OneToInf{Int64}}}
97+
end
98+
end
99+
100+
@testset "blockdiag" begin
101+
D = Diagonal(mortar(Fill.((-(0:∞) - (0:∞) .^ 2), 1:2:∞)))
102+
x = [randn(5); zeros(∞)]
103+
x̃ = BlockedArray(x, (axes(D, 1),))
104+
@test (D*x)[1:10] == (D*x̃)[1:10]
105+
end
106+
107+
@testset "sortedunion" begin
108+
a = cumsum(1:2:∞)
109+
@test BlockArrays.sortedunion(a, a) ≡ a
110+
@test BlockArrays.sortedunion([∞], a) ≡ BlockArrays.sortedunion(a, [∞]) ≡ a
111+
@test BlockArrays.sortedunion([∞], [∞]) == [∞]
112+
113+
b = Vcat([1, 2], 3:∞)
114+
c = Vcat(1, 3:∞)
115+
@test BlockArrays.sortedunion(b, b) ≡ b
116+
@test BlockArrays.sortedunion(c, c) ≡ c
117+
end
118+
119+
@testset "Algebra" begin
120+
@testset "Triangle OP recurrences" begin
121+
k = mortar(Base.OneTo.(1:∞))
122+
n = mortar(Fill.(1:∞, 1:∞))
123+
@test k[Block.(2:3)] isa BlockArray
124+
@test n[Block.(2:3)] isa BlockArray
125+
@test k[Block.(2:3)] == [1, 2, 1, 2, 3]
126+
@test n[Block.(2:3)] == [2, 2, 3, 3, 3]
127+
@test blocksize(BroadcastVector(exp, k)) == (ℵ₀,)
128+
@test BroadcastVector(exp, k)[Block.(2:3)] == exp.([1, 2, 1, 2, 3])
129+
# BroadcastVector(+,k,n)
130+
end
131+
# Multivariate OPs Corollary (3)
132+
# n = 5
133+
# BlockTridiagonal(Zeros.(1:∞,2:∞),
134+
# (n -> Diagonal(((n+2).+(0:n)))/ (2n + 2)).(0:∞),
135+
# Zeros.(2:∞,1:∞))
136+
end
137+
138+
@testset "findblock at +∞, HarmonicOrthogonalPolynomials#88" begin
139+
@test findblock(blockedrange(1:2:∞), RealInfinity()) == Block(ℵ₀)
140+
end
141+
142+
@testset "BlockTridiagonal Pert Toeplitz" begin
143+
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), ∞)),
144+
Vcat([zeros(1, 1)], Fill(zeros(2, 2), ∞)),
145+
Vcat([fill(1.0, 1, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)))
146+
147+
@test A isa BlockTriPertToeplitz
148+
@test A[Block.(1:2), Block(1)] == A[1:3, 1:1] == reshape([0.0, 1.0, 1.0], 3, 1)
149+
150+
@test (A-I)[1:100, 1:100] == A[1:100, 1:100] - I
151+
@test (A+I)[1:100, 1:100] == A[1:100, 1:100] + I
152+
@test (I+A)[1:100, 1:100] == I + A[1:100, 1:100]
153+
@test (I-A)[1:100, 1:100] == I - A[1:100, 1:100]
154+
155+
@test (A-im*I)[1:100, 1:100] == A[1:100, 1:100] - im * I
156+
@test (A+im*I)[1:100, 1:100] == A[1:100, 1:100] + im * I
157+
@test (im*I+A)[1:100, 1:100] == im * I + A[1:100, 1:100]
158+
@test (im*I-A)[1:100, 1:100] == im * I - A[1:100, 1:100]
159+
160+
@test (A'-I)[1:100, 1:100] == A'[1:100, 1:100] - I
161+
@test (A'+I)[1:100, 1:100] == A'[1:100, 1:100] + I
162+
@test (I+A')[1:100, 1:100] == I + A'[1:100, 1:100]
163+
@test (I-A')[1:100, 1:100] == I - A'[1:100, 1:100]
164+
165+
@test BlockTridiagonal(A')[Block.(1:10),Block.(1:10)] == A[Block.(1:10),Block.(1:10)]'
166+
end
167+
168+
@testset "Bi/Diagonal mortar" begin
169+
A = mortar(Diagonal(Fill([1 2; 3 4], ∞)))
170+
@test A[Block(1,1)] == [1 2; 3 4]
171+
172+
B = mortar(Bidiagonal(Fill([1 2; 3 4], ∞), Fill([1 2; 3 4], ∞), :U))
173+
@test B[Block(1,1)] == [1 2; 3 4]
174+
end
175+
end

‎test/test_infblockbanded.jl‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
2+
@testset "BlockBanded" begin
3+
a = b = c = 0.0
4+
n = mortar(Fill.(oneto(∞), oneto(∞)))
5+
k = mortar(Base.OneTo.(oneto(∞)))
6+
Dy = BlockBandedMatrices._BandedBlockBandedMatrix((k .+ (b + c))', axes(k, 1), (-1, 1), (-1, 1))
7+
N = 100
8+
@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))
9+
@test colsupport(Dy, axes(Dy,2)) == 1:∞
10+
@test rowsupport(Dy, axes(Dy,1)) == 2:∞
11+
end
12+
13+
14+
15+
@testset "BlockTridiagonal" begin
16+
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), ∞)),
17+
Vcat([zeros(1, 1)], Fill(zeros(2, 2), ∞)),
18+
Vcat([fill(1.0, 1, 2), Matrix(1.0I, 2, 2)], Fill(Matrix(1.0I, 2, 2), ∞)))
19+
20+
@test isblockbanded(A)
21+
@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]
22+
end

0 commit comments

Comments
 (0)