comparison src/LazyTensors/lazy_tensor_operations.jl @ 1888:eb70a0941cd6 allocation_testing

Merge default
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 03 Feb 2023 23:02:46 +0100
parents 2278730f9cee
children 6f51160c7ca7 f1c2a4fa0ee1
comparison
equal deleted inserted replaced
1887:24590890e124 1888:eb70a0941cd6
96 96
97 function apply_transpose(c::TensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D} 97 function apply_transpose(c::TensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D}
98 apply_transpose(c.t2, c.t1'*v, I...) 98 apply_transpose(c.t2, c.t1'*v, I...)
99 end 99 end
100 100
101
102 """ 101 """
103 TensorComposition(tm, tmi::IdentityTensor) 102 TensorComposition(tm, tmi::IdentityTensor)
104 TensorComposition(tmi::IdentityTensor, tm) 103 TensorComposition(tmi::IdentityTensor, tm)
105 104
106 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm` 105 Composes a `Tensormapping` `tm` with an `IdentityTensor` `tmi`, by returning `tm`
118 function TensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D} 117 function TensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D}
119 @boundscheck check_domain_size(tm, range_size(tmi)) 118 @boundscheck check_domain_size(tm, range_size(tmi))
120 return tmi 119 return tmi
121 end 120 end
122 121
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 124
124 """ 125 """
125 InflatedTensor{T,R,D} <: LazyTensor{T,R,D} 126 InflatedTensor{T,R,D} <: LazyTensor{T,R,D}
126 127
127 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions. 128 An inflated `LazyTensor` with dimensions added before and afer its actual dimensions.
266 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2) 267 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2)
267 268
268 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms) 269 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
269 270
270 271
272
273 """
274 inflate(tm::LazyTensor, sz, dir)
275
276 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 `tm`.
279
280 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 case we could have
283
284 ```julia
285 Dx = inflate(D, (10,10), 1)
286 Dy = inflate(D, (10,10), 2)
287 ```
288 """
289 function inflate(tm::LazyTensor, sz, dir)
290 Is = IdentityTensor{eltype(tm)}.(sz)
291 parts = Base.setindex(Is, tm, dir)
292 return foldl(⊗, parts)
293 end
294
271 function check_domain_size(tm::LazyTensor, sz) 295 function check_domain_size(tm::LazyTensor, sz)
272 if domain_size(tm) != sz 296 if domain_size(tm) != sz
273 throw(DomainSizeMismatch(tm,sz)) 297 throw(DomainSizeMismatch(tm,sz))
274 end 298 end
275 end 299 end