comparison test/LazyTensors/lazy_tensor_operations_test.jl @ 1100:157a78959e5d refactor/sbpoperators/inflation

Bring up to date
author Jonatan Werpers <jonatan@werpers.com>
date Tue, 10 May 2022 20:34:20 +0200
parents 2278730f9cee
children f1bb1b6d85dd
comparison
equal deleted inserted replaced
1099:05a25a5063bb 1100:157a78959e5d
34 @test range_size(m') == :domain_size 34 @test range_size(m') == :domain_size
35 @test domain_size(m') == :range_size 35 @test domain_size(m') == :range_size
36 end 36 end
37 37
38 38
39 @testset "LazyTensorApplication" begin 39 @testset "TensorApplication" begin
40 m = SizeDoublingMapping{Int, 1, 1}((3,)) 40 m = SizeDoublingMapping{Int, 1, 1}((3,))
41 mm = SizeDoublingMapping{Int, 1, 1}((6,)) 41 mm = SizeDoublingMapping{Int, 1, 1}((6,))
42 v = [0,1,2] 42 v = [0,1,2]
43 @test size(m*v) == 2 .*size(v) 43 @test size(m*v) == 2 .*size(v)
44 @test (m*v)[1] == (:apply,v,(1,)) 44 @test (m*v)[1] == (:apply,v,(1,))
157 @test_throws RangeSizeMismatch SizeDoublingMapping{Float64,1,1}((2,)) + ScalingTensor(2.0, (2,)) 157 @test_throws RangeSizeMismatch SizeDoublingMapping{Float64,1,1}((2,)) + ScalingTensor(2.0, (2,))
158 end 158 end
159 end 159 end
160 160
161 161
162 @testset "LazyTensorComposition" begin 162 @testset "TensorComposition" begin
163 A = rand(2,3) 163 A = rand(2,3)
164 B = rand(3,4) 164 B = rand(3,4)
165 165
166 Ã = LazyLinearMap(A, (1,), (2,)) 166 Ã = DenseTensor(A, (1,), (2,))
167 B̃ = LazyLinearMap(B, (1,), (2,)) 167 B̃ = DenseTensor(B, (1,), (2,))
168 168
169 @test Ã∘B̃ isa LazyTensorComposition 169 @test Ã∘B̃ isa TensorComposition
170 @test range_size(Ã∘B̃) == (2,) 170 @test range_size(Ã∘B̃) == (2,)
171 @test domain_size(Ã∘B̃) == (4,) 171 @test domain_size(Ã∘B̃) == (4,)
172 @test_throws DomainSizeMismatch B̃∘Ã 172 @test_throws DomainSizeMismatch B̃∘Ã
173 173
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) 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)
179 v = rand(2) 179 v = rand(2)
180 @test (Ã∘B̃)'*v ≈ B'*A'*v rtol=1e-14 180 @test (Ã∘B̃)'*v ≈ B'*A'*v rtol=1e-14
181 181
182 @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64 182 @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64
183 @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64 183 @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64
184 end 184
185 185 a = 2.
186 186 v = rand(3)
187 @testset "InflatedLazyTensor" begin 187 @test a*Ã isa TensorComposition
188 @test a*Ã == Ã*a
189 @test range_size(a*Ã) == range_size(Ã)
190 @test domain_size(a*Ã) == domain_size(Ã)
191 @test a*Ã*v == a.*A*v
192 end
193
194
195 @testset "InflatedTensor" begin
188 I(sz...) = IdentityTensor(sz...) 196 I(sz...) = IdentityTensor(sz...)
189 197
190 Ã = rand(4,2) 198 Ã = rand(4,2)
191 B̃ = rand(4,2,3) 199 B̃ = rand(4,2,3)
192 C̃ = rand(4,2,3) 200 C̃ = rand(4,2,3)
193 201
194 A = LazyLinearMap(Ã,(1,),(2,)) 202 A = DenseTensor(Ã,(1,),(2,))
195 B = LazyLinearMap(B̃,(1,2),(3,)) 203 B = DenseTensor(B̃,(1,2),(3,))
196 C = LazyLinearMap(C̃,(1,),(2,3)) 204 C = DenseTensor(C̃,(1,),(2,3))
197 205
198 @testset "Constructors" begin 206 @testset "Constructors" begin
199 @test InflatedLazyTensor(I(3,2), A, I(4)) isa LazyTensor{Float64, 4, 4} 207 @test InflatedTensor(I(3,2), A, I(4)) isa LazyTensor{Float64, 4, 4}
200 @test InflatedLazyTensor(I(3,2), B, I(4)) isa LazyTensor{Float64, 5, 4} 208 @test InflatedTensor(I(3,2), B, I(4)) isa LazyTensor{Float64, 5, 4}
201 @test InflatedLazyTensor(I(3), C, I(2,3)) isa LazyTensor{Float64, 4, 5} 209 @test InflatedTensor(I(3), C, I(2,3)) isa LazyTensor{Float64, 4, 5}
202 @test InflatedLazyTensor(C, I(2,3)) isa LazyTensor{Float64, 3, 4} 210 @test InflatedTensor(C, I(2,3)) isa LazyTensor{Float64, 3, 4}
203 @test InflatedLazyTensor(I(3), C) isa LazyTensor{Float64, 2, 3} 211 @test InflatedTensor(I(3), C) isa LazyTensor{Float64, 2, 3}
204 @test InflatedLazyTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3} 212 @test InflatedTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3}
205 end 213 end
206 214
207 @testset "Range and domain size" begin 215 @testset "Range and domain size" begin
208 @test range_size(InflatedLazyTensor(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)
209 @test domain_size(InflatedLazyTensor(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)
210 218
211 @test range_size(InflatedLazyTensor(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)
212 @test domain_size(InflatedLazyTensor(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)
213 221
214 @test range_size(InflatedLazyTensor(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)
215 @test domain_size(InflatedLazyTensor(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)
216 224
217 @inferred range_size(InflatedLazyTensor(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)
218 @inferred domain_size(InflatedLazyTensor(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)
219 end 227 end
220 228
221 @testset "Application" begin 229 @testset "Application" begin
222 # 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".
223 # 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.
224 cases = [ 232 cases = [
225 ( 233 (
226 InflatedLazyTensor(I(3,2), A, I(4)), 234 InflatedTensor(I(3,2), A, I(4)),
227 (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
228 (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
229 ), 237 ),
230 ( 238 (
231 InflatedLazyTensor(I(3,2), B, I(4)), 239 InflatedTensor(I(3,2), B, I(4)),
232 (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]),
233 (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]),
234 ), 242 ),
235 ( 243 (
236 InflatedLazyTensor(I(3,2), C, I(4)), 244 InflatedTensor(I(3,2), C, I(4)),
237 (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]),
238 (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]),
239 ), 247 ),
240 ( 248 (
241 InflatedLazyTensor(I(3,2), A), 249 InflatedTensor(I(3,2), A),
242 (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]),
243 (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]),
244 ), 252 ),
245 ( 253 (
246 InflatedLazyTensor(I(3,2), B), 254 InflatedTensor(I(3,2), B),
247 (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]),
248 (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]),
249 ), 257 ),
250 ( 258 (
251 InflatedLazyTensor(I(3,2), C), 259 InflatedTensor(I(3,2), C),
252 (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]),
253 (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]),
254 ), 262 ),
255 ( 263 (
256 InflatedLazyTensor(A,I(4)), 264 InflatedTensor(A,I(4)),
257 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]), 265 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]),
258 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]), 266 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]),
259 ), 267 ),
260 ( 268 (
261 InflatedLazyTensor(B,I(4)), 269 InflatedTensor(B,I(4)),
262 (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]),
263 (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]),
264 ), 272 ),
265 ( 273 (
266 InflatedLazyTensor(C,I(4)), 274 InflatedTensor(C,I(4)),
267 (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]),
268 (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]),
269 ), 277 ),
270 ] 278 ]
271 279
276 v = rand(range_size(tm)...) 284 v = rand(range_size(tm)...)
277 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14 285 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14
278 end 286 end
279 287
280 @testset "application to other type" begin 288 @testset "application to other type" begin
281 tm = InflatedLazyTensor(I(3,2), A, I(4)) 289 tm = InflatedTensor(I(3,2), A, I(4))
282 290
283 v = rand(ComplexF64, domain_size(tm)...) 291 v = rand(ComplexF64, domain_size(tm)...)
284 @test (tm*v)[1,2,3,1] isa ComplexF64 292 @test (tm*v)[1,2,3,1] isa ComplexF64
285 293
286 v = rand(ComplexF64, domain_size(tm')...) 294 v = rand(ComplexF64, domain_size(tm')...)
287 @test (tm'*v)[1,2,2,1] isa ComplexF64 295 @test (tm'*v)[1,2,2,1] isa ComplexF64
288 end 296 end
289 297
290 @testset "Inference of application" begin 298 @testset "Inference of application" begin
291 tm = InflatedLazyTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4)) 299 tm = InflatedTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4))
292 v = rand(domain_size(tm)...) 300 v = rand(domain_size(tm)...)
293 301
294 @inferred apply(tm,v,1,2,3,2,2,4) 302 @inferred apply(tm,v,1,2,3,2,2,4)
295 @inferred (tm*v)[1,2,3,2,2,4] 303 @inferred (tm*v)[1,2,3,2,2,4]
296 end 304 end
297 end 305 end
298 306
299 @testset "InflatedLazyTensor of InflatedLazyTensor" begin 307 @testset "InflatedTensor of InflatedTensor" begin
300 A = ScalingTensor(2.0,(2,3)) 308 A = ScalingTensor(2.0,(2,3))
301 itm = InflatedLazyTensor(I(3,2), A, I(4)) 309 itm = InflatedTensor(I(3,2), A, I(4))
302 @test InflatedLazyTensor(I(4), itm, I(2)) == InflatedLazyTensor(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))
303 @test InflatedLazyTensor(itm, I(2)) == InflatedLazyTensor(I(3,2), A, I(4,2)) 311 @test InflatedTensor(itm, I(2)) == InflatedTensor(I(3,2), A, I(4,2))
304 @test InflatedLazyTensor(I(4), itm) == InflatedLazyTensor(I(4,3,2), A, I(4)) 312 @test InflatedTensor(I(4), itm) == InflatedTensor(I(4,3,2), A, I(4))
305 313
306 @test InflatedLazyTensor(I(2), I(2), I(2)) isa InflatedLazyTensor # 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.
307 end 315 end
308 end 316 end
309 317
310 @testset "LazyOuterProduct" begin 318 @testset "LazyOuterProduct" begin
311 A = ScalingTensor(2.0, (5,)) 319 A = ScalingTensor(2.0, (5,))
333 B = rand(2,4,3) 341 B = rand(2,4,3)
334 342
335 v₁ = rand(2,4,3) 343 v₁ = rand(2,4,3)
336 v₂ = rand(4,3,2) 344 v₂ = rand(4,3,2)
337 345
338 Ã = LazyLinearMap(A,(1,),(2,)) 346 Ã = DenseTensor(A,(1,),(2,))
339 B̃ = LazyLinearMap(B,(1,),(2,3)) 347 B̃ = DenseTensor(B,(1,),(2,3))
340 348
341 ÃB̃ = LazyOuterProduct(Ã,B̃) 349 ÃB̃ = LazyOuterProduct(Ã,B̃)
342 @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]
343 @test ÃB̃*v₁ ≈ ABv 351 @test ÃB̃*v₁ ≈ ABv
344 352
347 @test B̃Ã*v₂ ≈ BAv 355 @test B̃Ã*v₂ ≈ BAv
348 356
349 @testset "Indentity mapping arguments" begin 357 @testset "Indentity mapping arguments" begin
350 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2) 358 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2)
351 359
352 Ã = LazyLinearMap(A,(1,),(2,)) 360 Ã = DenseTensor(A,(1,),(2,))
353 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedLazyTensor(IdentityTensor(3,2),Ã) 361 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedTensor(IdentityTensor(3,2),Ã)
354 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedLazyTensor(Ã,IdentityTensor(3,2)) 362 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedTensor(Ã,IdentityTensor(3,2))
355 363
356 I1 = IdentityTensor(3,2) 364 I1 = IdentityTensor(3,2)
357 I2 = IdentityTensor(4) 365 I2 = IdentityTensor(4)
358 @test I1⊗Ã⊗I2 == InflatedLazyTensor(I1, Ã, I2) 366 @test I1⊗Ã⊗I2 == InflatedTensor(I1, Ã, I2)
359 end 367 end
360 end 368 end
361 369
362 @testset "inflate" begin 370 @testset "inflate" begin
363 I = LazyTensors.inflate(IdentityTensor(),(3,4,5,6), 2) 371 I = LazyTensors.inflate(IdentityTensor(),(3,4,5,6), 2)
364 @test I isa LazyTensor{Float64, 3,3} 372 @test I isa LazyTensor{Float64, 3,3}
365 @test range_size(I) == (3,5,6) 373 @test range_size(I) == (3,5,6)
366 @test domain_size(I) == (3,5,6) 374 @test domain_size(I) == (3,5,6)
367 375
368 @test LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 1) == InflatedLazyTensor(IdentityTensor{Float64}(),ScalingTensor(1., (4,)),IdentityTensor(4,5,6)) 376 @test LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 1) == InflatedTensor(IdentityTensor{Float64}(),ScalingTensor(1., (4,)),IdentityTensor(4,5,6))
369 @test LazyTensors.inflate(ScalingTensor(2., (1,)),(3,4,5,6), 2) == InflatedLazyTensor(IdentityTensor(3),ScalingTensor(2., (1,)),IdentityTensor(5,6)) 377 @test LazyTensors.inflate(ScalingTensor(2., (1,)),(3,4,5,6), 2) == InflatedTensor(IdentityTensor(3),ScalingTensor(2., (1,)),IdentityTensor(5,6))
370 @test LazyTensors.inflate(ScalingTensor(3., (6,)),(3,4,5,6), 4) == InflatedLazyTensor(IdentityTensor(3,4,5),ScalingTensor(3., (6,)),IdentityTensor{Float64}()) 378 @test LazyTensors.inflate(ScalingTensor(3., (6,)),(3,4,5,6), 4) == InflatedTensor(IdentityTensor(3,4,5),ScalingTensor(3., (6,)),IdentityTensor{Float64}())
371 379
372 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 0) 380 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 0)
373 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 5) 381 @test_throws BoundsError LazyTensors.inflate(ScalingTensor(1., (4,)),(3,4,5,6), 5)
374 end 382 end