define BlockedTuple

Author: ogauthe

src/TensorAlgebra.jl

@@ -1,6 +1,7 @@
 module TensorAlgebra
 include("blockedpermutation.jl")
 include("BaseExtensions/BaseExtensions.jl")
+include("blockedtuple.jl")
 include("fusedims.jl")
 include("splitdims.jl")
 include("contract/contract.jl")

src/blockedtuple.jl

@@ -0,0 +1,76 @@
+# This file defines BlockedTuple, a Tuple of heterogeneous Tuple with a BlockArrays.jl
+# like interface
+
+using BlockArrays: Block, BlockArrays, BlockIndexRange, BlockRange, blockedrange
+
+struct BlockedTuple{Divs,Flat}
+  flat::Flat
+
+  function BlockedTuple{Divs}(flat::Tuple) where {Divs}
+    length(flat) != sum(Divs) && throw(DimensionMismatch("Invalid total length"))
+    return new{Divs,typeof(flat)}(flat)
+  end
+end
+
+# TensorAlgebra Interface
+BlockedTuple(tt::Vararg{Tuple}) = BlockedTuple{length.(tt)}(flatten_tuples(tt))
+BlockedTuple(bt::BlockedTuple) = bt
+flatten_tuples(bt::BlockedTuple) = Tuple(bt)
+
+# Base interface
+Base.Tuple(bt::BlockedTuple) = bt.flat
+
+Base.axes(bt::BlockedTuple) = (blockedrange([blocklengths(bt)...]),)
+
+Base.broadcastable(bt::BlockedTuple) = bt
+struct BlockedTupleBroadcastStyle{Divs} <: Broadcast.BroadcastStyle end
+function Base.BroadcastStyle(::Type{<:BlockedTuple{Divs}}) where {Divs}
+  return BlockedTupleBroadcastStyle{Divs}()
+end
+function Base.BroadcastStyle(::BlockedTupleBroadcastStyle, ::BlockedTupleBroadcastStyle)
+  throw(DimensionMismatch("Incompatible blocks"))
+end
+Base.BroadcastStyle(::T, ::T) where {T<:BlockedTupleBroadcastStyle} = T
+function Base.copy(bc::Broadcast.Broadcasted{BlockedTupleBroadcastStyle{Divs}}) where {Divs}
+  return BlockedTuple{Divs}(bc.f.((Tuple.(bc.args))...))
+end
+
+Base.copy(bt::BlockedTuple) = BlockedTuple{blocklengths(bt)}(copy.(Tuple(bt)))
+
+Base.deepcopy(bt::BlockedTuple) = BlockedTuple{blocklengths(bt)}(deepcopy.(Tuple(bt)))
+
+Base.firstindex(::BlockedTuple) = 1
+
+Base.getindex(bt::BlockedTuple, i::Integer) = Tuple(bt)[i]
+Base.getindex(bt::BlockedTuple, r::AbstractUnitRange) = Tuple(bt)[r]
+Base.getindex(bt::BlockedTuple, b::Block{1}) = blocks(bt)[Int(b)]
+Base.getindex(bt::BlockedTuple, br::BlockRange{1}) = blocks(bt)[Int.(br)]
+Base.getindex(bt::BlockedTuple, bi::BlockIndexRange{1}) = bt[Block(bi)][only(bi.indices)]
+
+Base.iterate(bt::BlockedTuple) = iterate(Tuple(bt))
+Base.iterate(bt::BlockedTuple, i::Int) = iterate(Tuple(bt), i)
+
+Base.lastindex(bt::BlockedTuple) = length(bt)
+
+Base.length(bt::BlockedTuple) = length(Tuple(bt))
+
+Base.map(f, bt::BlockedTuple) = BlockedTuple{blocklengths(bt)}(map(f, Tuple(bt)))
+
+# BlockArrays interface
+function BlockArrays.blockfirsts(bt::BlockedTuple)
+  return (0, cumsum(blocklengths(bt)[begin:(end - 1)])...) .+ 1
+end
+
+function BlockArrays.blocklasts(bt::BlockedTuple)
+  return cumsum(blocklengths(bt)[begin:end])
+end
+
+BlockArrays.blocklength(::BlockedTuple{Divs}) where {Divs} = length(Divs)
+
+BlockArrays.blocklengths(::BlockedTuple{Divs}) where {Divs} = Divs
+
+function BlockArrays.blocks(bt::BlockedTuple)
+  bf = blockfirsts(bt)
+  bl = blocklasts(bt)
+  return ntuple(i -> Tuple(bt)[bf[i]:bl[i]], blocklength(bt))
+end

test/test_blockedtuple.jl

@@ -0,0 +1,51 @@
+using Test: @test, @test_throws
+
+using BlockArrays: Block, blocklength, blocklengths, blockedrange, blockisequal, blocks
+
+using TensorAlgebra: BlockedTuple, flatten_tuples
+
+@testset "BlockedTuple" begin
+  flat = (1, 'a', 2, 'b', 3)
+  divs = (1, 2, 2)
+
+  bt = BlockedTuple{divs}(flat)
+
+  @test Tuple(bt) == flat
+  @test flatten_tuples(bt) == flat
+  @test bt == BlockedTuple((1,), ('a', 2), ('b', 3))
+  @test BlockedTuple(bt) == bt
+  @test blocklength(bt) == 3
+  @test blocklengths(bt) == (1, 2, 2)
+  @test blocks(bt) == ((1,), ('a', 2), ('b', 3))
+
+  @test bt[1] == 1
+  @test bt[2] == 'a'
+  @test bt[Block(1)] == blocks(bt)[1]
+  @test bt[Block(2)] == blocks(bt)[2]
+  @test bt[Block(1):Block(2)] == blocks(bt)[1:2]
+  @test bt[Block(2)[1:2]] == ('a', 2)
+
+  @test firstindex(bt) == 1
+  @test lastindex(bt) == 5
+  @test length(bt) == 5
+
+  @test iterate(bt) == (1, 2)
+  @test iterate(bt, 2) == ('a', 3)
+  @test blockisequal(only(axes(bt)), blockedrange([1, 2, 2]))
+
+  @test_throws DimensionMismatch BlockedTuple{(1, 2, 3)}(flat)
+
+  bt = BlockedTuple((1,), (4, 2), (5, 3))
+  @test Tuple(bt) == (1, 4, 2, 5, 3)
+  @test blocklengths(bt) == (1, 2, 2)
+  @test copy(bt) == bt
+  @test deepcopy(bt) == bt
+
+  @test map(n -> n + 1, bt) == BlockedTuple{blocklengths(bt)}(Tuple(bt) .+ 1)
+  @test bt .+ BlockedTuple((1,), (1, 1), (1, 1)) ==
+    BlockedTuple{blocklengths(bt)}(Tuple(bt) .+ 1)
+  @test_throws DimensionMismatch bt .+ BlockedTuple((1, 1), (1, 1), (1,))
+
+  bt = BlockedTuple((1:2, 1:2), (1:3,))
+  @test length.(bt) == BlockedTuple((2, 2), (3,))
+end