GradedArray (#6)

Author: mtfishman

Project.toml

@@ -5,19 +5,34 @@ version = "0.1.0"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
+
+[weakdeps]
+TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+
+[extensions]
+GradedArraysTensorAlgebraExt = "TensorAlgebra"
 
 [compat]
 BlockArrays = "1.5.0"
+BlockSparseArrays = "0.4.0"
 Compat = "4.16.0"
+DerivableInterfaces = "0.4.4"
 FillArrays = "1.13.0"
 HalfIntegers = "1.6.0"
+LinearAlgebra = "1.10.0"
 Random = "1.10.0"
 SplitApplyCombine = "1.2.3"
+TensorAlgebra = "0.2.7"
 TensorProducts = "0.1.3"
+TypeParameterAccessors = "0.3.9"
 julia = "1.10"

docs/Project.toml

@@ -1,4 +1,9 @@
 [deps]
-GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+
+[compat]
+Documenter = "1"
+GradedArrays = "0.1"
+Literate = "2"

examples/Project.toml

@@ -1,2 +1,5 @@
 [deps]
 GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
+
+[compat]
+GradedArrays = "0.1"

ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl

@@ -0,0 +1,109 @@
+module GradedArraysTensorAlgebraExt
+
+using BlockArrays: Block, BlockIndexRange, blockedrange, blocks
+using BlockSparseArrays:
+  BlockSparseArrays,
+  AbstractBlockSparseArray,
+  AbstractBlockSparseArrayInterface,
+  BlockSparseArray,
+  BlockSparseArrayInterface,
+  BlockSparseMatrix,
+  BlockSparseVector,
+  block_merge
+using DerivableInterfaces: @interface
+using GradedArrays.GradedUnitRanges:
+  GradedUnitRanges,
+  AbstractGradedUnitRange,
+  blockmergesortperm,
+  blocksortperm,
+  dual,
+  invblockperm,
+  nondual,
+  unmerged_tensor_product
+using LinearAlgebra: Adjoint, Transpose
+using TensorAlgebra:
+  TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims
+using TensorProducts: OneToOne
+
+#=
+    reducewhile(f, op, collection, state)
+
+reducewhile(x -> length(x) < 3, vcat, ["a", "b", "c", "d"], 2; init=String[]) ==
+  (["b", "c"], 4)
+=#
+function reducewhile(f, op, collection, state; init)
+  prev_result = init
+  prev_state = state
+  result = prev_result
+  while f(result)
+    prev_result = result
+    prev_state = state
+    value_and_state = iterate(collection, state)
+    isnothing(value_and_state) && break
+    value, state = value_and_state
+    result = op(result, value)
+  end
+  return prev_result, prev_state
+end
+
+#=
+    groupreducewhile(f, op, collection, ngroups)
+
+groupreducewhile((i, x) -> length(x) ≤ i, vcat, ["a", "b", "c", "d", "e", "f"], 3; init=String[]) ==
+  (["a"], ["b", "c"], ["d", "e", "f"])
+=#
+function groupreducewhile(f, op, collection, ngroups; init)
+  state = firstindex(collection)
+  return ntuple(ngroups) do group_number
+    result, state = reducewhile(x -> f(group_number, x), op, collection, state; init)
+    return result
+  end
+end
+
+TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion()
+
+# Sort the blocks by sector and then merge the common sectors.
+function block_mergesort(a::AbstractArray)
+  I = blockmergesortperm.(axes(a))
+  return a[I...]
+end
+
+function TensorAlgebra.fusedims(
+  ::SectorFusion, a::AbstractArray, merged_axes::AbstractUnitRange...
+)
+  # First perform a fusion using a block reshape.
+  # TODO avoid groupreducewhile. Require refactor of fusedims.
+  unmerged_axes = groupreducewhile(
+    unmerged_tensor_product, axes(a), length(merged_axes); init=OneToOne()
+  ) do i, axis
+    return length(axis) ≤ length(merged_axes[i])
+  end
+
+  a_reshaped = fusedims(BlockReshapeFusion(), a, unmerged_axes...)
+  # Sort the blocks by sector and merge the equivalent sectors.
+  return block_mergesort(a_reshaped)
+end
+
+function TensorAlgebra.splitdims(
+  ::SectorFusion, a::AbstractArray, split_axes::AbstractUnitRange...
+)
+  # First, fuse axes to get `blockmergesortperm`.
+  # Then unpermute the blocks.
+  axes_prod = groupreducewhile(
+    unmerged_tensor_product, split_axes, ndims(a); init=OneToOne()
+  ) do i, axis
+    return length(axis) ≤ length(axes(a, i))
+  end
+  blockperms = blocksortperm.(axes_prod)
+  sorted_axes = map((r, I) -> only(axes(r[I])), axes_prod, blockperms)
+
+  # TODO: This is doing extra copies of the blocks,
+  # use `@view a[axes_prod...]` instead.
+  # That will require implementing some reindexing logic
+  # for this combination of slicing.
+  a_unblocked = a[sorted_axes...]
+  a_blockpermed = a_unblocked[invblockperm.(blockperms)...]
+  return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...)
+end
+
+end

