Wrap a TensorKit TensorMap directly as an ITensor

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.10.4"
+version = "0.10.5"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -28,11 +28,13 @@ WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [extensions]
 ITensorBaseAdaptExt = "Adapt"
 ITensorBaseMooncakeExt = "Mooncake"
+ITensorBaseTensorKitExt = "TensorKit"
 ITensorBaseTensorOperationsExt = "TensorOperations"
 
 [compat]
@@ -50,10 +52,11 @@ OrderedCollections = "1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"
 TensorAlgebra = "0.16"
+TensorKit = "0.17"
 TensorOperations = "5.3.1"
 TermInterface = "2"
 TupleTools = "1.6"
 UUIDs = "1.10"
-VectorInterface = "0.5, 0.6"
+VectorInterface = "0.5"
 WrappedUnions = "0.3"
 julia = "1.10"

ext/ITensorBaseTensorKitExt.jl

@@ -0,0 +1,28 @@
+module ITensorBaseTensorKitExt
+
+using ITensorBase: ITensorBase, NamedUnitRange
+using TensorKit: ElementarySpace, dim
+
+# ================================  Index over a native space  ==============================
+# A native TensorKit space is stored directly as the axis value of a `NamedUnitRange`, so
+# `Index(V)` (and `dual`/`conj` of one) round-trips through the named layer with the space
+# intact. These terminal constructors short-circuit `to_range`; the element type is `Int`
+# (the flat index positions), matching how the space presents to TensorAlgebra as an axis.
+function ITensorBase.NamedUnitRange{Name}(unnamed::ElementarySpace, name) where {Name}
+    return ITensorBase.NamedUnitRange{Name, Int, typeof(unnamed)}(unnamed, name)
+end
+function ITensorBase.NamedUnitRange(unnamed::ElementarySpace, name)
+    return ITensorBase.NamedUnitRange{typeof(name), Int, typeof(unnamed)}(unnamed, name)
+end
+
+# `conj(index)` lowers to `named(conj(space), name)`, so `named` on a space must rebuild a
+# `NamedUnitRange` rather than fall through to the generic (array-shaped) `Named`.
+ITensorBase.named(r::ElementarySpace, name) = ITensorBase.NamedUnitRange(r, name)
+
+# The flat length of a space-backed axis is its total (dense) dimension, so `size`/`length`
+# of a `TensorMap`-backed ITensor report dense dimensions.
+function Base.length(r::NamedUnitRange{<:Any, <:Any, <:ElementarySpace})
+    return dim(ITensorBase.unnamed(r))
+end
+
+end

src/abstractnamedtensor.jl

@@ -157,9 +157,11 @@ end
 # Generic construction of named dims arrays.
 
 """
-    nameddims(a::AbstractArray, inds)
+    nameddims(a, dimnames)
 
-Construct a named dimensions array from an unnamed array `a` and named dimensions `inds`.
+Construct a named dimensions array from an unnamed parent `a` and named dimensions
+`dimnames`. The parent is usually an `AbstractArray`, but any object that a `NamedTensor`
+can wrap works (e.g. a TensorKit `TensorMap`).
 
 # Examples
 
@@ -173,17 +175,19 @@ named(Base.OneTo(2), :i)×named(Base.OneTo(3), :j) NamedTensor{Symbol}:
 
 See also [`NamedTensor`](@ref), [`named`](@ref).
 """
-function nameddims(a::AbstractArray, inds)
-    return NamedTensor(a, inds)
+function nameddims(a, dimnames)
+    return NamedTensor(a, dimnames)
 end
 
 #=
-    nameddimsof(a::AbstractNamedTensor, b::AbstractArray)
+    nameddimsof(a::AbstractNamedTensor, b)
 
 Construct a named dimensions array with the dimension names of `a`
-and with the data from `b`.
+and with the data from `b`. The parent `b` is usually an `AbstractArray` but may be any
+object a `NamedTensor` can wrap (e.g. a TensorKit `TensorMap`), so `copy`/`zero` of a
+named tensor round-trip through whatever backend `unnamed(a)` uses.
 =#
-function nameddimsof(a::AbstractNamedTensor, b::AbstractArray)
+function nameddimsof(a::AbstractNamedTensor, b)
     return nameddims(b, dimnames(a))
 end
 
@@ -283,8 +287,10 @@ function Base.AbstractArray{T, N}(a::AbstractNamedTensor) where {T, N}
     return dest
 end
 
+# Read the parent's axes through TensorAlgebra's interface (not `Base.axes`) so a non-array
+# backend like a TensorMap, whose axes are its native spaces, is supported.
 function Base.axes(a::AbstractNamedTensor)
