diff LazyTensors/src/lazy_operations.jl @ 210:2aa33d0eef90 boundary_conditions

Add ability to add and subtract TensorMappings
author Jonatan Werpers <jonatan@werpers.com>
date Tue, 25 Jun 2019 17:26:39 +0200
parents 7db145ce6e4d
children a20bb4fac23d
line wrap: on
line diff
--- a/LazyTensors/src/lazy_operations.jl	Tue Jun 25 16:37:35 2019 +0200
+++ b/LazyTensors/src/lazy_operations.jl	Tue Jun 25 17:26:39 2019 +0200
@@ -60,6 +60,7 @@
         return new{T,D,Op,T1,T2}(a,b)
     end
 end
+# TODO: Move Op to be the first parameter? Compare to Binary operations
 
 Base.size(v::LazyElementwiseOperation) = size(v.a)
 
@@ -143,6 +144,26 @@
 
 
 
+
+struct LazyTensorMappingBinaryOperation{Op,T,R,D,T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R}
+    A::T1
+    B::T2
+
+    @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}}
+        return new{Op,T,R,D,T1,T2}(A,B)
+    end
+end
+
+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...)
+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...)
+
+range_size(mp::LazyTensorMappingBinaryOperation{Op,T,R,D}, domain_size::NTuple{D,Integer}) where {Op,T,R,D} = range_size(mp.A, domain_size)
+domain_size(mp::LazyTensorMappingBinaryOperation{Op,T,R,D}, range_size::NTuple{R,Integer}) where {Op,T,R,D} = domain_size(mp.A, range_size)
+
+Base.:+(A::TensorMapping{T,R,D}, B::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(A,B)
+Base.:-(A::TensorMapping{T,R,D}, B::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(A,B)
+
+
 # TODO: Write tests and documentation for LazyTensorMappingComposition
 # struct LazyTensorMappingComposition{T,R,K,D} <: TensorMapping{T,R,D}
 #     t1::TensorMapping{T,R,K}