comparison test/LazyTensors/lazy_tensor_operations_test.jl @ 1023:52f07c77299d refactor/sbpoperators/inflation

Merge refactor/lazy_tensors
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 21 Mar 2022 09:51:07 +0100
parents bbbc31953367 56fe037641ef
children 5be17f647018
comparison
equal deleted inserted replaced
1022:bbbc31953367 1023:52f07c77299d
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
27 @testset "TensorApplication" begin 38
28 struct SizeDoublingMapping{T,R,D} <: TensorMapping{T,R,D} 39 @testset "LazyTensorApplication" begin
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 size(m*v) == 2 .*size(v) 43 @test size(m*v) == 2 .*size(v)
40 @test (m*v)[0] == (:apply,v,(0,)) 44 @test (m*v)[1] == (:apply,v,(1,))
41 @test (m*m*v)[1] == (:apply,m*v,(1,)) 45 @test (mm*m*v)[1] == (:apply,m*v,(1,))
42 @test (m*m*v)[3] == (:apply,m*v,(3,)) 46 @test (mm*m*v)[3] == (:apply,m*v,(3,))
43 @test (m*m*v)[6] == (:apply,m*v,(6,)) 47 @test (mm*m*v)[6] == (:apply,m*v,(6,))
44 @test_broken BoundsError == (m*m*v)[0]
45 @test_broken BoundsError == (m*m*v)[7]
46 @test_throws MethodError m*m 48 @test_throws MethodError m*m
47 49
48 @test (m*v)[CartesianIndex(2)] == (:apply,v,(2,)) 50 @test (m*v)[CartesianIndex(2)] == (:apply,v,(2,))
49 @test (m*m*v)[CartesianIndex(2)] == (:apply,m*v,(2,)) 51 @test (mm*m*v)[CartesianIndex(2)] == (:apply,m*v,(2,))
50
51 m = SizeDoublingMapping{Int, 2, 1}((3,))
52 @test_throws MethodError m*ones(Int,2,2)
53 @test_throws MethodError m*m*v
54 52
55 m = SizeDoublingMapping{Float64, 2, 2}((3,3)) 53 m = SizeDoublingMapping{Float64, 2, 2}((3,3))
54 mm = SizeDoublingMapping{Float64, 2, 2}((6,6))
56 v = ones(3,3) 55 v = ones(3,3)
57 @test size(m*v) == 2 .*size(v) 56 @test size(m*v) == 2 .*size(v)
58 @test (m*v)[1,2] == (:apply,v,(1,2)) 57 @test (m*v)[1,2] == (:apply,v,(1,2))
59 58
60 @test (m*v)[CartesianIndex(2,3)] == (:apply,v,(2,3)) 59 @test (m*v)[CartesianIndex(2,3)] == (:apply,v,(2,3))
61 @test (m*m*v)[CartesianIndex(4,3)] == (:apply,m*v,(4,3)) 60 @test (mm*m*v)[CartesianIndex(4,3)] == (:apply,m*v,(4,3))
62 61
63 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 62 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] 63 v = [1,2,3]
74 @test m*v isa AbstractVector 64 @test m*v isa AbstractVector
75 @test m*v == [2,4,6] 65 @test m*v == [2,4,6]
76 66
77 m = ScalingOperator{Int,2}(2,(2,2)) 67 m = ScalingTensor(2,(2,2))
78 v = [[1 2];[3 4]] 68 v = [[1 2];[3 4]]
79 @test m*v == [[2 4];[6 8]] 69 @test m*v == [[2 4];[6 8]]
80 @test (m*v)[2,1] == 6 70 @test (m*v)[2,1] == 6
81 71
72 @testset "Error on index out of bounds" begin
73 m = SizeDoublingMapping{Int, 1, 1}((3,))
74 v = [0,1,2]
75
76 @test_throws BoundsError (m*v)[0]
77 @test_throws BoundsError (m*v)[7]
78 end
79
80 @testset "Error on unmatched dimensions" begin
81 v = [0,1,2]
82 m = SizeDoublingMapping{Int, 2, 1}((3,))
83 @test_throws MethodError m*ones(Int,2,2)
84 @test_throws MethodError m*m*v
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
82 @testset "Type calculation" begin 93 @testset "Type calculation" begin
83 m = ScalingOperator{Int,1}(2,(3,)) 94 m = ScalingTensor(2,(3,))
84 v = [1.,2.,3.] 95 v = [1.,2.,3.]
85 @test m*v isa AbstractVector{Float64} 96 @test m*v isa AbstractVector{Float64}
86 @test m*v == [2.,4.,6.] 97 @test m*v == [2.,4.,6.]
87 @inferred m*v 98 @inferred m*v
88 @inferred (m*v)[1] 99 @inferred (m*v)[1]
89 100
90 m = ScalingOperator{Int,2}(2,(2,2)) 101 m = ScalingTensor(2,(2,2))
91 v = [[1. 2.];[3. 4.]] 102 v = [[1. 2.];[3. 4.]]
92 @test m*v == [[2. 4.];[6. 8.]] 103 @test m*v == [[2. 4.];[6. 8.]]
93 @test (m*v)[2,1] == 6. 104 @test (m*v)[2,1] == 6.
94 @inferred m*v 105 @inferred m*v
95 @inferred (m*v)[1] 106 @inferred (m*v)[1]
96 107
97 m = ScalingOperator{ComplexF64,1}(2. +2. *im,(3,)) 108 m = ScalingTensor(2. +2. *im,(3,))
98 v = [1.,2.,3.] 109 v = [1.,2.,3.]
99 @test m*v isa AbstractVector{ComplexF64} 110 @test m*v isa AbstractVector{ComplexF64}
100 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 111 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
101 @inferred m*v 112 @inferred m*v
102 @inferred (m*v)[1] 113 @inferred (m*v)[1]
103 114
104 m = ScalingOperator{ComplexF64,1}(1,(3,)) 115 m = ScalingTensor(1,(3,))
105 v = [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 116 v = [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
106 @test m*v isa AbstractVector{ComplexF64} 117 @test m*v isa AbstractVector{ComplexF64}
107 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im] 118 @test m*v == [2. + 2. *im, 4. + 4. *im, 6. + 6. *im]
108 @inferred m*v 119 @inferred m*v
109 @inferred (m*v)[1] 120 @inferred (m*v)[1]
110 121
111 m = ScalingOperator{Float64,1}(2., (3,)) 122 m = ScalingTensor(2., (3,))
112 v = [[1,2,3], [3,2,1],[1,3,1]] 123 v = [[1,2,3], [3,2,1],[1,3,1]]
113 @test m*v isa AbstractVector{Vector{Float64}} 124 @test m*v isa AbstractVector{Vector{Float64}}
114 @test m*v == [[2.,4.,6.], [6.,4.,2.],[2.,6.,2.]] 125 @test m*v == [[2.,4.,6.], [6.,4.,2.],[2.,6.,2.]]
115 @inferred m*v 126 @inferred m*v
116 @inferred (m*v)[1] 127 @inferred (m*v)[1]
117 end 128 end
118 end 129 end
119 130
120 @testset "TensorMapping binary operations" begin 131
121 struct ScalarMapping{T,R,D} <: TensorMapping{T,R,D} 132 @testset "LazyTensor binary operations" begin
122 λ::T 133 A = ScalingTensor(2.0, (3,))
123 range_size::NTuple{R,Int} 134 B = ScalingTensor(3.0, (3,))
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 135
134 v = [1.1,1.2,1.3] 136 v = [1.1,1.2,1.3]
135 for i ∈ eachindex(v) 137 for i ∈ eachindex(v)
136 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i] 138 @test ((A+B)*v)[i] == 2*v[i] + 3*v[i]
137 end 139 end
138 140
139 for i ∈ eachindex(v) 141 for i ∈ eachindex(v)
140 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i] 142 @test ((A-B)*v)[i] == 2*v[i] - 3*v[i]
141 end 143 end
142 144
145
143 @test range_size(A+B) == range_size(A) == range_size(B) 146 @test range_size(A+B) == range_size(A) == range_size(B)
144 @test domain_size(A+B) == domain_size(A) == domain_size(B) 147 @test domain_size(A+B) == domain_size(A) == domain_size(B)
145 148
146 @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64 149 @test ((A+B)*ComplexF64[1.1,1.2,1.3])[3] isa ComplexF64
147 end 150
148 151 @testset "Error on unmatched sizes" begin
149 152 @test_throws Union{DomainSizeMismatch, RangeSizeMismatch} ScalingTensor(2.0, (3,)) + ScalingTensor(2.0, (4,))
150 @testset "TensorMappingComposition" begin 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 "LazyTensorComposition" begin
151 A = rand(2,3) 163 A = rand(2,3)
152 B = rand(3,4) 164 B = rand(3,4)
153 165
154 Ã = LazyLinearMap(A, (1,), (2,)) 166 Ã = LazyLinearMap(A, (1,), (2,))
155 B̃ = LazyLinearMap(B, (1,), (2,)) 167 B̃ = LazyLinearMap(B, (1,), (2,))
156 168
157 @test Ã∘B̃ isa TensorMappingComposition 169 @test Ã∘B̃ isa LazyTensorComposition
158 @test range_size(Ã∘B̃) == (2,) 170 @test range_size(Ã∘B̃) == (2,)
159 @test domain_size(Ã∘B̃) == (4,) 171 @test domain_size(Ã∘B̃) == (4,)
160 @test_throws SizeMismatch B̃∘Ã 172 @test_throws DomainSizeMismatch B̃∘Ã
161 173
162 # @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)
163 175
164 v = rand(4) 176 v = rand(4)
165 @test Ã∘B̃*v ≈ A*B*v rtol=1e-14 177 @test Ã∘B̃*v ≈ A*B*v rtol=1e-14
169 181
170 @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64 182 @test (Ã∘B̃*ComplexF64[1.,2.,3.,4.])[1] isa ComplexF64
171 @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64 183 @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64
172 end 184 end
173 185
174 @testset "LazyLinearMap" begin 186
175 # Test a standard matrix-vector product 187 @testset "InflatedLazyTensor" begin
176 # mapping vectors of size 4 to vectors of size 3. 188 I(sz...) = IdentityTensor(sz...)
177 A = rand(3,4)
178 Ã = LazyLinearMap(A, (1,), (2,))
179 v = rand(4)
180 w = rand(3)
181
182 @test à isa LazyLinearMap{T,1,1} where T
183 @test à isa TensorMapping{T,1,1} where T
184 @test range_size(Ã) == (3,)
185 @test domain_size(Ã) == (4,)
186
187 @test Ã*ones(4) ≈ A*ones(4) atol=5e-13
188 @test Ã*v ≈ A*v atol=5e-13
189 @test Ã'*w ≈ A'*w
190
191 A = rand(2,3,4)
192 @test_throws DomainError LazyLinearMap(A, (3,1), (2,))
193
194 # Test more exotic mappings
195 B = rand(3,4,2)
196 # Map vectors of size 2 to matrices of size (3,4)
197 B̃ = LazyLinearMap(B, (1,2), (3,))
198 v = rand(2)
199
200 @test range_size(B̃) == (3,4)
201 @test domain_size(B̃) == (2,)
202 @test B̃ isa TensorMapping{T,2,1} where T
203 @test B̃*ones(2) ≈ B[:,:,1] + B[:,:,2] atol=5e-13
204 @test B̃*v ≈ B[:,:,1]*v[1] + B[:,:,2]*v[2] atol=5e-13
205
206 # Map matrices of size (3,2) to vectors of size 4
207 B̃ = LazyLinearMap(B, (2,), (1,3))
208 v = rand(3,2)
209
210 @test range_size(B̃) == (4,)
211 @test domain_size(B̃) == (3,2)
212 @test B̃ isa TensorMapping{T,1,2} where T
213 @test B̃*ones(3,2) ≈ B[1,:,1] + B[2,:,1] + B[3,:,1] +
214 B[1,:,2] + B[2,:,2] + B[3,:,2] atol=5e-13
215 @test B̃*v ≈ B[1,:,1]*v[1,1] + B[2,:,1]*v[2,1] + B[3,:,1]*v[3,1] +
216 B[1,:,2]v[1,2] + B[2,:,2]*v[2,2] + B[3,:,2]*v[3,2] atol=5e-13
217
218
219 # TODO:
220 # @inferred (B̃*v)[2]
221 end
222
223
224 @testset "IdentityMapping" begin
225 @test IdentityMapping{Float64}((4,5)) isa IdentityMapping{T,2} where T
226 @test IdentityMapping{Float64}((4,5)) isa TensorMapping{T,2,2} where T
227 @test IdentityMapping{Float64}((4,5)) == IdentityMapping{Float64}(4,5)
228
229 @test IdentityMapping(3,2) isa IdentityMapping{Float64,2}
230
231 for sz ∈ [(4,5),(3,),(5,6,4)]
232 I = IdentityMapping{Float64}(sz)
233 v = rand(sz...)
234 @test I*v == v
235 @test I'*v == v
236
237 v = rand(ComplexF64,sz...)
238 @test I*v == v
239 @test I'*v == v
240
241 @test range_size(I) == sz
242 @test domain_size(I) == sz
243 end
244
245 I = IdentityMapping{Float64}((4,5))
246 v = rand(4,5)
247 @inferred (I*v)[3,2]
248 @inferred (I'*v)[3,2]
249 @inferred range_size(I)
250
251 @inferred range_dim(I)
252 @inferred domain_dim(I)
253
254 Ã = rand(4,2)
255 A = LazyLinearMap(Ã,(1,),(2,))
256 I1 = IdentityMapping{Float64}(2)
257 I2 = IdentityMapping{Float64}(4)
258 @test A∘I1 == A
259 @test I2∘A == A
260 @test I1∘I1 == I1
261 @test_throws SizeMismatch I1∘A
262 @test_throws SizeMismatch A∘I2
263 @test_throws SizeMismatch I1∘I2
264 end
265
266 @testset "InflatedTensorMapping" begin
267 I(sz...) = IdentityMapping(sz...)
268 189
269 Ã = rand(4,2) 190 Ã = rand(4,2)
270 B̃ = rand(4,2,3) 191 B̃ = rand(4,2,3)
271 C̃ = rand(4,2,3) 192 C̃ = rand(4,2,3)
272 193
273 A = LazyLinearMap(Ã,(1,),(2,)) 194 A = LazyLinearMap(Ã,(1,),(2,))
274 B = LazyLinearMap(B̃,(1,2),(3,)) 195 B = LazyLinearMap(B̃,(1,2),(3,))
275 C = LazyLinearMap(C̃,(1,),(2,3)) 196 C = LazyLinearMap(C̃,(1,),(2,3))
276 197
277 @testset "Constructors" begin 198 @testset "Constructors" begin
278 @test InflatedTensorMapping(I(3,2), A, I(4)) isa TensorMapping{Float64, 4, 4} 199 @test InflatedLazyTensor(I(3,2), A, I(4)) isa LazyTensor{Float64, 4, 4}
279 @test InflatedTensorMapping(I(3,2), B, I(4)) isa TensorMapping{Float64, 5, 4} 200 @test InflatedLazyTensor(I(3,2), B, I(4)) isa LazyTensor{Float64, 5, 4}
280 @test InflatedTensorMapping(I(3), C, I(2,3)) isa TensorMapping{Float64, 4, 5} 201 @test InflatedLazyTensor(I(3), C, I(2,3)) isa LazyTensor{Float64, 4, 5}
281 @test InflatedTensorMapping(C, I(2,3)) isa TensorMapping{Float64, 3, 4} 202 @test InflatedLazyTensor(C, I(2,3)) isa LazyTensor{Float64, 3, 4}
282 @test InflatedTensorMapping(I(3), C) isa TensorMapping{Float64, 2, 3} 203 @test InflatedLazyTensor(I(3), C) isa LazyTensor{Float64, 2, 3}
283 @test InflatedTensorMapping(I(3), I(2,3)) isa TensorMapping{Float64, 3, 3} 204 @test InflatedLazyTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3}
284 end 205 end
285 206
286 @testset "Range and domain size" begin 207 @testset "Range and domain size" begin
287 @test range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 208 @test range_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,4,4)
288 @test domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 209 @test domain_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,2,4)
289 210
290 @test range_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,4,2,4) 211 @test range_size(InflatedLazyTensor(I(3,2), B, I(4))) == (3,2,4,2,4)
291 @test domain_size(InflatedTensorMapping(I(3,2), B, I(4))) == (3,2,3,4) 212 @test domain_size(InflatedLazyTensor(I(3,2), B, I(4))) == (3,2,3,4)
292 213
293 @test range_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,4,2,3) 214 @test range_size(InflatedLazyTensor(I(3), C, I(2,3))) == (3,4,2,3)
294 @test domain_size(InflatedTensorMapping(I(3), C, I(2,3))) == (3,2,3,2,3) 215 @test domain_size(InflatedLazyTensor(I(3), C, I(2,3))) == (3,2,3,2,3)
295 216
296 @inferred range_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,4,4) 217 @inferred range_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,4,4)
297 @inferred domain_size(InflatedTensorMapping(I(3,2), A, I(4))) == (3,2,2,4) 218 @inferred domain_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,2,4)
298 end 219 end
299 220
300 @testset "Application" begin 221 @testset "Application" begin
301 # Testing regular application and transposed application with inflation "before", "after" and "before and after". 222 # Testing regular application and transposed application with inflation "before", "after" and "before and after".
302 # The inflated tensor mappings are chosen to preserve, reduce and increase the dimension of the result compared to the input. 223 # The inflated tensor mappings are chosen to preserve, reduce and increase the dimension of the result compared to the input.
303 tests = [ 224 cases = [
304 ( 225 (
305 InflatedTensorMapping(I(3,2), A, I(4)), 226 InflatedLazyTensor(I(3,2), A, I(4)),
306 (v-> @tullio res[a,b,c,d] := Ã[c,i]*v[a,b,i,d]), # Expected result of apply 227 (v-> @tullio res[a,b,c,d] := Ã[c,i]*v[a,b,i,d]), # Expected result of apply
307 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose 228 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose
308 ), 229 ),
309 ( 230 (
310 InflatedTensorMapping(I(3,2), B, I(4)), 231 InflatedLazyTensor(I(3,2), B, I(4)),
311 (v-> @tullio res[a,b,c,d,e] := B̃[c,d,i]*v[a,b,i,e]), 232 (v-> @tullio res[a,b,c,d,e] := B̃[c,d,i]*v[a,b,i,e]),
312 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]), 233 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]),
313 ), 234 ),
314 ( 235 (
315 InflatedTensorMapping(I(3,2), C, I(4)), 236 InflatedLazyTensor(I(3,2), C, I(4)),
316 (v-> @tullio res[a,b,c,d] := C̃[c,i,j]*v[a,b,i,j,d]), 237 (v-> @tullio res[a,b,c,d] := C̃[c,i,j]*v[a,b,i,j,d]),
317 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]), 238 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]),
318 ), 239 ),
319 ( 240 (
320 InflatedTensorMapping(I(3,2), A), 241 InflatedLazyTensor(I(3,2), A),
321 (v-> @tullio res[a,b,c] := Ã[c,i]*v[a,b,i]), 242 (v-> @tullio res[a,b,c] := Ã[c,i]*v[a,b,i]),
322 (v-> @tullio res[a,b,c] := Ã[i,c]*v[a,b,i]), 243 (v-> @tullio res[a,b,c] := Ã[i,c]*v[a,b,i]),
323 ), 244 ),
324 ( 245 (
325 InflatedTensorMapping(I(3,2), B), 246 InflatedLazyTensor(I(3,2), B),
326 (v-> @tullio res[a,b,c,d] := B̃[c,d,i]*v[a,b,i]), 247 (v-> @tullio res[a,b,c,d] := B̃[c,d,i]*v[a,b,i]),
327 (v-> @tullio res[a,b,c] := B̃[i,j,c]*v[a,b,i,j]), 248 (v-> @tullio res[a,b,c] := B̃[i,j,c]*v[a,b,i,j]),
328 ), 249 ),
329 ( 250 (
330 InflatedTensorMapping(I(3,2), C), 251 InflatedLazyTensor(I(3,2), C),
331 (v-> @tullio res[a,b,c] := C̃[c,i,j]*v[a,b,i,j]), 252 (v-> @tullio res[a,b,c] := C̃[c,i,j]*v[a,b,i,j]),
332 (v-> @tullio res[a,b,c,d] := C̃[i,c,d]*v[a,b,i]), 253 (v-> @tullio res[a,b,c,d] := C̃[i,c,d]*v[a,b,i]),
333 ), 254 ),
334 ( 255 (
335 InflatedTensorMapping(A,I(4)), 256 InflatedLazyTensor(A,I(4)),
336 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]), 257 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]),
337 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]), 258 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]),
338 ), 259 ),
339 ( 260 (
340 InflatedTensorMapping(B,I(4)), 261 InflatedLazyTensor(B,I(4)),
341 (v-> @tullio res[a,b,c] := B̃[a,b,i]*v[i,c]), 262 (v-> @tullio res[a,b,c] := B̃[a,b,i]*v[i,c]),
342 (v-> @tullio res[a,b] := B̃[i,j,a]*v[i,j,b]), 263 (v-> @tullio res[a,b] := B̃[i,j,a]*v[i,j,b]),
343 ), 264 ),
344 ( 265 (
345 InflatedTensorMapping(C,I(4)), 266 InflatedLazyTensor(C,I(4)),
346 (v-> @tullio res[a,b] := C̃[a,i,j]*v[i,j,b]), 267 (v-> @tullio res[a,b] := C̃[a,i,j]*v[i,j,b]),
347 (v-> @tullio res[a,b,c] := C̃[i,a,b]*v[i,c]), 268 (v-> @tullio res[a,b,c] := C̃[i,a,b]*v[i,c]),
348 ), 269 ),
349 ] 270 ]
350 271
351 @testset "apply" begin 272 @testset "$tm" for (tm, true_apply, true_apply_transpose) ∈ cases
352 for i ∈ 1:length(tests) 273 v = rand(domain_size(tm)...)
353 tm = tests[i][1] 274 @test tm*v ≈ true_apply(v) rtol=1e-14
354 v = rand(domain_size(tm)...) 275
355 true_value = tests[i][2](v) 276 v = rand(range_size(tm)...)
356 @test tm*v ≈ true_value rtol=1e-14 277 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14
357 end
358 end 278 end
359 279
360 @testset "apply_transpose" begin
361 for i ∈ 1:length(tests)
362 tm = tests[i][1]
363 v = rand(range_size(tm)...)
364 true_value = tests[i][3](v)
365 @test tm'*v ≈ true_value rtol=1e-14
366 end
367 end
368
369 @testset "application to other type" begin 280 @testset "application to other type" begin
370 tm = InflatedTensorMapping(I(3,2), A, I(4)) 281 tm = InflatedLazyTensor(I(3,2), A, I(4))
371 282
372 v = rand(ComplexF64, domain_size(tm)...) 283 v = rand(ComplexF64, domain_size(tm)...)
373 @test (tm*v)[1,2,3,1] isa ComplexF64 284 @test (tm*v)[1,2,3,1] isa ComplexF64
374 285
375 v = rand(ComplexF64, domain_size(tm')...) 286 v = rand(ComplexF64, domain_size(tm')...)
376 @test (tm'*v)[1,2,2,1] isa ComplexF64 287 @test (tm'*v)[1,2,2,1] isa ComplexF64
377 end 288 end
378 289
379 @testset "Inference of application" begin 290 @testset "Inference of application" begin
380 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 291 tm = InflatedLazyTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4))
381 λ::T
382 size::NTuple{D,Int}
383 end
384
385 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
386 LazyTensors.range_size(m::ScalingOperator) = m.size
387 LazyTensors.domain_size(m::ScalingOperator) = m.size
388
389 tm = InflatedTensorMapping(I(2,3),ScalingOperator(2.0, (3,2)),I(3,4))
390 v = rand(domain_size(tm)...) 292 v = rand(domain_size(tm)...)
391 293
392 @inferred apply(tm,v,1,2,3,2,2,4) 294 @inferred apply(tm,v,1,2,3,2,2,4)
393 @inferred (tm*v)[1,2,3,2,2,4] 295 @inferred (tm*v)[1,2,3,2,2,4]
394 end 296 end
395 end 297 end
396 298
397 @testset "InflatedTensorMapping of InflatedTensorMapping" begin 299 @testset "InflatedLazyTensor of InflatedLazyTensor" begin
398 A = ScalingOperator(2.0,(2,3)) 300 A = ScalingTensor(2.0,(2,3))
399 itm = InflatedTensorMapping(I(3,2), A, I(4)) 301 itm = InflatedLazyTensor(I(3,2), A, I(4))
400 @test InflatedTensorMapping(I(4), itm, I(2)) == InflatedTensorMapping(I(4,3,2), A, I(4,2)) 302 @test InflatedLazyTensor(I(4), itm, I(2)) == InflatedLazyTensor(I(4,3,2), A, I(4,2))
401 @test InflatedTensorMapping(itm, I(2)) == InflatedTensorMapping(I(3,2), A, I(4,2)) 303 @test InflatedLazyTensor(itm, I(2)) == InflatedLazyTensor(I(3,2), A, I(4,2))
402 @test InflatedTensorMapping(I(4), itm) == InflatedTensorMapping(I(4,3,2), A, I(4)) 304 @test InflatedLazyTensor(I(4), itm) == InflatedLazyTensor(I(4,3,2), A, I(4))
403 305
404 @test InflatedTensorMapping(I(2), I(2), I(2)) isa InflatedTensorMapping # The constructor should always return its type. 306 @test InflatedLazyTensor(I(2), I(2), I(2)) isa InflatedLazyTensor # The constructor should always return its type.
405 end 307 end
406 end 308 end
407
408 @testset "split_index" begin
409 @test LazyTensors.split_index(Val(2),Val(1),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,5,6),(3,4))
410 @test LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,:,:,5,6),(3,4))
411 @test LazyTensors.split_index(Val(3),Val(1),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,5,6),(4,))
412 @test LazyTensors.split_index(Val(3),Val(2),Val(1),Val(2),1,2,3,4,5,6) == ((1,2,3,:,:,5,6),(4,))
413 @test LazyTensors.split_index(Val(1),Val(1),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,4,5,6),(2,3))
414 @test LazyTensors.split_index(Val(1),Val(2),Val(2),Val(3),1,2,3,4,5,6) == ((1,:,:,4,5,6),(2,3))
415
416 @test LazyTensors.split_index(Val(0),Val(1),Val(3),Val(3),1,2,3,4,5,6) == ((:,4,5,6),(1,2,3))
417 @test LazyTensors.split_index(Val(3),Val(1),Val(3),Val(0),1,2,3,4,5,6) == ((1,2,3,:),(4,5,6))
418
419 @inferred LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,2,2,4)
420 end
421
422 @testset "slice_tuple" begin
423 @test LazyTensors.slice_tuple((1,2,3),Val(1), Val(3)) == (1,2,3)
424 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(2), Val(5)) == (2,3,4,5)
425 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(1), Val(3)) == (1,2,3)
426 @test LazyTensors.slice_tuple((1,2,3,4,5,6),Val(4), Val(6)) == (4,5,6)
427 end
428
429 @testset "split_tuple" begin
430 @testset "2 parts" begin
431 @test LazyTensors.split_tuple((),Val(0)) == ((),())
432 @test LazyTensors.split_tuple((1,),Val(0)) == ((),(1,))
433 @test LazyTensors.split_tuple((1,),Val(1)) == ((1,),())
434
435 @test LazyTensors.split_tuple((1,2,3,4),Val(0)) == ((),(1,2,3,4))
436 @test LazyTensors.split_tuple((1,2,3,4),Val(1)) == ((1,),(2,3,4))
437 @test LazyTensors.split_tuple((1,2,3,4),Val(2)) == ((1,2),(3,4))
438 @test LazyTensors.split_tuple((1,2,3,4),Val(3)) == ((1,2,3),(4,))
439 @test LazyTensors.split_tuple((1,2,3,4),Val(4)) == ((1,2,3,4),())
440
441 @test LazyTensors.split_tuple((1,2,true,4),Val(3)) == ((1,2,true),(4,))
442
443 @inferred LazyTensors.split_tuple((1,2,3,4),Val(3))
444 @inferred LazyTensors.split_tuple((1,2,true,4),Val(3))
445 end
446
447 @testset "3 parts" begin
448 @test LazyTensors.split_tuple((),Val(0),Val(0)) == ((),(),())
449 @test LazyTensors.split_tuple((1,2,3),Val(1), Val(1)) == ((1,),(2,),(3,))
450 @test LazyTensors.split_tuple((1,true,3),Val(1), Val(1)) == ((1,),(true,),(3,))
451
452 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(1),Val(2)) == ((1,),(2,3),(4,5,6))
453 @test LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2)) == ((1,2,3),(4,5),(6,))
454
455 @inferred LazyTensors.split_tuple((1,2,3,4,5,6),Val(3),Val(2))
456 @inferred LazyTensors.split_tuple((1,true,3),Val(1), Val(1))
457 end
458 end
459
460 @testset "flatten_tuple" begin
461 @test LazyTensors.flatten_tuple((1,)) == (1,)
462 @test LazyTensors.flatten_tuple((1,2,3,4,5,6)) == (1,2,3,4,5,6)
463 @test LazyTensors.flatten_tuple((1,2,(3,4),5,6)) == (1,2,3,4,5,6)
464 @test LazyTensors.flatten_tuple((1,2,(3,(4,5)),6)) == (1,2,3,4,5,6)
465 @test LazyTensors.flatten_tuple(((1,2),(3,4),(5,),6)) == (1,2,3,4,5,6)
466 end
467
468 309
469 @testset "LazyOuterProduct" begin 310 @testset "LazyOuterProduct" begin
470 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 311 A = ScalingTensor(2.0, (5,))
471 λ::T 312 B = ScalingTensor(3.0, (3,))
472 size::NTuple{D,Int} 313 C = ScalingTensor(5.0, (3,2))
473 end
474
475 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
476 LazyTensors.range_size(m::ScalingOperator) = m.size
477 LazyTensors.domain_size(m::ScalingOperator) = m.size
478
479 A = ScalingOperator(2.0, (5,))
480 B = ScalingOperator(3.0, (3,))
481 C = ScalingOperator(5.0, (3,2))
482 314
483 AB = LazyOuterProduct(A,B) 315 AB = LazyOuterProduct(A,B)
484 @test AB isa TensorMapping{T,2,2} where T 316 @test AB isa LazyTensor{T,2,2} where T
485 @test range_size(AB) == (5,3) 317 @test range_size(AB) == (5,3)
486 @test domain_size(AB) == (5,3) 318 @test domain_size(AB) == (5,3)
487 319
488 v = rand(range_size(AB)...) 320 v = rand(range_size(AB)...)
489 @test AB*v == 6*v 321 @test AB*v == 6*v
490 322
491 ABC = LazyOuterProduct(A,B,C) 323 ABC = LazyOuterProduct(A,B,C)
492 324
493 @test ABC isa TensorMapping{T,4,4} where T 325 @test ABC isa LazyTensor{T,4,4} where T
494 @test range_size(ABC) == (5,3,3,2) 326 @test range_size(ABC) == (5,3,3,2)
495 @test domain_size(ABC) == (5,3,3,2) 327 @test domain_size(ABC) == (5,3,3,2)
496 328
497 @test A⊗B == AB 329 @test A⊗B == AB
498 @test A⊗B⊗C == ABC 330 @test A⊗B⊗C == ABC
513 B̃Ã = LazyOuterProduct(B̃,Ã) 345 B̃Ã = LazyOuterProduct(B̃,Ã)
514 @tullio BAv[k,i] := A[i,j]*B[k,l,m]*v₂[l,m,j] 346 @tullio BAv[k,i] := A[i,j]*B[k,l,m]*v₂[l,m,j]
515 @test B̃Ã*v₂ ≈ BAv 347 @test B̃Ã*v₂ ≈ BAv
516 348
517 @testset "Indentity mapping arguments" begin 349 @testset "Indentity mapping arguments" begin
518 @test LazyOuterProduct(IdentityMapping(3,2), IdentityMapping(1,2)) == IdentityMapping(3,2,1,2) 350 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2)
519 351
520 Ã = LazyLinearMap(A,(1,),(2,)) 352 Ã = LazyLinearMap(A,(1,),(2,))
521 @test LazyOuterProduct(IdentityMapping(3,2), Ã) == InflatedTensorMapping(IdentityMapping(3,2),Ã) 353 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedLazyTensor(IdentityTensor(3,2),Ã)
522 @test LazyOuterProduct(Ã, IdentityMapping(3,2)) == InflatedTensorMapping(Ã,IdentityMapping(3,2)) 354 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedLazyTensor(Ã,IdentityTensor(3,2))
523 355
524 I1 = IdentityMapping(3,2) 356 I1 = IdentityTensor(3,2)
525 I2 = IdentityMapping(4) 357 I2 = IdentityTensor(4)
526 @test I1⊗Ã⊗I2 == InflatedTensorMapping(I1, Ã, I2) 358 @test I1⊗Ã⊗I2 == InflatedLazyTensor(I1, Ã, I2)
527 end 359 end
528
529 end 360 end
530 361
531 @testset "inflate" begin 362 @testset "inflate" begin
532 struct ScalingOperator{T,D} <: TensorMapping{T,D,D} 363 I = LazyTensors.inflate(IdentityTensor(),(3,4,5,6), 2)
533 λ::T 364 @test I isa LazyTensor{Float64, 3,3}
534 size::NTuple{D,Int}
535 end
536
537 LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
538 LazyTensors.range_size(m::ScalingOperator) = m.size
539 LazyTensors.domain_size(m::ScalingOperator) = m.size
540
541
542 I = LazyTensors.inflate(IdentityMapping(),(3,4,5,6), 2)
543 @test I isa TensorMapping{Float64, 3,3}
544 @test range_size(I) == (3,5,6) 365 @test range_size(I) == (3,5,6)
545 @test domain_size(I) == (3,5,6) 366 @test domain_size(I) == (3,5,6)
546 367
547 # TODO: More tests 368 # TODO: More tests
548 369