define FusionTensorAxes (#58)

Author: ogauthe

Project.toml

@@ -1,7 +1,7 @@
 name = "FusionTensors"
 uuid = "e16ca583-1f51-4df0-8e12-57d32947d33e"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -21,7 +21,7 @@ WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 Accessors = "0.1.42"
 BlockArrays = "1.6"
 BlockSparseArrays = "0.7.4"
-GradedArrays = "0.4.7"
+GradedArrays = "0.4.13"
 HalfIntegers = "1.6"
 LRUCache = "1.6"
 LinearAlgebra = "1.10"

src/FusionTensors.jl

@@ -2,6 +2,7 @@ module FusionTensors
 
 include("fusion_trees/fusiontree.jl")
 include("fusion_trees/clebsch_gordan_tensors.jl")
+include("fusiontensor/fusiontensoraxes.jl")
 include("fusiontensor/fusiontensor.jl")
 include("fusiontensor/base_interface.jl")
 include("fusiontensor/array_cast.jl")

src/fusiontensor/base_interface.jl

@@ -44,7 +44,7 @@ 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))))
+  new_axes = FusionTensorAxes(dual.(domain_axes(ft)), dual.(codomain_axes(ft)))
   return FusionTensor(
     adjoint(data_matrix(ft)), new_axes, transpose_mapping(trees_block_mapping(ft))
   )
@@ -57,7 +57,7 @@ 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.(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)
@@ -66,8 +66,7 @@ function Base.deepcopy(ft::FusionTensor)
   )
 end
 
-# eachindex is automatically defined for AbstractArray. We do not want it.
-Base.eachindex(::FusionTensor) = error("eachindex not defined for FusionTensor")
+Base.eachindex(::FusionTensor) = throw(MethodError(eachindex, (FusionTensor,)))
 
 function Base.getindex(ft::FusionTensor, f1::SectorFusionTree, f2::SectorFusionTree)
   charge_matrix = data_matrix(ft)[trees_block_mapping(ft)[f1, f2]]

src/fusiontensor/fusiontensor.jl

@@ -19,136 +19,32 @@ using GradedArrays:
   sectormergesort,
   sectors,
   space_isequal
-using TensorAlgebra: BlockedTuple, tuplemortar
+using TensorAlgebra: BlockedTuple, trivial_axis, tuplemortar
 using TensorProducts: tensor_product
 using TypeParameterAccessors: type_parameters
 
-struct FusionTensor{T,N,Axes,Mat<:AbstractMatrix{T},Mapping} <: AbstractArray{T,N}
-  data_matrix::Mat
-  axes::Axes
-  trees_block_mapping::Mapping
-
-  # inner constructor to impose constraints on types
-  function FusionTensor{T,N,Axes,Mat,Mapping}(
-    mat, legs, trees_block_mapping
-  ) where {T,N,Axes,Mat,Mapping}
-    S = isempty(legs) ? TrivialSector : sector_type(first(legs))
-    @assert keytype(trees_block_mapping) <:
-      Tuple{<:SectorFusionTree{S},<:SectorFusionTree{S}}
-    @assert all(sector_type.(Tuple(legs)) .=== S)
-    return new{T,N,Axes,Mat,Mapping}(mat, legs, trees_block_mapping)
-  end
-end
-
-function FusionTensor(
-  mat::AbstractMatrix,
-  legs::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractGradedUnitRange}}},
-  trees_block_mapping::Dict{<:Tuple{<:SectorFusionTree,<:SectorFusionTree}},
-)
-  return FusionTensor{
-    eltype(mat),length(legs),typeof(legs),typeof(mat),typeof(trees_block_mapping)
-  }(
-    mat, legs, trees_block_mapping
-  )
-end
-
-# getters
-data_matrix(ft::FusionTensor) = ft.data_matrix
-trees_block_mapping(ft::FusionTensor) = ft.trees_block_mapping
-
-# misc access
-codomain_axes(ft::FusionTensor) = axes(ft)[Block(1)]
-domain_axes(ft::FusionTensor) = axes(ft)[Block(2)]
-ndims_codomain(ft::FusionTensor) = length(codomain_axes(ft))
-ndims_domain(ft::FusionTensor) = length(domain_axes(ft))
-
-dummy_axis(ft::FusionTensor) = dummy_axis(sector_type(ft))
-dummy_axis(::Type{S}) where {S<:AbstractSector} = gradedrange([trivial(S) => 1])
-
-function codomain_axis(ft::FusionTensor)
-  if ndims_codomain(ft) == 0
-    return dummy_axis(ft)
-  end
-  return ⊗(codomain_axes(ft)...)
-end
-function domain_axis(ft::FusionTensor)
-  if ndims_domain(ft) == 0
-    return dummy_axis(ft)
-  end
-  return dual(⊗(dual.(domain_axes(ft))...))
-end
-function charge_block_size(ft::FusionTensor, f1::SectorFusionTree, f2::SectorFusionTree)
-  b = Tuple(findblock(ft, f1, f2))
-  return ntuple(i -> sector_multiplicity(axes(ft, i)[b[i]]), ndims(ft))
-end
+# =======================================  Misc  ===========================================
 
