use BlockedTuple

Author: ogauthe

Project.toml

@@ -33,7 +33,7 @@ LinearAlgebra = "1.10.0"
 SparseArraysBase = "0.2.0"
 Strided = "2.2.0"
 SymmetrySectors = "0.1.1"
-TensorAlgebra = "0.1.0"
+TensorAlgebra = "0.1.5"
 TypeParameterAccessors = "0.2.0"
 WignerSymbols = "2.0.0"
 julia = "1.10"

src/fusiontensor/base_interface.jl

@@ -3,6 +3,7 @@
 using Accessors: @set
 
 using BlockSparseArrays: @view!
+using TensorAlgebra: BlockedTuple, tuplemortar
 
 set_data_matrix(ft::FusionTensor, data_matrix) = @set ft.data_matrix = data_matrix
 
@@ -43,35 +44,25 @@ function transpose_mapping(b::BlockIndexRange{2})
   return new_block[reverse(b.indices)...]
 end
 function Base.adjoint(ft::FusionTensor)
+  new_axes = tuplemortar((dual.(domain_axes(ft)), dual.(codomain_axes(ft))))
   return FusionTensor(
-    adjoint(data_matrix(ft)),
-    dual.(domain_axes(ft)),
-    dual.(codomain_axes(ft)),
-    transpose_mapping(trees_block_mapping(ft)),
+    adjoint(data_matrix(ft)), new_axes, transpose_mapping(trees_block_mapping(ft))
   )
 end
 
-Base.axes(ft::FusionTensor) = (codomain_axes(ft)..., domain_axes(ft)...)
+Base.axes(ft::FusionTensor) = ft.axes
 
 # conj is defined as coefficient wise complex conjugation, without axis dual
 Base.conj(ft::FusionTensor{<:Real}) = ft   # same object for real element type
 Base.conj(ft::FusionTensor) = set_data_matrix(ft, conj(data_matrix(ft)))
 
 function Base.copy(ft::FusionTensor)
-  return FusionTensor(
-    copy(data_matrix(ft)),
-    copy.(codomain_axes(ft)),
-    copy.(domain_axes(ft)),
-    copy(trees_block_mapping(ft)),
-  )
+  return FusionTensor(copy(data_matrix(ft)), copy.(axes(ft)), copy(trees_block_mapping(ft)))
 end
 
 function Base.deepcopy(ft::FusionTensor)
   return FusionTensor(
-    deepcopy(data_matrix(ft)),
-    deepcopy.(codomain_axes(ft)),
-    deepcopy.(domain_axes(ft)),
-    deepcopy(trees_block_mapping(ft)),
+    deepcopy(data_matrix(ft)), deepcopy(axes(ft)), deepcopy(trees_block_mapping(ft))
   )
 end
 
@@ -91,16 +82,33 @@ end
 
 Base.permutedims(ft::FusionTensor, args...) = fusiontensor_permutedims(ft, args...)
 
-function Base.similar(ft::FusionTensor, ::Type{T}) where {T}
+Base.similar(ft::FusionTensor) = similar(ft, eltype(ft))
+function Base.similar(ft::FusionTensor, T::Type)
+  # reuse trees_block_mapping
+
   # some fusion trees have Float64 eltype: need compatible type
   @assert promote_type(T, fusiontree_eltype(sector_type(ft))) === T
   mat = similar(data_matrix(ft), T)
   initialize_allowed_sectors!(mat)
-  return FusionTensor(mat, codomain_axes(ft), domain_axes(ft), trees_block_mapping(ft))
+  return FusionTensor(mat, axes(ft), trees_block_mapping(ft))
+end
+
+# trigger explicit error in TensorAlgebra.contract
+# TBD impose some convention? Remove?
+function Base.similar(
+  ft::FusionTensor, T::Type, new_axes::Tuple{Vararg{AbstractGradedUnitRange}}
+)
+  throw(DimensionMismatch("Need bituple of axes"))
+end
+function Base.similar(ft::FusionTensor, T::Type, new_axes::Tuple{})
+  throw(DimensionMismatch("Need bituple of axes"))
 end
 
