Mercurial > repos > public > sbplib_julia

comparison test/LazyTensors/lazy_tensor_operations_test.jl @ 1207:f1c2a4fa0ee1 performance/get_region_type_inference

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 03 Feb 2023 22:14:47 +0100
parents c94a12327737
children 4c0bc52e170f
comparison
equal deleted inserted replaced
919:b41180efb6c2 1207:f1c2a4fa0ee1
2 using Sbplib.LazyTensors 2 using Sbplib.LazyTensors
3 using Sbplib.RegionIndices 3 using Sbplib.RegionIndices
4 4
5 using Tullio 5 using Tullio
6 6
7 struct DummyMapping{T,R,D} <: LazyTensor{T,R,D} end
8
9 LazyTensors.apply(m::DummyMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = :apply
10 LazyTensors.apply_transpose(m::DummyMapping{T,R,D}, v, I::Vararg{Any,D}) where {T,R,D} = :apply_transpose
11
12 LazyTensors.range_size(m::DummyMapping) = :range_size
13 LazyTensors.domain_size(m::DummyMapping) = :domain_size
14
15
16 struct SizeDoublingMapping{T,R,D} <: LazyTensor{T,R,D}
17 domain_size::NTuple{D,Int}
18 end
19
20 LazyTensors.apply(m::SizeDoublingMapping{T,R}, v, i::Vararg{Any,R}) where {T,R} = (:apply,v,i)
21 LazyTensors.range_size(m::SizeDoublingMapping) = 2 .* m.domain_size
22 LazyTensors.domain_size(m::SizeDoublingMapping) = m.domain_size
23
24
25
7 @testset "Mapping transpose" begin 26 @testset "Mapping transpose" begin
8 struct DummyMapping{T,R,D} <: TensorMapping{T,R,D} end
9
10 LazyTensors.apply(m::DummyMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = :apply
11 LazyTensors.apply_transpose(m::DummyMapping{T,R,D}, v, I::Vararg{Any,D}) where {T,R,D} = :apply_transpose
12
13 LazyTensors.range_size(m::DummyMapping) = :range_size
14 LazyTensors.domain_size(m::DummyMapping) = :domain_size
15
16 m = DummyMapping{Float64,2,3}() 27 m = DummyMapping{Float64,2,3}()
17 @test m' isa TensorMapping{Float64, 3,2} 28 @test m' isa LazyTensor{Float64, 3,2}
18 @test m'' == m 29 @test m'' == m
19 @test apply(m',zeros(Float64,(0,0)), 0, 0, 0) == :apply_transpose 30 @test apply(m',zeros(Float64,(0,0)), 0, 0, 0) == :apply_transpose
20 @test apply(m'',zeros(Float64,(0,0,0)), 0, 0) == :apply 31 @test apply(m'',zeros(Float64,(0,0,0)), 0, 0) == :apply
21 @test apply_transpose(m', zeros(Float64,(0,0,0)), 0, 0) == :apply 32 @test apply_transpose(m', zeros(Float64,(0,0,0)), 0, 0) == :apply
22 33
23 @test range_size(m') == :domain_size 34 @test range_size(m') == :domain_size
24 @test domain_size(m') == :range_size 35 @test domain_size(m') == :range_size
25 end 36 end
26 37
38
27 @testset "TensorApplication" begin 39 @testset "TensorApplication" begin
28 struct SizeDoublingMapping{T,R,D} <: TensorMapping{T,R,D}
29 domain_size::NTuple{D,Int}
30 end
31
32 LazyTensors.apply(m::SizeDoublingMapping{T,R}, v, i::Vararg{Any,R}) where {T,R} = (:apply,v,i)
33 LazyTensors.range_size(m::SizeDoublingMapping) = 2 .* m.domain_size
34 LazyTensors.domain_size(m::SizeDoublingMapping) = m.domain_size
35
36
37 m = SizeDoublingMapping{Int, 1, 1}((3,)) 40 m = SizeDoublingMapping{Int, 1, 1}((3,))
41 mm = SizeDoublingMapping{Int, 1, 1}((6,))
38 v = [0,1,2] 42 v = [0,1,2]
39 @test m*v isa AbstractVector{Int}
40 @test size(m*v) == 2 .*size(v) 43 @test size(m*v) == 2 .*size(v)
41 @test (m*v)[0] == (:apply,v,(0,)) 44 @test (m*v)[1] == (:apply,v,(1,))
42 @test m*m*v isa AbstractVector{Int} 45 @test (mm*m*v)[1] == (:apply,m*v,(1,))
43 @test (m*m*v)[1] == (:apply,m*v,(1,)) 46 @test (mm*m*v)[3] == (:apply,m*v,(3,))
44 @test (m*m*v)[3] == (:apply,m*v,(3,)) 47 @test (mm*m*v)[6] == (:apply,m*v,(6,))
45 @test (m*m*v)[6] == (:apply,m*v,(6,))
46 @test_broken BoundsError == (m*m*v)[0]
47 @test_broken BoundsError == (m*m*v)[7]
48 @test_throws MethodError m*m 48 @test_throws MethodError m*m
49 49
50 m = SizeDoublingMapping{Int, 2, 1}((3,)) 50 @test (m*v)[CartesianIndex(2)] == (:apply,v,(2,))
51 @test_throws MethodError m*ones(Int,2,2) 51 @test (mm*m*v)[CartesianIndex(2)] == (:apply,m*v,(2,))
52 @test_throws MethodError m*m*v
53 52
54 m = SizeDoublingMapping{Float64, 2, 2}((3,3)) 53 m = SizeDoublingMapping{Float64, 2, 2}((3,3))
54 mm = SizeDoublingMapping{Float64, 2, 2}((6,6))
55 v = ones(3,3) 55 v = ones(3,3)
56 @test size(m*v) == 2 .*size(v) 56 @test size(m*v) == 2 .*size(v)
57 @test (m*v)[1,2] == (:apply,v,(1,2)) 57 @test (m*v)[1,2] == (:apply,v,(1,2))
58 58
59 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 59 @test (m*v)[CartesianIndex(2,3)] == (:apply,v,(2,3))
60 λ::T 60 @test (mm*m*v)[CartesianIndex(4,3)] == (:apply,m*v,(4,3))
61 size::NTuple{D,Int} 61
62 end 62 m = ScalingTensor(2,(3,))
63
64 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
65 LazyTensors.range_size(m::ScalingOperator) = m.size
66 LazyTensors.domain_size(m::ScalingOperator) = m.size
67
68 m = ScalingOperator{Int,1}(2,(3,))
69 v = [1,2,3] 63 v = [1,2,3]
70 @test m*v isa AbstractVector 64 @test m*v isa AbstractVector
71 @test m*v == [2,4,6] 65 @test m*v == [2,4,6]
72 66
73 m = ScalingOperator{Int,2}(2,(2,2)) 67 m = ScalingTensor(2,(2,2))
74 v = [[1 2];[3 4]] 68 v = [[1 2];[3 4]]
75 @test m*v == [[2 4];[6 8]] 69 @test m*v == [[2 4];[6 8]]
76 @test (m*v)[2,1] == 6 70 @test (m*v)[2,1] == 6
77 end 71
78 72 @testset "Error on index out of bounds" begin
79 @testset "TensorMapping binary operations" begin 73 m = SizeDoublingMapping{Int, 1, 1}((3,))
80 struct ScalarMapping{T,R,D} <: TensorMapping{T,R,D} 74 v = [0,1,2]
81 λ::T 75
82 range_size::NTuple{R,Int} 76 @test_throws BoundsError (m*v)[0]
83 domain_size::NTuple{D,Int} 77 @test_throws BoundsError (m*v)[7]
84 end 78 end
85 79
86 LazyTensors.apply(m::ScalarMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = m.λ*v[I...] 80 @testset "Error on unmatched dimensions" begin
87 LazyTensors.range_size(m::ScalarMapping) = m.domain_size 81 v = [0,1,2]
88 LazyTensors.domain_size(m::ScalarMapping) = m.range_size 82 m = SizeDoublingMapping{Int, 2, 1}((3,))
89 83 @test_throws MethodError m*ones(Int,2,2)
90 A = ScalarMapping{Float64,1,1}(2.0, (3,), (3,)) 84 @test_throws MethodError m*m*v
91 B = ScalarMapping{Float64,1,1}(3.0, (3,), (3,)) 85 end
86
87 @testset "Error on unmatched sizes" begin
88 @test_throws DomainSizeMismatch ScalingTensor(2,(2,))*ones(3)
89 @test_throws DomainSizeMismatch ScalingTensor(2,(2,))*ScalingTensor(2,(3,))*ones(3)
90 end
91
92
93 @testset "Type calculation" begin
94 m = ScalingTensor(2,(3,))
95 v = [1.,2.,3.]
96 @test m*v isa AbstractVector{Float64}
97 @test m*v == [2.,4.,6.]
98 @inferred m*v
99 @inferred (m*v)[1]
100
101 m = ScalingTensor(2,(2,2))
102 v = [[1. 2.];[3. 4.]]
103 @test m*v == [[2. 4.];[6. 8.]]
104 @test (m*v)[2,1] == 6.
105 @inferred m*v
106 @inferred (m*v)[1]
107
108 m = ScalingTensor(2. +2. *im,(3,))
109 v = [1.,2.,3.]
110 @test m*v isa AbstractVector{ComplexF64}
111 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
112 @inferred m*v
113 @inferred (m*v)[1]
114
115 m = ScalingTensor(1,(3,))
116 v = [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
117 @test m*v isa AbstractVector{ComplexF64}
118 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
119 @inferred m*v
120 @inferred (m*v)[1]
121
122 m = ScalingTensor(2., (3,))
123 v = [[1,2,3], [3,2,1],[1,3,1]]
124 @test m*v isa AbstractVector{Vector{Float64}}
125 @test m*v == [[2.,4.,6.], [6.,4.,2.],[2.,6.,2.]]
126 @inferred m*v
127 @inferred (m*v)[1]
128 end
129 end
130
131
132 @testset "LazyTensor binary operations" begin
133 A = ScalingTensor(2.0, (3,))
134 B = ScalingTensor(3.0, (3,))
92 135
93 v = [1.1,1.2,1.3] 136 v = [1.1,1.2,1.3]
94 for i ∈ eachindex(v) 137 for i ∈ eachindex(v)
95 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i] 138 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i]
96 end 139 end
97 140
98 for i ∈ eachindex(v) 141 for i ∈ eachindex(v)
99 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i] 142 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i]
100 end 143 end
101 144
145
102 @test range_size(A+B) == range_size(A) == range_size(B) 146 @test range_size(A+B) == range_size(A) == range_size(B)
103 @test domain_size(A+B) == domain_size(A) == domain_size(B) 147 @test domain_size(A+B) == domain_size(A) == domain_size(B)
104 end 148
105 149 @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64
106 150
107 @testset "TensorMappingComposition" begin 151 @testset "Error on unmatched sizes" begin
152 @test_throws Union{DomainSizeMismatch, RangeSizeMismatch} ScalingTensor(2.0, (3,)) + ScalingTensor(2.0, (4,))
153
154 @test_throws DomainSizeMismatch ScalingTensor(2.0, (4,)) + SizeDoublingMapping{Float64,1,1}((2,))
155 @test_throws DomainSizeMismatch SizeDoublingMapping{Float64,1,1}((2,)) + ScalingTensor(2.0, (4,))
156 @test_throws RangeSizeMismatch ScalingTensor(2.0, (2,)) + SizeDoublingMapping{Float64,1,1}((2,))
157 @test_throws RangeSizeMismatch SizeDoublingMapping{Float64,1,1}((2,)) + ScalingTensor(2.0, (2,))
158 end
159 end
160
161
162 @testset "TensorComposition" begin
108 A = rand(2,3) 163 A = rand(2,3)
109 B = rand(3,4) 164 B = rand(3,4)
110 165
111 Ã = LazyLinearMap(A, (1,), (2,)) 166 Ã = DenseTensor(A, (1,), (2,))
112 B̃ = LazyLinearMap(B, (1,), (2,)) 167 B̃ = DenseTensor(B, (1,), (2,))
113 168
114 @test Ã∘B̃ isa TensorMappingComposition 169 @test Ã∘B̃ isa TensorComposition
115 @test range_size(Ã∘B̃) == (2,) 170 @test range_size(Ã∘B̃) == (2,)
116 @test domain_size(Ã∘B̃) == (4,) 171 @test domain_size(Ã∘B̃) == (4,)
117 @test_throws SizeMismatch B̃∘Ã 172 @test_throws DomainSizeMismatch B̃∘Ã
118 173
119 # @test @inbounds B̃∘Ã # Should not error even though dimensions don't match. (Since ]test runs with forced boundschecking this is currently not testable 2020-10-16) 174 # @test @inbounds B̃∘Ã # Should not error even though dimensions don't match. (Since ]test runs with forced boundschecking this is currently not testable 2020-10-16)
120 175
121 v = rand(4) 176 v = rand(4)
122 @test Ã∘B̃*v ≈ A*B*v rtol=1e-14 177 @test Ã∘B̃*v ≈ A*B*v rtol=1e-14
123 178
124 v = rand(2) 179 v = rand(2)
125 @test (Ã∘B̃)'*v ≈ B'*A'*v rtol=1e-14 180 @test (Ã∘B̃)'*v ≈ B'*A'*v rtol=1e-14
126 end 181
127 182 @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64
128 @testset "LazyLinearMap" begin 183 @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64
129 # Test a standard matrix-vector product 184
130 # mapping vectors of size 4 to vectors of size 3. 185 a = 2.
131 A = rand(3,4) 186 v = rand(3)
132 Ã = LazyLinearMap(A, (1,), (2,)) 187 @test a*Ã isa TensorComposition
133 v = rand(4) 188 @test a*Ã == Ã*a
134 w = rand(3) 189 @test range_size(a*Ã) == range_size(Ã)
135 190 @test domain_size(a*Ã) == domain_size(Ã)
136 @test à isa LazyLinearMap{T,1,1} where T 191 @test a*Ã*v == a.*A*v
137 @test à isa TensorMapping{T,1,1} where T 192 end
138 @test range_size(Ã) == (3,) 193
139 @test domain_size(Ã) == (4,) 194
140 195 @testset "InflatedTensor" begin
141 @test Ã*ones(4) ≈ A*ones(4) atol=5e-13 196 I(sz...) = IdentityTensor(sz...)
142 @test Ã*v ≈ A*v atol=5e-13
143 @test Ã'*w ≈ A'*w
144
145 A = rand(2,3,4)
146 @test_throws DomainError LazyLinearMap(A, (3,1), (2,))
147
148 # Test more exotic mappings
149 B = rand(3,4,2)
150 # Map vectors of size 2 to matrices of size (3,4)
151 B̃ = LazyLinearMap(B, (1,2), (3,))
152 v = rand(2)
153
154 @test range_size(B̃) == (3,4)
155 @test domain_size(B̃) == (2,)
156 @test B̃ isa TensorMapping{T,2,1} where T
157 @test B̃*ones(2) ≈ B[:,:,1] + B[:,:,2] atol=5e-13
158 @test B̃*v ≈ B[:,:,1]*v[1] + B[:,:,2]*v[2] atol=5e-13
159
160 # Map matrices of size (3,2) to vectors of size 4
161 B̃ = LazyLinearMap(B, (2,), (1,3))
162 v = rand(3,2)
163
164 @test range_size(B̃) == (4,)
165 @test domain_size(B̃) == (3,2)
166 @test B̃ isa TensorMapping{T,1,2} where T
167 @test B̃*ones(3,2) ≈ B[1,:,1] + B[2,:,1] + B[3,:,1] +
168 B[1,:,2] + B[2,:,2] + B[3,:,2] atol=5e-13
169 @test B̃*v ≈ B[1,:,1]*v[1,1] + B[2,:,1]*v[2,1] + B[3,:,1]*v[3,1] +
170 B[1,:,2]v[1,2] + B[2,:,2]*v[2,2] + B[3,:,2]*v[3,2] atol=5e-13
171
172
173 # TODO:
174 # @inferred (B̃*v)[2]
175 end
176
177
178 @testset "IdentityMapping" begin
179 @test IdentityMapping{Float64}((4,5)) isa IdentityMapping{T,2} where T
180 @test IdentityMapping{Float64}((4,5)) isa TensorMapping{T,2,2} where T
181 @test IdentityMapping{Float64}((4,5)) == IdentityMapping{Float64}(4,5)
182
183 @test IdentityMapping(3,2) isa IdentityMapping{Float64,2}
184
185 for sz ∈ [(4,5),(3,),(5,6,4)]
186 I = IdentityMapping{Float64}(sz)
187 v = rand(sz...)
188 @test I*v == v
189 @test I'*v == v
190
191 @test range_size(I) == sz
192 @test domain_size(I) == sz
193 end
194
195 I = IdentityMapping{Float64}((4,5))
196 v = rand(4,5)
197 @inferred (I*v)[3,2]
198 @inferred (I'*v)[3,2]
199 @inferred range_size(I)
200
201 @inferred range_dim(I)
202 @inferred domain_dim(I)
203
204 Ã = rand(4,2)
205 A = LazyLinearMap(Ã,(1,),(2,))
206 I1 = IdentityMapping{Float64}(2)
207 I2 = IdentityMapping{Float64}(4)
208 @test A∘I1 == A
209 @test I2∘A == A
210 @test I1∘I1 == I1
211 @test_throws SizeMismatch I1∘A
212 @test_throws SizeMismatch A∘I2
213 @test_throws SizeMismatch I1∘I2
214 end
215
216 @testset "InflatedTensorMapping" begin
217 I(sz...) = IdentityMapping(sz...)
218 197
219 Ã = rand(4,2) 198 Ã = rand(4,2)
220 B̃ = rand(4,2,3) 199 B̃ = rand(4,2,3)
221 C̃ = rand(4,2,3) 200 C̃ = rand(4,2,3)
222 201
223 A = LazyLinearMap(Ã,(1,),(2,)) 202 A = DenseTensor(Ã,(1,),(2,))
224 B = LazyLinearMap(B̃,(1,2),(3,)) 203 B = DenseTensor(B̃,(1,2),(3,))
225 C = LazyLinearMap(C̃,(1,),(2,3)) 204 C = DenseTensor(C̃,(1,),(2,3))
226 205
227 @testset "Constructors" begin 206 @testset "Constructors" begin
228 @test InflatedTensorMapping(I(3,2), A, I(4)) isa TensorMapping{Float64, 4, 4} 207 @test InflatedTensor(I(3,2), A, I(4)) isa LazyTensor{Float64, 4, 4}
229 @test InflatedTensorMapping(I(3,2), B, I(4)) isa TensorMapping{Float64, 5, 4} 208 @test InflatedTensor(I(3,2), B, I(4)) isa LazyTensor{Float64, 5, 4}
230 @test InflatedTensorMapping(I(3), C, I(2,3)) isa TensorMapping{Float64, 4, 5} 209 @test InflatedTensor(I(3), C, I(2,3)) isa LazyTensor{Float64, 4, 5}
231 @test InflatedTensorMapping(C, I(2,3)) isa TensorMapping{Float64, 3, 4} 210 @test InflatedTensor(C, I(2,3)) isa LazyTensor{Float64, 3, 4}
232 @test InflatedTensorMapping(I(3), C) isa TensorMapping{Float64, 2, 3} 211 @test InflatedTensor(I(3), C) isa LazyTensor{Float64, 2, 3}
233 @test InflatedTensorMapping(I(3), I(2,3)) isa TensorMapping{Float64, 3, 3} 212 @test InflatedTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3}
234 end 213 end
235 214
236 @testset "Range and domain size" begin 215 @testset "Range and domain size" begin
237 @test range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 216 @test range_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,4,4)
238 @test domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 217 @test domain_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,2,4)
239 218
240 @test range_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,4,2,4) 219 @test range_size(InflatedTensor(I(3,2), B, I(4))) == (3,2,4,2,4)
241 @test domain_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,3,4) 220 @test domain_size(InflatedTensor(I(3,2), B, I(4))) == (3,2,3,4)
242 221
243 @test range_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,4,2,3) 222 @test range_size(InflatedTensor(I(3), C, I(2,3))) == (3,4,2,3)
244 @test domain_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,2,3,2,3) 223 @test domain_size(InflatedTensor(I(3), C, I(2,3))) == (3,2,3,2,3)
245 224
246 @inferred range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 225 @inferred range_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,4,4)
247 @inferred domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 226 @inferred domain_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,2,4)
248 end 227 end
249 228
250 @testset "Application" begin 229 @testset "Application" begin
251 # Testing regular application and transposed application with inflation "before", "after" and "before and after". 230 # Testing regular application and transposed application with inflation "before", "after" and "before and after".
252 # The inflated tensor mappings are chosen to preserve, reduce and increase the dimension of the result compared to the input. 231 # The inflated tensor mappings are chosen to preserve, reduce and increase the dimension of the result compared to the input.
253 tests = [ 232 cases = [
254 ( 233 (
255 InflatedTensorMapping(I(3,2), A, I(4)), 234 InflatedTensor(I(3,2), A, I(4)),
256 (v-> @tullio res[a,b,c,d] := Ã[c,i]*v[a,b,i,d]), # Expected result of apply 235 (v-> @tullio res[a,b,c,d] := Ã[c,i]*v[a,b,i,d]), # Expected result of apply
257 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose 236 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose
258 ), 237 ),
259 ( 238 (
260 InflatedTensorMapping(I(3,2), B, I(4)), 239 InflatedTensor(I(3,2), B, I(4)),
261 (v-> @tullio res[a,b,c,d,e] := B̃[c,d,i]*v[a,b,i,e]), 240 (v-> @tullio res[a,b,c,d,e] := B̃[c,d,i]*v[a,b,i,e]),
262 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]), 241 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]),
263 ), 242 ),
264 ( 243 (
265 InflatedTensorMapping(I(3,2), C, I(4)), 244 InflatedTensor(I(3,2), C, I(4)),
266 (v-> @tullio res[a,b,c,d] := C̃[c,i,j]*v[a,b,i,j,d]), 245 (v-> @tullio res[a,b,c,d] := C̃[c,i,j]*v[a,b,i,j,d]),
267 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]), 246 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]),
268 ), 247 ),
269 ( 248 (
270 InflatedTensorMapping(I(3,2), A), 249 InflatedTensor(I(3,2), A),
271 (v-> @tullio res[a,b,c] := Ã[c,i]*v[a,b,i]), 250 (v-> @tullio res[a,b,c] := Ã[c,i]*v[a,b,i]),
272 (v-> @tullio res[a,b,c] := Ã[i,c]*v[a,b,i]), 251 (v-> @tullio res[a,b,c] := Ã[i,c]*v[a,b,i]),
273 ), 252 ),
274 ( 253 (
275 InflatedTensorMapping(I(3,2), B), 254 InflatedTensor(I(3,2), B),
276 (v-> @tullio res[a,b,c,d] := B̃[c,d,i]*v[a,b,i]), 255 (v-> @tullio res[a,b,c,d] := B̃[c,d,i]*v[a,b,i]),
277 (v-> @tullio res[a,b,c] := B̃[i,j,c]*v[a,b,i,j]), 256 (v-> @tullio res[a,b,c] := B̃[i,j,c]*v[a,b,i,j]),
278 ), 257 ),
279 ( 258 (
280 InflatedTensorMapping(I(3,2), C), 259 InflatedTensor(I(3,2), C),
281 (v-> @tullio res[a,b,c] := C̃[c,i,j]*v[a,b,i,j]), 260 (v-> @tullio res[a,b,c] := C̃[c,i,j]*v[a,b,i,j]),
282 (v-> @tullio res[a,b,c,d] := C̃[i,c,d]*v[a,b,i]), 261 (v-> @tullio res[a,b,c,d] := C̃[i,c,d]*v[a,b,i]),
283 ), 262 ),
284 ( 263 (
285 InflatedTensorMapping(A,I(4)), 264 InflatedTensor(A,I(4)),
286 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]), 265 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]),
287 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]), 266 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]),
288 ), 267 ),
289 ( 268 (
290 InflatedTensorMapping(B,I(4)), 269 InflatedTensor(B,I(4)),
291 (v-> @tullio res[a,b,c] := B̃[a,b,i]*v[i,c]), 270 (v-> @tullio res[a,b,c] := B̃[a,b,i]*v[i,c]),
292 (v-> @tullio res[a,b] := B̃[i,j,a]*v[i,j,b]), 271 (v-> @tullio res[a,b] := B̃[i,j,a]*v[i,j,b]),
293 ), 272 ),
294 ( 273 (
295 InflatedTensorMapping(C,I(4)), 274 InflatedTensor(C,I(4)),
296 (v-> @tullio res[a,b] := C̃[a,i,j]*v[i,j,b]), 275 (v-> @tullio res[a,b] := C̃[a,i,j]*v[i,j,b]),
297 (v-> @tullio res[a,b,c] := C̃[i,a,b]*v[i,c]), 276 (v-> @tullio res[a,b,c] := C̃[i,a,b]*v[i,c]),
298 ), 277 ),
299 ] 278 ]
300 279
301 @testset "apply" begin 280 @testset "$tm" for (tm, true_apply, true_apply_transpose) ∈ cases
302 for i ∈ 1:length(tests) 281 v = rand(domain_size(tm)...)
303 tm = tests[i][1] 282 @test tm*v ≈ true_apply(v) rtol=1e-14
304 v = rand(domain_size(tm)...) 283
305 true_value = tests[i][2](v) 284 v = rand(range_size(tm)...)
306 @test tm*v ≈ true_value rtol=1e-14 285 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14
307 end
308 end 286 end
309 287
310 @testset "apply_transpose" begin 288 @testset "application to other type" begin
311 for i ∈ 1:length(tests) 289 tm = InflatedTensor(I(3,2), A, I(4))
312 tm = tests[i][1] 290
313 v = rand(range_size(tm)...) 291 v = rand(ComplexF64, domain_size(tm)...)
314 true_value = tests[i][3](v) 292 @test (tm*v)[1,2,3,1] isa ComplexF64
315 @test tm'*v ≈ true_value rtol=1e-14 293
316 end 294 v = rand(ComplexF64, domain_size(tm')...)
295 @test (tm'*v)[1,2,2,1] isa ComplexF64
317 end 296 end
318 297
319 @testset "Inference of application" begin 298 @testset "Inference of application" begin
320 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 299 tm = InflatedTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4))
321 λ::T
322 size::NTuple{D,Int}
323 end
324
325 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
326 LazyTensors.range_size(m::ScalingOperator) = m.size
327 LazyTensors.domain_size(m::ScalingOperator) = m.size
328
329 tm = InflatedTensorMapping(I(2,3),ScalingOperator(2.0, (3,2)),I(3,4))
330 v = rand(domain_size(tm)...) 300 v = rand(domain_size(tm)...)
331 301
332 @inferred apply(tm,v,1,2,3,2,2,4) 302 @inferred apply(tm,v,1,2,3,2,2,4)
333 @inferred (tm*v)[1,2,3,2,2,4] 303 @inferred (tm*v)[1,2,3,2,2,4]
334 end 304 end
335 end 305 end
336 306
337 @testset "InflatedTensorMapping of InflatedTensorMapping" begin 307 @testset "InflatedTensor of InflatedTensor" begin
338 A = ScalingOperator(2.0,(2,3)) 308 A = ScalingTensor(2.0,(2,3))
339 itm = InflatedTensorMapping(I(3,2), A, I(4)) 309 itm = InflatedTensor(I(3,2), A, I(4))
340 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2)) 310 @test InflatedTensor(I(4), itm, I(2)) == InflatedTensor(I(4,3,2), A, I(4,2))
341 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2)) 311 @test InflatedTensor(itm, I(2)) == InflatedTensor(I(3,2), A, I(4,2))
342 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4)) 312 @test InflatedTensor(I(4), itm) == InflatedTensor(I(4,3,2), A, I(4))
343 313
344 @test InflatedTensorMapping(I(2), I(2), I(2)) isa InflatedTensorMapping # The constructor should always return its type. 314 @test InflatedTensor(I(2), I(2), I(2)) isa InflatedTensor # The constructor should always return its type.
345 end 315 end
346 end 316 end
347
348 @testset "split_index" begin
349 @test LazyTensors.split_index(Val(2),Val(1),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,5,6),(3,4))
350 @test LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,:,:,5,6),(3,4))
351 @test LazyTensors.split_index(Val(3),Val(1),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,5,6),(4,))
352 @test LazyTensors.split_index(Val(3),Val(2),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,:,5,6),(4,))
353 @test LazyTensors.split_index(Val(1),Val(1),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,4,5,6),(2,3))
354 @test LazyTensors.split_index(Val(1),Val(2),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,:,4,5,6),(2,3))
355
356 @test LazyTensors.split_index(Val(0),Val(1),Val(3),Val(3),1,2,3,4,5,6) == ((:,4,5,6),(1,2,3))
357 @test LazyTensors.split_index(Val(3),Val(1),Val(3),Val(0),1,2,3,4,5,6) == ((1,2,3,:),(4,5,6))
358
359 @inferred LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,2,2,4)
360 end
361
362 @testset "slice_tuple" begin
363 @test LazyTensors.slice_tuple((1,2,3),Val(1), Val(3)) == (1,2,3)
364 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(2), Val(5)) == (2,3,4,5)
365 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(1), Val(3)) == (1,2,3)
366 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(4), Val(6)) == (4,5,6)
367 end
368
369 @testset "split_tuple" begin
370 @testset "2 parts" begin
371 @test LazyTensors.split_tuple((),Val(0)) == ((),())
372 @test LazyTensors.split_tuple((1,),Val(0)) == ((),(1,))
373 @test LazyTensors.split_tuple((1,),Val(1)) == ((1,),())
374
375 @test LazyTensors.split_tuple((1,2,3,4),Val(0)) == ((),(1,2,3,4))
376 @test LazyTensors.split_tuple((1,2,3,4),Val(1)) == ((1,),(2,3,4))
377 @test LazyTensors.split_tuple((1,2,3,4),Val(2)) == ((1,2),(3,4))
378 @test LazyTensors.split_tuple((1,2,3,4),Val(3)) == ((1,2,3),(4,))
379 @test LazyTensors.split_tuple((1,2,3,4),Val(4)) == ((1,2,3,4),())
380
381 @test LazyTensors.split_tuple((1,2,true,4),Val(3)) == ((1,2,true),(4,))
382
383 @inferred LazyTensors.split_tuple((1,2,3,4),Val(3))
384 @inferred LazyTensors.split_tuple((1,2,true,4),Val(3))
385 end
386
387 @testset "3 parts" begin
388 @test LazyTensors.split_tuple((),Val(0),Val(0)) == ((),(),())
389 @test LazyTensors.split_tuple((1,2,3),Val(1), Val(1)) == ((1,),(2,),(3,))
390 @test LazyTensors.split_tuple((1,true,3),Val(1), Val(1)) == ((1,),(true,),(3,))
391
392 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(1),Val(2)) == ((1,),(2,3),(4,5,6))
393 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2)) == ((1,2,3),(4,5),(6,))
394
395 @inferred LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2))
396 @inferred LazyTensors.split_tuple((1,true,3),Val(1), Val(1))
397 end
398 end
399
400 @testset "flatten_tuple" begin
401 @test LazyTensors.flatten_tuple((1,)) == (1,)
402 @test LazyTensors.flatten_tuple((1,2,3,4,5,6)) == (1,2,3,4,5,6)
403 @test LazyTensors.flatten_tuple((1,2,(3,4),5,6)) == (1,2,3,4,5,6)
404 @test LazyTensors.flatten_tuple((1,2,(3,(4,5)),6)) == (1,2,3,4,5,6)
405 @test LazyTensors.flatten_tuple(((1,2),(3,4),(5,),6)) == (1,2,3,4,5,6)
406 end
407
408 317
409 @testset "LazyOuterProduct" begin 318 @testset "LazyOuterProduct" begin
410 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 319 A = ScalingTensor(2.0, (5,))
411 λ::T 320 B = ScalingTensor(3.0, (3,))
412 size::NTuple{D,Int} 321 C = ScalingTensor(5.0, (3,2))
413 end
414
415 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
416 LazyTensors.range_size(m::ScalingOperator) = m.size
417 LazyTensors.domain_size(m::ScalingOperator) = m.size
418
419 A = ScalingOperator(2.0, (5,))
420 B = ScalingOperator(3.0, (3,))
421 C = ScalingOperator(5.0, (3,2))
422 322
423 AB = LazyOuterProduct(A,B) 323 AB = LazyOuterProduct(A,B)
424 @test AB isa TensorMapping{T,2,2} where T 324 @test AB isa LazyTensor{T,2,2} where T
425 @test range_size(AB) == (5,3) 325 @test range_size(AB) == (5,3)
426 @test domain_size(AB) == (5,3) 326 @test domain_size(AB) == (5,3)
427 327
428 v = rand(range_size(AB)...) 328 v = rand(range_size(AB)...)
429 @test AB*v == 6*v 329 @test AB*v == 6*v
430 330
431 ABC = LazyOuterProduct(A,B,C) 331 ABC = LazyOuterProduct(A,B,C)
432 332
433 @test ABC isa TensorMapping{T,4,4} where T 333 @test ABC isa LazyTensor{T,4,4} where T
434 @test range_size(ABC) == (5,3,3,2) 334 @test range_size(ABC) == (5,3,3,2)
435 @test domain_size(ABC) == (5,3,3,2) 335 @test domain_size(ABC) == (5,3,3,2)
436 336
437 @test A⊗B == AB 337 @test A⊗B == AB
438 @test A⊗B⊗C == ABC 338 @test A⊗B⊗C == ABC
441 B = rand(2,4,3) 341 B = rand(2,4,3)
442 342
443 v₁ = rand(2,4,3) 343 v₁ = rand(2,4,3)
444 v₂ = rand(4,3,2) 344 v₂ = rand(4,3,2)
445 345
446 Ã = LazyLinearMap(A,(1,),(2,)) 346 Ã = DenseTensor(A,(1,),(2,))
447 B̃ = LazyLinearMap(B,(1,),(2,3)) 347 B̃ = DenseTensor(B,(1,),(2,3))
448 348
449 ÃB̃ = LazyOuterProduct(Ã,B̃) 349 ÃB̃ = LazyOuterProduct(Ã,B̃)
450 @tullio ABv[i,k] := A[i,j]*B[k,l,m]*v₁[j,l,m] 350 @tullio ABv[i,k] := A[i,j]*B[k,l,m]*v₁[j,l,m]
451 @test ÃB̃*v₁ ≈ ABv 351 @test ÃB̃*v₁ ≈ ABv
452 352
453 B̃Ã = LazyOuterProduct(B̃,Ã) 353 B̃Ã = LazyOuterProduct(B̃,Ã)
454 @tullio BAv[k,i] := A[i,j]*B[k,l,m]*v₂[l,m,j] 354 @tullio BAv[k,i] := A[i,j]*B[k,l,m]*v₂[l,m,j]
455 @test B̃Ã*v₂ ≈ BAv 355 @test B̃Ã*v₂ ≈ BAv
456 356
457 @testset "Indentity mapping arguments" begin 357 @testset "Indentity mapping arguments" begin
458 @test LazyOuterProduct(IdentityMapping(3,2), IdentityMapping(1,2)) == IdentityMapping(3,2,1,2) 358 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2)
459 359
460 Ã = LazyLinearMap(A,(1,),(2,)) 360 Ã = DenseTensor(A,(1,),(2,))
461 @test LazyOuterProduct(IdentityMapping(3,2), Ã) == InflatedTensorMapping(IdentityMapping(3,2),Ã) 361 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedTensor(IdentityTensor(3,2),Ã)
462 @test LazyOuterProduct(Ã, IdentityMapping(3,2)) == InflatedTensorMapping(Ã,IdentityMapping(3,2)) 362 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedTensor(Ã,IdentityTensor(3,2))
463 363
464 I1 = IdentityMapping(3,2) 364 I1 = IdentityTensor(3,2)
465 I2 = IdentityMapping(4) 365 I2 = IdentityTensor(4)
466 @test I1⊗Ã⊗I2 == InflatedTensorMapping(I1, Ã, I2) 366 @test I1⊗Ã⊗I2 == InflatedTensor(I1, Ã, I2)
467 end 367 end
468 368 end
469 end 369
370 @testset "inflate" begin
371 I = LazyTensors.inflate(IdentityTensor(),(3,4,5,6), 2)
372 @test I isa LazyTensor{Float64, 3,3}
373 @test range_size(I) == (3,5,6)
374 @test domain_size(I) == (3,5,6)
375
376 @test LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 1) == InflatedTensor(IdentityTensor{Float64}(),ScalingTensor(1., (4,)),IdentityTensor(4,5,6))
377 @test LazyTensors.inflate(ScalingTensor(2., (1,)),(3,4,5,6), 2) == InflatedTensor(IdentityTensor(3),ScalingTensor(2., (1,)),IdentityTensor(5,6))
378 @test LazyTensors.inflate(ScalingTensor(3., (6,)),(3,4,5,6), 4) == InflatedTensor(IdentityTensor(3,4,5),ScalingTensor(3., (6,)),IdentityTensor{Float64}())
379
380 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 0)
381 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 5)
382 end