diff --git a/Project.toml b/Project.toml index b0a765e..383cc22 100644 --- a/Project.toml +++ b/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 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" diff --git a/ext/ITensorBaseTensorKitExt.jl b/ext/ITensorBaseTensorKitExt.jl new file mode 100644 index 0000000..924486f --- /dev/null +++ b/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 diff --git a/src/abstractnamedtensor.jl b/src/abstractnamedtensor.jl index 1a4a69f..abc89fc 100644 --- a/src/abstractnamedtensor.jl +++ b/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}) diff --git a/src/namedtensor.jl b/src/namedtensor.jl index 9d1fc50..c51c3d1 100644 --- a/src/namedtensor.jl +++ b/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,7 +35,7 @@ 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).")) @@ -38,7 +43,7 @@ struct NamedTensor{DimName} <: AbstractNamedTensor{DimName} 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)) diff --git a/src/namedunitrange.jl b/src/namedunitrange.jl index c2892f2..7e33d5b 100644 --- a/src/namedunitrange.jl +++ b/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) diff --git a/test/Project.toml b/test/Project.toml index 25a5b2a..882bd69 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl new file mode 100644 index 0000000..26930ff --- /dev/null +++ b/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