Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 995:1ba8a398af9c refactor/lazy_tensors

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Rename types
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 21:14:47 +0100
parents 55ab7801c45f
children 20c376dffe84
comparison
equal deleted inserted replaced
994:55ab7801c45f 995:1ba8a398af9c
1 # TBD: Is there a good way to split this file? 1 # TBD: Is there a good way to split this file?
2 2
3 """ 3 """
4 LazyTensorMappingApplication{T,R,D} <: LazyArray{T,R} 4 LazyTensorApplication{T,R,D} <: LazyArray{T,R}
5 5
6 Struct for lazy application of a TensorMapping. Created using `*`. 6 Struct for lazy application of a LazyTensor. Created using `*`.
7 7
8 Allows the result of a `TensorMapping` applied to a vector to be treated as an `AbstractArray`. 8 Allows the result of a `LazyTensor` applied to a vector to be treated as an `AbstractArray`.
9 With a mapping `m` and a vector `v` the LazyTensorMappingApplication object can be created by `m*v`. 9 With a mapping `m` and a vector `v` the LazyTensorApplication object can be created by `m*v`.
10 The actual result will be calcualted when indexing into `m*v`. 10 The actual result will be calcualted when indexing into `m*v`.
11 """ 11 """
12 struct LazyTensorMappingApplication{T,R,D, TM<:TensorMapping{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R} 12 struct LazyTensorApplication{T,R,D, TM<:LazyTensor{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R}
13 t::TM 13 t::TM
14 o::AA 14 o::AA
15 15
16 function LazyTensorMappingApplication(t::TensorMapping{<:Any,R,D}, o::AbstractArray{<:Any,D}) where {R,D} 16 function LazyTensorApplication(t::LazyTensor{<:Any,R,D}, o::AbstractArray{<:Any,D}) where {R,D}
17 I = ntuple(i->1, range_dim(t)) 17 I = ntuple(i->1, range_dim(t))
18 T = typeof(apply(t,o,I...)) 18 T = typeof(apply(t,o,I...))
19 return new{T,R,D,typeof(t), typeof(o)}(t,o) 19 return new{T,R,D,typeof(t), typeof(o)}(t,o)
20 end 20 end
21 end 21 end
22 # TODO: Do boundschecking on creation! 22 # TODO: Do boundschecking on creation!
23 23
24 Base.getindex(ta::LazyTensorMappingApplication{T,R}, I::Vararg{Any,R}) where {T,R} = apply(ta.t, ta.o, I...) 24 Base.getindex(ta::LazyTensorApplication{T,R}, I::Vararg{Any,R}) where {T,R} = apply(ta.t, ta.o, I...)
25 Base.getindex(ta::LazyTensorMappingApplication{T,1}, I::CartesianIndex{1}) where {T} = apply(ta.t, ta.o, I.I...) # Would otherwise be caught in the previous method. 25 Base.getindex(ta::LazyTensorApplication{T,1}, I::CartesianIndex{1}) where {T} = apply(ta.t, ta.o, I.I...) # Would otherwise be caught in the previous method.
26 Base.size(ta::LazyTensorMappingApplication) = range_size(ta.t) 26 Base.size(ta::LazyTensorApplication) = range_size(ta.t)
27 # TODO: What else is needed to implement the AbstractArray interface? 27 # TODO: What else is needed to implement the AbstractArray interface?
28 28
29 Base.:*(a::TensorMapping, v::AbstractArray) = LazyTensorMappingApplication(a,v) 29 Base.:*(a::LazyTensor, v::AbstractArray) = LazyTensorApplication(a,v)
30 Base.:*(a::TensorMapping, b::TensorMapping) = throw(MethodError(Base.:*,(a,b))) 30 Base.:*(a::LazyTensor, b::LazyTensor) = throw(MethodError(Base.:*,(a,b)))
31 Base.:*(a::TensorMapping, args::Union{TensorMapping, AbstractArray}...) = foldr(*,(a,args...)) 31 Base.:*(a::LazyTensor, args::Union{LazyTensor, AbstractArray}...) = foldr(*,(a,args...))
32 32
33 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→' 33 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→'
34 # # Should we overload some other infix binary opesrator? 34 # # Should we overload some other infix binary opesrator?
35 # →(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o) 35 # →(tm::LazyTensor{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorApplication(tm,o)
36 # TODO: We need to be really careful about good error messages. 36 # TODO: We need to be really careful about good error messages.
37 # For example what happens if you try to multiply LazyTensorMappingApplication with a TensorMapping(wrong order)? 37 # For example what happens if you try to multiply LazyTensorApplication with a LazyTensor(wrong order)?
38 38
39 """ 39 """
40 LazyTensorMappingTranspose{T,R,D} <: TensorMapping{T,D,R} 40 LazyTensorTranspose{T,R,D} <: LazyTensor{T,D,R}
41 41
42 Struct for lazy transpose of a TensorMapping. 42 Struct for lazy transpose of a LazyTensor.
43 43
44 If a mapping implements the the `apply_transpose` method this allows working with 44 If a mapping implements the the `apply_transpose` method this allows working with
45 the transpose of mapping `m` by using `m'`. `m'` will work as a regular TensorMapping lazily calling 45 the transpose of mapping `m` by using `m'`. `m'` will work as a regular LazyTensor lazily calling
46 the appropriate methods of `m`. 46 the appropriate methods of `m`.
47 """ 47 """
48 struct LazyTensorMappingTranspose{T,R,D, TM<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R} 48 struct LazyTensorTranspose{T,R,D, TM<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
49 tm::TM 49 tm::TM
50 end 50 end
51 51
52 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors? 52 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors?
53 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any TensorMapping even if it doesn't implement `apply_transpose`? 53 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any LazyTensor even if it doesn't implement `apply_transpose`?
54 Base.adjoint(tm::TensorMapping) = LazyTensorMappingTranspose(tm) 54 Base.adjoint(tm::LazyTensor) = LazyTensorTranspose(tm)
55 Base.adjoint(tmt::LazyTensorMappingTranspose) = tmt.tm 55 Base.adjoint(tmt::LazyTensorTranspose) = tmt.tm
56 56
57 apply(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...) 57 apply(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...)
58 apply_transpose(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...) 58 apply_transpose(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...)
59 59
60 range_size(tmt::LazyTensorMappingTranspose) = domain_size(tmt.tm) 60 range_size(tmt::LazyTensorTranspose) = domain_size(tmt.tm)
61 domain_size(tmt::LazyTensorMappingTranspose) = range_size(tmt.tm) 61 domain_size(tmt::LazyTensorTranspose) = range_size(tmt.tm)
62 62
63 63
64 struct LazyTensorMappingBinaryOperation{Op,T,R,D,T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R} 64 struct LazyLazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
65 tm1::T1 65 tm1::T1
66 tm2::T2 66 tm2::T2
67 67
68 @inline function LazyTensorMappingBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} 68 @inline function LazyLazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}}
69 return new{Op,T,R,D,T1,T2}(tm1,tm2) 69 return new{Op,T,R,D,T1,T2}(tm1,tm2)
70 end 70 end
71 end 71 end
72 # TODO: Boundschecking in constructor. 72 # TODO: Boundschecking in constructor.
73 73
74 apply(tmBinOp::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...) 74 apply(tmBinOp::LazyLazyTensorBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
75 apply(tmBinOp::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...) 75 apply(tmBinOp::LazyLazyTensorBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
76 76
77 range_size(tmBinOp::LazyTensorMappingBinaryOperation) = range_size(tmBinOp.tm1) 77 range_size(tmBinOp::LazyLazyTensorBinaryOperation) = range_size(tmBinOp.tm1)
78 domain_size(tmBinOp::LazyTensorMappingBinaryOperation) = domain_size(tmBinOp.tm1) 78 domain_size(tmBinOp::LazyLazyTensorBinaryOperation) = domain_size(tmBinOp.tm1)
79 79
80 Base.:+(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(tm1,tm2) 80 Base.:+(tm1::LazyTensor{T,R,D}, tm2::LazyTensor{T,R,D}) where {T,R,D} = LazyLazyTensorBinaryOperation{:+,T,R,D}(tm1,tm2)
81 Base.:-(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(tm1,tm2) 81 Base.:-(tm1::LazyTensor{T,R,D}, tm2::LazyTensor{T,R,D}) where {T,R,D} = LazyLazyTensorBinaryOperation{:-,T,R,D}(tm1,tm2)
82 82
83 """ 83 """
84 TensorMappingComposition{T,R,K,D} 84 LazyTensorComposition{T,R,K,D}
85 85
86 Lazily compose two `TensorMapping`s, so that they can be handled as a single `TensorMapping`. 86 Lazily compose two `LazyTensor`s, so that they can be handled as a single `LazyTensor`.
87 """ 87 """
88 struct TensorMappingComposition{T,R,K,D, TM1<:TensorMapping{T,R,K}, TM2<:TensorMapping{T,K,D}} <: TensorMapping{T,R,D} 88 struct LazyTensorComposition{T,R,K,D, TM1<:LazyTensor{T,R,K}, TM2<:LazyTensor{T,K,D}} <: LazyTensor{T,R,D}
89 t1::TM1 89 t1::TM1
90 t2::TM2 90 t2::TM2
91 91
92 @inline function TensorMappingComposition(t1::TensorMapping{T,R,K}, t2::TensorMapping{T,K,D}) where {T,R,K,D} 92 @inline function LazyTensorComposition(t1::LazyTensor{T,R,K}, t2::LazyTensor{T,K,D}) where {T,R,K,D}
93 @boundscheck check_domain_size(t1, range_size(t2)) 93 @boundscheck check_domain_size(t1, range_size(t2))
94 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2) 94 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2)
95 end 95 end
96 end 96 end
97 97
98 range_size(tm::TensorMappingComposition) = range_size(tm.t1) 98 range_size(tm::LazyTensorComposition) = range_size(tm.t1)
99 domain_size(tm::TensorMappingComposition) = domain_size(tm.t2) 99 domain_size(tm::LazyTensorComposition) = domain_size(tm.t2)
100 100
101 function apply(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,K,D} 101 function apply(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,K,D}
102 apply(c.t1, c.t2*v, I...) 102 apply(c.t1, c.t2*v, I...)
103 end 103 end
104 104
105 function apply_transpose(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D} 105 function apply_transpose(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D}
106 apply_transpose(c.t2, c.t1'*v, I...) 106 apply_transpose(c.t2, c.t1'*v, I...)
107 end 107 end
108 108
109 Base.@propagate_inbounds Base.:∘(s::TensorMapping, t::TensorMapping) = TensorMappingComposition(s,t) 109 Base.@propagate_inbounds Base.:∘(s::LazyTensor, t::LazyTensor) = LazyTensorComposition(s,t)
110 110
111 """ 111 """
112 LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies) 112 LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
113 113
114 TensorMapping defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should 114 LazyTensor defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
115 be considerd the range and domain of the TensorMapping. Each set of indices must be ordered in ascending order. 115 be considerd the range and domain of the LazyTensor. Each set of indices must be ordered in ascending order.
116 116
117 For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the 117 For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
118 standard matrix-vector product on vectors of size n. 118 standard matrix-vector product on vectors of size n.
119 """ 119 """
120 struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: TensorMapping{T,R,D} 120 struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: LazyTensor{T,R,D}
121 A::AA 121 A::AA
122 range_indicies::NTuple{R,Int} 122 range_indicies::NTuple{R,Int}
123 domain_indicies::NTuple{D,Int} 123 domain_indicies::NTuple{D,Int}
124 124
125 function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}} 125 function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
147 apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...) 147 apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
148 end 148 end
149 149
150 150
151 """ 151 """
152 IdentityMapping{T,D} <: TensorMapping{T,D,D} 152 IdentityTensor{T,D} <: LazyTensor{T,D,D}
153 153
154 The lazy identity TensorMapping for a given size. Usefull for building up higher dimensional tensor mappings from lower 154 The lazy identity LazyTensor for a given size. Usefull for building up higher dimensional tensor mappings from lower
155 dimensional ones through outer products. Also used in the Implementation for InflatedTensorMapping. 155 dimensional ones through outer products. Also used in the Implementation for InflatedLazyTensor.
156 """ 156 """
157 struct IdentityMapping{T,D} <: TensorMapping{T,D,D} 157 struct IdentityTensor{T,D} <: LazyTensor{T,D,D}
158 size::NTuple{D,Int} 158 size::NTuple{D,Int}
159 end 159 end
160 160
161 IdentityMapping{T}(size::NTuple{D,Int}) where {T,D} = IdentityMapping{T,D}(size) 161 IdentityTensor{T}(size::NTuple{D,Int}) where {T,D} = IdentityTensor{T,D}(size)
162 IdentityMapping{T}(size::Vararg{Int,D}) where {T,D} = IdentityMapping{T,D}(size) 162 IdentityTensor{T}(size::Vararg{Int,D}) where {T,D} = IdentityTensor{T,D}(size)
163 IdentityMapping(size::Vararg{Int,D}) where D = IdentityMapping{Float64,D}(size) 163 IdentityTensor(size::Vararg{Int,D}) where D = IdentityTensor{Float64,D}(size)
164 164
165 range_size(tmi::IdentityMapping) = tmi.size 165 range_size(tmi::IdentityTensor) = tmi.size
166 domain_size(tmi::IdentityMapping) = tmi.size 166 domain_size(tmi::IdentityTensor) = tmi.size
167 167
168 apply(tmi::IdentityMapping{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...] 168 apply(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
169 apply_transpose(tmi::IdentityMapping{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...] 169 apply_transpose(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
170 170
171 """ 171 """
172 Base.:∘(tm, tmi) 172 Base.:∘(tm, tmi)
173 Base.:∘(tmi, tm) 173 Base.:∘(tmi, tm)
174 174
175 Composes a `Tensormapping` `tm` with an `IdentityMapping` `tmi`, by returning `tm` 175 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm`
176 """ 176 """
177 @inline function Base.:∘(tm::TensorMapping{T,R,D}, tmi::IdentityMapping{T,D}) where {T,R,D} 177 @inline function Base.:∘(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D}
178 @boundscheck check_domain_size(tm, range_size(tmi)) 178 @boundscheck check_domain_size(tm, range_size(tmi))
179 return tm 179 return tm
180 end 180 end
181 181
182 @inline function Base.:∘(tmi::IdentityMapping{T,R}, tm::TensorMapping{T,R,D}) where {T,R,D} 182 @inline function Base.:∘(tmi::IdentityTensor{T,R}, tm::LazyTensor{T,R,D}) where {T,R,D}
183 @boundscheck check_domain_size(tmi, range_size(tm)) 183 @boundscheck check_domain_size(tmi, range_size(tm))
184 return tm 184 return tm
185 end 185 end
186 # Specialization for the case where tm is an IdentityMapping. Required to resolve ambiguity. 186 # Specialization for the case where tm is an IdentityTensor. Required to resolve ambiguity.
187 @inline function Base.:∘(tm::IdentityMapping{T,D}, tmi::IdentityMapping{T,D}) where {T,D} 187 @inline function Base.:∘(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D}
188 @boundscheck check_domain_size(tm, range_size(tmi)) 188 @boundscheck check_domain_size(tm, range_size(tmi))
189 return tmi 189 return tmi
190 end 190 end
191 # TODO: Implement the above as TensorMappingComposition instead 191 # TODO: Implement the above as LazyTensorComposition instead
192 # TODO: Move the operator definitions to one place 192 # TODO: Move the operator definitions to one place
193 193
194 """ 194 """
195 ScalingTensor{T,D} <: TensorMapping{T,D,D} 195 ScalingTensor{T,D} <: LazyTensor{T,D,D}
196 196
197 A lazy tensor that scales its input with `λ`. 197 A lazy tensor that scales its input with `λ`.
198 """ 198 """
199 struct ScalingTensor{T,D} <: TensorMapping{T,D,D} 199 struct ScalingTensor{T,D} <: LazyTensor{T,D,D}
200 λ::T 200 λ::T
201 size::NTuple{D,Int} 201 size::NTuple{D,Int}
202 end 202 end
203 203
204 LazyTensors.apply(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...] 204 LazyTensors.apply(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
209 209
210 # TODO: Rename everything with mapping 210 # TODO: Rename everything with mapping
211 # TODO: Remove ScalingOperator from tests 211 # TODO: Remove ScalingOperator from tests
212 212
213 """ 213 """
214 InflatedTensorMapping{T,R,D} <: TensorMapping{T,R,D} 214 InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D}
215 215
216 An inflated `TensorMapping` with dimensions added before and afer its actual dimensions. 216 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions.
217 """ 217 """
218 struct InflatedTensorMapping{T,R,D,D_before,R_middle,D_middle,D_after, TM<:TensorMapping{T,R_middle,D_middle}} <: TensorMapping{T,R,D} 218 struct InflatedLazyTensor{T,R,D,D_before,R_middle,D_middle,D_after, TM<:LazyTensor{T,R_middle,D_middle}} <: LazyTensor{T,R,D}
219 before::IdentityMapping{T,D_before} 219 before::IdentityTensor{T,D_before}
220 tm::TM 220 tm::TM
221 after::IdentityMapping{T,D_after} 221 after::IdentityTensor{T,D_after}
222 222
223 function InflatedTensorMapping(before, tm::TensorMapping{T}, after) where T 223 function InflatedLazyTensor(before, tm::LazyTensor{T}, after) where T
224 R_before = range_dim(before) 224 R_before = range_dim(before)
225 R_middle = range_dim(tm) 225 R_middle = range_dim(tm)
226 R_after = range_dim(after) 226 R_after = range_dim(after)
227 R = R_before+R_middle+R_after 227 R = R_before+R_middle+R_after
228 228
232 D = D_before+D_middle+D_after 232 D = D_before+D_middle+D_after
233 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after) 233 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
234 end 234 end
235 end 235 end
236 """ 236 """
237 InflatedTensorMapping(before, tm, after) 237 InflatedLazyTensor(before, tm, after)
238 InflatedTensorMapping(before,tm) 238 InflatedLazyTensor(before,tm)
239 InflatedTensorMapping(tm,after) 239 InflatedLazyTensor(tm,after)
240 240
241 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityMapping`s. 241 The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityTensor`s.
242 242
243 If one of `before` or `after` is left out, a 0-dimensional `IdentityMapping` is used as the default value. 243 If one of `before` or `after` is left out, a 0-dimensional `IdentityTensor` is used as the default value.
244 244
245 If `tm` already is an `InflatedTensorMapping`, `before` and `after` will be extended instead of 245 If `tm` already is an `InflatedLazyTensor`, `before` and `after` will be extended instead of
246 creating a nested `InflatedTensorMapping`. 246 creating a nested `InflatedLazyTensor`.
247 """ 247 """
248 InflatedTensorMapping(::IdentityMapping, ::TensorMapping, ::IdentityMapping) 248 InflatedLazyTensor(::IdentityTensor, ::LazyTensor, ::IdentityTensor)
249 249
250 function InflatedTensorMapping(before, itm::InflatedTensorMapping, after) 250 function InflatedLazyTensor(before, itm::InflatedLazyTensor, after)
251 return InflatedTensorMapping( 251 return InflatedLazyTensor(
252 IdentityMapping(before.size..., itm.before.size...), 252 IdentityTensor(before.size..., itm.before.size...),
253 itm.tm, 253 itm.tm,
254 IdentityMapping(itm.after.size..., after.size...), 254 IdentityTensor(itm.after.size..., after.size...),
255 ) 255 )
256 end 256 end
257 257
258 InflatedTensorMapping(before::IdentityMapping, tm::TensorMapping{T}) where T = InflatedTensorMapping(before,tm,IdentityMapping{T}()) 258 InflatedLazyTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedLazyTensor(before,tm,IdentityTensor{T}())
259 InflatedTensorMapping(tm::TensorMapping{T}, after::IdentityMapping) where T = InflatedTensorMapping(IdentityMapping{T}(),tm,after) 259 InflatedLazyTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedLazyTensor(IdentityTensor{T}(),tm,after)
260 # Resolve ambiguity between the two previous methods 260 # Resolve ambiguity between the two previous methods
261 InflatedTensorMapping(I1::IdentityMapping{T}, I2::IdentityMapping{T}) where T = InflatedTensorMapping(I1,I2,IdentityMapping{T}()) 261 InflatedLazyTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedLazyTensor(I1,I2,IdentityTensor{T}())
262 262
263 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensorMapping(I(3),B,I(2)) -> I(3)⊗B⊗I(2) 263 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedLazyTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
264 264
265 function range_size(itm::InflatedTensorMapping) 265 function range_size(itm::InflatedLazyTensor)
266 return flatten_tuple( 266 return flatten_tuple(
267 range_size(itm.before), 267 range_size(itm.before),
268 range_size(itm.tm), 268 range_size(itm.tm),
269 range_size(itm.after), 269 range_size(itm.after),
270 ) 270 )
271 end 271 end
272 272
273 function domain_size(itm::InflatedTensorMapping) 273 function domain_size(itm::InflatedLazyTensor)
274 return flatten_tuple( 274 return flatten_tuple(
275 domain_size(itm.before), 275 domain_size(itm.before),
276 domain_size(itm.tm), 276 domain_size(itm.tm),
277 domain_size(itm.after), 277 domain_size(itm.after),
278 ) 278 )
279 end 279 end
280 280
281 function apply(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} 281 function apply(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
282 dim_before = range_dim(itm.before) 282 dim_before = range_dim(itm.before)
283 dim_domain = domain_dim(itm.tm) 283 dim_domain = domain_dim(itm.tm)
284 dim_range = range_dim(itm.tm) 284 dim_range = range_dim(itm.tm)
285 dim_after = range_dim(itm.after) 285 dim_after = range_dim(itm.after)
286 286
288 288
289 v_inner = view(v, view_index...) 289 v_inner = view(v, view_index...)
290 return apply(itm.tm, v_inner, inner_index...) 290 return apply(itm.tm, v_inner, inner_index...)
291 end 291 end
292 292
293 function apply_transpose(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D} 293 function apply_transpose(itm::InflatedLazyTensor{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
294 dim_before = range_dim(itm.before) 294 dim_before = range_dim(itm.before)
295 dim_domain = domain_dim(itm.tm) 295 dim_domain = domain_dim(itm.tm)
296 dim_range = range_dim(itm.tm) 296 dim_range = range_dim(itm.tm)
297 dim_after = range_dim(itm.after) 297 dim_after = range_dim(itm.after)
298 298
381 flatten_tuple(ts::Vararg) = flatten_tuple(ts) 381 flatten_tuple(ts::Vararg) = flatten_tuple(ts)
382 382
383 @doc raw""" 383 @doc raw"""
384 LazyOuterProduct(tms...) 384 LazyOuterProduct(tms...)
385 385
386 Creates a `TensorMappingComposition` for the outerproduct of `tms...`. 386 Creates a `LazyTensorComposition` for the outerproduct of `tms...`.
387 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping. 387 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
388 388
389 First let 389 First let
390 ```math 390 ```math
391 \begin{aligned} 391 \begin{aligned}
417 (A⊗B⊗C)v = [(A⊗I_{|M|}⊗I_{|P|}) [(I_{|J|}⊗B⊗I_{|P|}) [(I_{|J|}⊗I_{|N|}⊗C)v]]] 417 (A⊗B⊗C)v = [(A⊗I_{|M|}⊗I_{|P|}) [(I_{|J|}⊗B⊗I_{|P|}) [(I_{|J|}⊗I_{|N|}⊗C)v]]]
418 ``` 418 ```
419 """ 419 """
420 function LazyOuterProduct end 420 function LazyOuterProduct end
421 421
422 function LazyOuterProduct(tm1::TensorMapping{T}, tm2::TensorMapping{T}) where T 422 function LazyOuterProduct(tm1::LazyTensor{T}, tm2::LazyTensor{T}) where T
423 itm1 = InflatedTensorMapping(tm1, IdentityMapping{T}(range_size(tm2))) 423 itm1 = InflatedLazyTensor(tm1, IdentityTensor{T}(range_size(tm2)))
424 itm2 = InflatedTensorMapping(IdentityMapping{T}(domain_size(tm1)),tm2) 424 itm2 = InflatedLazyTensor(IdentityTensor{T}(domain_size(tm1)),tm2)
425 425
426 return itm1∘itm2 426 return itm1∘itm2
427 end 427 end
428 428
429 LazyOuterProduct(t1::IdentityMapping{T}, t2::IdentityMapping{T}) where T = IdentityMapping{T}(t1.size...,t2.size...) 429 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...)
430 LazyOuterProduct(t1::TensorMapping, t2::IdentityMapping) = InflatedTensorMapping(t1, t2) 430 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedLazyTensor(t1, t2)
431 LazyOuterProduct(t1::IdentityMapping, t2::TensorMapping) = InflatedTensorMapping(t1, t2) 431 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedLazyTensor(t1, t2)
432 432
433 LazyOuterProduct(tms::Vararg{TensorMapping}) = foldl(LazyOuterProduct, tms) 433 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
434 434
435 ⊗(a::TensorMapping, b::TensorMapping) = LazyOuterProduct(a,b) 435 ⊗(a::LazyTensor, b::LazyTensor) = LazyOuterProduct(a,b)
436 436
437 437
438 function check_domain_size(tm::TensorMapping, sz) 438 function check_domain_size(tm::LazyTensor, sz)
439 if domain_size(tm) != sz 439 if domain_size(tm) != sz
440 throw(SizeMismatch(tm,sz)) 440 throw(SizeMismatch(tm,sz))
441 end 441 end
442 end 442 end
443 443
444 struct SizeMismatch <: Exception 444 struct SizeMismatch <: Exception
445 tm::TensorMapping 445 tm::LazyTensor
446 sz 446 sz
447 end 447 end
448 448
449 function Base.showerror(io::IO, err::SizeMismatch) 449 function Base.showerror(io::IO, err::SizeMismatch)
450 print(io, "SizeMismatch: ") 450 print(io, "SizeMismatch: ")
451 print(io, "domain size $(domain_size(err.tm)) of TensorMapping not matching size $(err.sz)") 451 print(io, "domain size $(domain_size(err.tm)) of LazyTensor not matching size $(err.sz)")
452 end 452 end