-function Base.similar(::FusionTensor, ::Type{T}, new_axes::Tuple{<:Tuple,<:Tuple}) where {T}
-  return FusionTensor(T, new_axes[1], new_axes[2])
+function Base.similar(ft::FusionTensor, T::Type, new_axes::Tuple{<:Tuple,<:Tuple})
+  return similar(ft, T, tuplemortar(new_axes))
+end
+function Base.similar(::FusionTensor, T::Type, new_axes::BlockedTuple{2})
+  return FusionTensor(T, new_axes)
 end
 
 Base.show(io::IO, ft::FusionTensor) = print(io, "$(ndims(ft))-dim FusionTensor")

src/fusiontensor/fusiontensor.jl

@@ -15,46 +15,39 @@ using GradedUnitRanges:
   sector_type,
   space_isequal
 using SymmetrySectors: SectorProduct, TrivialSector
+using TensorAlgebra: BlockedTuple, tuplemortar
 
-struct FusionTensor{T,N,CoDomainAxes,DomainAxes,Mat,Mapping} <: AbstractArray{T,N}
+struct FusionTensor{T,N,Axes,Mat,Mapping} <: AbstractArray{T,N}
   data_matrix::Mat
-  codomain_axes::CoDomainAxes
-  domain_axes::DomainAxes
+  axes::Axes
   trees_block_mapping::Mapping
 
   # inner constructor to impose constraints on types
   function FusionTensor(
     mat::AbstractMatrix,
-    codomain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
-    domain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
+    legs::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractGradedUnitRange}}},
     trees_block_mapping::Dict,
   )
     S = sector_type(axes(mat, 1))
     @assert sector_type(axes(mat, 2)) === S
     @assert keytype(trees_block_mapping) <:
       Tuple{<:SectorFusionTree{S},<:SectorFusionTree{S}}
-    @assert all(sector_type.(codomain_legs) .=== S)
-    @assert all(sector_type.(domain_legs) .=== S)
+    @assert all(sector_type.(Tuple(legs)) .=== S)
     return new{
-      eltype(mat),
-      length(codomain_legs) + length(domain_legs),
-      typeof(codomain_legs),
-      typeof(domain_legs),
-      typeof(mat),
-      typeof(trees_block_mapping),
+      eltype(mat),length(legs),typeof(legs),typeof(mat),typeof(trees_block_mapping)
     }(
-      mat, codomain_legs, domain_legs, trees_block_mapping
+      mat, legs, trees_block_mapping
     )
   end
 end
 
 # getters
 data_matrix(ft::FusionTensor) = ft.data_matrix
-codomain_axes(ft::FusionTensor) = ft.codomain_axes
-domain_axes(ft::FusionTensor) = ft.domain_axes
 trees_block_mapping(ft::FusionTensor) = ft.trees_block_mapping
 
 # misc access
+codomain_axes(ft::FusionTensor) = first(blocks(axes(ft)))
+domain_axes(ft::FusionTensor) = last(blocks(axes(ft)))
 ndims_codomain(ft::FusionTensor) = length(codomain_axes(ft))
 ndims_domain(ft::FusionTensor) = length(domain_axes(ft))
 
@@ -68,7 +61,7 @@ end
 
 # GradedUnitRanges interface
 function GradedUnitRanges.sector_type(
-  ::Type{<:FusionTensor{<:Any,<:Any,<:Any,<:Any,<:Any,<:Dict{<:Tuple{<:Any,F}}}}
+  ::Type{<:FusionTensor{<:Any,<:Any,<:Any,<:Any,<:Dict{<:Tuple{<:Any,F}}}}
 ) where {F}
   return sector_type(F)
 end
@@ -91,16 +84,10 @@ function sanitize_axes(raw_legs::Tuple{Vararg{AbstractGradedUnitRange}})
   @assert all(check_unique_blocklabels.(legs))
   return legs
 end
-sanitize_axes(::Tuple{}, ::Tuple{}) = TrivialSector, (), ()
-function sanitize_axes(
-  codomain_legs_raw::Tuple{Vararg{AbstractGradedUnitRange}},
-  domain_legs_raw::Tuple{Vararg{AbstractGradedUnitRange}},
-)
-  legs = sanitize_axes((codomain_legs_raw..., domain_legs_raw...))
-  S = sector_type(first(legs))
-  codomain_legs = legs[begin:length(codomain_legs_raw)]
-  domain_legs = legs[(length(codomain_legs_raw) + 1):end]
-  return S, domain_legs, codomain_legs
+sanitize_axes(legs::BlockedTuple{2,(0, 0)}) = TrivialSector, legs
+function sanitize_axes(raw_legs::BlockedTuple{2})
+  flat_legs = sanitize_axes(Tuple(raw_legs))
+  return sector_type(first(flat_legs)), BlockedTuple(flat_legs, Val(blocklengths(raw_legs)))
 end
 
 function check_unique_blocklabels(g::AbstractGradedUnitRange)