-    return named.(axes(unnamed(a)), Tuple(dimnames(a)))
+    return named.(TensorAlgebra.axes(unnamed(a)), Tuple(dimnames(a)))
 end
 function Base.size(a::AbstractNamedTensor)
     return length.(axes(a))
@@ -300,8 +306,9 @@ Base.axes(a::AbstractNamedTensor, d) = axes(a)[d]
 # Circumvent issue when ndims isn't known at compile time.
 Base.size(a::AbstractNamedTensor, d) = size(a)[d]
 
-# Circumvent issue when ndims isn't known at compile time.
-Base.ndims(a::AbstractNamedTensor) = ndims(unnamed(a))
+# Circumvent issue when ndims isn't known at compile time. Read through TensorAlgebra's
+# interface (not `Base.ndims`) so a non-array backend like a TensorMap is supported.
+Base.ndims(a::AbstractNamedTensor) = TensorAlgebra.ndims(unnamed(a))
 
 # Circumvent issue when eltype isn't known at compile time.
 Base.eltype(a::AbstractNamedTensor) = eltype(unnamed(a))
@@ -777,13 +784,16 @@ function ArrayLayouts.sub_materialize(::NamedTensorLayout, a, ax)
     return copy(a)
 end
 
+# Attaching names to a bare parent is not slicing, so this accepts any parent a `NamedTensor`
+# can wrap (an `AbstractArray`, or a non-array backend like a TensorKit `TensorMap`). `Name` is
+# an ITensorBase-owned index type, so the generic parent is not type piracy. The `AbstractArray`
+# methods carry the identical body and exist only to disambiguate against
+# `Base.getindex`/`view(::AbstractArray, I...)`, which would otherwise tie with the generic ones.
 # TODO: Should this be a view?
-function Base.getindex(a::AbstractArray, I1::Name, Irest::Name...)
-    return copy(view(a, I1, Irest...))
-end
-function Base.view(a::AbstractArray, I1::Name, Irest::Name...)
-    return nameddims(a, name.((I1, Irest...)))
-end
+Base.getindex(a, I1::Name, Irest::Name...) = copy(view(a, I1, Irest...))
+Base.getindex(a::AbstractArray, I1::Name, Irest::Name...) = copy(view(a, I1, Irest...))
+Base.view(a, I1::Name, Irest::Name...) = nameddims(a, name.((I1, Irest...)))
+Base.view(a::AbstractArray, I1::Name, Irest::Name...) = nameddims(a, name.((I1, Irest...)))
 
 function Base.getindex(a::AbstractArray, I1::NamedViewIndex, Irest::NamedViewIndex...)
     return copy(view(a, I1, Irest...))
@@ -974,33 +984,20 @@ end
 
 using Random: Random, AbstractRNG
 
-# Like `Base.rand` but supports axes, not just size.
-# TODO: Come up with a better name for this.
-_rand(args...) = Base.rand(args...)
-function _rand(
-        rng::AbstractRNG, elt::Type, dims::Tuple{Base.OneTo{Int}, Vararg{Base.OneTo{Int}}}
-    )
-    return Base.rand(rng, elt, length.(dims))
-end
-
-# Like `Base.randn` but supports axes, not just size.
-# TODO: Come up with a better name for this.
-_randn(args...) = Base.randn(args...)
-function _randn(
-        rng::AbstractRNG, elt::Type, dims::Tuple{Base.OneTo{Int}, Vararg{Base.OneTo{Int}}}
-    )
-    return Base.randn(rng, elt, length.(dims))
-end
-
+# Cold-start `rand`/`randn`/`zeros` over axes build an all-codomain map (trivial domain) with
+# `TensorAlgebra`'s map constructors. They dispatch on the axis type, so dense `Base.OneTo`
+# axes build an `Array`, graded axes a block-sparse array, and TensorKit spaces a `TensorMap`,
+# without ITensorBase choosing a backend or pirating `Base.rand`/`randn`/`zeros`.
 default_eltype() = Float64
-for (f, f′) in [(:rand, :_rand), (:randn, :_randn)]
+for f in [:rand, :randn]
+    f_map = Symbol(f, :_map)
     @eval begin
         function Base.$f(
                 rng::AbstractRNG,
                 elt::Type{<:Number},
                 ax::Tuple{NamedUnitRange, Vararg{NamedUnitRange}}
             )
