sbplib_julia: src/LazyTensors/lazy_tensor

comparison src/LazyTensors/lazy_tensor_operations.jl @ 509:b7e42384053a feature/boundary_ops

Merge w. default

author	Vidar Stiernström <vidar.stiernstrom@it.uu.se>
date	Sun, 08 Nov 2020 16:01:39 +0100
parents	4b9d124fe984
children	a5caa934b35f fe86ac896377

comparison

equal deleted inserted replaced

-:2ab687b1d221
+:b7e42384053a
 # TODO: Do boundschecking on creation!
 export LazyTensorMappingApplication
 # TODO: Go through and remove unneccerary type parameters on functions
-Base.:*(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o)
 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Index,R}) where {T,R,D} = apply(ta.t, ta.o, I...)
 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Int,R}) where {T,R,D} = apply(ta.t, ta.o, Index{Unknown}.(I)...)
 Base.size(ta::LazyTensorMappingApplication) = range_size(ta.t)
 # TODO: What else is needed to implement the AbstractArray interface?
+Base.:*(a::TensorMapping, v::AbstractArray) = LazyTensorMappingApplication(a,v)
+Base.:*(a::TensorMapping, b::TensorMapping) = throw(MethodError(Base.:*,(a,b)))
+Base.:*(a::TensorMapping, args::Union{TensorMapping, AbstractArray}...) = foldr(*,(a,args...))
 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→'