@@ -131,12 +118,16 @@ end
 
 # initialize with already computed data_matrix
 function FusionTensor(
-  mat::AbstractMatrix,
+  x,
   codomain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
   domain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
 )
+  return FusionTensor(x, tuplemortar((codomain_legs, domain_legs)))
+end
+
+function FusionTensor(mat::AbstractMatrix, legs::BlockedTuple{2})
   # init with empty data_matrix to construct trees_block_mapping
-  ft = FusionTensor(eltype(mat), codomain_legs, domain_legs)
+  ft = FusionTensor(eltype(mat), legs)
   @assert space_isequal(matrix_row_axis(ft), axes(mat, 1))
   @assert space_isequal(matrix_column_axis(ft), axes(mat, 2))
   for b in eachblockstoredindex(mat)
@@ -147,21 +138,17 @@ function FusionTensor(
 end
 
 # empty matrix
-function FusionTensor(
-  elt::Type,
-  codomain_legs_raw::Tuple{Vararg{AbstractGradedUnitRange}},
-  domain_legs_raw::Tuple{Vararg{AbstractGradedUnitRange}},
-)
-  S, domain_legs, codomain_legs = sanitize_axes(codomain_legs_raw, domain_legs_raw)
+function FusionTensor(elt::Type, raw_legs::BlockedTuple{2})
+  S, legs = sanitize_axes(raw_legs)
 
-  row_axis, codomain_trees_to_ranges_mapping = fuse_axes(S, codomain_legs)
-  nondual_col_axis, domain_trees_to_ranges_mapping = fuse_axes(S, dual.(domain_legs))
+  row_axis, codomain_trees_to_ranges_mapping = fuse_axes(S, first(blocks(legs)))
+  nondual_col_axis, domain_trees_to_ranges_mapping = fuse_axes(S, dual.(last(blocks(legs))))
 
   mat = initialize_data_matrix(elt, row_axis, nondual_col_axis)
   tree_to_block_mapping = intersect_codomain_domain(
     codomain_trees_to_ranges_mapping, domain_trees_to_ranges_mapping
   )
-  return FusionTensor(mat, codomain_legs, domain_legs, tree_to_block_mapping)
+  return FusionTensor(mat, legs, tree_to_block_mapping)
 end
 
 function fuse_axes(::Type{S}, ::Tuple{}) where {S<:AbstractSector}
@@ -178,14 +165,19 @@ end
 function fusion_trees_external_multiplicities(
   outer_legs::Tuple{Vararg{AbstractGradedUnitRange}}
 )
-  tree_arrows = isdual.(outer_legs)
   return mapreduce(vcat, CartesianIndices(blocklength.(outer_legs))) do it
-    block_sectors = map((g, i) -> blocklabels(g)[i], outer_legs, Tuple(it))
-    block_mult = mapreduce((g, i) -> blocklengths(g)[i], *, outer_legs, Tuple(it); init=1)
-    return build_trees(block_sectors, tree_arrows) .=> block_mult
+    return fusion_trees_external_multiplicities(outer_legs, Tuple(it))
   end
 end
 
+function fusion_trees_external_multiplicities(
+  outer_legs::NTuple{N,AbstractGradedUnitRange}, indices::NTuple{N,Int}
+) where {N}
+  block_sectors = map((g, i) -> blocklabels(g)[i], outer_legs, indices)
+  block_mult = mapreduce((g, i) -> blocklengths(g)[i], *, outer_legs, indices; init=1)
+  return build_trees(block_sectors, isdual.(outer_legs)) .=> block_mult
+end
+
 function compute_inner_ranges(
   fusion_trees_mult::AbstractVector{<:Pair{<:SectorFusionTree,<:Integer}}
 )

src/fusiontensor/tensor_algebra_interface.jl

@@ -6,38 +6,40 @@ using BlockArrays: Block
 
 using TensorAlgebra: BlockedPermutation, Matricize, TensorAlgebra
 
-# TBD how to deal with inner contraction = no ouput axis?
+# TODO how to deal with inner contraction = no ouput axis?
+# => currently biperm_dest is a BlockedPermutation{0}, change this
 function TensorAlgebra.allocate_output(
   ::typeof(contract),
   biperm_dest::BlockedPermutation{2},
-  a1::FusionTensor{T1,N},
-  biperm1::BlockedPermutation{2,N},
-  a2::FusionTensor{T2,M},
-  biperm2::BlockedPermutation{2,M},
+  a1::FusionTensor,
+  biperm1::BlockedPermutation{2},
+  a2::FusionTensor,
+  biperm2::BlockedPermutation{2},
   α::Number=true,
-) where {T1,T2,N,M}
+)
   axes_dest = (
-    (i -> axes(a1)[i]).(biperm1[Block(1)]), (i -> axes(a2)[i]).(biperm2[Block(2)])
+    map(i -> axes(a1)[i], first(blocks(biperm1))),
+    map(i -> axes(a2)[i], last(blocks(biperm2))),
   )
   return similar(a1, promote_type(eltype(a1), eltype(a2), typeof(α)), axes_dest)
 end
 
