Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 1002:271aa6ae1055 refactor/lazy_tensors

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Split out a file for tensor types
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 22:18:04 +0100
parents a3df203861d3
children becd95ba0fce
comparison
equal deleted inserted replaced
1001:a3df203861d3 1002:271aa6ae1055
1 # TBD: Is there a good way to split this file?
2 # TODO: Split out functions for composition
3 # TODO: We need to be really careful about good error messages. 1 # TODO: We need to be really careful about good error messages.
4 # TODO: Go over type parameters 2 # TODO: Go over type parameters
3
5 4
6 """ 5 """
7 LazyTensorApplication{T,R,D} <: LazyArray{T,R} 6 LazyTensorApplication{T,R,D} <: LazyArray{T,R}
8 7
9 Struct for lazy application of a LazyTensor. Created using `*`. 8 Struct for lazy application of a LazyTensor. Created using `*`.
52 apply_transpose(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...) 51 apply_transpose(tmt::LazyTensorTranspose{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmt.tm, v, I...)
53 52
54 range_size(tmt::LazyTensorTranspose) = domain_size(tmt.tm) 53 range_size(tmt::LazyTensorTranspose) = domain_size(tmt.tm)
55 domain_size(tmt::LazyTensorTranspose) = range_size(tmt.tm) 54 domain_size(tmt::LazyTensorTranspose) = range_size(tmt.tm)
56 55
57 # TODO: Rename this 56
58 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R} 57 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
59 tm1::T1 58 tm1::T1
60 tm2::T2 59 tm2::T2
61 60
62 function LazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} 61 function LazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}}
98 97
99 function apply_transpose(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D} 98 function apply_transpose(c::LazyTensorComposition{T,R,K,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,K,D}
100 apply_transpose(c.t2, c.t1'*v, I...) 99 apply_transpose(c.t2, c.t1'*v, I...)
101 end 100 end
102 101
103
104 """
105 LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
106
107 LazyTensor defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
108 be considerd the range and domain of the LazyTensor. Each set of indices must be ordered in ascending order.
109
110 For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
111 standard matrix-vector product on vectors of size n.
112 """
113 struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: LazyTensor{T,R,D}
114 A::AA
115 range_indicies::NTuple{R,Int}
116 domain_indicies::NTuple{D,Int}
117
118 function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
119 if !issorted(range_indicies) || !issorted(domain_indicies)
120 throw(DomainError("range_indicies and domain_indicies must be sorted in ascending order"))
121 end
122
123 return new{T,R,D,RD,AA}(A,range_indicies,domain_indicies)
124 end
125 end
126
127 range_size(llm::LazyLinearMap) = size(llm.A)[[llm.range_indicies...]]
128 domain_size(llm::LazyLinearMap) = size(llm.A)[[llm.domain_indicies...]]
129
130 function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
131 view_index = ntuple(i->:,ndims(llm.A))
132 for i ∈ 1:R
133 view_index = Base.setindex(view_index, Int(I[i]), llm.range_indicies[i])
134 end
135 A_view = @view llm.A[view_index...]
136 return sum(A_view.*v)
137 end
138
139 function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
140 apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
141 end
142
143
144 """
145 IdentityTensor{T,D} <: LazyTensor{T,D,D}
146
147 The lazy identity LazyTensor for a given size. Usefull for building up higher dimensional tensor mappings from lower
148 dimensional ones through outer products. Also used in the Implementation for InflatedLazyTensor.
149 """
150 struct IdentityTensor{T,D} <: LazyTensor{T,D,D}
151 size::NTuple{D,Int}
152 end
153
154 IdentityTensor{T}(size::NTuple{D,Int}) where {T,D} = IdentityTensor{T,D}(size)
155 IdentityTensor{T}(size::Vararg{Int,D}) where {T,D} = IdentityTensor{T,D}(size)
156 IdentityTensor(size::Vararg{Int,D}) where D = IdentityTensor{Float64,D}(size)
157
158 range_size(tmi::IdentityTensor) = tmi.size
159 domain_size(tmi::IdentityTensor) = tmi.size
160
161 apply(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
162 apply_transpose(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
163 102
164 """ 103 """
165 LazyTensorComposition(tm, tmi::IdentityTensor) 104 LazyTensorComposition(tm, tmi::IdentityTensor)
166 LazyTensorComposition(tmi::IdentityTensor, tm) 105 LazyTensorComposition(tmi::IdentityTensor, tm)
167 106
180 function LazyTensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D} 119 function LazyTensorComposition(tm::IdentityTensor{T,D}, tmi::IdentityTensor{T,D}) where {T,D}
181 @boundscheck check_domain_size(tm, range_size(tmi)) 120 @boundscheck check_domain_size(tm, range_size(tmi))
182 return tmi 121 return tmi
183 end 122 end
184 123
185 """
186 ScalingTensor{T,D} <: LazyTensor{T,D,D}
187
188 A lazy tensor that scales its input with `λ`.
189 """
190 struct ScalingTensor{T,D} <: LazyTensor{T,D,D}
191 λ::T
192 size::NTuple{D,Int}
193 end
194
195 LazyTensors.apply(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
196 LazyTensors.apply_transpose(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
197
198 LazyTensors.range_size(m::ScalingTensor) = m.size
199 LazyTensors.domain_size(m::ScalingTensor) = m.size
200 124
201 """ 125 """
202 InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D} 126 InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D}
203 127
204 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.
219 D_after = domain_dim(after) 143 D_after = domain_dim(after)
220 D = D_before+D_middle+D_after 144 D = D_before+D_middle+D_after
221 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after) 145 return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
222 end 146 end
223 end 147 end
148
224 """ 149 """
225 InflatedLazyTensor(before, tm, after) 150 InflatedLazyTensor(before, tm, after)
226 InflatedLazyTensor(before,tm) 151 InflatedLazyTensor(before,tm)
227 InflatedLazyTensor(tm,after) 152 InflatedLazyTensor(tm,after)
228 153