Move code from GradedUnitRanges.jl, LabelledNumbers.jl, SymmetrySectors.jl (#1)

Author: mtfishman

Project.toml

@@ -3,5 +3,21 @@ uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 authors = ["ITensor developers <[email protected]> and contributors"]
 version = "0.1.0"
 
+[deps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
+TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
+
 [compat]
+BlockArrays = "1.5.0"
+Compat = "4.16.0"
+FillArrays = "1.13.0"
+HalfIntegers = "1.6.0"
+Random = "1.10.0"
+SplitApplyCombine = "1.2.3"
+TensorProducts = "0.1.3"
 julia = "1.10"

docs/src/reference.md

@@ -1,5 +1,5 @@
 # Reference
 
 ```@autodocs
-Modules = [GradedArrays]
+Modules = [GradedArrays, GradedArrays.LabelledNumbers, GradedArrays.GradedUnitRanges, GradedArrays.SymmetrySectors]
 ```

src/GradedArrays.jl

@@ -1,5 +1,10 @@
 module GradedArrays
 
-# Write your package code here.
+include("lib/LabelledNumbers/LabelledNumbers.jl")
+using .LabelledNumbers: LabelledNumbers
+include("lib/GradedUnitRanges/GradedUnitRanges.jl")
+using .GradedUnitRanges: GradedUnitRanges
+include("lib/SymmetrySectors/SymmetrySectors.jl")
+using .SymmetrySectors: SymmetrySectors
 
 end

src/lib/GradedUnitRanges/GradedUnitRanges.jl

@@ -0,0 +1,12 @@
+module GradedUnitRanges
+
+export gradedrange
+
+include("blockedunitrange.jl")
+include("gradedunitrange.jl")
+include("dual.jl")
+include("labelledunitrangedual.jl")
+include("gradedunitrangedual.jl")
+include("fusion.jl")
+
+end

src/lib/GradedUnitRanges/blockedunitrange.jl

@@ -0,0 +1,184 @@
+using BlockArrays:
+  BlockArrays,
+  AbstractBlockedUnitRange,
+  AbstractBlockVector,
+  Block,
+  BlockIndex,
+  BlockIndexRange,
+  BlockRange,
+  BlockSlice,
+  BlockVector,
+  block,
+  blockindex,
+  findblock,
+  findblockindex,
+  mortar
+
+# Custom `BlockedUnitRange` constructor that takes a unit range
+# and a set of block lengths, similar to `BlockArray(::AbstractArray, blocklengths...)`.
+function blockedunitrange(a::AbstractUnitRange, blocklengths)
+  blocklengths_shifted = copy(blocklengths)
+  blocklengths_shifted[1] += (first(a) - 1)
+  blocklasts = cumsum(blocklengths_shifted)
+  return BlockArrays._BlockedUnitRange(first(a), blocklasts)
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+# TODO: Rename this. `BlockArrays.findblock(a, k)` finds the
+# block of the value `k`, while this finds the block of the index `k`.
+# This could make use of the `BlockIndices` object, i.e. `block(BlockIndices(a)[index])`.
+function blockedunitrange_findblock(a::AbstractBlockedUnitRange, index::Integer)
+  @boundscheck index in 1:length(a) || throw(BoundsError(a, index))
+  return @inbounds findblock(a, index + first(a) - 1)
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+# TODO: Rename this. `BlockArrays.findblockindex(a, k)` finds the
+# block index of the value `k`, while this finds the block index of the index `k`.
+# This could make use of the `BlockIndices` object, i.e. `BlockIndices(a)[index]`.
+function blockedunitrange_findblockindex(a::AbstractBlockedUnitRange, index::Integer)
+  @boundscheck index in 1:length(a) || throw(BoundsError())
+  return @inbounds findblockindex(a, index + first(a) - 1)
+end
+
+function blockedunitrange_getindices(a::AbstractUnitRange, indices)
+  return a[indices]
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+# Like `a[indices]` but preserves block structure.
+# TODO: Consider calling this something else, for example
+# `blocked_getindex`. See the discussion here:
+# https://github.com/JuliaArrays/BlockArrays.jl/issues/347
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::AbstractUnitRange{<:Integer}
+)
+  first_blockindex = blockedunitrange_findblockindex(a, first(indices))
+  last_blockindex = blockedunitrange_findblockindex(a, last(indices))
+  first_block = block(first_blockindex)
+  last_block = block(last_blockindex)
+  blocklengths = if first_block == last_block
+    [length(indices)]
+  else
+    map(first_block:last_block) do block
+      if block == first_block
+        return length(a[first_block]) - blockindex(first_blockindex) + 1
+      end
+      if block == last_block
+        return blockindex(last_blockindex)
+      end
+      return length(a[block])
+    end
+  end
+  return blockedunitrange(indices .+ (first(a) - 1), blocklengths)
+end
+
+# TODO: Make sure this handles block labels (AbstractGradedUnitRange) correctly.
+# TODO: Make a special case for `BlockedVector{<:Block{1},<:BlockRange{1}}`?
+# For example:
+# ```julia
+# blocklengths = map(bs -> sum(b -> length(a[b]), bs), blocks(indices))
+# return blockedrange(blocklengths)
+# ```
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::AbstractBlockVector{<:Block{1}}
+)
+  blks = map(bs -> mortar(map(b -> a[b], bs)), blocks(indices))
+  # We pass `length.(blks)` to `mortar` in order
+  # to pass block labels to the axes of the output,
+  # if they exist. This makes it so that
+  # `only(axes(a[indices])) isa `GradedUnitRange`
+  # if `a isa `GradedUnitRange`, for example.
+  # Note there is a more specialized definition:
+  # ```julia
+  # function blockedunitrange_getindices(
+  #   a::AbstractGradedUnitRange, indices::AbstractBlockVector{<:Block{1}}
+  # )
+  # ```
+  # that does a better job of preserving labels, since `length`
+  # may drop labels for certain block types.
+  return mortar(blks, length.(blks))
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockIndexRange)
+  return a[block(indices)][only(indices.indices)]
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockSlice)
+  # TODO: Is this a good definition? It ignores `indices.indices`.
+  return a[indices.block]
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::Vector{<:Integer}
+)
+  return map(index -> a[index], indices)
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+# TODO: Make a special definition for `BlockedVector{<:Block{1}}` in order
+# to merge blocks.
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
+)
+  # Without converting `indices` to `Vector`,
+  # mapping `indices` outputs a `BlockVector`
+  # which is harder to reason about.
+  blocks = map(index -> a[index], Vector(indices))
+  # We pass `length.(blocks)` to `mortar` in order
+  # to pass block labels to the axes of the output,
+  # if they exist. This makes it so that
+  # `only(axes(a[indices])) isa `GradedUnitRange`
+  # if `a isa `GradedUnitRange`, for example.
+  return mortar(blocks, length.(blocks))
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::Block{1})
+  return a[indices]
+end
+
+function blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
+)
+  return mortar(map(b -> a[b], blocks(indices)))
+end
+
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices)
+  return error("Not implemented.")
+end
+
+# The blocks of the corresponding slice.
+_blocks(a::AbstractUnitRange, indices) = error("Not implemented")
+function _blocks(a::AbstractUnitRange, indices::AbstractUnitRange)
+  return findblock(a, first(indices)):findblock(a, last(indices))
+end
+function _blocks(a::AbstractUnitRange, indices::BlockRange)
+  return indices
+end
+
+# Slice `a` by `I`, returning a:
+# `BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}`
+# with the `BlockIndex{1}` corresponding to each value of `I`.
+function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::UnitRange{<:Integer})
+  return mortar(
+    map(blocks(blockedunitrange_getindices(a, I))) do r
+      bi_first = findblockindex(a, first(r))
+      bi_last = findblockindex(a, last(r))
+      @assert block(bi_first) == block(bi_last)
+      return block(bi_first)[blockindex(bi_first):blockindex(bi_last)]
+    end,
+  )
+end
+
+# This handles non-blocked slices.
+# For example:
+# a = BlockSparseArray{Float64}([2, 2, 2, 2])
+# I = BlockedVector(Block.(1:4), [2, 2])
+# @views a[I][Block(1)]
+to_blockindices(a::Base.OneTo{<:Integer}, I::UnitRange{<:Integer}) = I

