Define BlockedTuple

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,83 @@
+# This file defines BlockedTuple, a Tuple of heterogeneous Tuple with a BlockArrays.jl
+# like interface
+
+using BlockArrays: Block, BlockArrays, BlockIndexRange, BlockRange, blockedrange
+
+abstract type AbstractBlockTuple end
+
+struct BlockedTuple{BlockLengths,Flat} <: AbstractBlockTuple
+  flat::Flat
+
+  function BlockedTuple{BlockLengths}(flat::Tuple) where {BlockLengths}
+    length(flat) != sum(BlockLengths) && throw(DimensionMismatch("Invalid total length"))
+    return new{BlockLengths,typeof(flat)}(flat)
+  end
+end
+
+# TensorAlgebra Interface
+BlockedTuple(tt::Vararg{Tuple}) = BlockedTuple{length.(tt)}(flatten_tuples(tt))
+BlockedTuple(bt::BlockedTuple) = bt
+
+# Base interface
+Base.Tuple(bt::BlockedTuple) = bt.flat
+
+Base.axes(bt::AbstractBlockTuple) = (blockedrange([blocklengths(bt)...]),)
+
+Base.broadcastable(bt::AbstractBlockTuple) = bt
+struct BlockedTupleBroadcastStyle{BlockLengths} <: Broadcast.BroadcastStyle end
+function Base.BroadcastStyle(type::Type{<:AbstractBlockTuple})
+  return BlockedTupleBroadcastStyle{blocklengths(type)}()
+end
+
+# BroadcastStyle is not called for two identical styles
+function Base.BroadcastStyle(::BlockedTupleBroadcastStyle, ::BlockedTupleBroadcastStyle)
+  throw(DimensionMismatch("Incompatible blocks"))
+end
+function Base.copy(
+  bc::Broadcast.Broadcasted{BlockedTupleBroadcastStyle{BlockLengths}}
+) where {BlockLengths}
+  return BlockedTuple{BlockLengths}(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(::AbstractBlockTuple) = 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::AbstractBlockTuple) = 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::AbstractBlockTuple)
+  return (0, cumsum(blocklengths(bt)[begin:(end - 1)])...) .+ 1
+end
+
+function BlockArrays.blocklasts(bt::AbstractBlockTuple)
+  return cumsum(blocklengths(bt)[begin:end])
+end
+
+BlockArrays.blocklength(bt::AbstractBlockTuple) = length(blocklengths(bt))
+
+BlockArrays.blocklengths(bt::AbstractBlockTuple) = blocklengths(typeof(bt))
+function BlockArrays.blocklengths(::Type{<:BlockedTuple{BlockLengths}}) where {BlockLengths}
+  return BlockLengths
+end
+
+function BlockArrays.blocks(bt::AbstractBlockTuple)
+  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
+
+@testset "BlockedTuple" begin
+  flat = (1, 'a', 2, 'b', 3)
+  divs = (1, 2, 2)
+
+  bt = BlockedTuple{divs}(flat)
+
+  @test Tuple(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 bt[2:4] == ('a', 2, 'b')
+
+  @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