-            a = $f′(rng, elt, unnamed.(ax))
+            a = TensorAlgebra.$f_map(rng, elt, unnamed.(ax), ())
             return a[Name.(name.(ax))...]
         end
         function Base.$f(
@@ -1030,12 +1027,21 @@ for (f, f′) in [(:rand, :_rand), (:randn, :_randn)]
         end
     end
 end
+# `zeros` routes through `TensorAlgebra.zeros_map` (all-codomain map, trivial domain), which
+# dispatches on the axis type to build a dense `Array`, a block-sparse array, or a TensorKit
+# `TensorMap`. `ones` has no map hook (an all-ones symmetric tensor is not well-defined), so it
+# stays on `Base.ones`, which already accepts axes.
 for f in [:zeros, :ones], dimtype in [:NamedInteger, :NamedUnitRange]
+    parent = if f === :zeros
+        :(TensorAlgebra.zeros_map(elt, unnamed.(ax), ()))
+    else
+        :(Base.ones(elt, unnamed.(ax)))
+    end
     @eval begin
         function Base.$f(
                 elt::Type{<:Number}, ax::Tuple{$dimtype, Vararg{$dimtype}}
             )
-            a = $f(elt, unnamed.(ax))
+            a = $parent
             return a[Name.(name.(ax))...]
         end
         function Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype})

src/namedtensor.jl

@@ -1,3 +1,5 @@
+using TensorAlgebra: TensorAlgebra
+
 """
     NamedTensor(array::AbstractArray, dims)
 
@@ -18,9 +20,12 @@ named(Base.OneTo(2), :i)×named(Base.OneTo(3), :j) NamedTensor{Symbol}:
 ```
 """
 struct NamedTensor{DimName} <: AbstractNamedTensor{DimName}
-    unnamed::AbstractArray
+    # The parent is usually an `AbstractArray`, but the field is left untyped so a non-array
+    # tensor backend (e.g. a TensorKit `TensorMap`, reached through TensorAlgebra's `ndims`/
+    # `axes`/algebra interface) can be the parent directly. See the TensorKit extension.
+    unnamed::Any
     dimnames::Vector{DimName}
-    function NamedTensor{DimName}(unnamed::AbstractArray, dimnames) where {DimName}
+    function NamedTensor{DimName}(unnamed, dimnames) where {DimName}
         dimnames = collect(DimName, dimnames)
         # Catch the common ITensors.jl-style mistake of passing indices as the names.
         any(dimname -> dimname isa NamedUnitRange, dimnames) && throw(
@@ -30,15 +35,15 @@ struct NamedTensor{DimName} <: AbstractNamedTensor{DimName}
                 array and indices, index the array instead, as in `array[i, j]`."
             )
         )
-        ndims(unnamed) == length(dimnames) ||
+        TensorAlgebra.ndims(unnamed) == length(dimnames) ||
             throw(ArgumentError("Number of named dims must match ndims."))
         allunique(dimnames) ||
             throw(ArgumentError("Dimension names must be distinct, got $(dimnames)."))
         return new{DimName}(unnamed, dimnames)
     end
 end
 
-NamedTensor(unnamed::AbstractArray, dims) = NamedTensor{eltype(dims)}(unnamed, dims)
+NamedTensor(unnamed, dims) = NamedTensor{eltype(dims)}(unnamed, dims)
 NamedTensor(a::AbstractNamedTensor, inds) = throw(ArgumentError("Already named."))
 NamedTensor(a::AbstractNamedTensor) = NamedTensor(unnamed(a), dimnames(a))
 

src/namedunitrange.jl

