comparison test/LazyTensors/lazy_tensor_operations_test.jl @ 993:2f9beee56a4c refactor/lazy_tensors

Use ScalingTensor instead of custom test type
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 20:41:31 +0100
parents bc384aaade30
children 1ba8a398af9c
comparison
equal deleted inserted replaced
992:bc384aaade30 993:2f9beee56a4c
58 @test (m*v)[1,2] == (:apply,v,(1,2)) 58 @test (m*v)[1,2] == (:apply,v,(1,2))
59 59
60 @test (m*v)[CartesianIndex(2,3)] == (:apply,v,(2,3)) 60 @test (m*v)[CartesianIndex(2,3)] == (:apply,v,(2,3))
61 @test (m*m*v)[CartesianIndex(4,3)] == (:apply,m*v,(4,3)) 61 @test (m*m*v)[CartesianIndex(4,3)] == (:apply,m*v,(4,3))
62 62
63 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 63 m = ScalingTensor(2,(3,))
64 λ::T
65 size::NTuple{D,Int}
66 end
67
68 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
69 LazyTensors.range_size(m::ScalingOperator) = m.size
70 LazyTensors.domain_size(m::ScalingOperator) = m.size
71
72 m = ScalingOperator{Int,1}(2,(3,))
73 v = [1,2,3] 64 v = [1,2,3]
74 @test m*v isa AbstractVector 65 @test m*v isa AbstractVector
75 @test m*v == [2,4,6] 66 @test m*v == [2,4,6]
76 67
77 m = ScalingOperator{Int,2}(2,(2,2)) 68 m = ScalingTensor(2,(2,2))
78 v = [[1 2];[3 4]] 69 v = [[1 2];[3 4]]
79 @test m*v == [[2 4];[6 8]] 70 @test m*v == [[2 4];[6 8]]
80 @test (m*v)[2,1] == 6 71 @test (m*v)[2,1] == 6
81 72
82 @testset "Type calculation" begin 73 @testset "Type calculation" begin
83 m = ScalingOperator{Int,1}(2,(3,)) 74 m = ScalingTensor(2,(3,))
84 v = [1.,2.,3.] 75 v = [1.,2.,3.]
85 @test m*v isa AbstractVector{Float64} 76 @test m*v isa AbstractVector{Float64}
86 @test m*v == [2.,4.,6.] 77 @test m*v == [2.,4.,6.]
87 @inferred m*v 78 @inferred m*v
88 @inferred (m*v)[1] 79 @inferred (m*v)[1]
89 80
90 m = ScalingOperator{Int,2}(2,(2,2)) 81 m = ScalingTensor(2,(2,2))
91 v = [[1. 2.];[3. 4.]] 82 v = [[1. 2.];[3. 4.]]
92 @test m*v == [[2. 4.];[6. 8.]] 83 @test m*v == [[2. 4.];[6. 8.]]
93 @test (m*v)[2,1] == 6. 84 @test (m*v)[2,1] == 6.
94 @inferred m*v 85 @inferred m*v
95 @inferred (m*v)[1] 86 @inferred (m*v)[1]
96 87
97 m = ScalingOperator{ComplexF64,1}(2. +2. *im,(3,)) 88 m = ScalingTensor(2. +2. *im,(3,))
98 v = [1.,2.,3.] 89 v = [1.,2.,3.]
99 @test m*v isa AbstractVector{ComplexF64} 90 @test m*v isa AbstractVector{ComplexF64}
100 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 91 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
101 @inferred m*v 92 @inferred m*v
102 @inferred (m*v)[1] 93 @inferred (m*v)[1]
103 94
104 m = ScalingOperator{ComplexF64,1}(1,(3,)) 95 m = ScalingTensor(1,(3,))
105 v = [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 96 v = [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
106 @test m*v isa AbstractVector{ComplexF64} 97 @test m*v isa AbstractVector{ComplexF64}
107 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 98 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
108 @inferred m*v 99 @inferred m*v
109 @inferred (m*v)[1] 100 @inferred (m*v)[1]
110 101
111 m = ScalingOperator{Float64,1}(2., (3,)) 102 m = ScalingTensor(2., (3,))
112 v = [[1,2,3], [3,2,1],[1,3,1]] 103 v = [[1,2,3], [3,2,1],[1,3,1]]
113 @test m*v isa AbstractVector{Vector{Float64}} 104 @test m*v isa AbstractVector{Vector{Float64}}
114 @test m*v == [[2.,4.,6.], [6.,4.,2.],[2.,6.,2.]] 105 @test m*v == [[2.,4.,6.], [6.,4.,2.],[2.,6.,2.]]
115 @inferred m*v 106 @inferred m*v
116 @inferred (m*v)[1] 107 @inferred (m*v)[1]
117 end 108 end
118 end 109 end
119 110
120 @testset "TensorMapping binary operations" begin 111 @testset "TensorMapping binary operations" begin
121 struct ScalarMapping{T,R,D} <: TensorMapping{T,R,D} 112 A = ScalingTensor(2.0, (3,))
122 λ::T 113 B = ScalingTensor(3.0, (3,))
123 range_size::NTuple{R,Int}
124 domain_size::NTuple{D,Int}
125 end
126
127 LazyTensors.apply(m::ScalarMapping{T,R}, v, I::Vararg{Any,R}) where {T,R} = m.λ*v[I...]
128 LazyTensors.range_size(m::ScalarMapping) = m.domain_size
129 LazyTensors.domain_size(m::ScalarMapping) = m.range_size
130
131 A = ScalarMapping{Float64,1,1}(2.0, (3,), (3,))
132 B = ScalarMapping{Float64,1,1}(3.0, (3,), (3,))
133 114
134 v = [1.1,1.2,1.3] 115 v = [1.1,1.2,1.3]
135 for i ∈ eachindex(v) 116 for i ∈ eachindex(v)
136 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i] 117 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i]
137 end 118 end
138 119
139 for i ∈ eachindex(v) 120 for i ∈ eachindex(v)
140 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i] 121 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i]
141 end 122 end
123
124 # TODO: Test with size changing tm
125 # TODO: Test for mismatch in dimensions (SizeMismatch?)
142 126
143 @test range_size(A+B) == range_size(A) == range_size(B) 127 @test range_size(A+B) == range_size(A) == range_size(B)
144 @test domain_size(A+B) == domain_size(A) == domain_size(B) 128 @test domain_size(A+B) == domain_size(A) == domain_size(B)
145 129
146 @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64 130 @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64
389 v = rand(ComplexF64, domain_size(tm')...) 373 v = rand(ComplexF64, domain_size(tm')...)
390 @test (tm'*v)[1,2,2,1] isa ComplexF64 374 @test (tm'*v)[1,2,2,1] isa ComplexF64
391 end 375 end
392 376
393 @testset "Inference of application" begin 377 @testset "Inference of application" begin
394 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 378 tm = InflatedTensorMapping(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4))
395 λ::T
396 size::NTuple{D,Int}
397 end
398
399 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
400 LazyTensors.range_size(m::ScalingOperator) = m.size
401 LazyTensors.domain_size(m::ScalingOperator) = m.size
402
403 tm = InflatedTensorMapping(I(2,3),ScalingOperator(2.0, (3,2)),I(3,4))
404 v = rand(domain_size(tm)...) 379 v = rand(domain_size(tm)...)
405 380
406 @inferred apply(tm,v,1,2,3,2,2,4) 381 @inferred apply(tm,v,1,2,3,2,2,4)
407 @inferred (tm*v)[1,2,3,2,2,4] 382 @inferred (tm*v)[1,2,3,2,2,4]
408 end 383 end
409 end 384 end
410 385
411 @testset "InflatedTensorMapping of InflatedTensorMapping" begin 386 @testset "InflatedTensorMapping of InflatedTensorMapping" begin
412 A = ScalingOperator(2.0,(2,3)) 387 A = ScalingTensor(2.0,(2,3))
413 itm = InflatedTensorMapping(I(3,2), A, I(4)) 388 itm = InflatedTensorMapping(I(3,2), A, I(4))
414 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2)) 389 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2))
415 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2)) 390 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2))
416 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4)) 391 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4))
417 392
479 @test LazyTensors.flatten_tuple(((1,2),(3,4),(5,),6)) == (1,2,3,4,5,6) 454 @test LazyTensors.flatten_tuple(((1,2),(3,4),(5,),6)) == (1,2,3,4,5,6)
480 end 455 end
481 456
482 457
483 @testset "LazyOuterProduct" begin 458 @testset "LazyOuterProduct" begin
484 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 459
485 λ::T 460 A = ScalingTensor(2.0, (5,))
486 size::NTuple{D,Int} 461 B = ScalingTensor(3.0, (3,))
487 end 462 C = ScalingTensor(5.0, (3,2))
488
489 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
490 LazyTensors.range_size(m::ScalingOperator) = m.size
491 LazyTensors.domain_size(m::ScalingOperator) = m.size
492
493 A = ScalingOperator(2.0, (5,))
494 B = ScalingOperator(3.0, (3,))
495 C = ScalingOperator(5.0, (3,2))
496 463
497 AB = LazyOuterProduct(A,B) 464 AB = LazyOuterProduct(A,B)
498 @test AB isa TensorMapping{T,2,2} where T 465 @test AB isa TensorMapping{T,2,2} where T
499 @test range_size(AB) == (5,3) 466 @test range_size(AB) == (5,3)
500 @test domain_size(AB) == (5,3) 467 @test domain_size(AB) == (5,3)