Define TensorAlgebra.matricize

ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl

@@ -21,55 +21,56 @@
 
 struct SectorFusion <: FusionStyle end
 
 TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion()
 
-# TODO consider heterogeneous sectors?
+# TBD consider heterogeneous sectors?
 TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = trivial(first(t))
 
 function matricize_axes(
   blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}
 )
+  @assert !isempty(blocked_axes)
+  default_axis = trivial_axis(Tuple(blocked_axes))
   codomain_axes, domain_axes = blocks(blocked_axes)
-  @assert !(isempty(codomain_axes) && isempty(domain_axes))
-  row_axis = unmerged_tensor_product(trivial_axes(domain_axes), codomain_axes...)
-  unflipped_col_axis = unmerged_tensor_product(trivial_axes(codomain_axes), domain_axes...)
+  row_axis = unmerged_tensor_product(default_axis, codomain_axes...)
+  unflipped_col_axis = unmerged_tensor_product(default_axis, domain_axes...)
   return row_axis, flip(unflipped_col_axis)
 end
 
 function TensorAlgebra.matricize(
   ::SectorFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2}
 )
   a_perm = permutedims(a, Tuple(biperm))
   row_axis, col_axis = matricize_axes(axes(a)[biperm])
   a_reshaped = blockreshape(a_perm, (row_axis, col_axis))
   # Sort the blocks by sector and merge the equivalent sectors.
   return sectormergesort(a_reshaped)
 end
 
 function TensorAlgebra.unmatricize(
   ::SectorFusion,
   m::AbstractMatrix,
   blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}},
 )
   # First, fuse axes to get `blockmergesortperm`.
   # Then unpermute the blocks.
-  row_col_axes = row_and_column_axes(blocked_axes)
+  row_col_axes = matricize_axes(blocked_axes)
 
   blockperms = blocksortperm.(row_col_axes)
   sorted_axes = map((r, I) -> only(axes(r[I])), row_col_axes, 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.
   m_unblocked = m[sorted_axes...]
   m_blockpermed = m_unblocked[invblockperm.(blockperms)...]
   return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes)
 end
 
 # Sort the blocks by sector and then merge the common sectors.
 function sectormergesort(a::AbstractArray)
   I = blockmergesortperm.(axes(a))
   return a[I...]
 end
 end