-# GradedArrays interface
-function GradedArrays.sector_type(
-  ::Type{<:FusionTensor{<:Any,<:Any,<:Any,<:Any,<:Dict{<:Tuple{<:Any,F}}}}
-) where {F}
-  return sector_type(F)
-end
-
-# BlockArrays interface
-function BlockArrays.findblock(ft::FusionTensor, f1::SectorFusionTree, f2::SectorFusionTree)
-  # find outer block corresponding to fusion trees
-  @assert typeof((f1, f2)) === keytype(trees_block_mapping(ft))
-  b1 = find_sector_block.(leaves(f1), codomain_axes(ft))
-  b2 = find_sector_block.(leaves(f2), domain_axes(ft))
-  return Block(b1..., b2...)
-end
 # TBD move to GradedArrays? rename findfirst_sector?
 function find_sector_block(s::AbstractSector, g::AbstractGradedUnitRange)
   return findfirst(==(s), sectors(flip_dual(g)))
 end
 
-function sanitize_axes(raw_legs::Tuple{Vararg{AbstractGradedUnitRange}})
-  legs = promote_sectors(raw_legs)
-  @assert all(check_unique_sectors.(legs))
-  return legs
-end
-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_sectors(g::AbstractGradedUnitRange)
-  return length(unique(sectors(g))) == blocklength(g)
-end
-
-promote_sectors(legs::NTuple{<:Any,<:AbstractGradedUnitRange}) = legs # nothing to do
-function promote_sectors(legs)
-  T = promote_sector_type(legs)
-  # fuse with trivial to insert all missing arguments inside each GradedAxis
-  # avoid depending on GradedArrays internals
-  s0 = trivial(T)
-  return map_sectors.(s -> only(sectors(to_gradedrange(tensor_product(s0, s)))), legs)
-end
-
-function promote_sector_type(legs)
-  # fuse trivial sectors to produce unified type
-  # avoid depending on GradedArrays internals
-  return sector_type(tensor_product(trivial.(legs)...))
+# TBD move to GradedArrays?
+function checkaxes(::Type{Bool}, axes1, axes2)
+  return length(axes1) == length(axes2) && all(space_isequal.(axes1, axes2))
 end
 
-# initialize with already computed data_matrix
-function FusionTensor(
-  x,
-  codomain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
-  domain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
-)
-  return FusionTensor(x, tuplemortar((codomain_legs, domain_legs)))
+# TBD move to GradedArrays?
+checkaxes_dual(axes1, axes2) = checkaxes(axes1, dual.(axes2))
+function checkaxes(ax1, ax2)
+  return checkaxes(Bool, ax1, ax2) ||
+         throw(DimensionMismatch(lazy"$ax1 does not match $ax2"))
 end
 
