Mercurial > repos > public > sbplib_julia

comparison test/testLazyTensors.jl @ 562:8f7919a9b398 feature/boundary_ops

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge with default
author Vidar Stiernström <vidar.stiernstrom@it.uu.se>
date Mon, 30 Nov 2020 18:30:24 +0100
parents 53828d3ed132
children 1c512e796c6d
comparison
equal deleted inserted replaced
544:884be64e82d9 562:8f7919a9b398
6 6
7 @testset "LazyTensors" begin 7 @testset "LazyTensors" begin
8 8
9 @testset "Generic Mapping methods" begin 9 @testset "Generic Mapping methods" begin
10 struct DummyMapping{T,R,D} <: TensorMapping{T,R,D} end 10 struct DummyMapping{T,R,D} <: TensorMapping{T,R,D} end
11 LazyTensors.apply(m::DummyMapping{T,R,D}, v, i::NTuple{R,Index{<:Region}}) where {T,R,D} = :apply 11 LazyTensors.apply(m::DummyMapping{T,R,D}, v, I::Vararg{Any,R}) where {T,R,D} = :apply
12 @test range_dim(DummyMapping{Int,2,3}()) == 2 12 @test range_dim(DummyMapping{Int,2,3}()) == 2
13 @test domain_dim(DummyMapping{Int,2,3}()) == 3 13 @test domain_dim(DummyMapping{Int,2,3}()) == 3
14 @test apply(DummyMapping{Int,2,3}(), zeros(Int, (0,0,0)),(Index{Unknown}(0),Index{Unknown}(0))) == :apply 14 @test apply(DummyMapping{Int,2,3}(), zeros(Int, (0,0,0)),0,0) == :apply
15 @test eltype(DummyMapping{Int,2,3}()) == Int 15 @test eltype(DummyMapping{Int,2,3}()) == Int
16 @test eltype(DummyMapping{Float64,2,3}()) == Float64 16 @test eltype(DummyMapping{Float64,2,3}()) == Float64
17 end 17 end
18 18
19 @testset "Mapping transpose" begin 19 @testset "Mapping transpose" begin
20 struct DummyMapping{T,R,D} <: TensorMapping{T,R,D} end 20 struct DummyMapping{T,R,D} <: TensorMapping{T,R,D} end
21 21
22 LazyTensors.apply(m::DummyMapping{T,R,D}, v, I::Vararg{Index{<:Region},R}) where {T,R,D} = :apply 22 LazyTensors.apply(m::DummyMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = :apply
23 LazyTensors.apply_transpose(m::DummyMapping{T,R,D}, v, I::Vararg{Index{<:Region},D}) where {T,R,D} = :apply_transpose 23 LazyTensors.apply_transpose(m::DummyMapping{T,R,D}, v, I::Vararg{Any,D}) where {T,R,D} = :apply_transpose
24 24
25 LazyTensors.range_size(m::DummyMapping{T,R,D}) where {T,R,D} = :range_size 25 LazyTensors.range_size(m::DummyMapping) = :range_size
26 LazyTensors.domain_size(m::DummyMapping{T,R,D}) where {T,R,D} = :domain_size 26 LazyTensors.domain_size(m::DummyMapping) = :domain_size
27 27
28 m = DummyMapping{Float64,2,3}() 28 m = DummyMapping{Float64,2,3}()
29 I = Index{Unknown}(0)
30 @test m' isa TensorMapping{Float64, 3,2} 29 @test m' isa TensorMapping{Float64, 3,2}
31 @test m'' == m 30 @test m'' == m
32 @test apply(m',zeros(Float64,(0,0)), I, I, I) == :apply_transpose 31 @test apply(m',zeros(Float64,(0,0)), 0, 0, 0) == :apply_transpose
33 @test apply(m'',zeros(Float64,(0,0,0)), I, I) == :apply 32 @test apply(m'',zeros(Float64,(0,0,0)), 0, 0) == :apply
34 @test apply_transpose(m', zeros(Float64,(0,0,0)), I, I) == :apply 33 @test apply_transpose(m', zeros(Float64,(0,0,0)), 0, 0) == :apply
35 34
36 @test range_size(m') == :domain_size 35 @test range_size(m') == :domain_size
37 @test domain_size(m') == :range_size 36 @test domain_size(m') == :range_size
38 end 37 end
39 38
40 @testset "TensorApplication" begin 39 @testset "TensorApplication" begin
41 struct SizeDoublingMapping{T,R,D} <: TensorMapping{T,R,D} 40 struct SizeDoublingMapping{T,R,D} <: TensorMapping{T,R,D}
42 domain_size::NTuple{D,Int} 41 domain_size::NTuple{D,Int}
43 end 42 end
44 43
45 LazyTensors.apply(m::SizeDoublingMapping{T,R,D}, v, i::Vararg{Index{<:Region},R}) where {T,R,D} = (:apply,v,i) 44 LazyTensors.apply(m::SizeDoublingMapping{T,R}, v, i::Vararg{Any,R}) where {T,R} = (:apply,v,i)
46 LazyTensors.range_size(m::SizeDoublingMapping) = 2 .* m.domain_size 45 LazyTensors.range_size(m::SizeDoublingMapping) = 2 .* m.domain_size
47 LazyTensors.domain_size(m::SizeDoublingMapping) = m.domain_size 46 LazyTensors.domain_size(m::SizeDoublingMapping) = m.domain_size
48 47
49 48
50 m = SizeDoublingMapping{Int, 1, 1}((3,)) 49 m = SizeDoublingMapping{Int, 1, 1}((3,))
51 v = [0,1,2] 50 v = [0,1,2]
52 @test m*v isa AbstractVector{Int} 51 @test m*v isa AbstractVector{Int}
53 @test size(m*v) == 2 .*size(v) 52 @test size(m*v) == 2 .*size(v)
54 @test (m*v)[Index{Upper}(0)] == (:apply,v,(Index{Upper}(0),)) 53 @test (m*v)[0] == (:apply,v,(0,))
55 @test (m*v)[0] == (:apply,v,(Index{Unknown}(0),))
56 @test m*m*v isa AbstractVector{Int} 54 @test m*m*v isa AbstractVector{Int}
57 @test (m*m*v)[Index{Upper}(1)] == (:apply,m*v,(Index{Upper}(1),)) 55 @test (m*m*v)[1] == (:apply,m*v,(1,))
58 @test (m*m*v)[1] == (:apply,m*v,(Index{Unknown}(1),)) 56 @test (m*m*v)[3] == (:apply,m*v,(3,))
59 @test (m*m*v)[Index{Interior}(3)] == (:apply,m*v,(Index{Interior}(3),)) 57 @test (m*m*v)[6] == (:apply,m*v,(6,))
60 @test (m*m*v)[3] == (:apply,m*v,(Index{Unknown}(3),))
61 @test (m*m*v)[Index{Lower}(6)] == (:apply,m*v,(Index{Lower}(6),))
62 @test (m*m*v)[6] == (:apply,m*v,(Index{Unknown}(6),))
63 @test_broken BoundsError == (m*m*v)[0] 58 @test_broken BoundsError == (m*m*v)[0]
64 @test_broken BoundsError == (m*m*v)[7] 59 @test_broken BoundsError == (m*m*v)[7]
65 @test_throws MethodError m*m 60 @test_throws MethodError m*m
66 61
67 m = SizeDoublingMapping{Int, 2, 1}((3,)) 62 m = SizeDoublingMapping{Int, 2, 1}((3,))
68 @test_throws MethodError m*ones(Int,2,2) 63 @test_throws MethodError m*ones(Int,2,2)
69 @test_throws MethodError m*m*v 64 @test_throws MethodError m*m*v
70 65
71 m = SizeDoublingMapping{Float64, 2, 2}((3,3)) 66 m = SizeDoublingMapping{Float64, 2, 2}((3,3))
72 v = ones(3,3) 67 v = ones(3,3)
73 I = (Index{Lower}(1),Index{Interior}(2));
74 @test size(m*v) == 2 .*size(v) 68 @test size(m*v) == 2 .*size(v)
75 @test (m*v)[I] == (:apply,v,I) 69 @test (m*v)[1,2] == (:apply,v,(1,2))
76 70
77 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 71 struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
78 λ::T 72 λ::T
79 size::NTuple{D,Int} 73 size::NTuple{D,Int}
80 end 74 end
81 75
82 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Index,D}) where {T,D} = m.λ*v[I] 76 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
83 LazyTensors.range_size(m::ScalingOperator) = m.size 77 LazyTensors.range_size(m::ScalingOperator) = m.size
84 LazyTensors.domain_size(m::ScalingOperator) = m.size 78 LazyTensors.domain_size(m::ScalingOperator) = m.size
85 79
86 m = ScalingOperator{Int,1}(2,(3,)) 80 m = ScalingOperator{Int,1}(2,(3,))
87 v = [1,2,3] 81 v = [1,2,3]
89 @test m*v == [2,4,6] 83 @test m*v == [2,4,6]
90 84
91 m = ScalingOperator{Int,2}(2,(2,2)) 85 m = ScalingOperator{Int,2}(2,(2,2))
92 v = [[1 2];[3 4]] 86 v = [[1 2];[3 4]]
93 @test m*v == [[2 4];[6 8]] 87 @test m*v == [[2 4];[6 8]]
94 I = (Index{Upper}(2),Index{Lower}(1)) 88 @test (m*v)[2,1] == 6
95 @test (m*v)[I] == 6
96 end 89 end
97 90
98 @testset "TensorMapping binary operations" begin 91 @testset "TensorMapping binary operations" begin
99 struct ScalarMapping{T,R,D} <: TensorMapping{T,R,D} 92 struct ScalarMapping{T,R,D} <: TensorMapping{T,R,D}
100 λ::T 93 λ::T
101 range_size::NTuple{R,Int} 94 range_size::NTuple{R,Int}
102 domain_size::NTuple{D,Int} 95 domain_size::NTuple{D,Int}
103 end 96 end
104 97
105 LazyTensors.apply(m::ScalarMapping{T,R,D}, v, I::Vararg{Index{<:Region}}) where {T,R,D} = m.λ*v[I...] 98 LazyTensors.apply(m::ScalarMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = m.λ*v[I...]
106 LazyTensors.range_size(m::ScalarMapping) = m.domain_size 99 LazyTensors.range_size(m::ScalarMapping) = m.domain_size
107 LazyTensors.domain_size(m::ScalarMapping) = m.range_size 100 LazyTensors.domain_size(m::ScalarMapping) = m.range_size
108 101
109 A = ScalarMapping{Float64,1,1}(2.0, (3,), (3,)) 102 A = ScalarMapping{Float64,1,1}(2.0, (3,), (3,))
110 B = ScalarMapping{Float64,1,1}(3.0, (3,), (3,)) 103 B = ScalarMapping{Float64,1,1}(3.0, (3,), (3,))
338 331
339 A = LazyLinearMap(Ã,(1,),(2,)) 332 A = LazyLinearMap(Ã,(1,),(2,))
340 B = LazyLinearMap(B̃,(1,2),(3,)) 333 B = LazyLinearMap(B̃,(1,2),(3,))
341 C = LazyLinearMap(C̃,(1,),(2,3)) 334 C = LazyLinearMap(C̃,(1,),(2,3))
342 335
343 @test InflatedTensorMapping(I(3,2), A, I(4)) isa TensorMapping{Float64, 4, 4} 336 @testset "Constructors" begin
344 @test InflatedTensorMapping(I(3,2), B, I(4)) isa TensorMapping{Float64, 5, 4} 337 @test InflatedTensorMapping(I(3,2), A, I(4)) isa TensorMapping{Float64, 4, 4}
345 @test InflatedTensorMapping(I(3), C, I(2,3)) isa TensorMapping{Float64, 4, 5} 338 @test InflatedTensorMapping(I(3,2), B, I(4)) isa TensorMapping{Float64, 5, 4}
346 @test InflatedTensorMapping(C, I(2,3)) isa TensorMapping{Float64, 3, 4} 339 @test InflatedTensorMapping(I(3), C, I(2,3)) isa TensorMapping{Float64, 4, 5}
347 @test InflatedTensorMapping(I(3), C) isa TensorMapping{Float64, 2, 3} 340 @test InflatedTensorMapping(C, I(2,3)) isa TensorMapping{Float64, 3, 4}
348 @test InflatedTensorMapping(I(3), I(2,3)) isa TensorMapping{Float64, 3, 3} 341 @test InflatedTensorMapping(I(3), C) isa TensorMapping{Float64, 2, 3}
349 342 @test InflatedTensorMapping(I(3), I(2,3)) isa TensorMapping{Float64, 3, 3}
350 @test range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 343 end
351 @test domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 344
352 345 @testset "Range and domain size" begin
353 @test range_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,4,2,4) 346 @test range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4)
354 @test domain_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,3,4) 347 @test domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4)
355 348
356 @test range_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,4,2,3) 349 @test range_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,4,2,4)
357 @test domain_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,2,3,2,3) 350 @test domain_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,3,4)
358 351
359 @inferred range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 352 @test range_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,4,2,3)
360 @inferred domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 353 @test domain_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,2,3,2,3)
361 354
362 # Test InflatedTensorMapping mapping w. before and after 355 @inferred range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4)
363 tm = InflatedTensorMapping(I(3,2), A, I(4)) 356 @inferred domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4)
364 v = rand(domain_size(tm)...) 357 end
365 @tullio IAIv[a,b,c,d] := Ã[c,i]*v[a,b,i,d] 358
366 @test tm*v ≈ IAIv rtol=1e-14 359 @testset "Application" begin
367 @inferred LazyTensors.split_index(tm,1,1,1,1) 360 # Testing regular application and transposed application with inflation "before", "after" and "before and after".
368 361 # The inflated tensor mappings are chosen to preserve, reduce and increase the dimension of the result compared to the input.
369 # Test InflatedTensorMapping mapping w. before 362 tests = [
370 tm = InflatedTensorMapping(I(3,2), A) 363 (
371 v = rand(domain_size(tm)...) 364 InflatedTensorMapping(I(3,2), A, I(4)),
372 @tullio IAIv[a,b,c] := Ã[c,i]*v[a,b,i] 365 (v-> @tullio res[a,b,c,d] := Ã[c,i]*v[a,b,i,d]), # Expected result of apply
373 @test tm*v ≈ IAIv rtol=1e-14 366 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose
374 @inferred LazyTensors.split_index(tm,1,1,1) 367 ),
375 368 (
376 # Test InflatedTensorMapping mapping w. after 369 InflatedTensorMapping(I(3,2), B, I(4)),
377 tm = InflatedTensorMapping(A,I(4)) 370 (v-> @tullio res[a,b,c,d,e] := B̃[c,d,i]*v[a,b,i,e]),
378 v = rand(domain_size(tm)...) 371 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]),
379 @tullio IAIv[c,d] := Ã[c,i]*v[i,d] 372 ),
380 @test tm*v ≈ IAIv rtol=1e-14 373 (
381 @inferred LazyTensors.split_index(tm,1,1) 374 InflatedTensorMapping(I(3,2), C, I(4)),
382 375 (v-> @tullio res[a,b,c,d] := C̃[c,i,j]*v[a,b,i,j,d]),
383 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 376 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]),
384 λ::T 377 ),
385 size::NTuple{D,Int} 378 (
386 end 379 InflatedTensorMapping(I(3,2), A),
387 380 (v-> @tullio res[a,b,c] := Ã[c,i]*v[a,b,i]),
388 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Index,D}) where {T,D} = m.λ*v[I] 381 (v-> @tullio res[a,b,c] := Ã[i,c]*v[a,b,i]),
389 LazyTensors.range_size(m::ScalingOperator) = m.size 382 ),
390 LazyTensors.domain_size(m::ScalingOperator) = m.size 383 (
391 384 InflatedTensorMapping(I(3,2), B),
392 tm = InflatedTensorMapping(I(2,3),ScalingOperator(2.0, (3,2)),I(3,4)) 385 (v-> @tullio res[a,b,c,d] := B̃[c,d,i]*v[a,b,i]),
393 v = rand(domain_size(tm)...) 386 (v-> @tullio res[a,b,c] := B̃[i,j,c]*v[a,b,i,j]),
394 387 ),
395 @inferred LazyTensors.split_index(tm,1,2,3,2,2,4) 388 (
396 @inferred apply(tm,v,Index{Unknown}.((1,2,3,2,2,4))...) 389 InflatedTensorMapping(I(3,2), C),
397 @inferred (tm*v)[1,2,3,2,2,4] 390 (v-> @tullio res[a,b,c] := C̃[c,i,j]*v[a,b,i,j]),
391 (v-> @tullio res[a,b,c,d] := C̃[i,c,d]*v[a,b,i]),
392 ),
393 (
394 InflatedTensorMapping(A,I(4)),
395 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]),
396 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]),
397 ),
398 (
399 InflatedTensorMapping(B,I(4)),
400 (v-> @tullio res[a,b,c] := B̃[a,b,i]*v[i,c]),
401 (v-> @tullio res[a,b] := B̃[i,j,a]*v[i,j,b]),
402 ),
403 (
404 InflatedTensorMapping(C,I(4)),
405 (v-> @tullio res[a,b] := C̃[a,i,j]*v[i,j,b]),
406 (v-> @tullio res[a,b,c] := C̃[i,a,b]*v[i,c]),
407 ),
408 ]
409
410 @testset "apply" begin
411 for i ∈ 1:length(tests)
412 tm = tests[i][1]
413 v = rand(domain_size(tm)...)
414 true_value = tests[i][2](v)
415 @test tm*v ≈ true_value rtol=1e-14
416 end
417 end
418
419 @testset "apply_transpose" begin
420 for i ∈ 1:length(tests)
421 tm = tests[i][1]
422 v = rand(range_size(tm)...)
423 true_value = tests[i][3](v)
424 @test tm'*v ≈ true_value rtol=1e-14
425 end
426 end
427
428 @testset "Inference of application" begin
429 struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
430 λ::T
431 size::NTuple{D,Int}
432 end
433
434 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
435 LazyTensors.range_size(m::ScalingOperator) = m.size
436 LazyTensors.domain_size(m::ScalingOperator) = m.size
437
438 tm = InflatedTensorMapping(I(2,3),ScalingOperator(2.0, (3,2)),I(3,4))
439 v = rand(domain_size(tm)...)
440
441 @inferred apply(tm,v,1,2,3,2,2,4)
442 @inferred (tm*v)[1,2,3,2,2,4]
443 end
444 end
398 445
399 @testset "InflatedTensorMapping of InflatedTensorMapping" begin 446 @testset "InflatedTensorMapping of InflatedTensorMapping" begin
400 A = ScalingOperator(2.0,(2,3)) 447 A = ScalingOperator(2.0,(2,3))
401 itm = InflatedTensorMapping(I(3,2), A, I(4)) 448 itm = InflatedTensorMapping(I(3,2), A, I(4))
402 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2)) 449 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2))
403 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2)) 450 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2))
404 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4)) 451 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4))
405 452
406 @test InflatedTensorMapping(I(2), I(2), I(2)) isa InflatedTensorMapping # The constructor should always return its type. 453 @test InflatedTensorMapping(I(2), I(2), I(2)) isa InflatedTensorMapping # The constructor should always return its type.
407 end 454 end
408 455 end
456
457 @testset "split_index" begin
458 @test LazyTensors.split_index(Val(2),Val(1),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,5,6),(3,4))
459 @test LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,:,:,5,6),(3,4))
460 @test LazyTensors.split_index(Val(3),Val(1),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,5,6),(4,))
461 @test LazyTensors.split_index(Val(3),Val(2),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,:,5,6),(4,))
462 @test LazyTensors.split_index(Val(1),Val(1),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,4,5,6),(2,3))
463 @test LazyTensors.split_index(Val(1),Val(2),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,:,4,5,6),(2,3))
464
465 @test LazyTensors.split_index(Val(0),Val(1),Val(3),Val(3),1,2,3,4,5,6) == ((:,4,5,6),(1,2,3))
466 @test LazyTensors.split_index(Val(3),Val(1),Val(3),Val(0),1,2,3,4,5,6) == ((1,2,3,:),(4,5,6))
467
468 @inferred LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,2,2,4)
409 end 469 end
410 470
411 @testset "slice_tuple" begin 471 @testset "slice_tuple" begin
412 @test LazyTensors.slice_tuple((1,2,3),Val(1), Val(3)) == (1,2,3) 472 @test LazyTensors.slice_tuple((1,2,3),Val(1), Val(3)) == (1,2,3)
413 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(2), Val(5)) == (2,3,4,5) 473 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(2), Val(5)) == (2,3,4,5)
414 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(1), Val(3)) == (1,2,3) 474 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(1), Val(3)) == (1,2,3)
415 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(4), Val(6)) == (4,5,6) 475 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(4), Val(6)) == (4,5,6)
476 end
477
478 @testset "split_tuple" begin
479 @testset "2 parts" begin
480 @test LazyTensors.split_tuple((),Val(0)) == ((),())
481 @test LazyTensors.split_tuple((1,),Val(0)) == ((),(1,))
482 @test LazyTensors.split_tuple((1,),Val(1)) == ((1,),())
483
484 @test LazyTensors.split_tuple((1,2,3,4),Val(0)) == ((),(1,2,3,4))
485 @test LazyTensors.split_tuple((1,2,3,4),Val(1)) == ((1,),(2,3,4))
486 @test LazyTensors.split_tuple((1,2,3,4),Val(2)) == ((1,2),(3,4))
487 @test LazyTensors.split_tuple((1,2,3,4),Val(3)) == ((1,2,3),(4,))
488 @test LazyTensors.split_tuple((1,2,3,4),Val(4)) == ((1,2,3,4),())
489
490 @inferred LazyTensors.split_tuple((1,2,3,4),Val(3))
491 end
492
493 @testset "3 parts" begin
494 @test LazyTensors.split_tuple((),Val(0),Val(0)) == ((),(),())
495 @test LazyTensors.split_tuple((1,2,3),Val(1), Val(1)) == ((1,),(2,),(3,))
496
497 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(1),Val(2)) == ((1,),(2,3),(4,5,6))
498 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2)) == ((1,2,3),(4,5),(6,))
499
500 @inferred LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2))
501 end
416 end 502 end
417 503
418 @testset "flatten_tuple" begin 504 @testset "flatten_tuple" begin
419 @test LazyTensors.flatten_tuple((1,)) == (1,) 505 @test LazyTensors.flatten_tuple((1,)) == (1,)
420 @test LazyTensors.flatten_tuple((1,2,3,4,5,6)) == (1,2,3,4,5,6) 506 @test LazyTensors.flatten_tuple((1,2,3,4,5,6)) == (1,2,3,4,5,6)
428 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 514 struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
429 λ::T 515 λ::T
430 size::NTuple{D,Int} 516 size::NTuple{D,Int}
431 end 517 end
432 518
433 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Index,D}) where {T,D} = m.λ*v[I] 519 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
434 LazyTensors.range_size(m::ScalingOperator) = m.size 520 LazyTensors.range_size(m::ScalingOperator) = m.size
435 LazyTensors.domain_size(m::ScalingOperator) = m.size 521 LazyTensors.domain_size(m::ScalingOperator) = m.size
436 522
437 A = ScalingOperator(2.0, (5,)) 523 A = ScalingOperator(2.0, (5,))
438 B = ScalingOperator(3.0, (3,)) 524 B = ScalingOperator(3.0, (3,))