comparison LazyTensors/src/lazy_operations.jl @ 231:fbabfd4e8f20

Merge in boundary_conditions
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 26 Jun 2019 15:07:47 +0200
parents 2aa33d0eef90
children a20bb4fac23d
comparison
equal deleted inserted replaced
144:ce56727e4232 231:fbabfd4e8f20
1 """
2 LazyArray{T,D} <: AbstractArray{T,D}
3
4 Array which is calcualted lazily when indexing.
5
6 A subtype of `LazyArray` will use lazy version of `+`, `-`, `*`, `/`.
7 """
8 abstract type LazyArray{T,D} <: AbstractArray{T,D} end
9 export LazyArray
10
11
12
13 """
14 LazyTensorMappingApplication{T,R,D} <: LazyArray{T,R}
15
16 Struct for lazy application of a TensorMapping. Created using `*`.
17
18 Allows the result of a `TensorMapping` applied to a vector to be treated as an `AbstractArray`.
19 With a mapping `m` and a vector `v` the LazyTensorMappingApplication object can be created by `m*v`.
20 The actual result will be calcualted when indexing into `m*v`.
21 """
22 struct LazyTensorMappingApplication{T,R,D} <: LazyArray{T,R}
23 t::TensorMapping{T,R,D}
24 o::AbstractArray{T,D}
25 end
26 export LazyTensorMappingApplication
27
28 Base.:*(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o)
29
30 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg) where {T,R,D} = apply(ta.t, ta.o, I...)
31 Base.size(ta::LazyTensorMappingApplication{T,R,D}) where {T,R,D} = range_size(ta.t,size(ta.o))
32 # TODO: What else is needed to implement the AbstractArray interface?
33
34 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→'
35 Base.:*(args::Union{TensorMapping{T}, AbstractArray{T}}...) where T = foldr(*,args)
36 # # Should we overload some other infix binary operator?
37 # →(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o)
38 # TODO: We need to be really careful about good error messages.
39 # For example what happens if you try to multiply LazyTensorMappingApplication with a TensorMapping(wrong order)?
40
41
42
43 """
44 LazyElementwiseOperation{T,D,Op, T1<:AbstractArray{T,D}, T2 <: AbstractArray{T,D}} <: AbstractArray{T,D}
45
46 Struct allowing for lazy evaluation of elementwise operations on AbstractArrays.
47
48 A LazyElementwiseOperation contains two AbstractArrays of equal size,
49 together with an operation. The operations are carried out when the
50 LazyElementwiseOperation is indexed.
51 """
52 struct LazyElementwiseOperation{T,D,Op, T1<:AbstractArray{T,D}, T2 <: AbstractArray{T,D}} <: LazyArray{T,D}
53 a::T1
54 b::T2
55
56 @inline function LazyElementwiseOperation{T,D,Op}(a::T1,b::T2) where {T,D,Op, T1<:AbstractArray{T,D}, T2<:AbstractArray{T,D}}
57 @boundscheck if size(a) != size(b)
58 throw(DimensionMismatch("dimensions must match"))
59 end
60 return new{T,D,Op,T1,T2}(a,b)
61 end
62 end
63 # TODO: Move Op to be the first parameter? Compare to Binary operations
64
65 Base.size(v::LazyElementwiseOperation) = size(v.a)
66
67 # TODO: Make sure boundschecking is done properly and that the lenght of the vectors are equal
68 # NOTE: Boundschecking in getindex functions now assumes that the size of the
69 # vectors in the LazyElementwiseOperation are the same size. If we remove the
70 # size assertion in the constructor we might have to handle
71 # boundschecking differently.
72 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:+}, I...) where {T,D}
73 @boundscheck if !checkbounds(Bool,leo.a,I...)
74 throw(BoundsError([leo],[I...]))
75 end
76 return leo.a[I...] + leo.b[I...]
77 end
78 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:-}, I...) where {T,D}
79 @boundscheck if !checkbounds(Bool,leo.a,I...)
80 throw(BoundsError([leo],[I...]))
81 end
82 return leo.a[I...] - leo.b[I...]
83 end
84 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:*}, I...) where {T,D}
85 @boundscheck if !checkbounds(Bool,leo.a,I...)
86 throw(BoundsError([leo],[I...]))
87 end
88 return leo.a[I...] * leo.b[I...]
89 end
90 Base.@propagate_inbounds @inline function Base.getindex(leo::LazyElementwiseOperation{T,D,:/}, I...) where {T,D}
91 @boundscheck if !checkbounds(Bool,leo.a,I...)
92 throw(BoundsError([leo],[I...]))
93 end
94 return leo.a[I...] / leo.b[I...]
95 end
96
97 # Define lazy operations for AbstractArrays. Operations constructs a LazyElementwiseOperation which
98 # can later be indexed into. Lazy operations are denoted by the usual operator followed by a tilde
99 Base.@propagate_inbounds +̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:+}(a,b)
100 Base.@propagate_inbounds -̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:-}(a,b)
101 Base.@propagate_inbounds *̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:*}(a,b)
102 Base.@propagate_inbounds /̃(a::AbstractArray{T,D}, b::AbstractArray{T,D}) where {T,D} = LazyElementwiseOperation{T,D,:/}(a,b)
103
104 # NOTE: Är det knas att vi har till exempel * istället för .* ??
105 # Oklart om det ens går att lösa..
106 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
107 Base.@propagate_inbounds Base.:+(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a +̃ b
108 Base.@propagate_inbounds Base.:+(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a +̃ b
109
110 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
111 Base.@propagate_inbounds Base.:-(a::LazyArray{T,D}, b::AbstractArray{T,D}) where {T,D} = a -̃ b
112 Base.@propagate_inbounds Base.:-(a::AbstractArray{T,D}, b::LazyArray{T,D}) where {T,D} = a -̃ b
113
114 # Element wise operation for `*` and `\` are not overloaded due to conflicts with the behavior
115 # of regular `*` and `/` for AbstractArrays. Use tilde versions instead.
116
117 export +̃, -̃, *̃, /̃
118
119
120
121 """
122 LazyTensorMappingTranspose{T,R,D} <: TensorMapping{T,D,R}
123
124 Struct for lazy transpose of a TensorMapping.
125
126 If a mapping implements the the `apply_transpose` method this allows working with
127 the transpose of mapping `m` by using `m'`. `m'` will work as a regular TensorMapping lazily calling
128 the appropriate methods of `m`.
129 """
130 struct LazyTensorMappingTranspose{T,R,D} <: TensorMapping{T,D,R}
131 tm::TensorMapping{T,R,D}
132 end
133 export LazyTensorMappingTranspose
134
135 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors?
136 Base.adjoint(t::TensorMapping) = LazyTensorMappingTranspose(t)
137 Base.adjoint(t::LazyTensorMappingTranspose) = t.tm
138
139 apply(tm::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,R}, I::Vararg) where {T,R,D} = apply_transpose(tm.tm, v, I...)
140 apply_transpose(tm::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,D}, I::Vararg) where {T,R,D} = apply(tm.tm, v, I...)
141
142 range_size(tmt::LazyTensorMappingTranspose{T,R,D}, d_size::NTuple{R,Integer}) where {T,R,D} = domain_size(tmt.tm, domain_size)
143 domain_size(tmt::LazyTensorMappingTranspose{T,R,D}, r_size::NTuple{D,Integer}) where {T,R,D} = range_size(tmt.tm, range_size)
144
145
146
147
148 struct LazyTensorMappingBinaryOperation{Op,T,R,D,T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R}
149 A::T1
150 B::T2
151
152 @inline function LazyTensorMappingBinaryOperation{Op,T,R,D}(A::T1,B::T2) where {Op,T,R,D, T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}}
153 return new{Op,T,R,D,T1,T2}(A,B)
154 end
155 end
156
157 apply(mb::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{T,D}, I::Vararg) where {T,R,D} = apply(mb.A, v, I...) + apply(mb.B,v,I...)
158 apply(mb::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{T,D}, I::Vararg) where {T,R,D} = apply(mb.A, v, I...) - apply(mb.B,v,I...)
159
160 range_size(mp::LazyTensorMappingBinaryOperation{Op,T,R,D}, domain_size::NTuple{D,Integer}) where {Op,T,R,D} = range_size(mp.A, domain_size)
161 domain_size(mp::LazyTensorMappingBinaryOperation{Op,T,R,D}, range_size::NTuple{R,Integer}) where {Op,T,R,D} = domain_size(mp.A, range_size)
162
163 Base.:+(A::TensorMapping{T,R,D}, B::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(A,B)
164 Base.:-(A::TensorMapping{T,R,D}, B::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(A,B)
165
166
167 # TODO: Write tests and documentation for LazyTensorMappingComposition
168 # struct LazyTensorMappingComposition{T,R,K,D} <: TensorMapping{T,R,D}
169 # t1::TensorMapping{T,R,K}
170 # t2::TensorMapping{T,K,D}
171 # end
172
173 # Base.:∘(s::TensorMapping{T,R,K}, t::TensorMapping{T,K,D}) where {T,R,K,D} = LazyTensorMappingComposition(s,t)
174
175 # function range_size(tm::LazyTensorMappingComposition{T,R,K,D}, domain_size::NTuple{D,Integer}) where {T,R,K,D}
176 # range_size(tm.t1, domain_size(tm.t2, domain_size))
177 # end
178
179 # function domain_size(tm::LazyTensorMappingComposition{T,R,K,D}, range_size::NTuple{R,Integer}) where {T,R,K,D}
180 # domain_size(tm.t1, domain_size(tm.t2, range_size))
181 # end
182
183 # function apply(c::LazyTensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::Vararg) where {T,R,K,D}
184 # apply(c.t1, LazyTensorMappingApplication(c.t2,v), I...)
185 # end
186
187 # function apply_transpose(c::LazyTensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::Vararg) where {T,R,K,D}
188 # apply_transpose(c.t2, LazyTensorMappingApplication(c.t1',v), I...)
189 # end
190
191 # # Have i gone too crazy with the type parameters? Maybe they aren't all needed?
192
193 # export →