Skip to content

Commit f67770f

Browse files
ogautheogauthe
authored
Test contract with BlockArrays and BlockedArray (#34)
1 parent 9407722 commit f67770f

File tree

3 files changed

+232
-3
lines changed

3 files changed

+232
-3
lines changed

‎Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "TensorAlgebra"
22
uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
33
authors = ["ITensor developers <[email protected]> and contributors"]
4-
version = "0.1.10"
4+
version = "0.1.11"
55

66
[deps]
77
ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

‎test/test_blockarrays_contract.jl

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
using BlockArrays: Block, BlockArray, BlockedArray, blockedrange, blocksize
2+
using BlockSparseArrays: BlockSparseArray
3+
using SparseArraysBase: densearray
4+
using TensorAlgebra: contract
5+
using Random: randn!
6+
using Test: @test, @test_broken, @testset
7+
8+
function randn_blockdiagonal(elt::Type, axes::Tuple)
9+
a = BlockSparseArray{elt}(axes)
10+
blockdiaglength = minimum(blocksize(a))
11+
for i in 1:blockdiaglength
12+
b = Block(ntuple(Returns(i), ndims(a)))
13+
a[b] = randn!(a[b])
14+
end
15+
return a
16+
end
17+
18+
const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
19+
@testset "`contract` blocked arrays (eltype=$elt)" for elt in elts
20+
d = blockedrange([2, 3])
21+
a1_sba = randn_blockdiagonal(elt, (d, d, d, d))
22+
a2_sba = randn_blockdiagonal(elt, (d, d, d, d))
23+
a3_sba = randn_blockdiagonal(elt, (d, d))
24+
a1_dense = densearray(a1_sba)
25+
a2_dense = densearray(a2_sba)
26+
a3_dense = densearray(a3_sba)
27+
28+
@testset "BlockArray" begin
29+
a1 = BlockArray(a1_sba)
30+
a2 = BlockArray(a2_sba)
31+
a3 = BlockArray(a3_sba)
32+
33+
# matrix matrix
34+
@test_broken a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4))
35+
#=
36+
a_dest_dense, dimnames_dest_dense = contract(
37+
a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4)
38+
)
39+
@test dimnames_dest == dimnames_dest_dense
40+
@test size(a_dest) == size(a_dest_dense)
41+
@test a_dest isa BlockArray
42+
@test a_dest ≈ a_dest_dense
43+
=#
44+
45+
# matrix vector
46+
@test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2))
47+
#=
48+
a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2))
49+
@test dimnames_dest == dimnames_dest_dense
50+
@test size(a_dest) == size(a_dest_dense)
51+
@test a_dest isa BlockArray
52+
@test a_dest ≈ a_dest_dense
53+
=#
54+
55+
# vector matrix
56+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1))
57+
#=
58+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1))
59+
@test dimnames_dest == dimnames_dest_dense
60+
@test size(a_dest) == size(a_dest_dense)
61+
@test a_dest isa BlockArray
62+
@test a_dest ≈ a_dest_dense
63+
=#
64+
65+
# vector vector
66+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1))
67+
a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1))
68+
@test dimnames_dest == dimnames_dest_dense
69+
@test size(a_dest) == size(a_dest_dense)
70+
@test_broken a_dest isa BlockArray # TBD relax to AbstractArray{elt,0}?
71+
@test a_dest ≈ a_dest_dense
72+
73+
# outer product
74+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4))
75+
#=
76+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4))
77+
@test dimnames_dest == dimnames_dest_dense
78+
@test size(a_dest) == size(a_dest_dense)
79+
@test a_dest isa BlockArray
80+
@test a_dest ≈ a_dest_dense
81+
=#
82+
end
83+
84+
@testset "BlockedArray" begin
85+
a1 = BlockedArray(a1_sba)
86+
a2 = BlockedArray(a2_sba)
87+
a3 = BlockedArray(a3_sba)
88+
89+
# matrix matrix
90+
a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4))
91+
a_dest_dense, dimnames_dest_dense = contract(
92+
a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4)
93+
)
94+
@test dimnames_dest == dimnames_dest_dense
95+
@test size(a_dest) == size(a_dest_dense)
96+
@test a_dest isa BlockedArray
97+
@test a_dest ≈ a_dest_dense
98+
99+
# matrix vector
100+
@test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2))
101+
#=
102+
a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2))
103+
@test dimnames_dest == dimnames_dest_dense
104+
@test size(a_dest) == size(a_dest_dense)
105+
@test a_dest isa BlockedArray
106+
@test a_dest ≈ a_dest_dense
107+
=#
108+
109+
# vector matrix
110+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1))
111+
#=
112+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1))
113+
@test dimnames_dest == dimnames_dest_dense
114+
@test size(a_dest) == size(a_dest_dense)
115+
@test a_dest isa BlockedArray
116+
@test a_dest ≈ a_dest_dense
117+
=#
118+
119+
# vector vector
120+
a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1))
121+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1))
122+
@test dimnames_dest == dimnames_dest_dense
123+
@test size(a_dest) == size(a_dest_dense)
124+
@test_broken a_dest isa BlockedArray # TBD relax to AbstractArray{elt,0}?
125+
@test a_dest ≈ a_dest_dense
126+
127+
# outer product
128+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4))
129+
a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4))
130+
@test dimnames_dest == dimnames_dest_dense
131+
@test size(a_dest) == size(a_dest_dense)
132+
@test a_dest isa BlockedArray
133+
@test a_dest ≈ a_dest_dense
134+
end
135+
136+
@testset "BlockSparseArray" begin
137+
a1, a2, a3 = a1_sba, a2_sba, a3_sba
138+
139+
# matrix matrix
140+
a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4))
141+
a_dest_dense, dimnames_dest_dense = contract(
142+
a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4)
143+
)
144+
@test dimnames_dest == dimnames_dest_dense
145+
@test size(a_dest) == size(a_dest_dense)
146+
@test a_dest isa BlockSparseArray
147+
@test a_dest ≈ a_dest_dense
148+
149+
# matrix vector
150+
@test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2))
151+
#=
152+
a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2))
153+
@test dimnames_dest == dimnames_dest_dense
154+
@test size(a_dest) == size(a_dest_dense)
155+
@test a_dest isa BlockSparseArray
156+
@test a_dest ≈ a_dest_dense
157+
=#
158+
159+
# vector matrix
160+
a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1))
161+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1))
162+
@test dimnames_dest == dimnames_dest_dense
163+
@test size(a_dest) == size(a_dest_dense)
164+
@test a_dest isa BlockSparseArray
165+
@test a_dest ≈ a_dest_dense
166+
167+
# vector vector
168+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1))
169+
a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1))
170+
@test dimnames_dest == dimnames_dest_dense
171+
@test size(a_dest) == size(a_dest_dense)
172+
@test a_dest isa BlockSparseArray
173+
@test a_dest ≈ a_dest_dense
174+
175+
# outer product
176+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4))
177+
a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4))
178+
@test dimnames_dest == dimnames_dest_dense
179+
@test size(a_dest) == size(a_dest_dense)
180+
@test a_dest isa BlockSparseArray
181+
@test a_dest ≈ a_dest_dense
182+
end
183+
end