src/lib/GradedUnitRanges/dual.jl

@@ -0,0 +1,39 @@
+using ..LabelledNumbers: LabelledInteger, label, labelled, unlabel
+
+"""
+    dual(x)
+
+Take the dual of the symmetry sector, graded unit range, etc.
+By default, it just returns `x`, i.e. it assumes the object
+is self-dual.
+"""
+dual(x) = x
+
+nondual(r::AbstractUnitRange) = r
+isdual(::AbstractUnitRange) = false
+
+dual_type(x) = dual_type(typeof(x))
+dual_type(T::Type) = T
+nondual_type(x) = nondual_type(typeof(x))
+nondual_type(T::Type) = T
+
+dual(i::LabelledInteger) = labelled(unlabel(i), dual(label(i)))
+flip(a::AbstractUnitRange) = dual(map_blocklabels(dual, a))
+
+"""
+    dag(r::AbstractUnitRange)
+
+Same as `dual(r)`.
+"""
+dag(r::AbstractUnitRange) = dual(r)
+
+"""
+    dag(a::AbstractArray)
+
+Complex conjugates `a` and takes the dual of the axes.
+"""
+function dag(a::AbstractArray)
+  a′ = similar(a, dual.(axes(a)))
+  a′ .= conj.(a)
+  return a′
+end