src/GradedArrays.jl

@@ -1,10 +1,32 @@
 module GradedArrays
 
-include("lib/LabelledNumbers/LabelledNumbers.jl")
+include("LabelledNumbers/LabelledNumbers.jl")
 using .LabelledNumbers: LabelledNumbers
-include("lib/GradedUnitRanges/GradedUnitRanges.jl")
-using .GradedUnitRanges: GradedUnitRanges
-include("lib/SymmetrySectors/SymmetrySectors.jl")
+include("GradedUnitRanges/GradedUnitRanges.jl")
+# This makes the following names accessible
+# as `GradedArrays.x`.
+using .GradedUnitRanges:
+  GradedUnitRanges,
+  GradedOneTo,
+  GradedUnitRange,
+  GradedUnitRangeDual,
+  LabelledUnitRangeDual,
+  blocklabels,
+  blockmergesortperm,
+  blocksortperm,
+  dag,
+  dual,
+  dual_type,
+  flip,
+  gradedrange,
+  isdual,
+  nondual,
+  nondual_type,
+  sector_type,
+  space_isequal,
+  unmerged_tensor_product
+include("SymmetrySectors/SymmetrySectors.jl")
 using .SymmetrySectors: SymmetrySectors
+include("gradedarray.jl")
 
 end

src/lib/GradedUnitRanges/GradedUnitRanges.jl renamed to src/GradedUnitRanges/GradedUnitRanges.jl

@@ -2,7 +2,6 @@ module GradedUnitRanges
 
 export gradedrange
 
-include("blockedunitrange.jl")
 include("gradedunitrange.jl")
 include("dual.jl")
 include("labelledunitrangedual.jl")

src/lib/GradedUnitRanges/dual.jl renamed to src/GradedUnitRanges/dual.jl

@@ -1,10 +1,13 @@
 using BlockArrays:
   BlockArrays,
+  AbstractBlockVector,
   AbstractBlockedUnitRange,
   Block,
+  BlockIndex,
   BlockIndexRange,
   BlockRange,
   BlockSlice,
+  BlockVector,
   BlockedOneTo,
   BlockedUnitRange,
   block,
@@ -15,7 +18,14 @@ using BlockArrays:
   blocks,
   blockindex,
   combine_blockaxes,
+  mortar,
   sortedunion
+using BlockSparseArrays:
+  BlockSparseArrays,
+  _blocks,
+  blockedunitrange_findblock,
+  blockedunitrange_findblockindex,
+  blockedunitrange_getindices
 using Compat: allequal
 using FillArrays: Fill
 using ..LabelledNumbers:
@@ -127,11 +137,15 @@ function BlockArrays.findblock(a::AbstractGradedUnitRange, index::Integer)
   return blockedunitrange_findblock(unlabel_blocks(a), index)
 end
 