-function FusionTensor(mat::AbstractMatrix, legs::BlockedTuple{2})
-  # init with empty data_matrix to construct trees_block_mapping
-  ft = FusionTensor(eltype(mat), legs)
-  for b in eachblockstoredindex(mat)
-    @assert b in eachblockstoredindex(data_matrix(ft))  # check matrix block is allowed
-    data_matrix(ft)[b] = mat[b]
-  end
-  return ft
+function to_blockindexrange(b1::BlockIndexRange{1}, b2::BlockIndexRange{1})
+  t = (b1, b2)
+  return Block(Block.(t))[to_block_indices.(t)...]
 end
 
 function flip_domain(nonflipped_col_axis, nonflipped_trees_to_ranges)
@@ -159,21 +55,8 @@ function flip_domain(nonflipped_col_axis, nonflipped_trees_to_ranges)
   return col_axis, domain_trees_to_ranges_mapping
 end
 
-# empty matrix
-function FusionTensor(elt::Type, raw_legs::BlockedTuple{2})
-  S, legs = sanitize_axes(raw_legs)
-  row_axis, codomain_trees_to_ranges = fuse_axes(S, legs[Block(1)])
-  col_axis, domain_trees_to_ranges = flip_domain(fuse_axes(S, dual.(legs[Block(2)]))...)
-
-  mat = initialize_data_matrix(elt, row_axis, col_axis)
-  tree_to_block_mapping = intersect_codomain_domain(
-    codomain_trees_to_ranges, domain_trees_to_ranges
-  )
-  return FusionTensor(mat, legs, tree_to_block_mapping)
-end
-
 function fuse_axes(::Type{S}, ::Tuple{}) where {S<:AbstractSector}
-  fused_axis = gradedrange([trivial(S) => 1])
+  fused_axis = trivial_axis(S)
   trees_to_ranges_mapping = Dict([SectorFusionTree{S}() => Block(1)[1:1]])
   return fused_axis, trees_to_ranges_mapping
 end
@@ -213,11 +96,6 @@ function compute_inner_ranges(fusion_trees_mult)
   return fused_leg, range_mapping
 end
 
