comparison test/LazyTensors/lazy_tensor_operations_test.jl @ 1017:6abbb2c6c3e4 refactor/lazy_tensors

Remove the Lazy prefix on some types
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 21 Mar 2022 15:22:22 +0100
parents 56fe037641ef
children f857057e61e6 2e606d4c0ab1
comparison
equal deleted inserted replaced
1016:5c8c148c56a3 1017:6abbb2c6c3e4
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)
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 end
185 185
186 186
187 @testset "InflatedLazyTensor" begin 187 @testset "InflatedTensor" begin
188 I(sz...) = IdentityTensor(sz...) 188 I(sz...) = IdentityTensor(sz...)
189 189
190 Ã = rand(4,2) 190 Ã = rand(4,2)
191 B̃ = rand(4,2,3) 191 B̃ = rand(4,2,3)
192 C̃ = rand(4,2,3) 192 C̃ = rand(4,2,3)
193 193
194 A = LazyLinearMap(Ã,(1,),(2,)) 194 A = DenseTensor(Ã,(1,),(2,))
195 B = LazyLinearMap(B̃,(1,2),(3,)) 195 B = DenseTensor(B̃,(1,2),(3,))
196 C = LazyLinearMap(C̃,(1,),(2,3)) 196 C = DenseTensor(C̃,(1,),(2,3))
197 197
198 @testset "Constructors" begin 198 @testset "Constructors" begin
199 @test InflatedLazyTensor(I(3,2), A, I(4)) isa LazyTensor{Float64, 4, 4} 199 @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} 200 @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} 201 @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} 202 @test InflatedTensor(C, I(2,3)) isa LazyTensor{Float64, 3, 4}
203 @test InflatedLazyTensor(I(3), C) isa LazyTensor{Float64, 2, 3} 203 @test InflatedTensor(I(3), C) isa LazyTensor{Float64, 2, 3}
204 @test InflatedLazyTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3} 204 @test InflatedTensor(I(3), I(2,3)) isa LazyTensor{Float64, 3, 3}
205 end 205 end
206 206
207 @testset "Range and domain size" begin 207 @testset "Range and domain size" begin
208 @test range_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,4,4) 208 @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) 209 @test domain_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,2,4)
210 210
211 @test range_size(InflatedLazyTensor(I(3,2), B, I(4))) == (3,2,4,2,4) 211 @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) 212 @test domain_size(InflatedTensor(I(3,2), B, I(4))) == (3,2,3,4)
213 213
214 @test range_size(InflatedLazyTensor(I(3), C, I(2,3))) == (3,4,2,3) 214 @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) 215 @test domain_size(InflatedTensor(I(3), C, I(2,3))) == (3,2,3,2,3)
216 216
217 @inferred range_size(InflatedLazyTensor(I(3,2), A, I(4))) == (3,2,4,4) 217 @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) 218 @inferred domain_size(InflatedTensor(I(3,2), A, I(4))) == (3,2,2,4)
219 end 219 end
220 220
221 @testset "Application" begin 221 @testset "Application" begin
222 # 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".
223 # 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.
224 cases = [ 224 cases = [
225 ( 225 (
226 InflatedLazyTensor(I(3,2), A, I(4)), 226 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 227 (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 228 (v-> @tullio res[a,b,c,d] := Ã[i,c]*v[a,b,i,d]), # Expected result of apply_transpose
229 ), 229 ),
230 ( 230 (
231 InflatedLazyTensor(I(3,2), B, I(4)), 231 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]), 232 (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]), 233 (v-> @tullio res[a,b,c,d] := B̃[i,j,c]*v[a,b,i,j,d]),
234 ), 234 ),
235 ( 235 (
236 InflatedLazyTensor(I(3,2), C, I(4)), 236 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]), 237 (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]), 238 (v-> @tullio res[a,b,c,d,e] := C̃[i,c,d]*v[a,b,i,e]),
239 ), 239 ),
240 ( 240 (
241 InflatedLazyTensor(I(3,2), A), 241 InflatedTensor(I(3,2), A),
242 (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]),
243 (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]),
244 ), 244 ),
245 ( 245 (
246 InflatedLazyTensor(I(3,2), B), 246 InflatedTensor(I(3,2), B),
247 (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]),
248 (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]),
249 ), 249 ),
250 ( 250 (
251 InflatedLazyTensor(I(3,2), C), 251 InflatedTensor(I(3,2), C),
252 (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]),
253 (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]),
254 ), 254 ),
255 ( 255 (
256 InflatedLazyTensor(A,I(4)), 256 InflatedTensor(A,I(4)),
257 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]), 257 (v-> @tullio res[a,b] := Ã[a,i]*v[i,b]),
258 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]), 258 (v-> @tullio res[a,b] := Ã[i,a]*v[i,b]),
259 ), 259 ),
260 ( 260 (
261 InflatedLazyTensor(B,I(4)), 261 InflatedTensor(B,I(4)),
262 (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]),
263 (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]),
264 ), 264 ),
265 ( 265 (
266 InflatedLazyTensor(C,I(4)), 266 InflatedTensor(C,I(4)),
267 (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]),
268 (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]),
269 ), 269 ),
270 ] 270 ]
271 271
276 v = rand(range_size(tm)...) 276 v = rand(range_size(tm)...)
277 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14 277 @test tm'*v ≈ true_apply_transpose(v) rtol=1e-14
278 end 278 end
279 279
280 @testset "application to other type" begin 280 @testset "application to other type" begin
281 tm = InflatedLazyTensor(I(3,2), A, I(4)) 281 tm = InflatedTensor(I(3,2), A, I(4))
282 282
283 v = rand(ComplexF64, domain_size(tm)...) 283 v = rand(ComplexF64, domain_size(tm)...)
284 @test (tm*v)[1,2,3,1] isa ComplexF64 284 @test (tm*v)[1,2,3,1] isa ComplexF64
285 285
286 v = rand(ComplexF64, domain_size(tm')...) 286 v = rand(ComplexF64, domain_size(tm')...)
287 @test (tm'*v)[1,2,2,1] isa ComplexF64 287 @test (tm'*v)[1,2,2,1] isa ComplexF64
288 end 288 end
289 289
290 @testset "Inference of application" begin 290 @testset "Inference of application" begin
291 tm = InflatedLazyTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4)) 291 tm = InflatedTensor(I(2,3),ScalingTensor(2.0, (3,2)),I(3,4))
292 v = rand(domain_size(tm)...) 292 v = rand(domain_size(tm)...)
293 293
294 @inferred apply(tm,v,1,2,3,2,2,4) 294 @inferred apply(tm,v,1,2,3,2,2,4)
295 @inferred (tm*v)[1,2,3,2,2,4] 295 @inferred (tm*v)[1,2,3,2,2,4]
296 end 296 end
297 end 297 end
298 298
299 @testset "InflatedLazyTensor of InflatedLazyTensor" begin 299 @testset "InflatedTensor of InflatedTensor" begin
300 A = ScalingTensor(2.0,(2,3)) 300 A = ScalingTensor(2.0,(2,3))
301 itm = InflatedLazyTensor(I(3,2), A, I(4)) 301 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)) 302 @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)) 303 @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)) 304 @test InflatedTensor(I(4), itm) == InflatedTensor(I(4,3,2), A, I(4))
305 305
306 @test InflatedLazyTensor(I(2), I(2), I(2)) isa InflatedLazyTensor # The constructor should always return its type. 306 @test InflatedTensor(I(2), I(2), I(2)) isa InflatedTensor # The constructor should always return its type.
307 end 307 end
308 end 308 end
309 309
310 @testset "LazyOuterProduct" begin 310 @testset "LazyOuterProduct" begin
311 A = ScalingTensor(2.0, (5,)) 311 A = ScalingTensor(2.0, (5,))
333 B = rand(2,4,3) 333 B = rand(2,4,3)
334 334
335 v₁ = rand(2,4,3) 335 v₁ = rand(2,4,3)
336 v₂ = rand(4,3,2) 336 v₂ = rand(4,3,2)
337 337
338 Ã = LazyLinearMap(A,(1,),(2,)) 338 Ã = DenseTensor(A,(1,),(2,))
339 B̃ = LazyLinearMap(B,(1,),(2,3)) 339 B̃ = DenseTensor(B,(1,),(2,3))
340 340
341 ÃB̃ = LazyOuterProduct(Ã,B̃) 341 ÃB̃ = LazyOuterProduct(Ã,B̃)
342 @tullio ABv[i,k] := A[i,j]*B[k,l,m]*v₁[j,l,m] 342 @tullio ABv[i,k] := A[i,j]*B[k,l,m]*v₁[j,l,m]
343 @test ÃB̃*v₁ ≈ ABv 343 @test ÃB̃*v₁ ≈ ABv
344 344
347 @test B̃Ã*v₂ ≈ BAv 347 @test B̃Ã*v₂ ≈ BAv
348 348
349 @testset "Indentity mapping arguments" begin 349 @testset "Indentity mapping arguments" begin
350 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2) 350 @test LazyOuterProduct(IdentityTensor(3,2), IdentityTensor(1,2)) == IdentityTensor(3,2,1,2)
351 351
352 Ã = LazyLinearMap(A,(1,),(2,)) 352 Ã = DenseTensor(A,(1,),(2,))
353 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedLazyTensor(IdentityTensor(3,2),Ã) 353 @test LazyOuterProduct(IdentityTensor(3,2), Ã) == InflatedTensor(IdentityTensor(3,2),Ã)
354 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedLazyTensor(Ã,IdentityTensor(3,2)) 354 @test LazyOuterProduct(Ã, IdentityTensor(3,2)) == InflatedTensor(Ã,IdentityTensor(3,2))
355 355
356 I1 = IdentityTensor(3,2) 356 I1 = IdentityTensor(3,2)
357 I2 = IdentityTensor(4) 357 I2 = IdentityTensor(4)
358 @test I1⊗Ã⊗I2 == InflatedLazyTensor(I1, Ã, I2) 358 @test I1⊗Ã⊗I2 == InflatedTensor(I1, Ã, I2)
359 end 359 end
360 end 360 end