Prevent construction of multifusion tensors with incompatible coloring

docs/src/Changelog.md

@@ -23,6 +23,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 ### Added
 
 ### Changed
+- For sector types with `GenericUnit` such that colorings are not unique, `GradedSpace`, `ProductSpace` and `HomSpace` now check for this compatibility. In particular, this prevents the construction of `TensorMap`s with incompatible colorings, which previously either errored or produced empty tensors inconsistently. ([#515](https://github.com/QuantumKitHub/TensorKit.jl/pull/515))
 
 ### Deprecated
 

src/spaces/gradedspace.jl

@@ -29,6 +29,15 @@ struct GradedSpace{I <: Sector, D} <: ElementarySpace
 end
 sectortype(::Type{<:GradedSpace{I}}) where {I <: Sector} = I
 
+# elementary spaces are homogeneously colored: all sectors share a left and a right unit.
+function _check_unit_homogeneity(::Type{I}, sectors) where {I <: Sector}
+    (UnitStyle(I) isa SimpleUnit || isempty(sectors)) && return nothing
+    l, r = leftunit(first(sectors)), rightunit(first(sectors))
+    all(c -> leftunit(c) == l && rightunit(c) == r, sectors) ||
+        throw(SpaceMismatch(lazy"sectors $(collect(sectors)) do not share a single left and right unit"))
+    return nothing
+end
+
 function GradedSpace{I, NTuple{N, Int}}(dims; dual::Bool = false) where {I, N}
     d = ntuple(n -> 0, N)
     isset = ntuple(n -> false, N)
@@ -40,6 +49,7 @@ function GradedSpace{I, NTuple{N, Int}}(dims; dual::Bool = false) where {I, N}
         isset = TupleTools.setindex(isset, true, i)
         d = TupleTools.setindex(d, dc, i)
     end
+    _check_unit_homogeneity(I, (values(I)[n] for n in 1:N if !iszero(d[n])))
     return GradedSpace{I, NTuple{N, Int}}(d, dual)
 end
 function GradedSpace{I, NTuple{N, Int}}(dims::Pair; dual::Bool = false) where {I, N}
@@ -54,6 +64,7 @@ function GradedSpace{I, SectorDict{I, Int}}(dims; dual::Bool = false) where {I <
         dc < 0 && throw(ArgumentError(lazy"Sector $k has negative dimension $dc"))
         !iszero(dc) && push!(d, k => dc)
     end
+    _check_unit_homogeneity(I, keys(d))
     return GradedSpace{I, SectorDict{I, Int}}(d, dual)
 end
 function GradedSpace{I, SectorDict{I, Int}}(dims::Pair; dual::Bool = false) where {I <: Sector}
@@ -122,6 +133,8 @@ function flip(V::GradedSpace{I}) where {I <: Sector}
 end
 
 function unitspace(S::Type{<:GradedSpace{I}}) where {I <: Sector}
+    UnitStyle(I) isa GenericUnit &&
+        throw(ArgumentError("Cannot construct unit space for sector types with semisimple unit structure."))
     return S(unit => 1 for unit in allunits(I))
 end
 zerospace(S::Type{<:GradedSpace}) = S()

src/spaces/homspace.jl

@@ -9,6 +9,38 @@ to denote categories and their objects, and keep `HomSpace` distinct.
 struct HomSpace{S <: ElementarySpace, P1 <: CompositeSpace{S}, P2 <: CompositeSpace{S}}
     codomain::P1
     domain::P2
+    function HomSpace{S, P1, P2}(codomain::P1, domain::P2) where {S <: ElementarySpace, P1 <: CompositeSpace{S}, P2 <: CompositeSpace{S}}
+        _check_unit_compatibility(codomain, domain)
+        return new{S, P1, P2}(codomain, domain)
+    end
+end
+function HomSpace(codomain::P1, domain::P2) where {S, P1 <: CompositeSpace{S}, P2 <: CompositeSpace{S}}
+    return HomSpace{S, P1, P2}(codomain, domain)
+end
+
+# check that the legs form a closed cycle of composable spaces:
+# codomain[1] … codomain[N₁], dual(domain[N₂]) … dual(domain[1]).
+function _check_unit_compatibility(
+        codomain::CompositeSpace{S}, domain::CompositeSpace{S}
+    ) where {S <: ElementarySpace}
+    UnitStyle(sectortype(S)) isa GenericUnit || return nothing
+    N₁, N₂ = length(codomain), length(domain)
+
+    if N₁ == 0 && N₂ == 0 # one() ← one(): empty cycle
+        return nothing
+    elseif N₁ == 0 # the domain segment closes onto itself
+        _matchunits(_leftunit(domain[1]), _rightunit(domain[N₂])) ||
+            throw(SpaceMismatch(lazy"domain $domain has incompatible left and right units"))
+    elseif N₂ == 0 # the codomain segment closes onto itself
+        _matchunits(_leftunit(codomain[1]), _rightunit(codomain[N₁])) ||
+            throw(SpaceMismatch(lazy"codomain $codomain has incompatible left and right units"))
+    else
+        _matchunits(_rightunit(codomain[N₁]), _rightunit(domain[N₂])) ||
+            throw(SpaceMismatch(lazy"HomSpace $codomain ← $domain has incompatible right units"))
+        _matchunits(_leftunit(codomain[1]), _leftunit(domain[1])) ||
+            throw(SpaceMismatch(lazy"HomSpace $codomain ← $domain has incompatible left units"))
+    end
+    return nothing
 end
 
 function HomSpace(codomain::S, domain::CompositeSpace{S}) where {S <: ElementarySpace}
@@ -270,6 +302,25 @@ function compose(W::HomSpace{S}, V::HomSpace{S}) where {S}
     return HomSpace(codomain(W), domain(V))
 end
 
+# workaround to permuting after composing intermediate spaces without constructing the latter
+function _contractedspace(
+        A::HomSpace{S}, (oindA, cindA)::Index2Tuple,
+        B::HomSpace{S}, (cindB, oindB)::Index2Tuple,
+        (p₁, p₂)::Index2Tuple{N₁, N₂}
+    ) where {S, N₁, N₂}
+    NA = length(oindA)
+
+    Acind = map(n -> dual(A[n]), cindA)
+    Bcind = map(n -> B[n], cindB)
+    Acind == Bcind || throw(SpaceMismatch(lazy"$(Acind) ≠ $(Bcind)"))
+
+    getopen(n) = n <= NA ? A[oindA[n]] : B[oindB[n - NA]]
+
+    cod = ProductSpace{S, N₁}(map(getopen, p₁))
+    dom = ProductSpace{S, N₂}(map(n -> dual(getopen(n)), p₂))
+    return cod ← dom
+end
+
 function TensorOperations.tensorcontract(
         A::HomSpace, pA::Index2Tuple, conjA::Bool,
         B::HomSpace, pB::Index2Tuple, conjB::Bool,
@@ -290,7 +341,7 @@ function TensorOperations.tensorcontract(
         pB′ = adjointtensorindices(B, pB)
         TensorOperations.tensorcontract(A, pA, false, B′, pB′, false, pAB)
     else
-        return permute(compose(permute(A, pA), permute(B, pB)), pAB)
+        _contractedspace(A, pA, B, pB, pAB)
     end
 end
 

src/spaces/productspace.jl

@@ -7,7 +7,23 @@ Only tensor products between [`ElementarySpace`](@ref) objects of the same type
 """
 struct ProductSpace{S <: ElementarySpace, N} <: CompositeSpace{S}
     spaces::NTuple{N, S}
-    ProductSpace{S, N}(spaces::NTuple{N, S}) where {S <: ElementarySpace, N} = new{S, N}(spaces)
+    function ProductSpace{S, N}(spaces::NTuple{N, S}) where {S <: ElementarySpace, N}
+        _check_unit_compatibility(spaces)
+        return new{S, N}(spaces)
+    end
+end
+
+# check that the factors form an open chain of composable spaces
+function _check_unit_compatibility(spaces::Tuple{Vararg{ElementarySpace}})
+    N = length(spaces)
+    N <= 1 && return nothing # no junctions to check
+    UnitStyle(sectortype(first(spaces))) isa GenericUnit || return nothing
+    @inbounds for i in 2:N
+        Vprev, V = spaces[i - 1], spaces[i]
+        _matchunits(_rightunit(Vprev), _leftunit(V)) ||
+            throw(SpaceMismatch(lazy"$Vprev and $V have incompatible coloring"))
+    end
+    return nothing
 end
 
 function ProductSpace{S, N}(spaces::Vararg{S, N}) where {S <: ElementarySpace, N}
@@ -64,7 +80,7 @@ Base.axes(P::ProductSpace) = map(axes, P)
 Base.axes(P::ProductSpace, n::Int) = axes(P[n])
 
 dual(P::ProductSpace{<:ElementarySpace, 0}) = P
-dual(P::ProductSpace) = ProductSpace(map(dual, reverse(P)))
+dual(P::ProductSpace) = ProductSpace(reverse(map(dual, P)))
 Base.conj(P::ProductSpace{<:ElementarySpace, 0}) = P
 Base.conj(P::ProductSpace) = ProductSpace(map(conj, P))
 

src/spaces/vectorspaces.jl

@@ -134,9 +134,8 @@ Always returns `false` for spaces where `V == conj(V)`, i.e. vector spaces over
 
 Return the corresponding vector space of type `S` that represents the trivial
 one-dimensional space, i.e. the space that is isomorphic to the corresponding field.
-For vector spaces where `I = sectortype(S)` has a semi-simple unit structure
-(`UnitStyle(I) == GenericUnit()`), this returns a multi-dimensional space corresponding to all unit sectors:
-`dim(unitspace(V), s) == 1` for all `s in allunits(I)`. 
+For vector spaces where `I = sectortype(S)` has a non-simple unit structure
+(`UnitStyle(I) == GenericUnit()`), this errors. 
 
 !!! note
     `unitspace(V)`is different from `one(V)`. The latter returns the empty product space
@@ -167,16 +166,10 @@ in the vector space.
 """
 function leftunitspace(V::ElementarySpace)
     I = sectortype(V)
-    if UnitStyle(I) isa SimpleUnit
-        return unitspace(typeof(V))
-    else
-        !isempty(sectors(V)) || throw(ArgumentError("Cannot determine the left unit of an empty space"))
-        _allequal(leftunit, sectors(V)) ||
-            throw(ArgumentError(lazy"sectors of $V do not have the same left unit"))
-
-        sector = leftunit(first(sectors(V)))
-        return spacetype(V)(sector => 1)
-    end
+    UnitStyle(I) isa SimpleUnit && return unitspace(typeof(V))
+    u = _leftunit(V)
+    isnothing(u) && throw(ArgumentError("Cannot determine the left unit of an empty space"))
+    return spacetype(V)(u => 1)
 end
 
 """
@@ -189,25 +182,38 @@ in the vector space.
 """
 function rightunitspace(V::ElementarySpace)
     I = sectortype(V)
-    if UnitStyle(I) isa SimpleUnit
-        return unitspace(typeof(V))
-    else
-        !isempty(sectors(V)) || throw(ArgumentError("Cannot determine the right unit of an empty space"))
-        _allequal(rightunit, sectors(V)) ||
-            throw(ArgumentError(lazy"sectors of $V do not have the same right unit"))
+    UnitStyle(I) isa SimpleUnit && return unitspace(typeof(V))
+    u = _rightunit(V)
+    isnothing(u) && throw(ArgumentError("Cannot determine the right unit of an empty space"))
+    return spacetype(V)(u => 1)
+end
 
-        sector = rightunit(first(sectors(V)))
-        return spacetype(V)(sector => 1)
-    end
+# Return the `(leftunit, rightunit)` of `V`, or `(nothing, nothing)` for the zero space,
+# whose coloring is unconstrained and thus acts as a wildcard. Elementary spaces are
+# homogeneously colored by construction, so any sector determines both.
+function _leftrightunit(V::ElementarySpace)
+    s = sectors(V)
+    isempty(s) && return (nothing, nothing)
+    c = first(s)
+    return (leftunit(c), rightunit(c))
 end
 
+_leftunit(V::ElementarySpace) = _leftrightunit(V)[1]
+_rightunit(V::ElementarySpace) = _leftrightunit(V)[2]
+
+# `nothing` acts as a wildcard, being compatible with any coloring
+_matchunits(::Nothing, ::Nothing) = true
+_matchunits(::Nothing, ::Sector) = true
+_matchunits(::Sector, ::Nothing) = true
+_matchunits(u₁::Sector, u₂::Sector) = u₁ == u₂
+
 """
     isunitspace(V::S) where {S <: ElementarySpace} -> Bool
 
 Return whether the elementary space `V` is a unit space, i.e. is isomorphic to the
 trivial one-dimensional space. For vector spaces of type `GradedSpace{I}` where `Sector` `I` has a
-semi-simple unit structure, this returns `true` if `V` is isomorphic to either the left, right or
-semi-simple unit space.
+semisimple unit structure, this returns `true` if `V` is isomorphic to the left or right unit
+space of its coloring.
 """
 function isunitspace(V::ElementarySpace)
     I = sectortype(V)

src/tensors/tensoroperations.jl

@@ -176,9 +176,9 @@ function TO.tensorcontract_structure(
         B::AbstractTensorMap, pB::Index2Tuple, conjB::Bool,
         pAB::Index2Tuple{N₁, N₂}
     ) where {N₁, N₂}
-    sA = TO.tensoradd_structure(A, pA, conjA)
-    sB = TO.tensoradd_structure(B, pB, conjB)
-    return permute(compose(sA, sB), pAB)
+    VA, pA′ = conjA ? (space(A)', adjointtensorindices(A, pA)) : (space(A), pA)
+    VB, pB′ = conjB ? (space(B)', adjointtensorindices(B, pB)) : (space(B), pB)
+    return _contractedspace(VA, pA′, VB, pB′, pAB)
 end
 
 function TO.checkcontractible(

test/chainrules/linalg.jl

@@ -155,7 +155,7 @@ for V in spacelist
                 test_rrule(inv, E; atol, rtol)
             end
 
-            A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5])
+            A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
             test_rrule(LinearAlgebra.adjoint, A; atol, rtol)
             test_rrule(LinearAlgebra.norm, A, 2; atol, rtol)
 

test/enzyme-vi-to/add.jl

@@ -20,7 +20,7 @@ fTβs = is_ci ? (Duplicated,) : (Duplicated, Const)
     α = randn(T)
     β = randn(T)
 
-    CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]
+    CV = V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])'
     C = randn(T, CV)
     A = randn(T, CV)
     for TC in (Duplicated,), TA in (Duplicated,)

test/enzyme-vi-to/inner.jl

@@ -17,7 +17,7 @@ fTs = is_ci ? (Duplicated,) : (Duplicated, Const)
         @testset for TC in (Duplicated,), TA in (Duplicated,), f in (identity, adjoint)
             atol = default_tol(T)
             rtol = default_tol(T)
-            CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]
+            CV = V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])'
             C = randn(T, CV)
             A = randn(T, CV)
             for RT in rTs

test/enzyme-vi-to/scale.jl

@@ -17,7 +17,7 @@ fTαs = is_ci ? (Duplicated,) : (Duplicated, Const)
         atol = default_tol(T)
         rtol = default_tol(T)
         α = randn(T)
-        CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]
+        CV = V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])'
         C = randn(T, CV)
         A = randn(T, CV)
         @testset for TC in (Duplicated,)

test/setup.jl

@@ -7,6 +7,7 @@ export random_fusion
 export sectorlist, fast_sectorlist
 # export dim_isapprox
 export default_spacelist, factorization_spacelist, ad_spacelist
+export VIBM, VIBMRepA4
 export test_ad_rrule
 export _isunitary, _isone