-# TBD do really I need to defined these as I cannot use them in contract! and has to redefine it?
-# TensorAlgebra.fusedims(ft::FusionTensor, perm::BlockedPermutation) = permutedims(ft, perm)
-# function TensorAlgebra.splitdims(ft1::FusionTensor, ft2::FusionTensor, blockedperm::BlockedPermutation)
-# function TensorAlgebra.splitdims!(ft1::FusionTensor, ft2::FusionTensor, blockedperm::BlockedPermutation)
+# TBD do really I need to define these as I cannot use them in contract! and has to redefine it?
+#TensorAlgebra.fusedims(ft::FusionTensor, perm::BlockedPermutation{2}) = permutedims(ft, perm)
+#function TensorAlgebra.splitdims(ft1::FusionTensor, ft2::FusionTensor, blockedperm::BlockedPermutation)
+#function TensorAlgebra.splitdims!(ft1::FusionTensor, ft2::FusionTensor, blockedperm::BlockedPermutation)
 
 # I cannot use contract! from TensorAlgebra/src/contract/contract_matricize/contract.jl
 # as it calls _mul!, which I should not overload.
-# TBD I can also overload higher up and do not allow use of different algorithms
+# TBD define fallback _mul!(::AbstractArray, ::AbstractArray, ::AbstractArray) in TensorAlgebra?
 function TensorAlgebra.contract!(
-  alg::Matricize,
+  ::Matricize,
   a_dest::FusionTensor,
-  biperm_dest::BlockedPermutation,
+  ::BlockedPermutation{2},
   a1::FusionTensor,
-  biperm1::BlockedPermutation,
+  biperm1::BlockedPermutation{2},
   a2::FusionTensor,
-  biperm2::BlockedPermutation,
+  biperm2::BlockedPermutation{2},
   α::Number,
   β::Number,
 )

test/basics/test_array_cast.jl

@@ -201,14 +201,13 @@ end
     if VERSION < v"1.11"
       @test_broken to_fusiontensor(zerodim, (), ()) isa FusionTensor  # https://github.com/JuliaLang/julia/issues/52615
     else
-      # TODO fix: add specialized method, maybe fix TensorAlgebra
-      @test_broken to_fusiontensor(zerodim, (), ())
-      #@test ft isa FusionTensor
-      #@test ndims(ft) == 0
-      #@test isnothing(check_sanity(ft))
-      #@test size(data_matrix(ft)) == (1, 1)
-      #@test data_matrix(ft)[1, 1] ≈ 1.0
-      #@test_broken Array(ft) ≈ zerodim  # cannot create zerodim BlockSparseArray
+      ft = to_fusiontensor(zerodim, (), ())
+      @test ft isa FusionTensor
+      @test ndims(ft) == 0
+      @test isnothing(check_sanity(ft))
+      @test size(data_matrix(ft)) == (1, 1)
+      @test data_matrix(ft)[1, 1] ≈ 1.0
+      @test_broken Array(ft) ≈ zerodim # https://github.com/ITensor/BlockSparseArrays.jl/issues/27
     end
   end
 end