Upgrade to TensorAlgebra v0.6 (#83)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.5.4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
+version = "0.5.5"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -16,7 +16,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
-TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
@@ -37,8 +36,7 @@ MatrixAlgebraKit = "0.6"
 Random = "1.10"
 SUNRepresentations = "0.3"
 SplitApplyCombine = "1.2.3"
-TensorAlgebra = "0.5"
+TensorAlgebra = "0.6.2"
 TensorKitSectors = "0.1, 0.2"
-TensorProducts = "0.1.3"
 TypeParameterAccessors = "0.4"
 julia = "1.10"

src/fusion.jl

@@ -1,22 +1,20 @@
 using BlockArrays: Block, blocks
 using SplitApplyCombine: groupcount
-using TensorProducts: TensorProducts, ⊗, OneToOne, tensor_product
 
 flip_dual(r::AbstractUnitRange) = isdual(r) ? flip(r) : r
 
-# TensorProducts interface
-function TensorProducts.tensor_product(sr1::SectorUnitRange, sr2::SectorUnitRange)
+function tensor_product(sr1::SectorUnitRange, sr2::SectorUnitRange)
     return tensor_product(combine_styles(SymmetryStyle(sr1), SymmetryStyle(sr2)), sr1, sr2)
 end
 
-function TensorProducts.tensor_product(
+function tensor_product(
         ::AbelianStyle, sr1::SectorUnitRange, sr2::SectorUnitRange
     )
     s = sector(flip_dual(sr1)) ⊗ sector(flip_dual(sr2))
     return sectorrange(s, sector_multiplicity(sr1) * sector_multiplicity(sr2))
 end
 
-function TensorProducts.tensor_product(
+function tensor_product(
         ::NotAbelianStyle, sr1::SectorUnitRange, sr2::SectorUnitRange
     )
     g0 = sector(flip_dual(sr1)) ⊗ sector(flip_dual(sr2))
@@ -27,41 +25,52 @@ function TensorProducts.tensor_product(
 end
 
 # allow to fuse a Sector with a GradedUnitRange
-function TensorProducts.tensor_product(
+function tensor_product(
         s::Union{SectorRange, SectorUnitRange}, g::AbstractGradedUnitRange
     )
     return to_gradedrange(s) ⊗ g
 end
 
-function TensorProducts.tensor_product(
+function tensor_product(
         g::AbstractGradedUnitRange, s::Union{SectorRange, SectorUnitRange}
     )
     return g ⊗ to_gradedrange(s)
 end
 
-function TensorProducts.tensor_product(sr::SectorUnitRange, s::SectorRange)
+function tensor_product(sr::SectorUnitRange, s::SectorRange)
     return sr ⊗ sectorrange(s, 1)
 end
 
-function TensorProducts.tensor_product(s::SectorRange, sr::SectorUnitRange)
+function tensor_product(s::SectorRange, sr::SectorUnitRange)
     return sectorrange(s, 1) ⊗ sr
 end
 
+function tensor_product(r1::AbstractUnitRange, r2::AbstractUnitRange)
+    (isone(first(r1)) && isone(first(r2))) ||
+        throw(ArgumentError("Only one-based axes are supported"))
+    return Base.OneTo(length(r1) * length(r2))
+end
+
+function tensor_product(
+        r1::AbstractUnitRange, r2::AbstractUnitRange, r3::AbstractUnitRange,
+        rs::AbstractUnitRange...,
+    )
+    return tensor_product(tensor_product(r1, r2), r3, rs...)
+end
+
 # unmerged_tensor_product is a private function needed in GradedArraysTensorAlgebraExt
 # to get block permutation
 # it is not aimed for generic use and does not support all tensor_product methods (no dispatch on SymmetryStyle)
-unmerged_tensor_product() = OneToOne()
+unmerged_tensor_product() = Base.OneTo(1)
 unmerged_tensor_product(a) = a
-unmerged_tensor_product(a, ::OneToOne) = a
-unmerged_tensor_product(::OneToOne, a) = a
-unmerged_tensor_product(::OneToOne, ::OneToOne) = OneToOne()
-function unmerged_tensor_product(a1, a2, as...)
-    return unmerged_tensor_product(unmerged_tensor_product(a1, a2), as...)
+function unmerged_tensor_product(a1, a2, a3, as...)
+    return unmerged_tensor_product(unmerged_tensor_product(a1, a2), a3, as...)
 end
 
 # default to tensor_product
 unmerged_tensor_product(a1, a2) = a1 ⊗ a2
 
+using BlockSparseArrays: mortar_axis
 function unmerged_tensor_product(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
     new_axes = map(splat(⊗), Iterators.flatten((Iterators.product(blocks(a1), blocks(a2)),)))
     return mortar_axis(new_axes)
@@ -103,9 +112,9 @@ end
 sectormergesort(g::AbstractUnitRange) = g
 
 # tensor_product produces a sorted, non-dual GradedUnitRange
-TensorProducts.tensor_product(g::AbstractGradedUnitRange) = sectormergesort(flip_dual(g))
+tensor_product(g::AbstractGradedUnitRange) = sectormergesort(flip_dual(g))
 
-function TensorProducts.tensor_product(
+function tensor_product(
         g1::AbstractGradedUnitRange, g2::AbstractGradedUnitRange
     )
     return sectormergesort(unmerged_tensor_product(g1, g2))

src/sectorrange.jl

@@ -1,6 +1,5 @@
 # This file defines the interface for type Sector
 # all fusion categories (Z{2}, SU2, Ising...) are subtypes of Sector
-using TensorProducts: TensorProducts, ⊗
 import TensorKitSectors as TKS
 
 """
@@ -122,17 +121,19 @@ function fusion_rule(r1::SectorRange, r2::SectorRange)
     )
 end
 
-# =============================  TensorProducts interface  =====--==========================
+# =============================  Tensor products  ==========================================
 
-TensorProducts.tensor_product(s::SectorRange) = s
-TensorProducts.tensor_product(c1::SectorRange, c2::SectorRange) = fusion_rule(c1, c2)
-function TensorProducts.tensor_product(c1::TKS.Sector, c2::TKS.Sector)
+function tensor_product end
+const ⊗ = tensor_product
+tensor_product(s::SectorRange) = s
+tensor_product(c1::SectorRange, c2::SectorRange) = fusion_rule(c1, c2)
+function tensor_product(c1::TKS.Sector, c2::TKS.Sector)
     return tensor_product(to_sector(c1), to_sector(c2))
 end
-function TensorProducts.tensor_product(c1::SectorRange, c2::TKS.Sector)
+function tensor_product(c1::SectorRange, c2::TKS.Sector)
     return tensor_product(c1, to_sector(c2))
 end
-function TensorProducts.tensor_product(c1::TKS.Sector, c2::SectorRange)
+function tensor_product(c1::TKS.Sector, c2::SectorRange)
     return tensor_product(to_sector(c1), c2)
 end
 

src/tensoralgebra.jl

@@ -1,58 +1,89 @@
-using BlockArrays: blocks
+using BlockArrays: blocks, eachblockaxes1
 using BlockSparseArrays: BlockSparseArray, blockreshape
-using GradedArrays:
-    AbstractGradedUnitRange,
-    SectorRange,
-    GradedArray,
-    flip,
-    gradedrange,
-    invblockperm,
-    sectormergesortperm,
-    sectorsortperm,
-    trivial,
-    unmerged_tensor_product
-using TensorAlgebra:
-    TensorAlgebra,
-    ⊗,
-    AbstractBlockPermutation,
-    BlockedTuple,
-    FusionStyle,
-    trivial_axis,
-    unmatricize
+using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockReshapeFusion,
+    BlockedTuple, FusionStyle, ReshapeFusion, matricize, matricize_axes,
+    tensor_product_axis, unmatricize
 
 struct SectorFusion <: FusionStyle end
 
 TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion()
 
-function TensorAlgebra.trivial_axis(t::Tuple{Vararg{G}}) where {G <: AbstractGradedUnitRange}
+function TensorAlgebra.trivial_axis(
+        ::BlockReshapeFusion,
+        ::Val{:codomain},
+        a::GradedArray,
+        axes_codomain::Tuple{Vararg{AbstractUnitRange}},
+        axes_domain::Tuple{Vararg{AbstractUnitRange}},
+    )
+    return trivial_gradedrange(axes(a))
+end
+function TensorAlgebra.trivial_axis(
+        ::BlockReshapeFusion,
+        ::Val{:domain},
+        a::GradedArray,
+        axes_codomain::Tuple{Vararg{AbstractUnitRange}},
+        axes_domain::Tuple{Vararg{AbstractUnitRange}},
+    )
+    return flip(trivial_gradedrange(axes(a)))
+end
+function trivial_gradedrange(t::Tuple{Vararg{G}}) where {G <: AbstractGradedUnitRange}
     return trivial(first(t))
 end
 # heterogeneous sectors
-TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = ⊗(trivial.(t)...)
+trivial_gradedrange(t::Tuple{Vararg{AbstractGradedUnitRange}}) = ⊗(trivial.(t)...)
 # trivial_axis from sector_type
-function TensorAlgebra.trivial_axis(::Type{S}) where {S <: SectorRange}
+function trivial_gradedrange(::Type{S}) where {S <: SectorRange}
     return gradedrange([trivial(S) => 1])
 end
 
-function matricize_axes(
-        blocked_axes::BlockedTuple{2, <:Any, <:Tuple{Vararg{AbstractUnitRange}}}
+function TensorAlgebra.tensor_product_axis(
+        ::ReshapeFusion, ::Val{:codomain}, r1::SectorUnitRange, r2::SectorUnitRange
     )
-    @assert !isempty(blocked_axes)
-    default_axis = trivial_axis(Tuple(blocked_axes))
-    codomain_axes, domain_axes = blocks(blocked_axes)
-    codomain_axis = unmerged_tensor_product(default_axis, codomain_axes...)
-    unflipped_domain_axis = unmerged_tensor_product(default_axis, domain_axes...)
-    return codomain_axis, flip(unflipped_domain_axis)
+    return r1 ⊗ r2
+end
+function TensorAlgebra.tensor_product_axis(
+        ::ReshapeFusion, ::Val{:domain}, r1::SectorUnitRange, r2::SectorUnitRange
+    )
+    return flip(r1 ⊗ r2)
+end
+function TensorAlgebra.tensor_product_axis(
+        style::BlockReshapeFusion,
+        side::Val{:codomain},
+        r1::AbstractGradedUnitRange,
+        r2::AbstractGradedUnitRange,
+    )
+    return tensor_product_gradedrange(style, side, r1, r2)
+end
+function TensorAlgebra.tensor_product_axis(
+        style::BlockReshapeFusion,
+        side::Val{:domain},
+        r1::AbstractGradedUnitRange,
+        r2::AbstractGradedUnitRange,
+    )
+    return tensor_product_gradedrange(style, side, r1, r2)
+end
+# TODO: Could this call out to a generic tensor_product_axis for AbstractBlockedUnitRange?
+function tensor_product_gradedrange(
+        ::BlockReshapeFusion,
+        side::Val,
+        r1::AbstractUnitRange,
+        r2::AbstractUnitRange,
+    )
+    (isone(first(r1)) && isone(first(r2))) ||
+        throw(ArgumentError("Only one-based axes are supported"))
+    blockaxpairs = Iterators.product(eachblockaxes1(r1), eachblockaxes1(r2))
+    blockaxs = map(blockaxpairs) do (b1, b2)
+        # TODO: Store a FusionStyle for the blocks in `BlockReshapeFusion`
+        # and use that here.
+        return tensor_product_axis(side, b1, b2)
+    end
+    return mortar_axis(vec(blockaxs))
 end
 
-using TensorAlgebra: blockedtrivialperm
 function TensorAlgebra.matricize(
-        ::SectorFusion, a::AbstractArray, codomain_length::Val, domain_length::Val
+        ::SectorFusion, a::AbstractArray, length_codomain::Val
     )
-    biperm = blockedtrivialperm((codomain_length, domain_length))
-    codomain_axis, domain_axis = matricize_axes(axes(a)[biperm])
-    a_reshaped = blockreshape(a, (codomain_axis, domain_axis))
-    # Sort the blocks by sector and merge the equivalent sectors.
+    a_reshaped = matricize(BlockReshapeFusion(), a, length_codomain)
     return sectormergesort(a_reshaped)
 end
 
@@ -74,7 +105,7 @@ function TensorAlgebra.unmatricize(
 
     # First, fuse axes to get `sectormergesortperm`.
     # Then unpermute the blocks.
-    fused_axes = matricize_axes(blocked_axes)
+    fused_axes = matricize_axes(BlockReshapeFusion(), m, codomain_axes, domain_axes)
 
     blockperms = sectorsortperm.(fused_axes)
     sorted_axes = map((r, I) -> only(axes(r[I])), fused_axes, blockperms)

test/Project.toml

@@ -12,7 +12,6 @@ SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
-TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 
@@ -31,8 +30,7 @@ SUNRepresentations = "0.3"
 SafeTestsets = "0.1"
 SparseArraysBase = "0.7"
 Suppressor = "0.2.8"
-TensorAlgebra = "0.5"
+TensorAlgebra = "0.6"
 TensorKitSectors = "0.1, 0.2"
-TensorProducts = "0.1.3"
 Test = "1.10"
 TestExtras = "0.3.1"

test/test_fusion_rule.jl

@@ -6,15 +6,16 @@ using GradedArrays:
     TrivialSector,
     U1,
     Z,
+    ⊗,
     dual,
     flip,
     gradedrange,
     nsymbol,
     quantum_dimension,
     space_isequal,
+    tensor_product,
     trivial,
     unmerged_tensor_product
-using TensorProducts: ⊗, tensor_product
 using SUNRepresentations: SUNIrrep
 using Test: @test, @test_throws, @testset
 using TestExtras: @constinferred

test/test_interface.jl

@@ -3,19 +3,8 @@ using BlockArrays: BlockedOneTo, blockedrange, blockisequal
 using GradedArrays:
     NoSector, dag, dual, flip, isdual, map_sectors, sectors, space_isequal, ungrade
 using Test: @test, @testset
-using TensorProducts: OneToOne
 
 @testset "GradedUnitRange interface for AbstractUnitRange" begin
-    a0 = OneToOne()
-    @test !isdual(a0)
-    @test dual(a0) isa OneToOne
-    @test space_isequal(a0, a0)
-    @test space_isequal(a0, dual(a0))
-    @test only(sectors(a0)) == NoSector()
-    @test ungrade(a0) === a0
-    @test map_sectors(identity, a0) === a0
-    @test dag(a0) === a0
-
     a = 1:3
     ad = dual(a)
     af = flip(a)

test/test_sectorproduct.jl

@@ -5,6 +5,7 @@ using GradedArrays:
     TrivialSector,
     U1,
     Z,
+    ⊗,
     ×,
     arguments,
     dual,
@@ -17,7 +18,6 @@ using GradedArrays:
     sectorrange,
     space_isequal,
     trivial
-using TensorProducts: ⊗
 using Test: @test, @test_broken, @test_throws, @testset
 using TestExtras: @constinferred
 using BlockArrays: blocklengths