‎test/test_gradedunitrangesext_contract.jl

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ using SparseArraysBase: densearray
55
using SymmetrySectors: U1
66
using TensorAlgebra: contract
77
using Random: randn!
8-
using Test: @test, @testset
8+
using Test: @test, @test_broken, @testset
99

1010
function randn_blockdiagonal(elt::Type, axes::Tuple)
1111
a = BlockSparseArray{elt}(axes)
@@ -22,12 +22,58 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
2222
d = gradedrange([U1(0) => 2, U1(1) => 3])
2323
a1 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d)))
2424
a2 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d)))
25-
a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4))
25+
a3 = randn_blockdiagonal(elt, (d, dual(d)))
2626
a1_dense = densearray(a1)
2727
a2_dense = densearray(a2)
28+
a3_dense = densearray(a3)
29+
30+
# matrix matrix
31+
a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4))
2832
a_dest_dense, dimnames_dest_dense = contract(
2933
a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4)
3034
)
3135
@test dimnames_dest == dimnames_dest_dense
36+
@test size(a_dest) == size(a_dest_dense)
37+
@test a_dest isa BlockSparseArray
38+
@test a_dest ≈ a_dest_dense
39+
40+
# matrix vector
41+
@test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2))
42+
#=
43+
a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2))
44+
@test dimnames_dest == dimnames_dest_dense
45+
@test size(a_dest) == size(a_dest_dense)
46+
@test a_dest isa BlockSparseArray
47+
@test a_dest ≈ a_dest_dense
48+
=#
49+
50+
# vector matrix
51+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1))
52+
#=
53+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1))
54+
@test dimnames_dest == dimnames_dest_dense
55+
@test size(a_dest) == size(a_dest_dense)
56+
@test a_dest isa BlockSparseArray
57+
@test a_dest ≈ a_dest_dense
58+
=#
59+
60+
# vector vector
61+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1))
62+
#=
63+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1))
64+
@test dimnames_dest == dimnames_dest_dense
65+
@test size(a_dest) == size(a_dest_dense)
66+
@test a_dest isa BlockSparseArray
67+
@test a_dest ≈ a_dest_dense
68+
=#
69+
70+
# outer product
71+
@test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4))
72+
#=
73+
a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4))
74+
@test dimnames_dest == dimnames_dest_dense
75+
@test size(a_dest) == size(a_dest_dense)
76+
@test a_dest isa BlockSparseArray
3277
@test a_dest ≈ a_dest_dense
78+
=#
3379
end

0 commit comments

Comments
 (0)