-function to_blockindexrange(b1::BlockIndexRange{1}, b2::BlockIndexRange{1})
-  t = (b1, b2)
-  return Block(Block.(t))[to_block_indices.(t)...]
-end
-
 function intersect_codomain_domain(
   codomain_trees_to_ranges_mapping::Dict{<:SectorFusionTree,<:BlockIndexRange{1}},
   domain_trees_to_ranges_mapping::Dict{<:SectorFusionTree,<:BlockIndexRange{1}},
@@ -257,11 +135,112 @@ function initialize_data_matrix(
   return mat
 end
 
-checkaxes_dual(axes1, axes2) = checkaxes(axes1, dual.(axes2))
-function checkaxes(ax1, ax2)
-  return checkaxes(Bool, ax1, ax2) ||
-         throw(DimensionMismatch(lazy"$ax1 does not match $ax2"))
+# ====================================  Definitions  =======================================
+
+struct FusionTensor{T,N,Axes<:FusionTensorAxes,Mat<:AbstractMatrix{T},Mapping} <:
+       AbstractArray{T,N}
+  data_matrix::Mat
+  axes::Axes
+  trees_block_mapping::Mapping
+
+  # inner constructor to impose constraints on types
+  function FusionTensor{T,N,Axes,Mat,Mapping}(
+    mat, legs, trees_block_mapping
+  ) where {T,N,Axes,Mat,Mapping}
+    S = sector_type(legs)
+    @assert keytype(trees_block_mapping) <:
+      Tuple{<:SectorFusionTree{S},<:SectorFusionTree{S}}
+    return new{T,N,Axes,Mat,Mapping}(mat, legs, trees_block_mapping)
+  end
 end
-function checkaxes(::Type{Bool}, axes1, axes2)
-  return length(axes1) == length(axes2) && all(space_isequal.(axes1, axes2))
+
+# =====================================  Accessors  ========================================
+
+data_matrix(ft::FusionTensor) = ft.data_matrix
+trees_block_mapping(ft::FusionTensor) = ft.trees_block_mapping
+
+# ====================================  Constructors  ======================================
+
+function FusionTensor(
+  mat::AbstractMatrix,
+  legs::FusionTensorAxes,
+  trees_block_mapping::Dict{<:Tuple{<:SectorFusionTree,<:SectorFusionTree}},
+)
+  return FusionTensor{
+    eltype(mat),length(legs),typeof(legs),typeof(mat),typeof(trees_block_mapping)
+  }(
+    mat, legs, trees_block_mapping
+  )
+end
+
+# empty matrix
+function FusionTensor(elt::Type, legs::FusionTensorAxes)
+  S = sector_type(legs)
+  row_axis, codomain_trees_to_ranges = fuse_axes(S, codomain(legs))
+  col_axis, domain_trees_to_ranges = flip_domain(fuse_axes(S, dual.(domain(legs)))...)
+
+  mat = initialize_data_matrix(elt, row_axis, col_axis)
+  tree_to_block_mapping = intersect_codomain_domain(
+    codomain_trees_to_ranges, domain_trees_to_ranges
+  )
+  return FusionTensor(mat, legs, tree_to_block_mapping)
+end
+
+#constructor from precomputed data_matrix
+function FusionTensor(mat::AbstractMatrix, legs::FusionTensorAxes)
+  # init with empty data_matrix to construct trees_block_mapping
+  ft = FusionTensor(eltype(mat), legs)
+  for b in eachblockstoredindex(mat)
+    b in eachblockstoredindex(data_matrix(ft)) ||
+      throw(ArgumentError("matrix block $b is not allowed"))
+    data_matrix(ft)[b] = mat[b]
+  end
+  return ft
+end
+
+FusionTensor(x, legs::BlockedTuple{2}) = FusionTensor(x, FusionTensorAxes(legs))
+
+# constructor from split axes
+function FusionTensor(
+  x,
+  codomain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
+  domain_legs::Tuple{Vararg{AbstractGradedUnitRange}},
+)
+  return FusionTensor(x, tuplemortar((codomain_legs, domain_legs)))
+end
+
+# ================================  BlockArrays interface  =================================
+
+function BlockArrays.findblock(ft::FusionTensor, f1::SectorFusionTree, f2::SectorFusionTree)
+  # find outer block corresponding to fusion trees
+  @assert typeof((f1, f2)) === keytype(trees_block_mapping(ft))
+  b1 = find_sector_block.(leaves(f1), codomain_axes(ft))
+  b2 = find_sector_block.(leaves(f2), domain_axes(ft))
+  return Block(b1..., b2...)
+end
+
+# ==============================  GradedArrays interface  ==================================
+
+function GradedArrays.sector_type(::Type{FT}) where {FT<:FusionTensor}
+  return sector_type(type_parameters(FT, 3))
+end
+
+# ==============================  FusionTensor interface  ==================================
+
+# misc access
+codomain_axes(ft::FusionTensor) = codomain(axes(ft))
+
+domain_axes(ft::FusionTensor) = domain(axes(ft))
+
+codomain_axis(ft::FusionTensor) = fused_codomain(axes(ft))
+
+domain_axis(ft::FusionTensor) = fused_domain(axes(ft))
+
+ndims_codomain(ft::FusionTensor) = length_codomain(axes(ft))
+
+ndims_domain(ft::FusionTensor) = length_domain(axes(ft))
+
+function charge_block_size(ft::FusionTensor, f1::SectorFusionTree, f2::SectorFusionTree)
+  b = Tuple(findblock(ft, f1, f2))
+  return ntuple(i -> sector_multiplicity(axes(ft, i)[b[i]]), ndims(ft))
 end