Mercurial > repos > public > sbplib_julia
diff 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 |
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl Mon Oct 19 19:02:48 2020 +0200 +++ b/src/LazyTensors/lazy_tensor_operations.jl Sun Nov 08 16:01:39 2020 +0100 @@ -16,14 +16,16 @@ # 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. @@ -38,8 +40,8 @@ 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} - tm::TensorMapping{T,R,D} +struct LazyTensorMappingTranspose{T,R,D, TM<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R} + tm::TM end export LazyTensorMappingTranspose @@ -84,12 +86,9 @@ 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) - throw(DimensionMismatch("the first argument has domain size $(domain_size(t1)) while the second has range size $(range_size(t2)) ")) - end + @boundscheck check_domain_size(t1, range_size(t2)) return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2) end - # Add check for matching sizes as a boundscheck end export TensorMappingComposition @@ -145,3 +144,238 @@ 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