Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 1954:b0915f43b122 feature/sbp_operators/laplace_curvilinear

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Merge feature/grids/geometry_functions
author Jonatan Werpers <jonatan@werpers.com>
date Sat, 08 Feb 2025 09:38:58 +0100
parents ed50eec18365
children
comparison
equal deleted inserted replaced
1953:835b1dcee38e 1954:b0915f43b122
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,TT<:NTuple{N, LazyTensor{T,R,D}} where N} <: LazyTensor{T,R,D} 55 """
56 TensorNegation{T,R,D,...} <: LazyTensor{T,R,D}
57
58 The negation of a LazyTensor.
59 """
60 struct TensorNegation{T,R,D,TM<:LazyTensor{T,R,D}} <: LazyTensor{T,R,D}
61 tm::TM
62 end
63
64 apply(tm::TensorNegation, v, I...) = -apply(tm.tm, v, I...)
65 apply_transpose(tm::TensorNegation, v, I...) = -apply_transpose(tm.tm, v, I...)
66
67 range_size(tm::TensorNegation) = range_size(tm.tm)
68 domain_size(tm::TensorNegation) = domain_size(tm.tm)
69
70
71 """
72 TensorSum{T,R,D,...} <: LazyTensor{T,R,D}
73
74 The lazy sum of 2 or more lazy tensors.
75 """
76 struct TensorSum{T,R,D,TT<:NTuple{N, LazyTensor{T,R,D}} where N} <: LazyTensor{T,R,D}
56 tms::TT 77 tms::TT
57 78
58 function ElementwiseTensorOperation{Op,T,R,D}(tms::TT) where {Op,T,R,D, TT<:NTuple{N, LazyTensor{T,R,D}} where N} 79 function TensorSum{T,R,D}(tms::TT) where {T,R,D, TT<:NTuple{N, LazyTensor{T,R,D}} where N}
59 @boundscheck map(tms) do tm 80 @boundscheck map(tms) do tm
60 check_domain_size(tm, domain_size(tms[1])) 81 check_domain_size(tm, domain_size(tms[1]))
61 check_range_size(tm, range_size(tms[1])) 82 check_range_size(tm, range_size(tms[1]))
62 end 83 end
63 84
64 return new{Op,T,R,D,TT}(tms) 85 return new{T,R,D,TT}(tms)
65 end 86 end
66 end 87 end
67 # TBD: Can we introduce negation of LazyTensors? It could be done generically 88
68 # with a ScalingTensor but also using specializations for specific tensor 89 """
69 # types. This would allow simplification of ElementwiseTensorOperation to 90 TensorSum(ts::Vararg{LazyTensor})
70 # TensorSum. The implementation of `-` can be done using negation and the 91
71 # TensorSum type. We should make sure this doesn't impact the efficiency of 92 The lazy sum of the tensors `ts`.
72 # for example SATs. 93 """
73 94 function TensorSum(ts::Vararg{LazyTensor})
74 95 T = eltype(ts[1])
75 function ElementwiseTensorOperation{:+}(ts::Vararg{LazyTensor}) 96 R = range_dim(ts[1])
76 return ElementwiseTensorOperation{:+,eltype(ts[1]), range_dim(ts[1]), domain_dim(ts[1])}(ts) 97 D = domain_dim(ts[1])
77 end 98 return TensorSum{T,R,D}(ts)
78 99 end
79 # The following methods for :+ are intended to reduce the depth of the tree of operations in some caes 100
80 function ElementwiseTensorOperation{:+}(t1::ElementwiseTensorOperation{:+}, t2::ElementwiseTensorOperation{:+}) 101 function apply(tmBinOp::TensorSum{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
81 ElementwiseTensorOperation{:+}(t1.tms..., t2.tms...) 102 return sum(tmBinOp.tms) do tm
82 end
83
84 function ElementwiseTensorOperation{:+}(t1::ElementwiseTensorOperation{:+}, t2::LazyTensor)
85 ElementwiseTensorOperation{:+}(t1.tms..., t2)
86 end
87
88 function ElementwiseTensorOperation{:+}(t1::LazyTensor, t2::ElementwiseTensorOperation{:+})
89 ElementwiseTensorOperation{:+}(t1, t2.tms...)
90 end
91
92 function ElementwiseTensorOperation{:-}(t1::LazyTensor, t2::LazyTensor)
93 return ElementwiseTensorOperation{:-,eltype(t1), range_dim(t1), domain_dim(t1)}((t1,t2))
94 end
95
96 function apply(tmBinOp::ElementwiseTensorOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
97 vs = map(tmBinOp.tms) do tm
98 apply(tm,v,I...) 103 apply(tm,v,I...)
99 end 104 end
100 105 end
101 return +(vs...) 106
102 end 107 function apply_transpose(tmBinOp::TensorSum{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
103 function apply(tmBinOp::ElementwiseTensorOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} 108 return sum(tmBinOp.tms) do tm
104 apply(tmBinOp.tms[1], v, I...) - apply(tmBinOp.tms[2], v, I...) 109 apply_transpose(tm,v,I...)
105 end 110 end
106 111 end
107 range_size(tmBinOp::ElementwiseTensorOperation) = range_size(tmBinOp.tms[1]) 112
108 domain_size(tmBinOp::ElementwiseTensorOperation) = domain_size(tmBinOp.tms[1]) 113 range_size(tmBinOp::TensorSum) = range_size(tmBinOp.tms[1])
114 domain_size(tmBinOp::TensorSum) = domain_size(tmBinOp.tms[1])
109 115
110 116
111 """ 117 """
112 TensorComposition{T,R,K,D} 118 TensorComposition{T,R,K,D}
113 119
155 return tmi 161 return tmi
156 end 162 end
157 163
158 Base.:*(a::T, tm::LazyTensor{T}) where T = TensorComposition(ScalingTensor{T,range_dim(tm)}(a,range_size(tm)), tm) 164 Base.:*(a::T, tm::LazyTensor{T}) where T = TensorComposition(ScalingTensor{T,range_dim(tm)}(a,range_size(tm)), tm)
159 Base.:*(tm::LazyTensor{T}, a::T) where T = a*tm 165 Base.:*(tm::LazyTensor{T}, a::T) where T = a*tm
160 Base.:-(tm::LazyTensor) = (-one(eltype(tm)))*tm
161 166
162 """ 167 """
163 InflatedTensor{T,R,D} <: LazyTensor{T,R,D} 168 InflatedTensor{T,R,D} <: LazyTensor{T,R,D}
164 169
165 An inflated `LazyTensor` with dimensions added before and after its actual dimensions. 170 An inflated `LazyTensor` with dimensions added before and after its actual dimensions.
203 itm.tm, 208 itm.tm,
204 IdentityTensor(itm.after.size..., after.size...), 209 IdentityTensor(itm.after.size..., after.size...),
205 ) 210 )
206 end 211 end
207 212
208 InflatedTensor(before::IdentityTensor, tm::LazyTensor{T}) where T = InflatedTensor(before,tm,IdentityTensor{T}()) 213 InflatedTensor(before::IdentityTensor, tm::LazyTensor) = InflatedTensor(before,tm,IdentityTensor{eltype(tm)}())
209 InflatedTensor(tm::LazyTensor{T}, after::IdentityTensor) where T = InflatedTensor(IdentityTensor{T}(),tm,after) 214 InflatedTensor(tm::LazyTensor, after::IdentityTensor) = InflatedTensor(IdentityTensor{eltype(tm)}(),tm,after)
210 # Resolve ambiguity between the two previous methods 215 # Resolve ambiguity between the two previous methods
211 InflatedTensor(I1::IdentityTensor{T}, I2::IdentityTensor{T}) where T = InflatedTensor(I1,I2,IdentityTensor{T}()) 216 InflatedTensor(I1::IdentityTensor, I2::IdentityTensor) = InflatedTensor(I1,I2,IdentityTensor{promote_type(eltype(I1), eltype(I2))}())
212 217
213 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2) 218 # TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensor(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
214 219
215 function range_size(itm::InflatedTensor) 220 function range_size(itm::InflatedTensor)
216 return concatenate_tuples( 221 return concatenate_tuples(
297 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2) 302 itm2 = InflatedTensor(IdentityTensor{T}(domain_size(tm1)),tm2)
298 303
299 return itm1∘itm2 304 return itm1∘itm2
300 end 305 end
301 306
302 LazyOuterProduct(t1::IdentityTensor{T}, t2::IdentityTensor{T}) where T = IdentityTensor{T}(t1.size...,t2.size...) 307 LazyOuterProduct(t1::IdentityTensor, t2::IdentityTensor) = IdentityTensor{promote_type(eltype(t1),eltype(t2))}(t1.size...,t2.size...)
303 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2) 308 LazyOuterProduct(t1::LazyTensor, t2::IdentityTensor) = InflatedTensor(t1, t2)
304 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2) 309 LazyOuterProduct(t1::IdentityTensor, t2::LazyTensor) = InflatedTensor(t1, t2)
305 310
306 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms) 311 LazyOuterProduct(tms::Vararg{LazyTensor}) = foldl(LazyOuterProduct, tms)
307 312