-Base.:*(a::TensorMapping{T,R,D}, b::TensorMapping{T,D,K}, args::Union{TensorMapping{T}, AbstractArray{T}}...) where {T,R,D,K} = foldr(*,(a,b,args...))
 # # Should we overload some other infix binary opesrator?
 # →(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o)
 # TODO: We need to be really careful about good error messages.
 # For example what happens if you try to multiply LazyTensorMappingApplication with a TensorMapping(wrong order)?
 If a mapping implements the the `apply_transpose` method this allows working with
 the transpose of mapping `m` by using `m'`. `m'` will work as a regular TensorMapping lazily calling
 the appropriate methods of `m`.
 """
-struct LazyTensorMappingTranspose{T,R,D} <: TensorMapping{T,D,R}
+struct LazyTensorMappingTranspose{T,R,D, TM<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R}
-tm::TensorMapping{T,R,D}
+tm::TM
 end
 export LazyTensorMappingTranspose
 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors?
 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any TensorMapping even if it doesn't implement `apply_transpose`?
 struct TensorMappingComposition{T,R,K,D, TM1<:TensorMapping{T,R,K}, TM2<:TensorMapping{T,K,D}} <: TensorMapping{T,R,D}
 t1::TM1
 t2::TM2
 @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)
+@boundscheck check_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
 range_size(tm::TensorMappingComposition) = range_size(tm.t1)
 domain_size(tm::TensorMappingComposition) = domain_size(tm.t2)
 end
 function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Index,D}) where {T,R,D}
 apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
 end
+"""
+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 IdentityMapping{T,D} <: TensorMapping{T,D,D}
+size::NTuple{D,Int}
+end
+export IdentityMapping
+IdentityMapping{T}(size::NTuple{D,Int}) where {T,D} = IdentityMapping{T,D}(size)
+IdentityMapping{T}(size::Vararg{Int,D}) where {T,D} = IdentityMapping{T,D}(size)
+IdentityMapping(size::Vararg{Int,D}) where D = IdentityMapping{Float64,D}(size)
+range_size(tmi::IdentityMapping) = tmi.size
+domain_size(tmi::IdentityMapping) = tmi.size
+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...]
+"""
+Base.:∘(tm, tmi)
+Base.:∘(tmi, tm)
+Composes a `Tensormapping` `tm` with an `IdentityMapping` `tmi`, by returning `tm`
+"""
+@inline function Base.:∘(tm::TensorMapping{T,R,D}, tmi::IdentityMapping{T,D}) where {T,R,D}
+@boundscheck check_domain_size(tm, range_size(tmi))
+return tm
+end
+@inline function Base.:∘(tmi::IdentityMapping{T,R}, tm::TensorMapping{T,R,D}) where {T,R,D}
+@boundscheck check_domain_size(tmi, range_size(tm))
+return tm
+end
+# Specialization for the case where tm is an IdentityMapping. Required to resolve ambiguity.
+@inline function Base.:∘(tm::IdentityMapping{T,D}, tmi::IdentityMapping{T,D}) where {T,D}
+@boundscheck check_domain_size(tm, range_size(tmi))
+return tmi
+end
+"""
+InflatedTensorMapping{T,R,D} <: TensorMapping{T,R,D}
+An inflated `TensorMapping` with dimensions added before and afer its actual dimensions.
+"""
+struct InflatedTensorMapping{T,R,D,D_before,R_middle,D_middle,D_after, TM<:TensorMapping{T,R_middle,D_middle}} <: TensorMapping{T,R,D}
+before::IdentityMapping{T,D_before}
+tm::TM
+after::IdentityMapping{T,D_after}
+function InflatedTensorMapping(before, tm::TensorMapping{T}, after) where T
+R_before = range_dim(before)
+R_middle = range_dim(tm)
+R_after = range_dim(after)
+R = R_before+R_middle+R_after
+D_before = domain_dim(before)
+D_middle = domain_dim(tm)
+D_after = domain_dim(after)
+D = D_before+D_middle+D_after
+return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
+end
+end
+export InflatedTensorMapping
+"""
+InflatedTensorMapping(before, tm, after)
+InflatedTensorMapping(before,tm)
+InflatedTensorMapping(tm,after)
+The outer product of `before`, `tm` and `after`, where `before` and `after` are `IdentityMapping`s.
+If one of `before` or `after` is left out, a 0-dimensional `IdentityMapping` is used as the default value.
+If `tm` already is an `InflatedTensorMapping`, `before` and `after` will be extended instead of
+creating a nested `InflatedTensorMapping`.
+"""
+InflatedTensorMapping(::IdentityMapping, ::TensorMapping, ::IdentityMapping)
+function InflatedTensorMapping(before, itm::InflatedTensorMapping, after)
+return InflatedTensorMapping(
+IdentityMapping(before.size...,  itm.before.size...),
+itm.tm,
+IdentityMapping(itm.after.size..., after.size...),
+)
+end
+InflatedTensorMapping(before::IdentityMapping, tm::TensorMapping{T}) where T = InflatedTensorMapping(before,tm,IdentityMapping{T}())
+InflatedTensorMapping(tm::TensorMapping{T}, after::IdentityMapping) where T = InflatedTensorMapping(IdentityMapping{T}(),tm,after)
+# Resolve ambiguity between the two previous methods
+InflatedTensorMapping(I1::IdentityMapping{T}, I2::IdentityMapping{T}) where T = InflatedTensorMapping(I1,I2,IdentityMapping{T}())
+# TODO: Implement syntax and constructors for products of different combinations of InflatedTensorMapping and IdentityMapping
+# TODO: Implement some pretty printing in terms of ⊗. E.g InflatedTensorMapping(I(3),B,I(2)) -> I(3)⊗B⊗I(2)
+function range_size(itm::InflatedTensorMapping)
+return flatten_tuple(
+range_size(itm.before),
+range_size(itm.tm),
+range_size(itm.after),
+)
+end
+function domain_size(itm::InflatedTensorMapping)
+return flatten_tuple(
+domain_size(itm.before),
+domain_size(itm.tm),
+domain_size(itm.after),
+)
+end
+function apply(itm::InflatedTensorMapping{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Any,R}) where {T,R,D}
+view_index, inner_index = split_index(itm, I...)
+v_inner = view(v, view_index...)
+return apply(itm.tm, v_inner, inner_index...)
+end
+"""
+split_index(...)
+Splits the multi-index into two parts. One part for the view that the inner TensorMapping acts on, and one part for indexing the result
+Eg.
+```
+(1,2,3,4) -> (1,:,:,4), (2,3)
+```
+"""
+function split_index(itm::InflatedTensorMapping{T,R,D}, I::Vararg{Any,R}) where {T,R,D}
+I_before = slice_tuple(I, Val(1), Val(range_dim(itm.before)))
+I_after = slice_tuple(I, Val(R-range_dim(itm.after)+1), Val(R))
+view_index = (I_before..., ntuple((i)->:,domain_dim(itm.tm))..., I_after...)
+inner_index = slice_tuple(I, Val(range_dim(itm.before)+1), Val(R-range_dim(itm.after)))
+return (view_index, inner_index)
+end
+# TODO: Can this be replaced by something more elegant while still being type stable? 2020-10-21
+# See:
+# https://github.com/JuliaLang/julia/issues/34884
+# https://github.com/JuliaLang/julia/issues/30386
+"""
+slice_tuple(t, Val(l), Val(u))
+Get a slice of a tuple in a type stable way.
+Equivalent to t[l:u] but type stable.
+"""
+function slice_tuple(t,::Val{L},::Val{U}) where {L,U}
+return ntuple(i->t[i+L-1], U-L+1)
+end
+"""
+flatten_tuple(t)
+Takes a nested tuple and flattens the whole structure
+"""
+flatten_tuple(t::NTuple{N, Number} where N) = t
+flatten_tuple(t::Tuple) = ((flatten_tuple.(t)...)...,) # simplify?
+flatten_tuple(ts::Vararg) = flatten_tuple(ts)
+"""
+LazyOuterProduct(tms...)
+Creates a `TensorMappingComposition` for the outerproduct of `tms...`.
+This is done by separating the outer product into regular products of outer products involving only identity mappings and one non-identity mapping.
+First let
+```math
+A = A_{I,J}
+B = B_{M,N}
+C = C_{P,Q}
+```
+where ``I``, ``M``, ``P`` are  multi-indexes for the ranges of ``A``, ``B``, ``C``, and ``J``, ``N``, ``Q`` are multi-indexes of the domains.
+We use ``⊗`` to denote the outer product
+```math
+(A⊗B)_{IM,JN} = A_{I,J}B_{M,N}
+```
+We note that
+```math
+A⊗B⊗C = (A⊗B⊗C)_{IMP,JNQ} = A_{I,J}B_{M,N}C_{P,Q}
+```
+And that
+```math
+A⊗B⊗C = (A⊗I_{|M|}⊗I_{|P|})(I_{|J|}⊗B⊗I_{|P|})(I_{|J|}⊗I_{|N|}⊗C)
+```
+where |.| of a multi-index is a vector of sizes for each dimension. ``I_v`` denotes the identity tensor of size ``v[i]`` in each direction
+To apply ``A⊗B⊗C`` we evaluate
+(A⊗B⊗C)v = [(A⊗I_{|M|}⊗I_{|P|})  [(I_{|J|}⊗B⊗I_{|P|}) [(I_{|J|}⊗I_{|N|}⊗C)v]]]
+"""
+function LazyOuterProduct end
+export LazyOuterProduct
+function LazyOuterProduct(tm1::TensorMapping{T}, tm2::TensorMapping{T}) where T
+itm1 = InflatedTensorMapping(tm1, IdentityMapping{T}(range_size(tm2)))
+itm2 = InflatedTensorMapping(IdentityMapping{T}(domain_size(tm1)),tm2)
+return itm1∘itm2
+end
+LazyOuterProduct(t1::IdentityMapping{T}, t2::IdentityMapping{T}) where T = IdentityMapping{T}(t1.size...,t2.size...)
+LazyOuterProduct(t1::TensorMapping, t2::IdentityMapping) = InflatedTensorMapping(t1, t2)
+LazyOuterProduct(t1::IdentityMapping, t2::TensorMapping) = InflatedTensorMapping(t1, t2)
+LazyOuterProduct(tms::Vararg{TensorMapping}) = foldl(LazyOuterProduct, tms)
+⊗(a::TensorMapping, b::TensorMapping) = LazyOuterProduct(a,b)
+export ⊗
+function check_domain_size(tm::TensorMapping, sz)
+if domain_size(tm) != sz
+throw(SizeMismatch(tm,sz))
+end
+end
+struct SizeMismatch <: Exception
+tm::TensorMapping
+sz
+end
+export SizeMismatch
+function Base.showerror(io::IO, err::SizeMismatch)
+print(io, "SizeMismatch: ")
+print(io, "domain size $(domain_size(err.tm)) of TensorMapping not matching size $(err.sz)")
+end

Mercurial > repos > public > sbplib_julia

comparison src/LazyTensors/lazy_tensor_operations.jl @ 509:b7e42384053a feature/boundary_ops