src/lib/GradedUnitRanges/fusion.jl

@@ -0,0 +1,81 @@
+using BlockArrays: AbstractBlockedUnitRange, blocklengths
+using ..LabelledNumbers: LabelledInteger, label, labelled
+using SplitApplyCombine: groupcount
+using TensorProducts: TensorProducts, OneToOne, tensor_product
+
+flip_dual(r::AbstractUnitRange) = r
+flip_dual(r::GradedUnitRangeDual) = flip(r)
+
+function fuse_labels(x, y)
+  return error(
+    "`fuse_labels` not implemented for object of type `$(typeof(x))` and `$(typeof(y))`."
+  )
+end
+
+function fuse_blocklengths(x::LabelledInteger, y::LabelledInteger)
+  # return blocked unit range to keep non-abelian interface
+  return blockedrange([labelled(x * y, fuse_labels(label(x), label(y)))])
+end
+
+unmerged_tensor_product() = OneToOne()
+unmerged_tensor_product(a) = a
+unmerged_tensor_product(a, ::OneToOne) = a
+unmerged_tensor_product(::OneToOne, a) = a
+unmerged_tensor_product(::OneToOne, ::OneToOne) = OneToOne()
+unmerged_tensor_product(a1, a2) = tensor_product(a1, a2)
+function unmerged_tensor_product(a1, a2, as...)
+  return unmerged_tensor_product(unmerged_tensor_product(a1, a2), as...)
+end
+
+function unmerged_tensor_product(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
+  nested = map(Iterators.flatten((Iterators.product(blocks(a1), blocks(a2)),))) do it
+    return mapreduce(length, fuse_blocklengths, it)
+  end
+  new_blocklengths = mapreduce(blocklengths, vcat, nested)
+  return blockedrange(new_blocklengths)
+end
+
+# convention: sort GradedUnitRangeDual according to nondual blocks
+function blocksortperm(a::AbstractUnitRange)
+  return Block.(sortperm(blocklabels(nondual(a))))
+end
+
+# Get the permutation for sorting, then group by common elements.
+# groupsortperm([2, 1, 2, 3]) == [[2], [1, 3], [4]]
+function groupsortperm(v; kwargs...)
+  perm = sortperm(v; kwargs...)
+  v_sorted = @view v[perm]
+  group_lengths = collect(groupcount(identity, v_sorted))
+  return BlockVector(perm, group_lengths)
+end
+
+# Used by `TensorAlgebra.splitdims` in `BlockSparseArraysGradedUnitRangesExt`.
+# Get the permutation for sorting, then group by common elements.
+# groupsortperm([2, 1, 2, 3]) == [[2], [1, 3], [4]]
+function blockmergesortperm(a::AbstractUnitRange)
+  return Block.(groupsortperm(blocklabels(nondual(a))))
+end
+
+# Used by `TensorAlgebra.splitdims` in `BlockSparseArraysGradedUnitRangesExt`.
+invblockperm(a::Vector{<:Block{1}}) = Block.(invperm(Int.(a)))
+
+function blockmergesort(g::AbstractGradedUnitRange)
+  glabels = blocklabels(g)
+  gblocklengths = blocklengths(g)
+  new_blocklengths = map(sort(unique(glabels))) do la
+    return labelled(sum(gblocklengths[findall(==(la), glabels)]; init=0), la)
+  end
+  return gradedrange(new_blocklengths)
+end
+
+blockmergesort(g::GradedUnitRangeDual) = flip(blockmergesort(flip(g)))
+blockmergesort(g::AbstractUnitRange) = g
+
+# tensor_product produces a sorted, non-dual GradedUnitRange
+TensorProducts.tensor_product(g::AbstractGradedUnitRange) = blockmergesort(flip_dual(g))
+
+function TensorProducts.tensor_product(
+  g1::AbstractGradedUnitRange, g2::AbstractGradedUnitRange
+)
+  return blockmergesort(unmerged_tensor_product(g1, g2))
+end