Mercurial > repos > public > sbplib_julia
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}