-function blockedunitrange_findblock(a::AbstractGradedUnitRange, index::Integer)
+function BlockSparseArrays.blockedunitrange_findblock(
+  a::AbstractGradedUnitRange, index::Integer
+)
   return blockedunitrange_findblock(unlabel_blocks(a), index)
 end
 
-function blockedunitrange_findblockindex(a::AbstractGradedUnitRange, index::Integer)
+function BlockSparseArrays.blockedunitrange_findblockindex(
+  a::AbstractGradedUnitRange, index::Integer
+)
   return blockedunitrange_findblockindex(unlabel_blocks(a), index)
 end
 
@@ -221,32 +235,36 @@ function firstblockindices(a::AbstractGradedUnitRange)
   return labelled.(firstblockindices(unlabel_blocks(a)), blocklabels(a))
 end
 
-function blockedunitrange_getindices(a::AbstractGradedUnitRange, index::Block{1})
+function BlockSparseArrays.blockedunitrange_getindices(
+  a::AbstractGradedUnitRange, index::Block{1}
+)
   return labelled(unlabel_blocks(a)[index], get_label(a, index))
 end
 
-function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::Vector{<:Integer})
+function BlockSparseArrays.blockedunitrange_getindices(
+  a::AbstractGradedUnitRange, indices::Vector{<:Integer}
+)
   return map(index -> a[index], indices)
 end
 
-function blockedunitrange_getindices(
+function BlockSparseArrays.blockedunitrange_getindices(
   a::AbstractGradedUnitRange,
   indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
 )
   return mortar(map(b -> a[b], blocks(indices)))
 end
 
-function blockedunitrange_getindices(a::AbstractGradedUnitRange, index)
+function BlockSparseArrays.blockedunitrange_getindices(a::AbstractGradedUnitRange, index)
   return labelled(unlabel_blocks(a)[index], get_label(a, index))
 end
 
-function blockedunitrange_getindices(
+function BlockSparseArrays.blockedunitrange_getindices(
   a::AbstractGradedUnitRange, indices::BlockIndexRange{1}
 )
   return a[block(indices)][only(indices.indices)]
 end
 
-function blockedunitrange_getindices(
+function BlockSparseArrays.blockedunitrange_getindices(
   a::AbstractGradedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
 )
   # Without converting `indices` to `Vector`,
@@ -268,22 +286,28 @@ function blocklabels(a::AbstractUnitRange, indices)
   end
 end
 
-function blockedunitrange_getindices(
+function BlockSparseArrays.blockedunitrange_getindices(
   ga::AbstractGradedUnitRange, indices::AbstractUnitRange{<:Integer}
 )
   a_indices = blockedunitrange_getindices(unlabel_blocks(ga), indices)
   return labelled_blocks(a_indices, blocklabels(ga, indices))
 end
 
-function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::BlockSlice)
+function BlockSparseArrays.blockedunitrange_getindices(
+  a::AbstractGradedUnitRange, indices::BlockSlice
+)
   return a[indices.block]
 end
 
-function blockedunitrange_getindices(ga::AbstractGradedUnitRange, indices::BlockRange)
+function BlockSparseArrays.blockedunitrange_getindices(
+  ga::AbstractGradedUnitRange, indices::BlockRange
+)
   return labelled_blocks(unlabel_blocks(ga)[indices], blocklabels(ga, indices))
 end
 
-function blockedunitrange_getindices(a::AbstractGradedUnitRange, indices::BlockIndex{1})
+function BlockSparseArrays.blockedunitrange_getindices(
+  a::AbstractGradedUnitRange, indices::BlockIndex{1}
+)
   return a[block(indices)][blockindex(indices)]
 end
 
@@ -398,7 +422,7 @@ end
 # blocklengths = map(bs -> sum(b -> length(a[b]), bs), blocks(indices))
 # return blockedrange(blocklengths)
 # ```
-function blockedunitrange_getindices(
+function BlockSparseArrays.blockedunitrange_getindices(
   a::AbstractGradedUnitRange, indices::AbstractBlockVector{<:Block{1}}
 )
   blks = map(bs -> mortar(map(b -> a[b], bs)), blocks(indices))