changeset 503:fbbb3733650c

Merge in feature/outer_product
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 06 Nov 2020 21:37:10 +0100
parents d8075fb14418 (current diff) f4c245feb273 (diff)
children 21fba50cb5b0 27e64b3d3efa b7e42384053a
files
diffstat 2 files changed, 116 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Thu Nov 05 15:18:11 2020 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Fri Nov 06 21:37:10 2020 +0100
@@ -310,6 +310,59 @@
 flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify?
 flatten_tuple(ts::Vararg) = flatten_tuple(ts)
 
+"""
+   LazyOuterProduct(tms...)
+
+Creates a `TensorMappingComposition` for the outerproduct of `tms...`.
+This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
+
+First let
+```math
+A = A_{I,J}
+B = B_{M,N}
+C = C_{P,Q}
+```
+
+where ``I``, ``M``, ``P`` are  multi-indexes for the ranges of ``A``, ``B``, ``C``, and ``J``, ``N``, ``Q`` are multi-indexes of the domains.
+
+We use ``⊗`` to denote the outer product
+```math
+(A⊗B)_{IM,JN} = A_{I,J}B_{M,N}
+```
+
+We note that
+```math
+A⊗B⊗C = (A⊗B⊗C)_{IMP,JNQ} = A_{I,J}B_{M,N}C_{P,Q}
+```
+And that
+```math
+A⊗B⊗C = (A⊗I_{|M|}⊗I_{|P|})(I_{|J|}⊗B⊗I_{|P|})(I_{|J|}⊗I_{|N|}⊗C)
+```
+where |.| of a multi-index is a vector of sizes for each dimension. ``I_v`` denotes the identity tensor of size ``v[i]`` in each direction
+To apply ``A⊗B⊗C`` we evaluate
+
+(A⊗B⊗C)v = [(A⊗I_{|M|}⊗I_{|P|})  [(I_{|J|}⊗B⊗I_{|P|}) [(I_{|J|}⊗I_{|N|}⊗C)v]]]
+"""
+function LazyOuterProduct end
+export LazyOuterProduct
+
+function LazyOuterProduct(tm1::TensorMapping{T}, tm2::TensorMapping{T}) where T
+    itm1 = InflatedTensorMapping(tm1, IdentityMapping{T}(range_size(tm2)))
+    itm2 = InflatedTensorMapping(IdentityMapping{T}(domain_size(tm1)),tm2)
+
+    return itm1∘itm2
+end
+
+LazyOuterProduct(t1::IdentityMapping{T}, t2::IdentityMapping{T}) where T = IdentityMapping{T}(t1.size...,t2.size...)
+LazyOuterProduct(t1::TensorMapping, t2::IdentityMapping) = InflatedTensorMapping(t1, t2)
+LazyOuterProduct(t1::IdentityMapping, t2::TensorMapping) = InflatedTensorMapping(t1, t2)
+
+LazyOuterProduct(tms::Vararg{TensorMapping}) = foldl(LazyOuterProduct, tms)
+
+⊗(a::TensorMapping, b::TensorMapping) = LazyOuterProduct(a,b)
+export ⊗
+
+
 function check_domain_size(tm::TensorMapping, sz)
     if domain_size(tm) != sz
         throw(SizeMismatch(tm,sz))
--- a/test/testLazyTensors.jl	Thu Nov 05 15:18:11 2020 +0100
+++ b/test/testLazyTensors.jl	Fri Nov 06 21:37:10 2020 +0100
@@ -423,4 +423,67 @@
     @test LazyTensors.flatten_tuple(((1,2),(3,4),(5,),6)) == (1,2,3,4,5,6)
 end
 
+
+@testset "LazyOuterProduct" begin
+    struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
+        λ::T
+        size::NTuple{D,Int}
+    end
+
+    LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Index,D}) where {T,D} = m.λ*v[I]
+    LazyTensors.range_size(m::ScalingOperator) = m.size
+    LazyTensors.domain_size(m::ScalingOperator) = m.size
+
+    A = ScalingOperator(2.0, (5,))
+    B = ScalingOperator(3.0, (3,))
+    C = ScalingOperator(5.0, (3,2))
+
+    AB = LazyOuterProduct(A,B)
+    @test AB isa TensorMapping{T,2,2} where T
+    @test range_size(AB) == (5,3)
+    @test domain_size(AB) == (5,3)
+
+    v = rand(range_size(AB)...)
+    @test AB*v == 6*v
+
+    ABC = LazyOuterProduct(A,B,C)
+
+    @test ABC isa TensorMapping{T,4,4} where T
+    @test range_size(ABC) == (5,3,3,2)
+    @test domain_size(ABC) == (5,3,3,2)
+
+    @test A⊗B == AB
+    @test A⊗B⊗C == ABC
+
+    A = rand(3,2)
+    B = rand(2,4,3)
+
+    v₁ = rand(2,4,3)
+    v₂ = rand(4,3,2)
+
+    Ã = LazyLinearMap(A,(1,),(2,))
+    B̃ = LazyLinearMap(B,(1,),(2,3))
+
+    ÃB̃ = LazyOuterProduct(Ã,B̃)
+    @tullio ABv[i,k] := A[i,j]*B[k,l,m]*v₁[j,l,m]
+    @test ÃB̃*v₁ ≈ ABv
+
+    B̃Ã = LazyOuterProduct(B̃,Ã)
+    @tullio BAv[k,i] := A[i,j]*B[k,l,m]*v₂[l,m,j]
+    @test B̃Ã*v₂ ≈ BAv
+
+    @testset "Indentity mapping arguments" begin
+        @test LazyOuterProduct(IdentityMapping(3,2), IdentityMapping(1,2)) == IdentityMapping(3,2,1,2)
+
+        Ã = LazyLinearMap(A,(1,),(2,))
+        @test LazyOuterProduct(IdentityMapping(3,2), Ã) == InflatedTensorMapping(IdentityMapping(3,2),Ã)
+        @test LazyOuterProduct(Ã, IdentityMapping(3,2)) == InflatedTensorMapping(Ã,IdentityMapping(3,2))
+
+        I1 = IdentityMapping(3,2)
+        I2 = IdentityMapping(4)
+        @test I1⊗Ã⊗I2 == InflatedTensorMapping(I1, Ã, I2)
+    end
+
 end
+
+end