diff src/LazyTensors/lazy_tensor_operations.jl @ 447:27e0e256e5d9 feature/inflated_tensormapping

Merge in feature/lazy_identity
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 19 Oct 2020 20:59:08 +0200
parents 904aae1899df 648a36ebac99
children 912ae510dec9
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Mon Oct 19 08:37:35 2020 +0200
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Mon Oct 19 20:59:08 2020 +0200
@@ -74,34 +74,38 @@
 Base.:+(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(tm1,tm2)
 Base.:-(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(tm1,tm2)
 
+"""
+    TensorMappingComposition{T,R,K,D}
 
-# TODO: Write tests and documentation for LazyTensorMappingComposition
-# struct LazyTensorMappingComposition{T,R,K,D} <: TensorMapping{T,R,D}
-#     t1::TensorMapping{T,R,K}
-#     t2::TensorMapping{T,K,D}
-# end
-
-# Base.:∘(s::TensorMapping{T,R,K}, t::TensorMapping{T,K,D}) where {T,R,K,D} = LazyTensorMappingComposition(s,t)
-
-# function range_size(tm::LazyTensorMappingComposition{T,R,K,D}, domain_size::NTuple{D,Integer}) where {T,R,K,D}
-#     range_size(tm.t1, domain_size(tm.t2, domain_size))
-# end
+Lazily compose two TensorMappings, so that they can be handled as a single TensorMapping.
+"""
+struct TensorMappingComposition{T,R,K,D, TM1<:TensorMapping{T,R,K}, TM2<:TensorMapping{T,K,D}} <: TensorMapping{T,R,D}
+    t1::TM1
+    t2::TM2
 
-# function domain_size(tm::LazyTensorMappingComposition{T,R,K,D}, range_size::NTuple{R,Integer}) where {T,R,K,D}
-#     domain_size(tm.t1, domain_size(tm.t2, range_size))
-# end
-
-# function apply(c::LazyTensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::NTuple{R,Int}) where {T,R,K,D}
-#     apply(c.t1, LazyTensorMappingApplication(c.t2,v), I...)
-# end
+    @inline function TensorMappingComposition(t1::TensorMapping{T,R,K}, t2::TensorMapping{T,K,D}) where {T,R,K,D}
+        @boundscheck if domain_size(t1) != range_size(t2)
+            throw(DimensionMismatch("the first argument has domain size $(domain_size(t1)) while the second has range size $(range_size(t2)) "))
+        end
+        return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2)
+    end
+    # Add check for matching sizes as a boundscheck
+end
+export TensorMappingComposition
 
-# function apply_transpose(c::LazyTensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::NTuple{D,Int}) where {T,R,K,D}
-#     apply_transpose(c.t2, LazyTensorMappingApplication(c.t1',v), I...)
-# end
+range_size(tm::TensorMappingComposition) = range_size(tm.t1)
+domain_size(tm::TensorMappingComposition) = domain_size(tm.t2)
+
+function apply(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{T,D}, I::Vararg{S,R} where S) where {T,R,K,D}
+    apply(c.t1, c.t2*v, I...)
+end
 
-# # Have i gone too crazy with the type parameters? Maybe they aren't all needed?
+function apply_transpose(c::TensorMappingComposition{T,R,K,D}, v::AbstractArray{T,R}, I::Vararg{S,D} where S) where {T,R,K,D}
+    apply_transpose(c.t2, c.t1'*v, I...)
+end
 
-# export →
+Base.@propagate_inbounds Base.:∘(s::TensorMapping, t::TensorMapping) = TensorMappingComposition(s,t)
+
 """
     LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
 
@@ -144,23 +148,23 @@
 
 
 """
-    LazyIdentity{T,D} <: TensorMapping{T,D,D}
+    IdentityMapping{T,D} <: TensorMapping{T,D,D}
 
 The lazy identity TensorMapping for a given size. Usefull for building up higher dimensional tensor mappings from lower
 dimensional ones through outer products. Also used in the Implementation for InflatedTensorMapping.
 """
-struct LazyIdentity{T,D} <: TensorMapping{T,D,D}
+struct IdentityMapping{T,D} <: TensorMapping{T,D,D}
     size::NTuple{D,Int}
 end
-export LazyIdentity
+export IdentityMapping
 
-LazyIdentity{T}(size::NTuple{D,Int}) where {T,D} = LazyIdentity{T,D}(size)
+IdentityMapping{T}(size::NTuple{D,Int}) where {T,D} = IdentityMapping{T,D}(size)
 
-range_size(tmi::LazyIdentity) = tmi.size
-domain_size(tmi::LazyIdentity) = tmi.size
+range_size(tmi::IdentityMapping) = tmi.size
+domain_size(tmi::IdentityMapping) = tmi.size
 
-apply(tmi::LazyIdentity{T,D}, v::AbstractArray{T,D}, I::Vararg{Index,D}) where {T,D} = v[Int.(I)...]
-apply_transpose(tmi::LazyIdentity{T,D}, v::AbstractArray{T,D}, I::Vararg{Index,D}) where {T,D} = v[Int.(I)...]
+apply(tmi::IdentityMapping{T,D}, v::AbstractArray{T,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
+apply_transpose(tmi::IdentityMapping{T,D}, v::AbstractArray{T,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
 
 struct InflatedTensorMapping{T,R,D,D_before,R_middle,D_middle,D_after} <: TensorMapping{T,R,D}
     before::LazyIdentity{T,D_before}