[GradedAxes] Introduce GradedUnitRangeDual

[ogauthe]

This PR aims to simplify the use of GradedAxes.dual, especially to be used in a BlockSparseArray. It proposes to remove the UnitRangeDual type:

• a generic AbstractUnitRange is now self dual
• a labelled GradedUnitRange or GradedOneTo are cast to GradedUnitRangeDual{<:LabelledInteger,<:AbstractGradedUnitRange} <: AbstractGradedUnitRange

This should simplify slicing a BlockSparseArray with dual axes and allows to preserve labels or take their dual.

This branch is similar to #1529 but is based on main.

Currently, this branche solves:

• preserving labels in blocklengths(dual(...))

g = gradedrange([U1(1) => 2, U1(2) => 2])
@test GradedAxes.blocklengths(dual(g)) isa Vector{<:LabelledInteger}  #    # broken in main and in non-abelian_fusion

• preserving GradedUnitRange behavior when slicing:

blocklengths(dual(gradedrange([U1(0) => 2, U1(1) => 2])[1:3])) isa Vector  # broken in main and in non-abelian_fusion

• manipulating arrays with DualGradedUnitRange

r = gradedrange([U1(0) => 2, U1(1) => 2])[1:3]
a = BlockSparseArray{Float64}(r, r)
b = BlockSparseArray{Float64}(dual(r),dual(r))
@test a[:, :] isa BlockSparseArray   # broken in main and in non-abelian_fusion
@test b[:, :] isa BlockSparseArray   # broken in main and in non-abelian_fusion

Meaning one can create and manipulate BlockSparseArray with GradedUnitRange or its dual (not just BlockedOneTo). However axes(a[:,:]) are still Base.IdentityUnitRange(GradedUnitRangeDual{...}), which creates other issues.

It also adds many tests, including broken ones.

• The current plan is to have multiple dual unit range types, such as UnitRangeDual <: UnitRange, GradedUnitRangeDual <: AbstractGradedUnitRange, and maybe BlockedUnitRangeDual <: AbstractBlockedUnitRange. In that case, it would be useful to define an IsDual trait based on the isdual function (say using SimpleTraits.jl). See https://github.com/ITensor/ITensors.jl/pull/1363/files#r1779660676 and https://github.com/ITensor/ITensors.jl/pull/1363/files#r1779708817 for examples of where that would be useful, say for generalizing blocklabels(a::UnitRangeDual) and flip_dual(r::UnitRangeDual).

[ogauthe]

Overall I still find there are too many kind of UnitRange. I wonder if we could avoid the doublet GradedOneTo and GradedUnitRange. gradedrange creates a GradedOneTo, but when slicing it one gets a GradedUnitRange.

For instance, blockfirsts(gradedrange(["x" => 2, "y" => 3])[1:5]) currently fails due to ambiguity.

[mtfishman]

For dispatch it is helpful to have special ranges that start at one since that is used as axes by most arrays, like how Base has UnitRange and Base.OneTo. Handling both shouldn't be too difficult if we write good generic code.

[mtfishman]

To be more specific, if we write most code in terms of AbstractGradedUnitRange and AbstractUnitRange, that should allow us to have a lot of subtypes (such as GradedOneTo, GradedUnitRage, GradedUnitRangeDual, etc.) without adding too much code complexity. Though perhaps one question is, is there a need for a GradedOneToDual? That does start to sound excessive and it is best to try to design it without that for now.