@@ -17,8 +17,10 @@ named(1:3, :i)
 
 See also [`Index`](@ref), [`named`](@ref).
 """
-struct NamedUnitRange{Name, UnnamedT <: Integer, Unnamed <: AbstractUnitRange{UnnamedT}} <:
-    AbstractNamedVector{Name, UnnamedT}
+struct NamedUnitRange{Name, UnnamedT, Unnamed} <: AbstractNamedVector{Name, UnnamedT}
+    # The `value` is usually an integer `AbstractUnitRange`, but the bound is left open so a
+    # backend can store a richer axis object directly, e.g. a native TensorKit space, which is
+    # not an `AbstractUnitRange` but is its own axis (see the TensorKit extension).
     value::Unnamed
     name::Name
 end
@@ -47,6 +49,13 @@ end
 function NamedUnitRange{Name}(unnamed::AbstractUnitRange, name) where {Name}
     return NamedUnitRange{Name, eltype(unnamed), typeof(unnamed)}(unnamed, name)
 end
+# Base case for the name-inferred path: a ready-made range and a name. Loosening the struct's
+# `Unnamed <: AbstractUnitRange` bound removed the compiler-synthesized 2-arg constructor that
+# used to terminate here, so it is spelled out explicitly (a backend storing a non-range axis,
+# e.g. a TensorKit space, adds its own terminal — see the TensorKit extension).
+function NamedUnitRange(unnamed::AbstractUnitRange, name)
+    return NamedUnitRange{typeof(name), eltype(unnamed), typeof(unnamed)}(unnamed, name)
+end
 # A space and an explicit name, name type inferred from `name`.
 function NamedUnitRange(space, name)
     return NamedUnitRange(to_range(space), name)

test/Project.toml

@@ -14,6 +14,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -40,9 +41,10 @@ SafeTestsets = "0.1"
 StableRNGs = "1"
 Suppressor = "0.2"
 TensorAlgebra = "0.16"
+TensorKit = "0.17"
 TensorOperations = "5.3.1"
 TermInterface = "2"
 Test = "1.10"
 UUIDs = "1.10"
-VectorInterface = "0.5, 0.6"
+VectorInterface = "0.5"
 WrappedUnions = "0.3"

test/test_tensorkitext.jl

@@ -0,0 +1,71 @@
+using ITensorBase: ITensorBase, Index, dimnames, name, unnamed
+using LinearAlgebra: norm
+using MatrixAlgebraKit: qr_compact, svd_compact
+using StableRNGs: StableRNG
+using TensorKit: TensorKit, @tensor, AbstractTensorMap, SU2Irrep, U1Irrep, Vect, dim, dual,
+    scalar, space, ⊗
+using Test: @test, @test_throws, @testset
+
+# A native TensorKit space flows into `Index`, so an `ITensor` wraps a `TensorMap` directly.
+# Cover an abelian (U₁) and a non-abelian (SU₂) symmetry; the non-abelian case is the point
+# — no block-sparse abelian backend can represent it.
+@testset "TensorKitExt (eltype = $elt)" for elt in (Float64, ComplexF64)
+    rng = StableRNG(1234)
+    @testset "$label" for (label, Vi, Vj, Vk) in (
+            (
+                "U₁",
+                Vect[U1Irrep](0 => 2, 1 => 3),
+                Vect[U1Irrep](0 => 1, 1 => 2),
+                Vect[U1Irrep](-1 => 1, 0 => 2),
+            ),
+            (
+                "SU₂",
+                Vect[SU2Irrep](0 => 2, 1 // 2 => 1),
+                Vect[SU2Irrep](1 // 2 => 1, 1 => 1),
+                Vect[SU2Irrep](0 => 1, 1 // 2 => 2),
+            ),
+        )
+        i, j, k = Index(Vi), Index(Vj), Index(Vk)
+
+        # `Index` stores the native space directly.
+        @test unnamed(i) === Vi
+
+        # `conj(index)` round-trips to an `Index` carrying the dual space, same name.
+        @test conj(i) isa Index
+        @test unnamed(conj(i)) == dual(Vi)
+        @test name(conj(i)) == name(i)
+
+        # Cold-start construction wraps a `TensorMap`; size/eltype report dense values.
+        a = randn(rng, elt, i, j)
+        @test unnamed(a) isa AbstractTensorMap
+        @test size(a) == (dim(Vi), dim(Vj))
+        @test eltype(a) == elt
+        @test norm(unnamed(zeros(elt, i, j))) == 0
+
+        # Contraction over the shared (dualized) leg matches a direct TensorKit reference.
+        b = randn(rng, elt, conj(j), k)
+        c = a * b
+        @test Set(dimnames(c)) == Set(name.((i, k)))
+        ta, tb, gc = unnamed(a), unnamed(b), unnamed(c)
+        @tensor ref[vi; vk] := ta[vi, vj] * tb[vj, vk]
+        @test space(ref) == space(gc)
+        @test ref ≈ gc
+
+        # Linear-combination broadcast lowers to `bipermutedimsopadd!`; element-wise errors.
+        b2 = randn(rng, elt, i, j)
+        @test unnamed(a + b2) ≈ unnamed(a) + unnamed(b2)
+        @test unnamed(2 * a) ≈ 2 * unnamed(a)
+        @test unnamed(a .- 3 .* b2) ≈ unnamed(a) - 3 * unnamed(b2)
+        @test_throws ErrorException sin.(a)
+
+        # Factorizations reconstruct the tensor (lowered through matricize / MatrixAlgebraKit).
+        # Checked by full contraction to a scalar, which is bipartition-independent.
+        a3 = randn(rng, elt, i, j, k)
+        w = randn(rng, elt, conj(i), conj(j), conj(k))
+        sca(x) = scalar(unnamed(x))
+        u, s, v = svd_compact(a3, (i,), (j, k))
+        @test sca((u * s * v) * w) ≈ sca(a3 * w)
+        q, r = qr_compact(a3, (i,), (j, k))
+        @test sca((q * r) * w) ≈ sca(a3 * w)
+    end
+end