comparison src/LazyTensors/lazy_tensor_operations.jl @ 1854:654a2b7e6824 tooling/benchmarks

Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Sat, 11 Jan 2025 10:19:47 +0100
parents 164e26a6cf79
children 21e5fe1545c0
comparison
equal deleted inserted replaced
1378:2b5480e2d4bf 1854:654a2b7e6824
3 3
4 Struct for lazy application of a LazyTensor. Created using `*`. 4 Struct for lazy application of a LazyTensor. Created using `*`.
5 5
6 Allows the result of a `LazyTensor` applied to a vector to be treated as an `AbstractArray`. 6 Allows the result of a `LazyTensor` applied to a vector to be treated as an `AbstractArray`.
7 With a mapping `m` and a vector `v` the TensorApplication object can be created by `m*v`. 7 With a mapping `m` and a vector `v` the TensorApplication object can be created by `m*v`.
8 The actual result will be calcualted when indexing into `m*v`. 8 The actual result will be calculated when indexing into `m*v`.
9 """ 9 """
10 struct TensorApplication{T,R,D, TM<:LazyTensor{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R} 10 struct TensorApplication{T,R,D, TM<:LazyTensor{<:Any,R,D}, AA<:AbstractArray{<:Any,D}} <: LazyArray{T,R}
11 t::TM 11 t::TM
12 o::AA 12 o::AA
13 13
50 50
51 range_size(tmt::TensorTranspose) = domain_size(tmt.tm) 51 range_size(tmt::TensorTranspose) = domain_size(tmt.tm)
52 domain_size(tmt::TensorTranspose) = range_size(tmt.tm) 52 domain_size(tmt::TensorTranspose) = range_size(tmt.tm)
53 53
54 54
55 struct ElementwiseTensorOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R} 55 struct ElementwiseTensorOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,R,D}
56 tm1::T1 56 tm1::T1
57 tm2::T2 57 tm2::T2
58 58
59 function ElementwiseTensorOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} 59 function ElementwiseTensorOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}}
60 @boundscheck check_domain_size(tm2, domain_size(tm1)) 60 @boundscheck check_domain_size(tm2, domain_size(tm1))
100 100
101 """ 101 """
102 TensorComposition(tm, tmi::IdentityTensor) 102 TensorComposition(tm, tmi::IdentityTensor)
103 TensorComposition(tmi::IdentityTensor, tm) 103 TensorComposition(tmi::IdentityTensor, tm)
104 104
105 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm` 105 Composes a `LazyTensor` `tm` with an `IdentityTensor` `tmi`, by returning `tm`
106 """ 106 """
107 function TensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D} 107 function TensorComposition(tm::LazyTensor{T,R,D}, tmi::IdentityTensor{T,D}) where {T,R,D}
108 @boundscheck check_domain_size(tm, range_size(tmi)) 108 @boundscheck check_domain_size(tm, range_size(tmi))
109 return tm 109 return tm
110 end 110 end
119 return tmi 119 return tmi
120 end 120 end
121 121
122 Base.:*(a::T, tm::LazyTensor{T}) where T = TensorComposition(ScalingTensor{T,range_dim(tm)}(a,range_size(tm)), tm) 122 Base.:*(a::T, tm::LazyTensor{T}) where T = TensorComposition(ScalingTensor{T,range_dim(tm)}(a,range_size(tm)), tm)
123 Base.:*(tm::LazyTensor{T}, a::T) where T = a*tm 123 Base.:*(tm::LazyTensor{T}, a::T) where T = a*tm
124 Base.:-(tm::LazyTensor) = (-one(eltype(tm)))*tm
124 125
125 """ 126 """
126 InflatedTensor{T,R,D} <: LazyTensor{T,R,D} 127 InflatedTensor{T,R,D} <: LazyTensor{T,R,D}
127 128
128 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions. 129 An inflated `LazyTensor` with dimensions added before and after its actual dimensions.
129 """ 130 """
130 struct InflatedTensor{T,R,D,D_before,R_middle,D_middle,D_after, TM<:LazyTensor{T,R_middle,D_middle}} <: LazyTensor{T,R,D} 131 struct InflatedTensor{T,R,D,D_before,R_middle,D_middle,D_after, TM<:LazyTensor{T,R_middle,D_middle}} <: LazyTensor{T,R,D}
131 before::IdentityTensor{T,D_before} 132 before::IdentityTensor{T,D_before}
132 tm::TM 133 tm::TM
133 after::IdentityTensor{T,D_after} 134 after::IdentityTensor{T,D_after}
166 itm.tm, 167 itm.tm,
167 IdentityTensor(itm.after.size..., after.size...), 168 IdentityTensor(itm.after.size..., after.size...),
168 ) 169 )
169 end 170 end
170 171
171 InflatedTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedTensor(before,tm,IdentityTensor{T}()) 172 InflatedTensor(before::IdentityTensor, tm::LazyTensor) = InflatedTensor(before,tm,IdentityTensor{eltype(tm)}())
172 InflatedTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedTensor(IdentityTensor{T}(),tm,after) 173 InflatedTensor(tm::LazyTensor, after::IdentityTensor) = InflatedTensor(IdentityTensor{eltype(tm)}(),tm,after)
173 # Resolve ambiguity between the two previous methods 174 # Resolve ambiguity between the two previous methods
174 InflatedTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedTensor(I1,I2,IdentityTensor{T}()) 175 InflatedTensor(I1::IdentityTensor, I2::IdentityTensor) = InflatedTensor(I1,I2,IdentityTensor{promote_type(eltype(I1), eltype(I2))}())
175 176
176 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2) 177 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
177 178
178 function range_size(itm::InflatedTensor) 179 function range_size(itm::InflatedTensor)
179 return concatenate_tuples( 180 return concatenate_tuples(
217 218
218 219
219 @doc raw""" 220 @doc raw"""
220 LazyOuterProduct(tms...) 221 LazyOuterProduct(tms...)
221 222
222 Creates a `TensorComposition` for the outerproduct of `tms...`. 223 Creates a `TensorComposition` for the outer product of `tms...`.
223 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping. 224 This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
224 225
225 First let 226 First let
226 ```math 227 ```math
227 \begin{aligned} 228 \begin{aligned}
260 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2) 261 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2)
261 262
262 return itm1∘itm2 263 return itm1∘itm2
263 end 264 end
264 265
265 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...) 266 LazyOuterProduct(t1::IdentityTensor, t2::IdentityTensor) = IdentityTensor{promote_type(eltype(t1),eltype(t2))}(t1.size...,t2.size...)
266 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2) 267 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2)
267 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2) 268 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2)
268 269
269 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms) 270 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
270 271
276 Inflate `tm` such that it gets the size `sz` in all directions except `dir`. 277 Inflate `tm` such that it gets the size `sz` in all directions except `dir`.
277 Here `sz[dir]` is ignored and replaced with the range and domains size of 278 Here `sz[dir]` is ignored and replaced with the range and domains size of
278 `tm`. 279 `tm`.
279 280
280 An example of when this operation is useful is when extending a one 281 An example of when this operation is useful is when extending a one
281 dimensional difference operator `D` to a 2D grid of a ceratin size. In that 282 dimensional difference operator `D` to a 2D grid of a certain size. In that
282 case we could have 283 case we could have
283 284
284 ```julia 285 ```julia
285 Dx = inflate(D, (10,10), 1) 286 Dx = inflate(D, (10,10), 1)
286 Dy = inflate(D, (10,10), 2) 287 Dy = inflate(D, (10,10), 2)