One thing to consider is that dual only really makes sense for graded unit ranges, so maybe we could consider dropping UnitRangeDual, and just define dual(r::AbstractUnitRange) = r as a generic fallback (i.e. if the unit range isn't graded, the dual is itself).

[ogauthe]

I have no strong opinion whether there should be a UnitRangeDual type or not. This is mostly a design question for BlockSparseArray, about how much complexity you want there. In the specific case of FusionTensor, I will always use a graded axis. My objective here is to make dual(GradedUnitRange) an AbstractGradedUnitRange to deal with labels in a generic way.

However I believe dual and not dual axes should be handled on equal footing: deciding who is dual and who is not is a pure convention and they should always appear with a 1:1 ratio for contraction to be possible. If there is some gain in having both GradedOneTo and GradedUnitRangeDual, then I think there should be a GradedOneToDual.

[ogauthe]

I realize that with the current design, there are already GradedUnitRangeDual{<:LabelledInteger, BlockArrays.BlockedOneTo and GradedUnitRangeDual{<:LabelledInteger, BlockArrays.BlockedUnitRange}, so there should be no need for an explicit GradedOneToDual.

I will remove UnitRangeDual and use self-dual axes as the default.

[mtfishman]

I realize that with the current design, there are already GradedUnitRangeDual{<:LabelledInteger, BlockArrays.BlockedOneTo and GradedUnitRangeDual{<:LabelledInteger, BlockArrays.BlockedUnitRange}, so there should be no need for an explicit GradedOneToDual.

I will remove UnitRangeDual and use self-dual axes as the default.

Sounds good, that's what I was going to suggest (i.e. GradedOneToDual can be a type alias, if needed).

[ogauthe]

I think this is now ready for GradedAxes. I will focus on BlockSparseArrays in another PR.

[mtfishman]

@ogauthe I realized that one reason to keep UnitRangeDual is that it may be a reasonable return type when slicing blocks of a GradedUnitRangeDual, i.e. it might make sense for dual(gradedrange([U1(0) => 2, U1(1) => 2]))[Block(1)] to return a UnitRangeDual, to preserve the fact that the block was taken from a dual space.

That could be helpful since then it could make it easier to create a dual graded unit range out of the blocks of another dual graded unit range (say if you have a dual graded unit range and want to remove one of the blocks, it could make it easier to write that kind of code).

[ogauthe]

Indeed slicing a GradedUnitRangeDual with a Block was an issue I spotted (this is associated with the slicing issues in #1336). Currently it yields a LabelledUnitRange and the dual information is lost.

Now, how generic UnitRangeDual should be? I guess it should be a subtype of AbstractUnitRange and should be valid axis for a BlockSparseArray. Should we come back to "any AbstractUnitRange can be dualed and become a UnitRangeDual" or should it be more specific, as "obtained from slicing a GradedUnitRangeDual"?

I favor the second as it would make it simpler to use a BlockSparseArray without caring about mathematical details, but I think both options carry the same amount of code complexity (mostly coming from being valid BlockSparseArray axes).

[mtfishman]

I agree with your assessment that UnitRangeDual should be a subtype of AbstractUnitRange, and also that we should be thoughtful about which operations create UnitRangeDual. I'll try to summarize that view:

I think it should be valid to make a BlockSparseArray with any kind of unit range as long as it satisfies the BlockArrays.jl axis interface, and maybe some additional interface functions we require in BlockSparseArrays, so axes may be dual or non-dual of whatever kind.

The question to me more comes down to how and when dual axes get constructed, depending on the unit range type and operation. So to be concrete, if you have a block sparse matrix a, in general a' should have axes: axes(a') == dual.(reverse(axes(a))), so that operations like a' * a make sense (i.e. involve a contraction of a dual axis with a non-dual axis).

However, I think it would be surprising and unnecessarily complicated if a started with non-graded axes, say OneTo or BlockedOneTo, and a' had dual axes, say UnitRangeDual{<:Any,<:OneTo} or UnitRangeDual{<:Any,<:BlockedOneTo}, since they are self-dual. So that leads me to think that we should define dual(r::OneTo) = r and dual(r::BlockedOneTo) = r (since adjoint(::BlockSparseMatrix) will be defined in terms of dual to take the dual of the axes where appropriate), and if a user wants to create a BlockSparseArray with dual axes of type UnitRangeDual{<:Any,<:BlockedOneTo}, they'll have to do that manually in another way besides the dual(...) function.

EDIT: One idea for an interface for constructing the dual type of a unit range, even if it is defined to be self-dual with dual(r::AbstractUnitRange) = r, would be a generic function Dual(r::AbstractUnitRange) = UnitRangeDual(r). Of course, someone can use UnitRangeDual(r), but that isn't generic, i.e. that should always output an object of type UnitRangeDual, while Dual(r) could output UnitRangeDual, BlockedUnitRangeDual, GradedUnitRangeDual, etc. depending on the type of r. So schematically, Dual could be defined as:

function Dual(r::AbstractUnitRange)
  r_dual = dual(r)
  !isdual(r_dual) && return UnitRangeDual(r)
  return r_dual
end

though maybe would choose another implementation in practice. An interesting side comment is that this could be a new design for array wrappers in Julia, for example instead of a single Transpose wrapper there could be ArrayTranspose, BlockSparseArrayTranpose, SparseArrayTranspose, etc. which could all be constructed from the same generic function transpose but provide more specific type information about what is being wrapped. Though that would be a big change to Base and would require introducing an IsTransposed trait in order to implement generic code on all of those transpose types.

We may want to define that dual(r::LabelledUnitRange) outputs a UnitRangeDual(r), though in that case it may be better to rename LabelledUnitRange to something else to indicate that it is more specifically the block of a GradedUnitRange.

[mtfishman]

This is copied from this comment: #1363 (comment):

It would be nice to think through how functions like:

flip(g::AbstractGradedUnitRange) = dual(gradedrange(label_dual.(blocklengths(g))))

could be written in a simpler way.

One proposal would be to make this work as an implementation:

flip(g::AbstractGradedUnitRange) = gradedrange(flip.(blocklengths(g)))

In order to make that work, we would need to have flip work on sectors, define the concept of dual sectors, and make it so that gradedrange outputs a dual unit range when constructed from dual sectors.

An implementation could be the following:

flip(l::LabelledInteger) = labelled(unlabel(l), flip(label(l)))

struct CategoryDual{NondualCategory<:AbstractCategory} <: AbstractCategory
  nondual_category::NondualCategory
end
flip(c::AbstractCategory) = CategoryDual(dual(c))
nondual(c::AbstractCategory) = c
nondual(c::CategoryDual) = c.nondual_category
dual(c::CategoryDual) = nondual(c)

In general, dual(flip(x)) can be the canonical/generic way of instantiating a dual sector/space/range, and could supersede label_dual (which I never liked but didn't have a better idea for at the time).

Additionally, gradedrange would have to be able to detect that all of the labels are CategoryDual and in that case construct a UnitRangeDual/GradedUnitRangeDual.

EDIT: I don't think this is needed right now and may not be a good idea, just something to keep in mind for future work.

[ogauthe]

The other problems where due to the change blockedunitrange_getindices/gradedunitrange_getindices. Maybe I should avoid defining gradedunitrange_getindices and just specialize blockedunitrange_getindices for AbstractGradedUnitRange.

[mtfishman]

The other problems where due to the change blockedunitrange_getindices/gradedunitrange_getindices. Maybe I should avoid defining gradedunitrange_getindices and just specialize blockedunitrange_getindices for AbstractGradedUnitRange.

That seems simpler, for context the primary purpose of blockedunitrange_getindices was to circumvent issues in BlockArrays.jl.

[ogauthe]

I fixed all the tests for GradedOneTo and GradedUnitRange. I only define blockedunitrange_getindex. I am not very satisfied with the amount of code duplication for blockedunitrange_getindex but most versions are needed to fix ambiguities.

Now all errors boil down to some version of

g = dual(gradedrange([U1(-1) => 1, U1(1) => 2, U1(2) => 3]))
@test_broken isdual(g[Block(1)])

i.e. slicing a GradedUnitRangeDual with a block returns a non-dual object. I hope to fix them all with a UnitRangeDual, or maybe LabelledUnitRangeDual.

[mtfishman]

@ogauthe this PR looks good to me, given the current type hierarchy that we are aiming for right now. I definitely get your point about there being a lot of code that needs to be defined to make all of the different unit range types work. I think there are a few potential solutions, some of which are more benign like reassessing what code can be made more generic and can get shared across unit range types or even simplified by making upstream changes to BlockArrays.jl, and some more extreme like making an isdual flag or moving the information about duality into the sectors. As I said let's stick with this design for now and see how far it takes us and reassess once more of the FusionTensors and BlockSparseArrays code is fleshed out.

From what I can tell, at least with the current code design and organization a lot of the complexity for dealing with the different unit range types is living in GradedAxes while BlockSparseArrays and FusionTensors can remain relatively agnostic about how the graded unit range types are designed, so that should make it easier to try out different designs within GradedAxes in the future without affecting those libraries too much.

Regarding slicing GradedUnitRangeDual, do you plan to fix that in this PR or in a follow-up? I'm fine either way, though this PR has gotten pretty long already so maybe it would be nicer to merge this and work on that in a future PR, for the sake of making reviewing easier.

[ogauthe]

Thanks for the review. I already started working on a slicing GradedUnitRangeDual, but this is non trivial.

I am pretty happy that many issues appearing in #1336 actually come from GradedAxes and could be fixed here. I added many tests such that test failure would appear clearly in GradedAxes test suit instead of buried deep below a BlockSparseArrays error. So I agree it would be best to merge this PR as it is now.

[mtfishman]

Oops, it looks like #1562 and #1563 caused test failures, I'll revert those. They passed separately but together they allowed an upgrade to Functors.jl v0.5 which breaks something in NDTensors.jl.

[ogauthe]

All tests are passing locally. The failures come from unchanged UnallocatedArrays, I suppose they are not related to this PR.

[mtfishman]

I am pretty happy that many issues appearing in #1336 actually come from GradedAxes and could be fixed here. I added many tests such that test failure would appear clearly in GradedAxes test suit instead of buried deep below a BlockSparseArrays error. So I agree it would be best to merge this PR as it is now.

That's a good sign about the design, my goal with designing BlockSparseArrays was to make it as generic as possible about the axes (and also is the design goal of BlockArrays.jl) so I'm glad to see that seems to be working out so far. Those tests will be helpful, thanks!

[mtfishman]

@ogauthe looks like the test failures are fixed now, is this ready to merge once the rest of the tests pass?

[ogauthe]

@ogauthe looks like the test failures are fixed now, is this ready to merge once the rest of the tests pass?

yes it is

[mtfishman]

Thanks for the PR @ogauthe. Can you check off which issues listed in #1336 that this fixes?