diff --git a/ext/ITensorBaseAbstractTreesExt/ITensorBaseAbstractTreesExt.jl b/ext/ITensorBaseAbstractTreesExt/ITensorBaseAbstractTreesExt.jl index aacfe91..75886ef 100644 --- a/ext/ITensorBaseAbstractTreesExt/ITensorBaseAbstractTreesExt.jl +++ b/ext/ITensorBaseAbstractTreesExt/ITensorBaseAbstractTreesExt.jl @@ -1,10 +1,10 @@ module ITensorBaseAbstractTreesExt using AbstractTrees: AbstractTrees -using ITensorBase: AbstractNamedDimsArray, dimnames +using ITensorBase: AbstractITensor, dimnames # Only print the dimension names when printing with `AbstractTrees.print_tree`. -function AbstractTrees.printnode(io::IO, a::AbstractNamedDimsArray) +function AbstractTrees.printnode(io::IO, a::AbstractITensor) dimnames_a = "{" * join(map(s -> "\"$s\"", dimnames(a)), ", ") * "}" print(io, dimnames_a) return nothing diff --git a/ext/ITensorBaseAdaptExt/ITensorBaseAdaptExt.jl b/ext/ITensorBaseAdaptExt/ITensorBaseAdaptExt.jl index a688bc2..dc8dc0d 100644 --- a/ext/ITensorBaseAdaptExt/ITensorBaseAdaptExt.jl +++ b/ext/ITensorBaseAdaptExt/ITensorBaseAdaptExt.jl @@ -1,9 +1,9 @@ module ITensorBaseAdaptExt using Adapt: Adapt, adapt -using ITensorBase: AbstractNamedDimsArray, denamed, dimnames, nameddims +using ITensorBase: AbstractITensor, denamed, dimnames, nameddims -function Adapt.adapt_structure(to, a::AbstractNamedDimsArray) +function Adapt.adapt_structure(to, a::AbstractITensor) return nameddims(adapt(to, denamed(a)), dimnames(a)) end diff --git a/ext/ITensorBaseBlockArraysExt/ITensorBaseBlockArraysExt.jl b/ext/ITensorBaseBlockArraysExt/ITensorBaseBlockArraysExt.jl index d919afd..a2d7f6c 100644 --- a/ext/ITensorBaseBlockArraysExt/ITensorBaseBlockArraysExt.jl +++ b/ext/ITensorBaseBlockArraysExt/ITensorBaseBlockArraysExt.jl @@ -1,8 +1,7 @@ module ITensorBaseBlockArraysExt using ArrayLayouts: ArrayLayouts using BlockArrays: Block, BlockRange -using ITensorBase: AbstractNamedDimsArray, AbstractNamedDimsMatrix, AbstractNamedUnitRange, - getindex_named, view_nameddims +using ITensorBase: AbstractITensor, AbstractNamedUnitRange, getindex_named, view_nameddims function Base.getindex(r::AbstractNamedUnitRange{<:Integer}, I::Block{1}) # TODO: Use `Derive.@interface NamedArrayInterface() r[I]` instead. @@ -16,39 +15,39 @@ end const BlockIndex{N} = Union{Block{N}, BlockRange{N}, AbstractVector{<:Block{N}}} -function Base.view(a::AbstractNamedDimsArray, I1::Block{1}, Irest::BlockIndex{1}...) - # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead. +function Base.view(a::AbstractITensor, I1::Block{1}, Irest::BlockIndex{1}...) + # TODO: Use `Derive.@interface ITensorInterface() r[I]` instead. return view_nameddims(a, I1, Irest...) end -function Base.view(a::AbstractNamedDimsArray, I::Block) - # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead. +function Base.view(a::AbstractITensor, I::Block) + # TODO: Use `Derive.@interface ITensorInterface() r[I]` instead. return view_nameddims(a, Tuple(I)...) end -function Base.view(a::AbstractNamedDimsArray, I1::BlockIndex{1}, Irest::BlockIndex{1}...) - # TODO: Use `Derive.@interface NamedDimsArrayInterface() r[I]` instead. +function Base.view(a::AbstractITensor, I1::BlockIndex{1}, Irest::BlockIndex{1}...) + # TODO: Use `Derive.@interface ITensorInterface() r[I]` instead. return view_nameddims(a, I1, Irest...) end # Fix ambiguity error. function Base.getindex( - a::AbstractNamedDimsArray, I1::BlockRange{1}, Irest::BlockRange{1}... + a::AbstractITensor, I1::BlockRange{1}, Irest::BlockRange{1}... ) return ArrayLayouts.layout_getindex(a, I1, Irest...) end # Fix ambiguity errors. -function Base.getindex(a::AbstractNamedDimsArray, I1::Block{1}, Irest...) +function Base.getindex(a::AbstractITensor, I1::Block{1}, Irest...) return copy(view(a, I1, Irest...)) end -function Base.getindex(a::AbstractNamedDimsMatrix, I1::AbstractVector, I2::Block{1}) +function Base.getindex(a::AbstractITensor, I1::AbstractVector, I2::Block{1}) return copy(view(a, I1, I2)) end -function Base.getindex(a::AbstractNamedDimsMatrix, I1::Block{1}, I2::AbstractVector) +function Base.getindex(a::AbstractITensor, I1::Block{1}, I2::AbstractVector) return copy(view(a, I1, I2)) end -function Base.getindex(a::AbstractNamedDimsArray{<:Any, N}, I::Block{N}) where {N} +function Base.getindex(a::AbstractITensor, I::Block{N}) where {N} return copy(view(a, I)) end diff --git a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl index f0080db..6cb2ebd 100644 --- a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl +++ b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl @@ -1,6 +1,6 @@ module ITensorBaseMooncakeExt -using ITensorBase: AbstractNamedDimsArray, AbstractNamedUnitRange, blockedperm_nameddims, +using ITensorBase: AbstractITensor, AbstractNamedUnitRange, blockedperm_nameddims, combine_nameddimsconstructors, dimnames, dimnames_setdiff, inds, name, nameddimsconstructorof, randname, to_inds using Mooncake: Mooncake, @zero_derivative, DefaultCtx @@ -8,7 +8,7 @@ using TensorAlgebra: blockedperm Mooncake.tangent_type(::Type{<:AbstractNamedUnitRange}) = Mooncake.NoTangent -@zero_derivative DefaultCtx Tuple{typeof(blockedperm), AbstractNamedDimsArray, Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(blockedperm), AbstractITensor, Any, Any} @zero_derivative DefaultCtx Tuple{typeof(blockedperm_nameddims), Any, Any, Any} @zero_derivative DefaultCtx Tuple{typeof(combine_nameddimsconstructors), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(dimnames), Any} @@ -22,9 +22,9 @@ Mooncake.tangent_type(::Type{<:AbstractNamedUnitRange}) = Mooncake.NoTangent @zero_derivative DefaultCtx Tuple{typeof(randname), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(to_inds), Any, Any} -using ITensorBase: AbstractNamedDimsArray, NamedDimsArray, denamed +using ITensorBase: AbstractITensor, ITensor, denamed using Mooncake: Tangent -function Base.copyto!(dest::NamedDimsArray, src::Tangent) +function Base.copyto!(dest::ITensor, src::Tangent) # TODO: Account for the `inds` of the Tangent? In other words, is the tangent data # aligned with the `dest` data? copyto!(denamed(dest), src.fields.parent) diff --git a/src/ITensorBase.jl b/src/ITensorBase.jl index 4bdb536..a0e93c5 100644 --- a/src/ITensorBase.jl +++ b/src/ITensorBase.jl @@ -1,6 +1,6 @@ module ITensorBase -export ITensor, Index, NamedDimsArray, aligndims, dimnametype, named, nameddims, +export ITensor, Index, aligndims, dimnametype, named, nameddims, operator, similar_operator using Compat: @compat @compat public to_inds @@ -16,16 +16,15 @@ include("abstractnamedarray.jl") include("namedarray.jl") include("abstractnamedunitrange.jl") include("namedunitrange.jl") -include("abstractnameddimsarray.jl") +include("abstractitensor.jl") include("broadcast.jl") include("tensoralgebra.jl") include("linearalgebra.jl") -include("nameddimsarray.jl") -include("nameddimsoperator.jl") +include("itensor.jl") +include("itensoroperator.jl") -# ITensor layer built on the named-array machinery. +# `IndexName` dimname flavor and the `Index` named unit range. include("index.jl") -include("abstractitensor.jl") include("quirks.jl") end diff --git a/src/abstractitensor.jl b/src/abstractitensor.jl index 656eb80..6e08c95 100644 --- a/src/abstractitensor.jl +++ b/src/abstractitensor.jl @@ -1,21 +1,892 @@ -abstract type AbstractITensor <: AbstractNamedDimsArray{Any, Any} end +using LinearAlgebra: LinearAlgebra +using Random: Random +using TensorAlgebra: permuteddims +using TypeParameterAccessors: unspecify_type_parameters -nameddimsconstructor(::Type{<:IndexName}) = ITensor -dimnametype(::Type{<:AbstractITensor}) = IndexName +# Some of the interface is inspired by: +# https://github.com/ITensor/ITensors.jl +# https://github.com/invenia/NamedDims.jl +# https://github.com/mcabbott/NamedPlus.jl +# https://pytorch.org/docs/stable/named_tensor.html +abstract type AbstractITensor{DimName} <: AbstractArray{Any, Any} end + +# Rank and element type live in the data, not the type. The `<: AbstractArray` +# subtyping is kept for now to reuse generic array functionality, with `Any` for +# both parameters since neither is fixed by the type. Base.ndims(::Type{<:AbstractITensor}) = Any -struct ITensor <: AbstractITensor - denamed::AbstractArray - dimnames::Any - function ITensor(denamed::AbstractArray, dimnames) - ndims(denamed) == length(dimnames) || - throw(ArgumentError("Number of named dims must match ndims.")) - all(dimname -> dimname isa IndexName, dimnames) || - throw(ArgumentError("All dimnames must be of type IndexName.")) - return new(denamed, dimnames) +dimnames(a::AbstractITensor) = throw(MethodError(dimnames, a)) +function dimnames(a::AbstractITensor, dim::Int) + return dimnames(a)[dim] +end + +""" + dimnametype(a::AbstractITensor) + dimnametype(type::Type{<:AbstractITensor}) + +The type of an individual dimension name of `a`. The primary method dispatches +on the array type, and `dimnametype(a)` forwards to `dimnametype(typeof(a))`. A +type that does not fix its dimname flavor (such as the unparameterized `ITensor`) +returns `Any`, the same way `eltype(Array)` is `Any`. + +# Examples + +```jldoctest +julia> a = nameddims(zeros(2, 3), (:i, :j)); + +julia> dimnametype(a) +Symbol + +julia> dimnametype(typeof(a)) +Symbol +``` +""" +function dimnametype end +dimnametype(a::AbstractITensor) = dimnametype(typeof(a)) +dimnametype(type::Type{<:AbstractITensor}) = Any + +# Unwrapping the names (named-array interface). +# TODO: Use `IsNamed` trait? +denamed(a::AbstractITensor) = throw(MethodError(denamed, a)) +denamed(a::AbstractITensor, inds) = denamed(aligneddims(a, inds)) +dename(a::AbstractITensor, inds) = denamed(aligndims(a, inds)) + +# Output the named axes/indices of the named dims array. +inds(a::AbstractArray) = LittleSet(named.(axes(denamed(a)), dimnames(a))) +inds(a::AbstractArray, dim::Int) = inds(a)[dim] + +isnamed(::Type{<:AbstractITensor}) = true + +function dim(a::AbstractArray, n) + return findfirst(==(name(n)), dimnames(a)) +end +dims(a::AbstractArray, ns) = Base.Fix1(dim, a).(ns) + +dimname_isequal(x) = Base.Fix1(dimname_isequal, x) +dimname_isequal(x, y) = isequal(x, y) + +dimname_isequal(r1::AbstractNamedArray, r2::AbstractNamedArray) = isequal(r1, r2) +dimname_isequal(r1::AbstractNamedArray, r2) = name(r1) == r2 +dimname_isequal(r1, r2::AbstractNamedArray) = r1 == name(r2) + +dimname_isequal(r1::AbstractNamedArray, r2::Name) = name(r1) == name(r2) +dimname_isequal(r1::Name, r2::AbstractNamedArray) = name(r1) == name(r2) + +dimname_isequal(r1::AbstractNamedUnitRange, r2::AbstractNamedUnitRange) = isequal(r1, r2) +dimname_isequal(r1::AbstractNamedUnitRange, r2) = name(r1) == r2 +dimname_isequal(r1, r2::AbstractNamedUnitRange) = r1 == name(r2) + +dimname_isequal(r1::AbstractNamedUnitRange, r2::Name) = name(r1) == name(r2) +dimname_isequal(r1::Name, r2::AbstractNamedUnitRange) = name(r1) == name(r2) + +function to_inds(a::AbstractArray, dims) + is = Base.Fix1(dim, a).(name.(dims)) + return Base.Fix1(inds, a).(is) +end + +# Generic construction of named dims arrays. + +""" + nameddims(a::AbstractArray, inds) + +Construct a named dimensions array from an denamed array `a` and named dimensions `inds`. +""" +function nameddims(a::AbstractArray, inds) + return nameddimsconstructor(a, inds)(a, inds) +end + +#= + nameddimsof(a::AbstractITensor, b::AbstractArray) + +Construct a named dimensions array with the dimension names of `a` +and based on the type of `a` but with the data from `b`. +=# +function nameddimsof(a::AbstractITensor, b::AbstractArray) + return nameddimsconstructorof(a)(b, dimnames(a)) +end + +# Interface inspired by +# [ConstructionBase.constructorof](https://github.com/JuliaObjects/ConstructionBase.jl). +nameddimsconstructorof(a::AbstractITensor) = nameddimsconstructorof(typeof(a)) +function nameddimsconstructorof(type::Type{<:AbstractITensor}) + return unspecify_type_parameters(type) +end + +# Output a constructor for a named dims array (that should accept and denamed array and +# a set of named dimensions/axes/indices) based on the dimension names. +function nameddimsconstructor(a::AbstractArray, dims) + dimnames = name.(dims) + isempty(dimnames) && return ITensor + return mapreduce(nameddimsconstructor, combine_nameddimsconstructors, dimnames) +end + +nameddimsconstructor(nameddim) = nameddimsconstructor(typeof(nameddim)) +# Can overload this to get custom named dims array wrapper +# depending on the dimension name types, for example +# output an `ITensor` if the dimension names are `IndexName`s. +nameddimsconstructor(nameddimtype::Type) = ITensor +function nameddimsconstructor(nameddimtype::Type{<:AbstractNamedUnitRange}) + return nameddimsconstructor(nametype(nameddimtype)) +end +function combine_nameddimsconstructors( + ::Type{<:AbstractITensor}, ::Type{<:AbstractITensor} + ) + return ITensor +end +combine_nameddimsconstructors(::Type{T}, ::Type{T}) where {T <: AbstractITensor} = T + +# TODO: Move to `utils.jl` file. +# TODO: Use `Base.indexin`? +function getperm(x, y; isequal = isequal) + return map(yᵢ -> findfirst(isequal(yᵢ), x), y) +end + +# TODO: Move to `utils.jl` file. +function checked_indexin(x, y) + I = indexin(x, y) + return something.(I) +end + +function checked_indexin(x::Number, y) + return findfirst(==(x), y) +end + +function checked_indexin(x::AbstractUnitRange, y::AbstractUnitRange) + return findfirst(==(first(x)), y):findfirst(==(last(x)), y) +end + +Base.copy(a::AbstractITensor) = nameddimsof(a, copy(denamed(a))) + +# Forward `conj` to the underlying so that graded axes flip their sector arrows. +# The default `AbstractArray` fallback would broadcast `conj` over elements without +# touching the axes, which silently changes the contraction convention for tensors +# with graded (dual-tagged) axes. +Base.conj(a::AbstractITensor) = nameddimsof(a, conj(denamed(a))) + +# `LinearAlgebra.normalize` infers result eltype via `typeof(first(a)/nrm)`, which +# scalar-indexes block-structured storage. `a / norm(a, p)` already preserves names. +function LinearAlgebra.normalize(a::AbstractITensor, p::Real = 2) + return a / LinearAlgebra.norm(a, p) +end + +# Forward `Random.randn!` / `Random.rand!` to the concrete storage so they +# see the runtime eltype. +function Random.randn!(rng::Random.AbstractRNG, a::AbstractITensor) + Random.randn!(rng, denamed(a)) + return a +end +function Random.rand!(rng::Random.AbstractRNG, a::AbstractITensor) + Random.rand!(rng, denamed(a)) + return a +end + +function Base.copyto!(a_dest::AbstractITensor, a_src::AbstractITensor) + a′_dest = denamed(a_dest) + # TODO: Use `denamed` to do the permutations lazily. + a′_src = dename(a_src, inds(a_dest)) + copyto!(a′_dest, a′_src) + return a_dest +end + +# Conversion + +# Copied from `Base` (defined in abstractarray.jl). +@noinline function _checkaxs(axd, axs) + axd == axs || throw(DimensionMismatch("axes must agree, got $axd and $axs")) + return nothing +end +function copyto_axcheck!(dest, src) + _checkaxs(axes(dest), axes(src)) + copyto!(dest, src) + return dest +end + +# These are defined since the Base versions assume the eltype and ndims are known +# at compile time, which isn't true for ITensors. +Base.Array(a::AbstractITensor) = Array(denamed(a)) +Base.Array{T}(a::AbstractITensor) where {T} = Array{T}(denamed(a)) +Base.Array{T, N}(a::AbstractITensor) where {T, N} = Array{T, N}(denamed(a)) +Base.AbstractArray{T}(a::AbstractITensor) where {T} = AbstractArray{T, ndims(a)}(a) +function Base.AbstractArray{T, N}(a::AbstractITensor) where {T, N} + dest = similar(a, T) + copyto_axcheck!(denamed(dest), denamed(a)) + return dest +end + +function Base.axes(a::AbstractITensor) + return inds(a) +end +function Base.size(a::AbstractITensor) + return length.(axes(a)) +end + +function Base.length(a::AbstractITensor) + return prod(size(a); init = 1) +end + +# Circumvent issue when ndims isn't known at compile time. +Base.axes(a::AbstractITensor, d) = axes(a)[d] + +# Circumvent issue when ndims isn't known at compile time. +Base.size(a::AbstractITensor, d) = size(a)[d] + +# Circumvent issue when ndims isn't known at compile time. +Base.ndims(a::AbstractITensor) = ndims(denamed(a)) + +# Circumvent issue when eltype isn't known at compile time. +Base.eltype(a::AbstractITensor) = eltype(denamed(a)) + +using VectorInterface: VectorInterface, scalartype +# Circumvent issue when eltype isn't known at compile time. +VectorInterface.scalartype(a::AbstractITensor) = scalartype(denamed(a)) + +Base.axes(a::AbstractITensor, dimname::Name) = axes(a, dim(a, dimname)) +Base.size(a::AbstractITensor, dimname::Name) = size(a, dim(a, dimname)) + +function similar_nameddims(a::AbstractITensor, elt::Type, ax) + return nameddimsconstructorof(a)( + similar(denamed(a), elt, denamed.(Tuple(ax))), + name.(ax) + ) +end +function similar_nameddims(a::AbstractArray, elt::Type, ax) + return nameddims(similar(a, elt, denamed.(Tuple(ax))), name.(ax)) +end + +# Base.similar gets the eltype at compile time. +Base.similar(a::AbstractITensor) = similar(a, eltype(a)) +function Base.similar(a::AbstractITensor, elt::Type) + return similar_nameddims(a, elt) +end +function similar_nameddims(a::AbstractITensor, elt::Type) + return nameddimsof(a, similar(denamed(a), elt)) +end + +# This is defined explicitly since the Base version expects the eltype +# to be known at compile time, which isn't true for ITensors. +function Base.similar( + a::AbstractArray, inds::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} + ) + return similar(a, eltype(a), inds) +end + +function Base.similar( + a::AbstractArray, elt::Type, + inds::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} + ) + return similar_nameddims(a, elt, inds) +end + +function Base.similar(a::AbstractArray, inds::LittleSet) + return similar_nameddims(a, eltype(a), inds) +end + +function Base.similar(a::AbstractArray, elt::Type, inds::LittleSet) + return similar_nameddims(a, elt, inds) +end + +function setinds(a::AbstractITensor, inds) + return nameddimsconstructorof(a)(denamed(a), inds) +end + +function replacedimnames(a::AbstractITensor, replacements::Pair...) + new_dimnames = replace(dimnames(a), replacements...) + return nameddims(denamed(a), new_dimnames) +end +function replacedimnames(f, a::AbstractITensor) + new_dimnames = replace(f, dimnames(a)) + return nameddims(denamed(a), new_dimnames) +end +mapdimnames(f, a::AbstractITensor) = replacedimnames(f, a) + +function replaceinds(a::AbstractITensor, replacements::Pair...) + new_inds = replace(inds(a), replacements...) + return denamed(a)[new_inds...] +end +function replaceinds(f, a::AbstractITensor) + new_inds = replace(f, inds(a)) + return denamed(a)[new_inds...] +end +mapinds(f, a::AbstractITensor) = replaceinds(f, a) + +# `Base.isempty(a::AbstractArray)` is defined as `length(a) == 0`, +# which involves comparing a named integer to an denamed integer +# which isn't well defined. +Base.isempty(a::AbstractITensor) = isempty(denamed(a)) + +# Define this on objects rather than types in case the wrapper type +# isn't known at compile time, like for the ITensor type. +Base.IndexStyle(a::AbstractITensor) = IndexStyle(denamed(a)) +Base.eachindex(a::AbstractITensor) = eachindex(denamed(a)) + +# Cartesian indices with names attached. +struct NamedIndexCartesian <: IndexStyle end + +# When multiple named dims arrays are involved, use the named +# dimensions. +function Base.IndexStyle(a1::AbstractITensor, a2::AbstractITensor) + return NamedIndexCartesian() +end +# Define promotion of index styles. +Base.IndexStyle(s1::NamedIndexCartesian, s2::NamedIndexCartesian) = NamedIndexCartesian() +Base.IndexStyle(s1::IndexStyle, s2::NamedIndexCartesian) = NamedIndexCartesian() +Base.IndexStyle(s1::NamedIndexCartesian, s2::IndexStyle) = NamedIndexCartesian() + +# Like CartesianIndex but with named dimensions. +struct NamedDimsCartesianIndex{N, Index <: Tuple{Vararg{AbstractNamedInteger, N}}} <: + Base.AbstractCartesianIndex{N} + I::Index +end +NamedDimsCartesianIndex(I::AbstractNamedInteger...) = NamedDimsCartesianIndex(I) +Base.Tuple(I::NamedDimsCartesianIndex) = I.I +function Base.show(io::IO, I::NamedDimsCartesianIndex) + print(io, "NamedDimsCartesianIndex") + show(io, Tuple(I)) + return nothing +end + +# Like CartesianIndices but with named dimensions. +struct NamedDimsCartesianIndices{ + N, + DimName, + Indices <: Tuple{Vararg{AbstractNamedUnitRange, N}}, + Index <: Tuple{Vararg{AbstractNamedInteger, N}}, + } <: AbstractITensor{DimName} + indices::Indices + function NamedDimsCartesianIndices(indices::Tuple{Vararg{AbstractNamedUnitRange}}) + dimname = eltype(name.(indices)) + return new{length(indices), dimname, typeof(indices), Tuple{eltype.(indices)...}}( + indices + ) + end +end +function NamedDimsCartesianIndices(indices::LittleSet) + return NamedDimsCartesianIndices(Tuple(indices)) +end + +# The element type is no longer carried by the (rank-erased) supertype, so recover +# it from the stored index-tuple parameter. +function Base.eltype( + ::Type{<:NamedDimsCartesianIndices{N, <:Any, <:Any, Index}} + ) where {N, Index} + return NamedDimsCartesianIndex{N, Index} +end +Base.eltype(I::NamedDimsCartesianIndices) = eltype(typeof(I)) +Base.axes(I::NamedDimsCartesianIndices) = (only ∘ axes ∘ denamed).(I.indices) +Base.size(I::NamedDimsCartesianIndices) = (length ∘ denamed).(I.indices) + +function Base.getindex(a::NamedDimsCartesianIndices{N}, I::Vararg{Int, N}) where {N} + index = map(a.indices, I) do r, i + return r[i] end + return NamedDimsCartesianIndex(index) +end + +function denamed(I::NamedDimsCartesianIndices) + return CartesianIndices(denamed.(I.indices)) +end + +function Base.eachindex(::NamedIndexCartesian, a1::AbstractArray, a_rest::AbstractArray...) + all(a -> issetequal(inds(a1), inds(a)), a_rest) || + throw(NameMismatch("Dimension name mismatch $(inds.((a1, a_rest...))).")) + # TODO: Check the shapes match. + return NamedDimsCartesianIndices(inds(a1)) +end + +# Base version ignores dimension names. +# TODO: Use `mapreduce(isequal, &&, a1, a2)`? +function Base.isequal(a1::AbstractITensor, a2::AbstractITensor) + (inds(a1) == inds(a2)) || return false + return isequal(denamed(a1), denamed(a2, inds(a1))) +end + +# Base version ignores dimension names. +# TODO: Use `mapreduce(==, &&, a1, a2)`? +# TODO: Handle `missing` values properly. +function Base.:(==)(a1::AbstractITensor, a2::AbstractITensor) + (inds(a1) == inds(a2)) || return false + return denamed(a1) == denamed(a2, inds(a1)) +end + +# Generalization of `Base.sort` to Tuples for Julia v1.10 compatibility. +# TODO: Remove when we drop support for Julia v1.10. +_sort(x; kwargs...) = sort(x; kwargs...) +_sort(x::NTuple{N}; kwargs...) where {N} = NTuple{N}(sort(collect(x); kwargs...)) + +function Base.hash(a::AbstractITensor, h::UInt64) + h = hash(:ITensor, h) + a′ = aligneddims(a, _sort(dimnames(a))) + h = hash(denamed(a′), h) + for i in inds(a′) + h = hash(i, h) + end + return h +end + +# Indexing. + +# Scalar indexing + +Base.firstindex(a::AbstractITensor) = firstindex(denamed(a)) +Base.lastindex(a::AbstractITensor) = lastindex(denamed(a)) + +function Base.firstindex(a::AbstractITensor, d) + return FirstIndex(a, d) +end + +function Base.lastindex(a::AbstractITensor, d) + return LastIndex(a, d) +end + +# Redefine generic definition which expects `axes(a)` to +# return a Tuple. +function Base.to_indices(a::AbstractITensor, I::Tuple) + return to_indices(a, Tuple(axes(a)), I) +end +# Fix ambiguity error with Base. +function Base.to_indices( + a::AbstractITensor, + I::Tuple{Union{Integer, CartesianIndex}} + ) + return to_indices(a, Tuple(axes(a)), I) +end +function Base.checkbounds(::Type{Bool}, a::AbstractITensor, I::Int...) + return checkbounds(Bool, denamed(a), I...) +end + +function Base.to_indices( + a::AbstractITensor, I::Tuple{AbstractNamedInteger, Vararg{AbstractNamedInteger}} + ) + perm = getperm(name.(I), dimnames(a)) + # TODO: Throw a `NameMismatch` error. + @assert isperm(perm) + I = map(p -> I[p], perm) + return map(inds(a), I) do dimname, i + return checked_indexin(denamed(i), denamed(dimname)) + end +end +function Base.to_indices( + a::AbstractITensor, I::Tuple{Pair{<:Any, Int}, Vararg{Pair{<:Any, Int}}} + ) + inds = to_inds(a, first.(I)) + return to_indices(a, map((i, name) -> name[i], last.(I), inds)) +end +function Base.to_indices(a::AbstractITensor, I::Tuple{Pair, Vararg{Pair}}) + inds = to_inds(a, first.(I)) + return map((i, name) -> name[i], last.(I), inds) + return to_indices(a, named.(last.(I), first.(I))) +end + +function Base.to_indices(a::AbstractITensor, I::Tuple{NamedDimsCartesianIndex}) + return to_indices(a, Tuple(only(I))) +end + +function Base.getindex(a::AbstractITensor, I...) + return getindex(a, to_indices(a, I)...) +end + +function Base.getindex(a::AbstractITensor, I1::Int, Irest::Int...) + return getindex(denamed(a), I1, Irest...) +end +function Base.getindex( + a::AbstractITensor, I1::AbstractNamedInteger, Irest::AbstractNamedInteger... + ) + return getindex(a, to_indices(a, (I1, Irest...))...) +end +function Base.getindex(a::AbstractITensor) + return getindex(denamed(a)) +end +# Linear indexing. +function Base.getindex(a::AbstractITensor, I::Int) + return getindex(denamed(a), I) +end + +function Base.setindex!(a::AbstractITensor, value, I1::Int, Irest::Int...) + setindex!(denamed(a), value, I1, Irest...) + return a +end +function Base.setindex!(a::AbstractITensor, value, I::CartesianIndex) + setindex!(a, value, to_indices(a, (I,))...) + return a +end + +function Base.setindex!( + a::AbstractITensor, value, I1::AbstractNamedInteger, + Irest::AbstractNamedInteger... + ) + setindex!(a, value, to_indices(a, (I1, Irest...))...) + return a +end +function Base.setindex!(a::AbstractITensor, value, I::NamedDimsCartesianIndex) + setindex!(a, value, to_indices(a, (I,))...) + return a +end +function Base.setindex!(a::AbstractITensor, value, I1::Pair, Irest::Pair...) + setindex!(a, value, to_indices(a, (I1, Irest...))...) + return a +end +function Base.setindex!(a::AbstractITensor, value) + setindex!(denamed(a), value) + return a +end +# Linear indexing. +function Base.setindex!(a::AbstractITensor, value, I::Int) + setindex!(denamed(a), value, I) + return a +end + +function Base.isassigned(a::AbstractITensor, I::Int...) + return isassigned(denamed(a), I...) +end + +# Slicing + +# Like `const ViewIndex = Union{Real,AbstractArray}`. +const NamedViewIndex = + Union{AbstractNamedInteger, AbstractNamedUnitRange, AbstractNamedArray} + +using ArrayLayouts: ArrayLayouts, MemoryLayout + +abstract type AbstractITensorLayout <: MemoryLayout end +struct ITensorLayout{ParentLayout} <: AbstractITensorLayout end + +function ArrayLayouts.MemoryLayout(arrtype::Type{<:AbstractITensor}) + return ITensorLayout{typeof(MemoryLayout(parenttype(arrtype)))}() +end + +function ArrayLayouts.sub_materialize(::ITensorLayout, a, ax) + return copy(a) +end + +# 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 + +function Base.getindex(a::AbstractArray, I1::NamedViewIndex, Irest::NamedViewIndex...) + return copy(view(a, I1, Irest...)) +end +# Disambiguate from `Base.getindex(A::Array, I::AbstractUnitRange{<:Integer})`. +function Base.getindex(a::Array, I1::AbstractNamedUnitRange{<:Integer}) + return copy(view(a, I1)) +end +function Base.view(a::AbstractArray, I1::NamedViewIndex, Irest::NamedViewIndex...) + I = (I1, Irest...) + return nameddims(view(a, denamed.(I)...), name.(I)) +end + +# TODO: Should this be a view? +function Base.getindex(a::AbstractITensor, I1::Name, Irest::Name...) + return copy(view(a, I1, Irest...)) +end +function Base.view(a::AbstractITensor, I1::Name, Irest::Name...) + issetequal(dimnames(a), name.((I1, Irest...))) || + throw( + NameMismatch( + "Dimension name mismatch $(dimnames(a)), $(name.((I1, Irest...)))." + ) + ) + return a +end + +function Base.getindex(a::AbstractITensor, I1::Pair, Irest::Pair...) + return getindex(a, to_indices(a, (I1, Irest...))...) +end +function Base.view(a::AbstractITensor, I1::Pair, Irest::Pair...) + I = (I1, Irest...) + inds = to_inds(a, first.(I)) + return view(a, map((i, name) -> name[i], last.(I), inds)...) +end + +function Base.getindex( + a::AbstractITensor, I1::NamedViewIndex, Irest::NamedViewIndex... + ) + return copy(view(a, I1, Irest...)) +end +function Base.view(a::AbstractITensor, I1::NamedViewIndex, Irest::NamedViewIndex...) + I = (I1, Irest...) + perm = getperm(name.(I), dimnames(a)) + isperm(perm) || throw( + NameMismatch( + "Dimension name mismatch $(dimnames(a)), $(name.(I))." + ) + ) + Ip = map(p -> denamed(I[p]), perm) + return view_nameddims(a, Ip...) +end + +# Repeated definition of `Base.ViewIndex`. +const ViewIndex = Union{Real, AbstractArray} + +# Like `Base.ScalarIndex` but as a trait function. +# This catches cases like `Colon`, `BlockArrays.Block`, etc. which are not AbstractArray +# indices but also aren't scalar indices. +isscalarindex(I) = false +isscalarindex(I::Real) = true + +# Slicing with unnamed indices, such as: +# a = ITensor(rand(3,4), (:x, :y)) +# b = view(a, 1:2, 2) +function view_nameddims(a::AbstractITensor, I...) + nonscalar_dims = filter(dim -> !isscalarindex(I[dim]), ntuple(identity, ndims(a))) + nonscalar_dimnames = map(dim -> dimnames(a, dim), nonscalar_dims) + return nameddimsconstructorof(a)(view(denamed(a), I...), nonscalar_dimnames) +end + +function Base.view(a::AbstractITensor, I::ViewIndex...) + return view_nameddims(a, I...) +end + +function getindex_nameddims(a::AbstractArray, I...) + return copy(view(a, I...)) +end + +function Base.getindex(a::AbstractITensor, I::ViewIndex...) + return getindex_nameddims(a, I...) +end + +function Base.setindex!( + a::AbstractITensor, + value::AbstractITensor, + I1::NamedViewIndex, + Irest::NamedViewIndex... + ) + view(a, I1, Irest...) .= value + return a +end +function Base.setindex!( + a::AbstractITensor, + value::AbstractArray, + I1::NamedViewIndex, + Irest::NamedViewIndex... + ) + I = (I1, Irest...) + a[I...] = nameddimsconstructorof(a)(value, name.(I)) + return a +end +function Base.setindex!( + a::AbstractITensor, + value::AbstractITensor, + I1::ViewIndex, + Irest::ViewIndex... + ) + view(a, I1, Irest...) .= value + return a +end +function Base.setindex!( + a::AbstractITensor, value::AbstractArray, I1::ViewIndex, Irest::ViewIndex... + ) + setindex!(denamed(a), value, I1, Irest...) + return a +end + +# Permute/align dimensions + +function aligndims(a::AbstractArray, dims) + new_dimnames = name.(dims) + perm = Tuple(getperm(dimnames(a), new_dimnames)) + isperm(perm) || throw( + NameMismatch( + "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." + ) + ) + return nameddimsconstructorof(a)(permutedims(denamed(a), perm), new_dimnames) +end + +function aligneddims(a::AbstractArray, dims) + new_dimnames = name.(dims) + perm = getperm(dimnames(a), new_dimnames) + isperm(perm) || throw( + NameMismatch( + "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." + ) + ) + return nameddimsconstructorof(a)( + permuteddims(denamed(a), perm), new_dimnames + ) +end + +# Convenient constructors + +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 + +default_eltype() = Float64 +for (f, f′) in [(:rand, :_rand), (:randn, :_randn)] + @eval begin + function Base.$f( + rng::AbstractRNG, + elt::Type{<:Number}, + ax::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} + ) + a = $f′(rng, elt, denamed.(ax)) + return a[Name.(name.(ax))...] + end + function Base.$f( + rng::AbstractRNG, + elt::Type{<:Number}, + dims::Tuple{AbstractNamedInteger, Vararg{AbstractNamedInteger}} + ) + return $f(rng, elt, Base.oneto.(dims)) + end + end + for dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] + @eval begin + function Base.$f( + rng::AbstractRNG, elt::Type{<:Number}, dim1::$dimtype, + dims::Vararg{$dimtype} + ) + return $f(rng, elt, (dim1, dims...)) + end + Base.$f(elt::Type{<:Number}, dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f( + Random.default_rng(), elt, dims + ) + Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype}) = $f( + elt, (dim1, dims...) + ) + Base.$f(dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(default_eltype(), dims) + Base.$f(dim1::$dimtype, dims::Vararg{$dimtype}) = $f((dim1, dims...)) + end + end +end +for f in [:zeros, :ones], dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] + @eval begin + function Base.$f( + elt::Type{<:Number}, ax::Tuple{$dimtype, Vararg{$dimtype}} + ) + a = $f(elt, denamed.(ax)) + return a[Name.(name.(ax))...] + end + function Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype}) + return $f(elt, (dim1, dims...)) + end + Base.$f(dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(default_eltype(), dims) + Base.$f(dim1::$dimtype, dims::Vararg{$dimtype}) = $f((dim1, dims...)) + end +end +for dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] + @eval begin + function Base.fill(value, ax::Tuple{$dimtype, Vararg{$dimtype}}) + a = fill(value, denamed.(ax)) + return a[Name.(name.(ax))...] + end + function Base.fill(value, dim1::$dimtype, dims::Vararg{$dimtype}) + return fill(value, (dim1, dims...)) + end + end +end + +function Base.fill!(a::AbstractITensor, v) + fill!(denamed(a), v) + return a +end + +function Base.map!(f, a_dest::AbstractITensor, a_srcs::AbstractITensor...) + a′_dest = denamed(a_dest) + # TODO: Use `denamed` to do the permutations lazily. + # TODO: Define `dename[d](dimnames) = Base.Fix1(dename[d], dimnames)` and use it here? + a′_srcs = Base.Fix2(dename, dimnames(a_dest)).(a_srcs) + map!(f, a′_dest, a′_srcs...) + return a_dest +end + +function Base.map(f, a_srcs::AbstractITensor...) + # copy(mapped(f, a_srcs...)) + return f.(a_srcs...) +end + +function Base.mapreduce(f, op, a::AbstractITensor; kwargs...) + return mapreduce(f, op, denamed(a); kwargs...) +end + +# `sum` is routed to the underlying data rather than left to fall back on the +# `mapreduce` method above because some array types (such as graded arrays) define +# `Base.sum` directly but not the general `mapreduce`, so the denamed `sum` is the +# path that works for them. +function Base.sum(a::AbstractITensor; kwargs...) + return sum(denamed(a); kwargs...) +end + +function LinearAlgebra.promote_leaf_eltypes(a::AbstractITensor) + return LinearAlgebra.promote_leaf_eltypes(denamed(a)) +end + +# Printing + +# Copy of `Base.dims2string` defined in `show.jl`. +function dims_to_string(d) + isempty(d) && return "0-dimensional" + length(d) == 1 && return "$(d[1])-element" + return join(map(string, d), '×') +end + +using TypeParameterAccessors: type_parameters, unspecify_type_parameters +function concretetype_to_string_truncated( + type::Type; + param_truncation_length = typemax(Int) + ) + isconcretetype(type) || throw(ArgumentError("Type must be concrete.")) + alias = Base.make_typealias(type) + base_type, params = if isnothing(alias) + unspecify_type_parameters(type), type_parameters(type) + else + base_type_globalref, params_svec = alias + base_type_globalref.name, params_svec + end + str = string(base_type) + if isempty(params) + return str + end + str *= '{' + param_strings = map(params) do param + param_string = string(param) + if length(param_string) > param_truncation_length + return "…" + end + return param_string + end + str *= join(param_strings, ", ") + str *= '}' + return str +end + +function Base.summary(io::IO, a::AbstractITensor) + print(io, dims_to_string(inds(a))) + print(io, ' ') + print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length = 40)) + return nothing +end + +function Base.show(io::IO, mime::MIME"text/plain", a::AbstractITensor) + summary(io, a) + println(io, ":") + show(io, mime, denamed(a)) + return nothing +end + +function Base.show(io::IO, a::AbstractITensor) + show(io, denamed(a)) + print(io, "[", join(inds(a), ", "), "]") + return nothing end -dimnames(a::ITensor) = LittleSet(a.dimnames) -denamed(a::ITensor) = a.denamed -Base.parent(a::ITensor) = denamed(a) diff --git a/src/abstractnameddimsarray.jl b/src/abstractnameddimsarray.jl deleted file mode 100644 index a91f33f..0000000 --- a/src/abstractnameddimsarray.jl +++ /dev/null @@ -1,877 +0,0 @@ -using LinearAlgebra: LinearAlgebra -using Random: Random -using TensorAlgebra: permuteddims -using TypeParameterAccessors: unspecify_type_parameters - -# Some of the interface is inspired by: -# https://github.com/ITensor/ITensors.jl -# https://github.com/invenia/NamedDims.jl -# https://github.com/mcabbott/NamedPlus.jl -# https://pytorch.org/docs/stable/named_tensor.html - -abstract type AbstractNamedDimsArray{T, N} <: AbstractArray{T, N} end - -const AbstractNamedDimsVector{T} = AbstractNamedDimsArray{T, 1} -const AbstractNamedDimsMatrix{T} = AbstractNamedDimsArray{T, 2} - -dimnames(a::AbstractNamedDimsArray) = throw(MethodError(dimnames, a)) -function dimnames(a::AbstractNamedDimsArray, dim::Int) - return dimnames(a)[dim] -end - -""" - dimnametype(a::AbstractNamedDimsArray) - dimnametype(type::Type{<:AbstractNamedDimsArray}) - -The type of an individual dimension name of `a`. The primary method dispatches -on the array type, and `dimnametype(a)` forwards to `dimnametype(typeof(a))`. - -# Examples - -```jldoctest -julia> a = nameddims(zeros(2, 3), (:i, :j)); - -julia> dimnametype(a) -Symbol - -julia> dimnametype(typeof(a)) -Symbol -``` -""" -function dimnametype end -dimnametype(a::AbstractNamedDimsArray) = dimnametype(typeof(a)) -dimnametype(type::Type{<:AbstractNamedDimsArray}) = throw(MethodError(dimnametype, type)) - -# Unwrapping the names (named-array interface). -# TODO: Use `IsNamed` trait? -denamed(a::AbstractNamedDimsArray) = throw(MethodError(denamed, a)) -denamed(a::AbstractNamedDimsArray, inds) = denamed(aligneddims(a, inds)) -dename(a::AbstractNamedDimsArray, inds) = denamed(aligndims(a, inds)) - -# Output the named axes/indices of the named dims array. -inds(a::AbstractArray) = LittleSet(named.(axes(denamed(a)), dimnames(a))) -inds(a::AbstractArray, dim::Int) = inds(a)[dim] - -isnamed(::Type{<:AbstractNamedDimsArray}) = true - -function dim(a::AbstractArray, n) - return findfirst(==(name(n)), dimnames(a)) -end -dims(a::AbstractArray, ns) = Base.Fix1(dim, a).(ns) - -dimname_isequal(x) = Base.Fix1(dimname_isequal, x) -dimname_isequal(x, y) = isequal(x, y) - -dimname_isequal(r1::AbstractNamedArray, r2::AbstractNamedArray) = isequal(r1, r2) -dimname_isequal(r1::AbstractNamedArray, r2) = name(r1) == r2 -dimname_isequal(r1, r2::AbstractNamedArray) = r1 == name(r2) - -dimname_isequal(r1::AbstractNamedArray, r2::Name) = name(r1) == name(r2) -dimname_isequal(r1::Name, r2::AbstractNamedArray) = name(r1) == name(r2) - -dimname_isequal(r1::AbstractNamedUnitRange, r2::AbstractNamedUnitRange) = isequal(r1, r2) -dimname_isequal(r1::AbstractNamedUnitRange, r2) = name(r1) == r2 -dimname_isequal(r1, r2::AbstractNamedUnitRange) = r1 == name(r2) - -dimname_isequal(r1::AbstractNamedUnitRange, r2::Name) = name(r1) == name(r2) -dimname_isequal(r1::Name, r2::AbstractNamedUnitRange) = name(r1) == name(r2) - -function to_inds(a::AbstractArray, dims) - is = Base.Fix1(dim, a).(name.(dims)) - return Base.Fix1(inds, a).(is) -end - -# Generic construction of named dims arrays. - -""" - nameddims(a::AbstractArray, inds) - -Construct a named dimensions array from an denamed array `a` and named dimensions `inds`. -""" -function nameddims(a::AbstractArray, inds) - return nameddimsconstructor(a, inds)(a, inds) -end - -#= - nameddimsof(a::AbstractNamedDimsArray, b::AbstractArray) - -Construct a named dimensions array with the dimension names of `a` -and based on the type of `a` but with the data from `b`. -=# -function nameddimsof(a::AbstractNamedDimsArray, b::AbstractArray) - return nameddimsconstructorof(a)(b, dimnames(a)) -end - -# Interface inspired by -# [ConstructionBase.constructorof](https://github.com/JuliaObjects/ConstructionBase.jl). -nameddimsconstructorof(a::AbstractNamedDimsArray) = nameddimsconstructorof(typeof(a)) -function nameddimsconstructorof(type::Type{<:AbstractNamedDimsArray}) - return unspecify_type_parameters(type) -end - -# Output a constructor for a named dims array (that should accept and denamed array and -# a set of named dimensions/axes/indices) based on the dimension names. -function nameddimsconstructor(a::AbstractArray, dims) - dimnames = name.(dims) - isempty(dimnames) && return NamedDimsArray - return mapreduce(nameddimsconstructor, combine_nameddimsconstructors, dimnames) -end - -nameddimsconstructor(nameddim) = nameddimsconstructor(typeof(nameddim)) -# Can overload this to get custom named dims array wrapper -# depending on the dimension name types, for example -# output an `ITensor` if the dimension names are `IndexName`s. -nameddimsconstructor(nameddimtype::Type) = NamedDimsArray -function nameddimsconstructor(nameddimtype::Type{<:AbstractNamedUnitRange}) - return nameddimsconstructor(nametype(nameddimtype)) -end -function combine_nameddimsconstructors( - ::Type{<:AbstractNamedDimsArray}, ::Type{<:AbstractNamedDimsArray} - ) - return NamedDimsArray -end -combine_nameddimsconstructors(::Type{T}, ::Type{T}) where {T <: AbstractNamedDimsArray} = T - -# TODO: Move to `utils.jl` file. -# TODO: Use `Base.indexin`? -function getperm(x, y; isequal = isequal) - return map(yᵢ -> findfirst(isequal(yᵢ), x), y) -end - -# TODO: Move to `utils.jl` file. -function checked_indexin(x, y) - I = indexin(x, y) - return something.(I) -end - -function checked_indexin(x::Number, y) - return findfirst(==(x), y) -end - -function checked_indexin(x::AbstractUnitRange, y::AbstractUnitRange) - return findfirst(==(first(x)), y):findfirst(==(last(x)), y) -end - -Base.copy(a::AbstractNamedDimsArray) = nameddimsof(a, copy(denamed(a))) - -# Forward `conj` to the underlying so that graded axes flip their sector arrows. -# The default `AbstractArray` fallback would broadcast `conj` over elements without -# touching the axes, which silently changes the contraction convention for tensors -# with graded (dual-tagged) axes. -Base.conj(a::AbstractNamedDimsArray) = nameddimsof(a, conj(denamed(a))) - -# `LinearAlgebra.normalize` infers result eltype via `typeof(first(a)/nrm)`, which -# scalar-indexes block-structured storage. `a / norm(a, p)` already preserves names. -function LinearAlgebra.normalize(a::AbstractNamedDimsArray, p::Real = 2) - return a / LinearAlgebra.norm(a, p) -end - -# Forward `Random.randn!` / `Random.rand!` to the concrete storage so they -# see the runtime eltype. -function Random.randn!(rng::Random.AbstractRNG, a::AbstractNamedDimsArray) - Random.randn!(rng, denamed(a)) - return a -end -function Random.rand!(rng::Random.AbstractRNG, a::AbstractNamedDimsArray) - Random.rand!(rng, denamed(a)) - return a -end - -function Base.copyto!(a_dest::AbstractNamedDimsArray, a_src::AbstractNamedDimsArray) - a′_dest = denamed(a_dest) - # TODO: Use `denamed` to do the permutations lazily. - a′_src = dename(a_src, inds(a_dest)) - copyto!(a′_dest, a′_src) - return a_dest -end - -# Conversion - -# Copied from `Base` (defined in abstractarray.jl). -@noinline function _checkaxs(axd, axs) - axd == axs || throw(DimensionMismatch("axes must agree, got $axd and $axs")) - return nothing -end -function copyto_axcheck!(dest, src) - _checkaxs(axes(dest), axes(src)) - copyto!(dest, src) - return dest -end - -# These are defined since the Base versions assume the eltype and ndims are known -# at compile time, which isn't true for ITensors. -Base.Array(a::AbstractNamedDimsArray) = Array(denamed(a)) -Base.Array{T}(a::AbstractNamedDimsArray) where {T} = Array{T}(denamed(a)) -Base.Array{T, N}(a::AbstractNamedDimsArray) where {T, N} = Array{T, N}(denamed(a)) -Base.AbstractArray{T}(a::AbstractNamedDimsArray) where {T} = AbstractArray{T, ndims(a)}(a) -function Base.AbstractArray{T, N}(a::AbstractNamedDimsArray) where {T, N} - dest = similar(a, T) - copyto_axcheck!(denamed(dest), denamed(a)) - return dest -end - -function Base.axes(a::AbstractNamedDimsArray) - return inds(a) -end -function Base.size(a::AbstractNamedDimsArray) - return length.(axes(a)) -end - -function Base.length(a::AbstractNamedDimsArray) - return prod(size(a); init = 1) -end - -# Circumvent issue when ndims isn't known at compile time. -Base.axes(a::AbstractNamedDimsArray, d) = axes(a)[d] - -# Circumvent issue when ndims isn't known at compile time. -Base.size(a::AbstractNamedDimsArray, d) = size(a)[d] - -# Circumvent issue when ndims isn't known at compile time. -Base.ndims(a::AbstractNamedDimsArray) = ndims(denamed(a)) - -# Circumvent issue when eltype isn't known at compile time. -Base.eltype(a::AbstractNamedDimsArray) = eltype(denamed(a)) - -using VectorInterface: VectorInterface, scalartype -# Circumvent issue when eltype isn't known at compile time. -VectorInterface.scalartype(a::AbstractNamedDimsArray) = scalartype(denamed(a)) - -Base.axes(a::AbstractNamedDimsArray, dimname::Name) = axes(a, dim(a, dimname)) -Base.size(a::AbstractNamedDimsArray, dimname::Name) = size(a, dim(a, dimname)) - -function similar_nameddims(a::AbstractNamedDimsArray, elt::Type, ax) - return nameddimsconstructorof(a)( - similar(denamed(a), elt, denamed.(Tuple(ax))), - name.(ax) - ) -end -function similar_nameddims(a::AbstractArray, elt::Type, ax) - return nameddims(similar(a, elt, denamed.(Tuple(ax))), name.(ax)) -end - -# Base.similar gets the eltype at compile time. -Base.similar(a::AbstractNamedDimsArray) = similar(a, eltype(a)) -function Base.similar(a::AbstractNamedDimsArray, elt::Type) - return similar_nameddims(a, elt) -end -function similar_nameddims(a::AbstractNamedDimsArray, elt::Type) - return nameddimsof(a, similar(denamed(a), elt)) -end - -# This is defined explicitly since the Base version expects the eltype -# to be known at compile time, which isn't true for ITensors. -function Base.similar( - a::AbstractArray, inds::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} - ) - return similar(a, eltype(a), inds) -end - -function Base.similar( - a::AbstractArray, elt::Type, - inds::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} - ) - return similar_nameddims(a, elt, inds) -end - -function Base.similar(a::AbstractArray, inds::LittleSet) - return similar_nameddims(a, eltype(a), inds) -end - -function Base.similar(a::AbstractArray, elt::Type, inds::LittleSet) - return similar_nameddims(a, elt, inds) -end - -function setinds(a::AbstractNamedDimsArray, inds) - return nameddimsconstructorof(a)(denamed(a), inds) -end - -function replacedimnames(a::AbstractNamedDimsArray, replacements::Pair...) - new_dimnames = replace(dimnames(a), replacements...) - return nameddims(denamed(a), new_dimnames) -end -function replacedimnames(f, a::AbstractNamedDimsArray) - new_dimnames = replace(f, dimnames(a)) - return nameddims(denamed(a), new_dimnames) -end -mapdimnames(f, a::AbstractNamedDimsArray) = replacedimnames(f, a) - -function replaceinds(a::AbstractNamedDimsArray, replacements::Pair...) - new_inds = replace(inds(a), replacements...) - return denamed(a)[new_inds...] -end -function replaceinds(f, a::AbstractNamedDimsArray) - new_inds = replace(f, inds(a)) - return denamed(a)[new_inds...] -end -mapinds(f, a::AbstractNamedDimsArray) = replaceinds(f, a) - -# `Base.isempty(a::AbstractArray)` is defined as `length(a) == 0`, -# which involves comparing a named integer to an denamed integer -# which isn't well defined. -Base.isempty(a::AbstractNamedDimsArray) = isempty(denamed(a)) - -# Define this on objects rather than types in case the wrapper type -# isn't known at compile time, like for the ITensor type. -Base.IndexStyle(a::AbstractNamedDimsArray) = IndexStyle(denamed(a)) -Base.eachindex(a::AbstractNamedDimsArray) = eachindex(denamed(a)) - -# Cartesian indices with names attached. -struct NamedIndexCartesian <: IndexStyle end - -# When multiple named dims arrays are involved, use the named -# dimensions. -function Base.IndexStyle(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) - return NamedIndexCartesian() -end -# Define promotion of index styles. -Base.IndexStyle(s1::NamedIndexCartesian, s2::NamedIndexCartesian) = NamedIndexCartesian() -Base.IndexStyle(s1::IndexStyle, s2::NamedIndexCartesian) = NamedIndexCartesian() -Base.IndexStyle(s1::NamedIndexCartesian, s2::IndexStyle) = NamedIndexCartesian() - -# Like CartesianIndex but with named dimensions. -struct NamedDimsCartesianIndex{N, Index <: Tuple{Vararg{AbstractNamedInteger, N}}} <: - Base.AbstractCartesianIndex{N} - I::Index -end -NamedDimsCartesianIndex(I::AbstractNamedInteger...) = NamedDimsCartesianIndex(I) -Base.Tuple(I::NamedDimsCartesianIndex) = I.I -function Base.show(io::IO, I::NamedDimsCartesianIndex) - print(io, "NamedDimsCartesianIndex") - show(io, Tuple(I)) - return nothing -end - -# Like CartesianIndices but with named dimensions. -struct NamedDimsCartesianIndices{ - N, - Indices <: Tuple{Vararg{AbstractNamedUnitRange, N}}, - Index <: Tuple{Vararg{AbstractNamedInteger, N}}, - } <: AbstractNamedDimsArray{NamedDimsCartesianIndex{N, Index}, N} - indices::Indices - function NamedDimsCartesianIndices(indices::Tuple{Vararg{AbstractNamedUnitRange}}) - return new{length(indices), typeof(indices), Tuple{eltype.(indices)...}}(indices) - end -end -function NamedDimsCartesianIndices(indices::LittleSet) - return NamedDimsCartesianIndices(Tuple(indices)) -end - -Base.eltype(I::NamedDimsCartesianIndices) = eltype(typeof(I)) -Base.axes(I::NamedDimsCartesianIndices) = (only ∘ axes ∘ denamed).(I.indices) -Base.size(I::NamedDimsCartesianIndices) = (length ∘ denamed).(I.indices) - -function Base.getindex(a::NamedDimsCartesianIndices{N}, I::Vararg{Int, N}) where {N} - index = map(a.indices, I) do r, i - return r[i] - end - return NamedDimsCartesianIndex(index) -end - -function denamed(I::NamedDimsCartesianIndices) - return CartesianIndices(denamed.(I.indices)) -end - -function Base.eachindex(::NamedIndexCartesian, a1::AbstractArray, a_rest::AbstractArray...) - all(a -> issetequal(inds(a1), inds(a)), a_rest) || - throw(NameMismatch("Dimension name mismatch $(inds.((a1, a_rest...))).")) - # TODO: Check the shapes match. - return NamedDimsCartesianIndices(inds(a1)) -end - -# Base version ignores dimension names. -# TODO: Use `mapreduce(isequal, &&, a1, a2)`? -function Base.isequal(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) - (inds(a1) == inds(a2)) || return false - return isequal(denamed(a1), denamed(a2, inds(a1))) -end - -# Base version ignores dimension names. -# TODO: Use `mapreduce(==, &&, a1, a2)`? -# TODO: Handle `missing` values properly. -function Base.:(==)(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) - (inds(a1) == inds(a2)) || return false - return denamed(a1) == denamed(a2, inds(a1)) -end - -# Generalization of `Base.sort` to Tuples for Julia v1.10 compatibility. -# TODO: Remove when we drop support for Julia v1.10. -_sort(x; kwargs...) = sort(x; kwargs...) -_sort(x::NTuple{N}; kwargs...) where {N} = NTuple{N}(sort(collect(x); kwargs...)) - -function Base.hash(a::AbstractNamedDimsArray, h::UInt64) - h = hash(:NamedDimsArray, h) - a′ = aligneddims(a, _sort(dimnames(a))) - h = hash(denamed(a′), h) - for i in inds(a′) - h = hash(i, h) - end - return h -end - -# Indexing. - -# Scalar indexing - -Base.firstindex(a::AbstractNamedDimsArray) = firstindex(denamed(a)) -Base.lastindex(a::AbstractNamedDimsArray) = lastindex(denamed(a)) - -function Base.firstindex(a::AbstractNamedDimsArray, d) - return FirstIndex(a, d) -end - -function Base.lastindex(a::AbstractNamedDimsArray, d) - return LastIndex(a, d) -end - -# Redefine generic definition which expects `axes(a)` to -# return a Tuple. -function Base.to_indices(a::AbstractNamedDimsArray, I::Tuple) - return to_indices(a, Tuple(axes(a)), I) -end -# Fix ambiguity error with Base. -function Base.to_indices( - a::AbstractNamedDimsArray, - I::Tuple{Union{Integer, CartesianIndex}} - ) - return to_indices(a, Tuple(axes(a)), I) -end -function Base.checkbounds(::Type{Bool}, a::AbstractNamedDimsArray, I::Int...) - return checkbounds(Bool, denamed(a), I...) -end - -function Base.to_indices( - a::AbstractNamedDimsArray, I::Tuple{AbstractNamedInteger, Vararg{AbstractNamedInteger}} - ) - perm = getperm(name.(I), dimnames(a)) - # TODO: Throw a `NameMismatch` error. - @assert isperm(perm) - I = map(p -> I[p], perm) - return map(inds(a), I) do dimname, i - return checked_indexin(denamed(i), denamed(dimname)) - end -end -function Base.to_indices( - a::AbstractNamedDimsArray, I::Tuple{Pair{<:Any, Int}, Vararg{Pair{<:Any, Int}}} - ) - inds = to_inds(a, first.(I)) - return to_indices(a, map((i, name) -> name[i], last.(I), inds)) -end -function Base.to_indices(a::AbstractNamedDimsArray, I::Tuple{Pair, Vararg{Pair}}) - inds = to_inds(a, first.(I)) - return map((i, name) -> name[i], last.(I), inds) - return to_indices(a, named.(last.(I), first.(I))) -end - -function Base.to_indices(a::AbstractNamedDimsArray, I::Tuple{NamedDimsCartesianIndex}) - return to_indices(a, Tuple(only(I))) -end - -function Base.getindex(a::AbstractNamedDimsArray, I...) - return getindex(a, to_indices(a, I)...) -end - -function Base.getindex(a::AbstractNamedDimsArray, I1::Int, Irest::Int...) - return getindex(denamed(a), I1, Irest...) -end -function Base.getindex( - a::AbstractNamedDimsArray, I1::AbstractNamedInteger, Irest::AbstractNamedInteger... - ) - return getindex(a, to_indices(a, (I1, Irest...))...) -end -function Base.getindex(a::AbstractNamedDimsArray) - return getindex(denamed(a)) -end -# Linear indexing. -function Base.getindex(a::AbstractNamedDimsArray, I::Int) - return getindex(denamed(a), I) -end - -function Base.setindex!(a::AbstractNamedDimsArray, value, I1::Int, Irest::Int...) - setindex!(denamed(a), value, I1, Irest...) - return a -end -function Base.setindex!(a::AbstractNamedDimsArray, value, I::CartesianIndex) - setindex!(a, value, to_indices(a, (I,))...) - return a -end - -function Base.setindex!( - a::AbstractNamedDimsArray, value, I1::AbstractNamedInteger, - Irest::AbstractNamedInteger... - ) - setindex!(a, value, to_indices(a, (I1, Irest...))...) - return a -end -function Base.setindex!(a::AbstractNamedDimsArray, value, I::NamedDimsCartesianIndex) - setindex!(a, value, to_indices(a, (I,))...) - return a -end -function Base.setindex!(a::AbstractNamedDimsArray, value, I1::Pair, Irest::Pair...) - setindex!(a, value, to_indices(a, (I1, Irest...))...) - return a -end -function Base.setindex!(a::AbstractNamedDimsArray, value) - setindex!(denamed(a), value) - return a -end -# Linear indexing. -function Base.setindex!(a::AbstractNamedDimsArray, value, I::Int) - setindex!(denamed(a), value, I) - return a -end - -function Base.isassigned(a::AbstractNamedDimsArray, I::Int...) - return isassigned(denamed(a), I...) -end - -# Slicing - -# Like `const ViewIndex = Union{Real,AbstractArray}`. -const NamedViewIndex = - Union{AbstractNamedInteger, AbstractNamedUnitRange, AbstractNamedArray} - -using ArrayLayouts: ArrayLayouts, MemoryLayout - -abstract type AbstractNamedDimsArrayLayout <: MemoryLayout end -struct NamedDimsArrayLayout{ParentLayout} <: AbstractNamedDimsArrayLayout end - -function ArrayLayouts.MemoryLayout(arrtype::Type{<:AbstractNamedDimsArray}) - return NamedDimsArrayLayout{typeof(MemoryLayout(parenttype(arrtype)))}() -end - -function ArrayLayouts.sub_materialize(::NamedDimsArrayLayout, a, ax) - return copy(a) -end - -# 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 - -function Base.getindex(a::AbstractArray, I1::NamedViewIndex, Irest::NamedViewIndex...) - return copy(view(a, I1, Irest...)) -end -# Disambiguate from `Base.getindex(A::Array, I::AbstractUnitRange{<:Integer})`. -function Base.getindex(a::Array, I1::AbstractNamedUnitRange{<:Integer}) - return copy(view(a, I1)) -end -function Base.view(a::AbstractArray, I1::NamedViewIndex, Irest::NamedViewIndex...) - I = (I1, Irest...) - return nameddims(view(a, denamed.(I)...), name.(I)) -end - -# TODO: Should this be a view? -function Base.getindex(a::AbstractNamedDimsArray, I1::Name, Irest::Name...) - return copy(view(a, I1, Irest...)) -end -function Base.view(a::AbstractNamedDimsArray, I1::Name, Irest::Name...) - issetequal(dimnames(a), name.((I1, Irest...))) || - throw( - NameMismatch( - "Dimension name mismatch $(dimnames(a)), $(name.((I1, Irest...)))." - ) - ) - return a -end - -function Base.getindex(a::AbstractNamedDimsArray, I1::Pair, Irest::Pair...) - return getindex(a, to_indices(a, (I1, Irest...))...) -end -function Base.view(a::AbstractNamedDimsArray, I1::Pair, Irest::Pair...) - I = (I1, Irest...) - inds = to_inds(a, first.(I)) - return view(a, map((i, name) -> name[i], last.(I), inds)...) -end - -function Base.getindex( - a::AbstractNamedDimsArray, I1::NamedViewIndex, Irest::NamedViewIndex... - ) - return copy(view(a, I1, Irest...)) -end -function Base.view(a::AbstractNamedDimsArray, I1::NamedViewIndex, Irest::NamedViewIndex...) - I = (I1, Irest...) - perm = getperm(name.(I), dimnames(a)) - isperm(perm) || throw( - NameMismatch( - "Dimension name mismatch $(dimnames(a)), $(name.(I))." - ) - ) - Ip = map(p -> denamed(I[p]), perm) - return view_nameddims(a, Ip...) -end - -# Repeated definition of `Base.ViewIndex`. -const ViewIndex = Union{Real, AbstractArray} - -# Like `Base.ScalarIndex` but as a trait function. -# This catches cases like `Colon`, `BlockArrays.Block`, etc. which are not AbstractArray -# indices but also aren't scalar indices. -isscalarindex(I) = false -isscalarindex(I::Real) = true - -# Slicing with unnamed indices, such as: -# a = NamedDimsArray(rand(3,4), (:x, :y)) -# b = view(a, 1:2, 2) -function view_nameddims(a::AbstractNamedDimsArray, I...) - nonscalar_dims = filter(dim -> !isscalarindex(I[dim]), ntuple(identity, ndims(a))) - nonscalar_dimnames = map(dim -> dimnames(a, dim), nonscalar_dims) - return nameddimsconstructorof(a)(view(denamed(a), I...), nonscalar_dimnames) -end - -function Base.view(a::AbstractNamedDimsArray, I::ViewIndex...) - return view_nameddims(a, I...) -end - -function getindex_nameddims(a::AbstractArray, I...) - return copy(view(a, I...)) -end - -function Base.getindex(a::AbstractNamedDimsArray, I::ViewIndex...) - return getindex_nameddims(a, I...) -end - -function Base.setindex!( - a::AbstractNamedDimsArray, - value::AbstractNamedDimsArray, - I1::NamedViewIndex, - Irest::NamedViewIndex... - ) - view(a, I1, Irest...) .= value - return a -end -function Base.setindex!( - a::AbstractNamedDimsArray, - value::AbstractArray, - I1::NamedViewIndex, - Irest::NamedViewIndex... - ) - I = (I1, Irest...) - a[I...] = nameddimsconstructorof(a)(value, name.(I)) - return a -end -function Base.setindex!( - a::AbstractNamedDimsArray, - value::AbstractNamedDimsArray, - I1::ViewIndex, - Irest::ViewIndex... - ) - view(a, I1, Irest...) .= value - return a -end -function Base.setindex!( - a::AbstractNamedDimsArray, value::AbstractArray, I1::ViewIndex, Irest::ViewIndex... - ) - setindex!(denamed(a), value, I1, Irest...) - return a -end - -# Permute/align dimensions - -function aligndims(a::AbstractArray, dims) - new_dimnames = name.(dims) - perm = Tuple(getperm(dimnames(a), new_dimnames)) - isperm(perm) || throw( - NameMismatch( - "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." - ) - ) - return nameddimsconstructorof(a)(permutedims(denamed(a), perm), new_dimnames) -end - -function aligneddims(a::AbstractArray, dims) - new_dimnames = name.(dims) - perm = getperm(dimnames(a), new_dimnames) - isperm(perm) || throw( - NameMismatch( - "Dimension name mismatch $(dimnames(a)), $(new_dimnames)." - ) - ) - return nameddimsconstructorof(a)( - permuteddims(denamed(a), perm), new_dimnames - ) -end - -# Convenient constructors - -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 - -default_eltype() = Float64 -for (f, f′) in [(:rand, :_rand), (:randn, :_randn)] - @eval begin - function Base.$f( - rng::AbstractRNG, - elt::Type{<:Number}, - ax::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} - ) - a = $f′(rng, elt, denamed.(ax)) - return a[Name.(name.(ax))...] - end - function Base.$f( - rng::AbstractRNG, - elt::Type{<:Number}, - dims::Tuple{AbstractNamedInteger, Vararg{AbstractNamedInteger}} - ) - return $f(rng, elt, Base.oneto.(dims)) - end - end - for dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] - @eval begin - function Base.$f( - rng::AbstractRNG, elt::Type{<:Number}, dim1::$dimtype, - dims::Vararg{$dimtype} - ) - return $f(rng, elt, (dim1, dims...)) - end - Base.$f(elt::Type{<:Number}, dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f( - Random.default_rng(), elt, dims - ) - Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype}) = $f( - elt, (dim1, dims...) - ) - Base.$f(dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(default_eltype(), dims) - Base.$f(dim1::$dimtype, dims::Vararg{$dimtype}) = $f((dim1, dims...)) - end - end -end -for f in [:zeros, :ones], dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] - @eval begin - function Base.$f( - elt::Type{<:Number}, ax::Tuple{$dimtype, Vararg{$dimtype}} - ) - a = $f(elt, denamed.(ax)) - return a[Name.(name.(ax))...] - end - function Base.$f(elt::Type{<:Number}, dim1::$dimtype, dims::Vararg{$dimtype}) - return $f(elt, (dim1, dims...)) - end - Base.$f(dims::Tuple{$dimtype, Vararg{$dimtype}}) = $f(default_eltype(), dims) - Base.$f(dim1::$dimtype, dims::Vararg{$dimtype}) = $f((dim1, dims...)) - end -end -for dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange] - @eval begin - function Base.fill(value, ax::Tuple{$dimtype, Vararg{$dimtype}}) - a = fill(value, denamed.(ax)) - return a[Name.(name.(ax))...] - end - function Base.fill(value, dim1::$dimtype, dims::Vararg{$dimtype}) - return fill(value, (dim1, dims...)) - end - end -end - -function Base.fill!(a::AbstractNamedDimsArray, v) - fill!(denamed(a), v) - return a -end - -function Base.map!(f, a_dest::AbstractNamedDimsArray, a_srcs::AbstractNamedDimsArray...) - a′_dest = denamed(a_dest) - # TODO: Use `denamed` to do the permutations lazily. - # TODO: Define `dename[d](dimnames) = Base.Fix1(dename[d], dimnames)` and use it here? - a′_srcs = Base.Fix2(dename, dimnames(a_dest)).(a_srcs) - map!(f, a′_dest, a′_srcs...) - return a_dest -end - -function Base.map(f, a_srcs::AbstractNamedDimsArray...) - # copy(mapped(f, a_srcs...)) - return f.(a_srcs...) -end - -function Base.mapreduce(f, op, a::AbstractNamedDimsArray; kwargs...) - return mapreduce(f, op, denamed(a); kwargs...) -end - -# `sum` is routed to the underlying data rather than left to fall back on the -# `mapreduce` method above because some array types (such as graded arrays) define -# `Base.sum` directly but not the general `mapreduce`, so the denamed `sum` is the -# path that works for them. -function Base.sum(a::AbstractNamedDimsArray; kwargs...) - return sum(denamed(a); kwargs...) -end - -function LinearAlgebra.promote_leaf_eltypes(a::AbstractNamedDimsArray) - return LinearAlgebra.promote_leaf_eltypes(denamed(a)) -end - -# Printing - -# Copy of `Base.dims2string` defined in `show.jl`. -function dims_to_string(d) - isempty(d) && return "0-dimensional" - length(d) == 1 && return "$(d[1])-element" - return join(map(string, d), '×') -end - -using TypeParameterAccessors: type_parameters, unspecify_type_parameters -function concretetype_to_string_truncated( - type::Type; - param_truncation_length = typemax(Int) - ) - isconcretetype(type) || throw(ArgumentError("Type must be concrete.")) - alias = Base.make_typealias(type) - base_type, params = if isnothing(alias) - unspecify_type_parameters(type), type_parameters(type) - else - base_type_globalref, params_svec = alias - base_type_globalref.name, params_svec - end - str = string(base_type) - if isempty(params) - return str - end - str *= '{' - param_strings = map(params) do param - param_string = string(param) - if length(param_string) > param_truncation_length - return "…" - end - return param_string - end - str *= join(param_strings, ", ") - str *= '}' - return str -end - -function Base.summary(io::IO, a::AbstractNamedDimsArray) - print(io, dims_to_string(inds(a))) - print(io, ' ') - print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length = 40)) - return nothing -end - -function Base.show(io::IO, mime::MIME"text/plain", a::AbstractNamedDimsArray) - summary(io, a) - println(io, ":") - show(io, mime, denamed(a)) - return nothing -end - -function Base.show(io::IO, a::AbstractNamedDimsArray) - show(io, denamed(a)) - print(io, "[", join(inds(a), ", "), "]") - return nothing -end diff --git a/src/broadcast.jl b/src/broadcast.jl index 6b1feb9..4743a05 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -1,33 +1,33 @@ -using ..ITensorBase: AbstractNamedDimsArray, AbstractNamedUnitRange, ITensorBase, LittleSet, +using ..ITensorBase: AbstractITensor, AbstractNamedUnitRange, ITensorBase, LittleSet, dename, denamed, getperm, inds, name, named, nameddimsconstructorof using Base.Broadcast: Broadcast as BC, Broadcasted, broadcast_shape, broadcasted, check_broadcast_shape, combine_axes using TensorAlgebra: TensorAlgebra as TA -abstract type AbstractNamedDimsArrayStyle{N} <: BC.AbstractArrayStyle{N} end +abstract type AbstractITensorStyle{N} <: BC.AbstractArrayStyle{N} end -struct NamedDimsArrayStyle{N, NDA} <: AbstractNamedDimsArrayStyle{N} end -NamedDimsArrayStyle(::Val{N}) where {N} = NamedDimsArrayStyle{N, NamedDimsArray}() -NamedDimsArrayStyle{M}(::Val{N}) where {M, N} = NamedDimsArrayStyle{N, NamedDimsArray}() -NamedDimsArrayStyle{M, NDA}(::Val{N}) where {M, N, NDA} = NamedDimsArrayStyle{N, NDA}() +struct ITensorStyle{N, NDA} <: AbstractITensorStyle{N} end +ITensorStyle(::Val{N}) where {N} = ITensorStyle{N, ITensor}() +ITensorStyle{M}(::Val{N}) where {M, N} = ITensorStyle{N, ITensor}() +ITensorStyle{M, NDA}(::Val{N}) where {M, N, NDA} = ITensorStyle{N, NDA}() -function nameddimstype(style::NamedDimsArrayStyle{<:Any, NDA}) where {NDA} +function nameddimstype(style::ITensorStyle{<:Any, NDA}) where {NDA} return NDA end -function BC.BroadcastStyle(arraytype::Type{<:AbstractNamedDimsArray}) - return NamedDimsArrayStyle{ndims(arraytype), nameddimsconstructorof(arraytype)}() +function BC.BroadcastStyle(arraytype::Type{<:AbstractITensor}) + return ITensorStyle{ndims(arraytype), nameddimsconstructorof(arraytype)}() end function BC.combine_axes( - a1::AbstractNamedDimsArray, a_rest::AbstractNamedDimsArray... + a1::AbstractITensor, a_rest::AbstractITensor... ) return broadcast_shape(axes(a1), combine_axes(a_rest...)) end -function BC.combine_axes(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) +function BC.combine_axes(a1::AbstractITensor, a2::AbstractITensor) return broadcast_shape(axes(a1), axes(a2)) end -BC.combine_axes(a::AbstractNamedDimsArray) = axes(a) +BC.combine_axes(a::AbstractITensor) = axes(a) function BC.broadcast_shape( ax1::LittleSet, ax2::LittleSet, ax_rest::LittleSet... @@ -62,7 +62,7 @@ function set_promote_shape( end # Handle operations like `randn() + randn(2, 2)[i, j]``. -# TODO: Decide if this should be a general definition for `AbstractNamedDimsArray`, +# TODO: Decide if this should be a general definition for `AbstractITensor`, # or just for `AbstractITensor`. function set_promote_shape( ax1::Tuple{}, ax2::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}} @@ -71,7 +71,7 @@ function set_promote_shape( end # Handle operations like `randn(2, 2)[i, j] + randn()`. -# TODO: Decide if this should be a general definition for `AbstractNamedDimsArray`, +# TODO: Decide if this should be a general definition for `AbstractITensor`, # or just for `AbstractITensor`. function set_promote_shape( ax1::Tuple{AbstractNamedUnitRange, Vararg{AbstractNamedUnitRange}}, ax2::Tuple{} @@ -104,11 +104,11 @@ end # result inherits the operands' backend (e.g. graded) rather than a lazy permuted # wrapper's `similar` (which can drop the backend). denamed_prototype(bc::Broadcasted) = denamed_prototype(bc.args...) -denamed_prototype(arg::AbstractNamedDimsArray, args...) = denamed(arg) +denamed_prototype(arg::AbstractITensor, args...) = denamed(arg) denamed_prototype(arg::Broadcasted, args...) = denamed_prototype(arg.args..., args...) denamed_prototype(arg, args...) = denamed_prototype(args...) -function Base.similar(bc::Broadcasted{<:AbstractNamedDimsArrayStyle}, elt::Type, ax) +function Base.similar(bc::Broadcasted{<:AbstractITensorStyle}, elt::Type, ax) inds_a = name.(ax) bc_denamed = broadcasted_denamed(bc, inds_a) a_denamed = similar(bc_denamed, elt) @@ -116,7 +116,7 @@ function Base.similar(bc::Broadcasted{<:AbstractNamedDimsArrayStyle}, elt::Type, end inds(bc::Broadcasted) = name.(axes(bc)) -function Base.copy(bc::Broadcasted{<:AbstractNamedDimsArrayStyle}) +function Base.copy(bc::Broadcasted{<:AbstractITensorStyle}) # We could use: # ```julia # elt = combine_eltypes(bc.f, bc.args) @@ -142,7 +142,7 @@ function Base.copy(bc::Broadcasted{<:AbstractNamedDimsArrayStyle}) return nameddimstype(bc.style)(dest_denamed, inds_dest) end -function Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractNamedDimsArrayStyle}) +function Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractITensorStyle}) dest_denamed = denamed(dest) inds_dest = inds(dest) bc_denamed = broadcasted_denamed(bc, inds_dest) diff --git a/src/itensor.jl b/src/itensor.jl new file mode 100644 index 0000000..e21cc34 --- /dev/null +++ b/src/itensor.jl @@ -0,0 +1,33 @@ +using TypeParameterAccessors: TypeParameterAccessors, parenttype + +# TODO: Check `allunique(dimnames)`? +struct ITensor{DimName} <: AbstractITensor{DimName} + denamed::AbstractArray + dimnames::Vector{DimName} +end + +# TODO: Check `allunique(dimnames)`? +function ITensor(denamed::AbstractArray, dims) + ndims(denamed) == length(dims) || + throw(ArgumentError("Number of named dims must match ndims.")) + dimnames = collect(dims) + return ITensor{eltype(dimnames)}(denamed, dimnames) +end +ITensor(a::AbstractITensor, inds) = throw(ArgumentError("Already named.")) +ITensor(a::AbstractITensor) = ITensor(denamed(a), dimnames(a)) + +# Minimal interface. The dimnames are stored as a `Vector{DimName}`, but the +# accessor still returns a `LittleSet` over a `Tuple` (unchanged public behavior). +dimnames(a::ITensor) = LittleSet(Tuple(a.dimnames)) +denamed(a::ITensor) = a.denamed +Base.parent(a::ITensor) = denamed(a) + +dimnametype(::Type{<:ITensor{DimName}}) where {DimName} = DimName + +# The parent array is erased at the field level, so its concrete type is not part +# of `ITensor`'s signature. An instance still carries the parent, so the instance +# methods recover the concrete type while the type methods report `AbstractArray`. +denamedtype(a::ITensor) = typeof(denamed(a)) +denamedtype(::Type{<:ITensor}) = AbstractArray +TypeParameterAccessors.parenttype(a::ITensor) = typeof(parent(a)) +TypeParameterAccessors.parenttype(::Type{<:ITensor}) = AbstractArray diff --git a/src/nameddimsoperator.jl b/src/itensoroperator.jl similarity index 79% rename from src/nameddimsoperator.jl rename to src/itensoroperator.jl index bd3c9b7..a9af8ef 100644 --- a/src/nameddimsoperator.jl +++ b/src/itensoroperator.jl @@ -21,14 +21,14 @@ get_codomain_name(a, i) = throw(MethodError(get_codomain_name, (a, i))) # If it doesn't exist, return the index itself. get_domain_name(a, i) = throw(MethodError(get_domain_name, (a, i))) -function apply(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) +function apply(x::AbstractITensor, y::AbstractITensor) xy = x * y return mapdimnames(xy) do i return get_domain_name(x, i) end end -function apply_dag(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) +function apply_dag(x::AbstractITensor, y::AbstractITensor) xy = x * y return mapdimnames(xy) do i return get_codomain_name(y, i) @@ -37,7 +37,7 @@ end # TODO: Define versions that accept codomain and domain names, # i.e. `transpose(a, codomain, domain)` and `adjoint(a, codomain, domain)` (?). -function Base.transpose(a::AbstractNamedDimsArray) +function Base.transpose(a::AbstractITensor) c = codomainnames(a) d = domainnames(a) a_map = merge(Dict(c .=> d), Dict(d .=> c)) @@ -46,11 +46,11 @@ function Base.transpose(a::AbstractNamedDimsArray) end return operator(a′, c, d) end -function Base.adjoint(a::AbstractNamedDimsArray) +function Base.adjoint(a::AbstractITensor) return transpose(conj(a)) end -function product(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) +function product(x::AbstractITensor, y::AbstractITensor) c = codomainnames(x) d = domainnames(x) c′ = randname.(c) @@ -97,25 +97,25 @@ Base.iterate(b::Bijection) = iterate(b.domain_to_codomain) Base.iterate(b::Bijection, state) = iterate(b.domain_to_codomain, state) Base.length(b::Bijection) = length(b.domain_to_codomain) -abstract type AbstractNamedDimsOperator{T, N} <: AbstractNamedDimsArray{T, N} end +abstract type AbstractITensorOperator{DimName} <: AbstractITensor{DimName} end -state(a::AbstractNamedDimsArray) = a -dimnames(a::AbstractNamedDimsOperator) = dimnames(state(a)) +state(a::AbstractITensor) = a +dimnames(a::AbstractITensorOperator) = dimnames(state(a)) # TODO: Unify these two functions. function operator(a::AbstractArray, codomain, domain) na = nameddims(a, (codomain..., domain...)) return operator(na, codomain, domain) end -function operator(a::AbstractNamedDimsArray, codomain, domain) - return NamedDimsOperator(a, codomain, domain) +function operator(a::AbstractITensor, codomain, domain) + return ITensorOperator(a, codomain, domain) end -# This helps preserve the NamedDimsArray type when multiplying, -# for example when a NamedDimsOperator wraps an ITensor. -Base.:*(a::AbstractNamedDimsOperator, b::AbstractNamedDimsOperator) = state(a) * state(b) -Base.:*(a::AbstractNamedDimsOperator, b::AbstractNamedDimsArray) = state(a) * state(b) -Base.:*(a::AbstractNamedDimsArray, b::AbstractNamedDimsOperator) = state(a) * state(b) +# This helps preserve the ITensor type when multiplying, +# for example when a ITensorOperator wraps an ITensor. +Base.:*(a::AbstractITensorOperator, b::AbstractITensorOperator) = state(a) * state(b) +Base.:*(a::AbstractITensorOperator, b::AbstractITensor) = state(a) * state(b) +Base.:*(a::AbstractITensor, b::AbstractITensorOperator) = state(a) * state(b) # Broadcasting peels an operator to its `state` before the broadcast style is # computed. This covers the array operations that lower to broadcasting too (`+`, @@ -125,11 +125,11 @@ Base.:*(a::AbstractNamedDimsArray, b::AbstractNamedDimsOperator) = state(a) * st # `*`; keeping it an operator (carrying the codomain/domain bijection) is future # work that first needs the mixed cases settled (`state + operator`, and # `operator + operator` where names match but the codomain/domain split does not). -Base.broadcastable(a::AbstractNamedDimsOperator) = state(a) +Base.broadcastable(a::AbstractITensorOperator) = state(a) for f in MATRIX_FUNCTIONS @eval begin - function Base.$f(a::AbstractNamedDimsOperator) + function Base.$f(a::AbstractITensorOperator) c = codomainnames(a) d = domainnames(a) return operator($f(state(a), c, d), c, d) @@ -138,7 +138,7 @@ for f in MATRIX_FUNCTIONS end # Operator entries for the gram factorizations defined in `tensoralgebra.jl`. -# Placed here because `AbstractNamedDimsOperator` is defined below +# Placed here because `AbstractITensorOperator` is defined below # `tensoralgebra.jl` in the include order. # # Per-method docstrings are factored out into `const` strings and attached @@ -147,7 +147,7 @@ end # don't share enough structure to warrant `$($f)`-interpolation. const _gram_eigh_full_operator_docstring = """ - TensorAlgebra.gram_eigh_full(a::AbstractNamedDimsOperator; kwargs...) -> x + TensorAlgebra.gram_eigh_full(a::AbstractITensorOperator; kwargs...) -> x Gram factorization of a Hermitian positive semi-definite named operator `a`, returning `x` such that `x * x_cod ≈ state(a)`, where `x_cod` is @@ -180,7 +180,7 @@ true """ const _gram_eigh_full_with_pinv_operator_docstring = """ - TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractNamedDimsOperator; kwargs...) -> x, y + TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractITensorOperator; kwargs...) -> x, y Like `TensorAlgebra.gram_eigh_full`, but additionally returns a named array `y` that is a left inverse of `x`: `y * x ≈ I` on the @@ -216,14 +216,14 @@ true for f in (:gram_eigh_full, :gram_eigh_full_with_pinv) doc_sym = Symbol("_", f, "_operator_docstring") @eval begin - @doc $doc_sym function TA.$f(a::AbstractNamedDimsOperator; kwargs...) + @doc $doc_sym function TA.$f(a::AbstractITensorOperator; kwargs...) return TA.$f(state(a), codomainnames(a), domainnames(a); kwargs...) end end end """ - Base.one(op::AbstractNamedDimsOperator) -> Id + Base.one(op::AbstractITensorOperator) -> Id Return the identity operator with the same codomain/domain names and shape as `op`. `op` is treated as a shape prototype and is not mutated. @@ -249,7 +249,7 @@ julia> apply(Id, v) ≈ v true ``` """ -function Base.one(op::AbstractNamedDimsOperator) +function Base.one(op::AbstractITensorOperator) co, dom = codomainnames(op), domainnames(op) return operator(one(state(op), co, dom), co, dom) end @@ -308,52 +308,52 @@ function similar_operator(prototype, named_domain_axes) end # Forward `Random.randn!` / `Random.rand!` to the operator's state, which -# itself peels to the concrete storage via the generic AbstractNamedDimsArray +# itself peels to the concrete storage via the generic AbstractITensor # method. -function Random.randn!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator) +function Random.randn!(rng::Random.AbstractRNG, op::AbstractITensorOperator) Random.randn!(rng, state(op)) return op end -function Random.rand!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator) +function Random.rand!(rng::Random.AbstractRNG, op::AbstractITensorOperator) Random.rand!(rng, state(op)) return op end -struct NamedDimsOperator{T, N, P <: AbstractNamedDimsArray{T, N}, D, C} <: - AbstractNamedDimsOperator{T, N} +struct ITensorOperator{DimName, P <: AbstractITensor{DimName}, D, C} <: + AbstractITensorOperator{DimName} parent::P dimnames_bijection::Bijection{D, C} end -state(a::NamedDimsOperator) = a.parent -Base.parent(a::NamedDimsOperator) = state(a) -denamed(a::NamedDimsOperator) = denamed(state(a)) +state(a::ITensorOperator) = a.parent +Base.parent(a::ITensorOperator) = state(a) +denamed(a::ITensorOperator) = denamed(state(a)) -function NamedDimsOperator(a::AbstractNamedDimsArray, codomainnames, domainnames) - return NamedDimsOperator(a, Bijection(domainnames, codomainnames)) +function ITensorOperator(a::AbstractITensor, codomainnames, domainnames) + return ITensorOperator(a, Bijection(domainnames, codomainnames)) end using TypeParameterAccessors: TypeParameterAccessors, parenttype -function TypeParameterAccessors.parenttype(type::Type{<:NamedDimsOperator}) +function TypeParameterAccessors.parenttype(type::Type{<:ITensorOperator}) return fieldtype(type, :parent) end -statetype(type::Type{<:NamedDimsOperator}) = parenttype(type) +statetype(type::Type{<:ITensorOperator}) = parenttype(type) -function nameddimsof(a::NamedDimsOperator, b::AbstractArray) - return NamedDimsOperator(nameddimsof(state(a), b), a.dimnames_bijection) +function nameddimsof(a::ITensorOperator, b::AbstractArray) + return ITensorOperator(nameddimsof(state(a), b), a.dimnames_bijection) end -function nameddimsconstructorof(type::Type{<:NamedDimsOperator}) +function nameddimsconstructorof(type::Type{<:ITensorOperator}) return nameddimsconstructorof(statetype(type)) end -codomainnames(a::NamedDimsOperator) = codomain(a.dimnames_bijection) -domainnames(a::NamedDimsOperator) = domain(a.dimnames_bijection) +codomainnames(a::ITensorOperator) = codomain(a.dimnames_bijection) +domainnames(a::ITensorOperator) = domain(a.dimnames_bijection) -function get_codomain_name(a::NamedDimsOperator, i) +function get_codomain_name(a::ITensorOperator, i) return get(a.dimnames_bijection, i, i) end -function get_domain_name(a::NamedDimsOperator, i) +function get_domain_name(a::ITensorOperator, i) return get(inverse(a.dimnames_bijection), i, i) end diff --git a/src/linearalgebra.jl b/src/linearalgebra.jl index d36464b..7584e5f 100644 --- a/src/linearalgebra.jl +++ b/src/linearalgebra.jl @@ -6,7 +6,7 @@ using LinearAlgebra: LinearAlgebra as LA # which isn't friendly for named arrays wrapping GPU arrays. # This implicitly helps with defining `LA.normalize[!]` as well (though note that # uses `LinearAlgebra.rmul!` as well). -function LA.norm(a::AbstractNamedDimsArray, p::Real = 2; kwargs...) +function LA.norm(a::AbstractITensor, p::Real = 2; kwargs...) return LA.norm(denamed(a), p; kwargs...) end @@ -18,11 +18,11 @@ for f! in [:mul!, :div!] lf! = Symbol(:l, f!) rf! = Symbol(:r, f!) @eval begin - function LA.$rf!(a::AbstractNamedDimsArray, α::Number) + function LA.$rf!(a::AbstractITensor, α::Number) LA.$rf!(denamed(a), α) return a end - function LA.$lf!(α::Number, a::AbstractNamedDimsArray) + function LA.$lf!(α::Number, a::AbstractITensor) LA.$lf!(α, denamed(a)) return a end @@ -33,6 +33,6 @@ end # uses scalar indexing: # https://github.com/JuliaLang/LinearAlgebra.jl/blob/3a4fdad7f608928ecb4b41e76b1e9ecacd058444/src/generic.jl#L919-L1009 # which isn't friendly for named arrays wrapping GPU arrays. -function LA.dot(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) +function LA.dot(a1::AbstractITensor, a2::AbstractITensor) return (conj(a1) * a2)[] end diff --git a/src/littleset.jl b/src/littleset.jl index 1db3759..c1da480 100644 --- a/src/littleset.jl +++ b/src/littleset.jl @@ -30,8 +30,8 @@ Base.Broadcast.BroadcastStyle(s1::AbstractArrayStyle, s2::Style{LittleSet}) = s1 Base.Broadcast.broadcastable(s::LittleSet) = s Base.to_shape(s::LittleSet) = s -# Needed for functionality such as `CartesianIndices(::AbstractNamedDimsArray)`, -# `pairs(::AbstractNamedDimsArray)`, etc. +# Needed for functionality such as `CartesianIndices(::AbstractITensor)`, +# `pairs(::AbstractITensor)`, etc. Base.CartesianIndices(s::LittleSet) = CartesianIndices(s.values) function Base.copy( diff --git a/src/nameddimsarray.jl b/src/nameddimsarray.jl deleted file mode 100644 index 5b99dae..0000000 --- a/src/nameddimsarray.jl +++ /dev/null @@ -1,45 +0,0 @@ -using TypeParameterAccessors: - TypeParameterAccessors, Position, get_type_parameters, parenttype - -# TODO: Check `allunique(dimnames)`? -struct NamedDimsArray{ - T, - N, - Denamed <: AbstractArray{T, N}, - DimNames <: Tuple{Vararg{Any, N}}, - } <: - AbstractNamedDimsArray{T, N} - denamed::Denamed - dimnames::DimNames -end - -const NamedDimsVector{T, Denamed <: AbstractVector{T}, Inds <: Tuple{Any}} = NamedDimsArray{ - T, 1, Denamed, Inds, -} -const NamedDimsMatrix{T, Denamed <: AbstractMatrix{T}, Inds <: Tuple{Any, Any}} = - NamedDimsArray{ - T, 2, Denamed, Inds, -} - -# TODO: Check `allunique(dimnames)`? -function NamedDimsArray(denamed::AbstractArray, dims) - ndims(denamed) == length(dims) || - throw(ArgumentError("Number of named dims must match ndims.")) - return NamedDimsArray(denamed, Tuple{Vararg{Any, ndims(denamed)}}(dims)) -end -NamedDimsArray(a::AbstractNamedDimsArray, inds) = throw(ArgumentError("Already named.")) -NamedDimsArray(a::AbstractNamedDimsArray) = NamedDimsArray(denamed(a), dimnames(a)) - -# Minimal interface. -dimnames(a::NamedDimsArray) = LittleSet(a.dimnames) -denamed(a::NamedDimsArray) = a.denamed -Base.parent(a::NamedDimsArray) = denamed(a) - -denamedtype(T::Type{<:NamedDimsArray}) = get_type_parameters(T, Position(3)) -dimnametype(T::Type{<:NamedDimsArray}) = eltype(get_type_parameters(T, Position(4))) - -function TypeParameterAccessors.position( - ::Type{<:AbstractNamedDimsArray}, ::typeof(parenttype) - ) - return Position(3) -end diff --git a/src/quirks.jl b/src/quirks.jl index 294fe3f..3094ac2 100644 --- a/src/quirks.jl +++ b/src/quirks.jl @@ -1,39 +1,3 @@ # This seems to be needed to get broadcasting working. # TODO: Investigate this and see if we can get rid of it. Base.Broadcast.extrude(a::AbstractITensor) = a - -# See: https://github.com/JuliaLang/julia/blob/v1.11.4/base/namedtuple.jl#L269 -# `filter(f, ::NamedTuple)` is available in Julia v1.11, delete once -# we drop support for Julia v1.10. -filter_namedtuple(f, xs::NamedTuple) = xs[filter(k -> f(xs[k]), keys(xs))] - -function translate_factorize_kwargs(; - # MatrixAlgebraKit.jl/TensorAlgebra.jl kwargs. - orth = nothing, - rtol = nothing, - maxrank = nothing, - # ITensors.jl kwargs. - ortho = nothing, - cutoff = nothing, - maxdim = nothing, - kwargs... - ) - orth = Symbol(@something orth ortho :left) - rtol = @something rtol cutoff Some(nothing) - maxrank = @something maxrank maxdim Some(nothing) - trunc = (; rtol, maxrank) - # !isnothing(maxrank) && error("`maxrank` not supported yet.") - return filter_namedtuple(!isnothing, (; orth, trunc, kwargs...)) -end - -using TensorAlgebra: TensorAlgebra, factorize -function TensorAlgebra.factorize(a::AbstractITensor, codomain_inds, domain_inds; kwargs...) - return invoke( - factorize, - Tuple{AbstractNamedDimsArray, Any, Any}, - a, - codomain_inds, - domain_inds; - translate_factorize_kwargs(; kwargs...)... - ) -end diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 64200d6..4a4a9fb 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -5,7 +5,7 @@ using TupleTools: TupleTools # This layer is used to define derivative rules (to skip differentiating `setdiff`). dimnames_setdiff(s1, s2) = setdiff(s1, s2) -Base.:*(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) = mul_nameddims(a1, a2) +Base.:*(a1::AbstractITensor, a2::AbstractITensor) = mul_nameddims(a1, a2) function mul_nameddims(a1::AbstractArray, a2::AbstractArray) a_dest, dimnames_dest = TA.contract( denamed(a1), dimnames(a1), denamed(a2), dimnames(a2) @@ -24,8 +24,8 @@ end # ``` # that optimize matrix multiplication sequence. function Base.:*( - a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray, - a3::AbstractNamedDimsArray, a_rest::AbstractNamedDimsArray... + a1::AbstractITensor, a2::AbstractITensor, + a3::AbstractITensor, a_rest::AbstractITensor... ) return mul_nameddims(a1, a2, a3, a_rest...) end @@ -37,8 +37,8 @@ function mul_nameddims( end function LA.mul!( - a_dest::AbstractNamedDimsArray, - a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray, + a_dest::AbstractITensor, + a1::AbstractITensor, a2::AbstractITensor, α::Number, β::Number ) return mul!_nameddims(a_dest, a1, a2, α, β) @@ -58,8 +58,8 @@ function mul!_nameddims( end function LA.mul!( - a_dest::AbstractNamedDimsArray, - a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray + a_dest::AbstractITensor, + a1::AbstractITensor, a2::AbstractITensor ) return mul!_nameddims(a_dest, a1, a2) end @@ -75,7 +75,7 @@ function mul!_nameddims( return a_dest end -function TA.blockedperm(na::AbstractNamedDimsArray, nameddim_blocks::Tuple...) +function TA.blockedperm(na::AbstractITensor, nameddim_blocks::Tuple...) return blockedperm_nameddims(na, nameddim_blocks...) end function blockedperm_nameddims(a::AbstractArray, nameddim_blocks::Tuple...) @@ -92,7 +92,7 @@ end # matricize(a, (i, k) => "a") # matricize(a, (i, k) => "a", (j, l) => "b") # TODO: Rewrite in terms of `matricize(a, .., (1, 3))` interface. -function TA.matricize(a::AbstractNamedDimsArray, fusions::Vararg{Pair, 2}) +function TA.matricize(a::AbstractITensor, fusions::Vararg{Pair, 2}) return matricize_nameddims(a, fusions...) end function matricize_nameddims(na::AbstractArray, fusions::Vararg{Pair, 2}) @@ -113,7 +113,7 @@ function matricize_nameddims(na::AbstractArray, fusions::Vararg{Pair, 2}) return nameddims(a_fused, dimnames_fused) end -function TA.unmatricize(na::AbstractNamedDimsArray, splitters::Vararg{Pair, 2}) +function TA.unmatricize(na::AbstractITensor, splitters::Vararg{Pair, 2}) return unmatricize_nameddims(na, splitters...) end function unmatricize_nameddims(na::AbstractArray, splitters::Vararg{Pair, 2}) @@ -138,6 +138,39 @@ function unmatricize_nameddims(na::AbstractArray, splitters::Vararg{Pair, 2}) return nameddims(a_split, names_split) end +# See: https://github.com/JuliaLang/julia/blob/v1.11.4/base/namedtuple.jl#L269 +# `filter(f, ::NamedTuple)` is available in Julia v1.11, delete once +# we drop support for Julia v1.10. +filter_namedtuple(f, xs::NamedTuple) = xs[filter(k -> f(xs[k]), keys(xs))] + +# `factorize` additionally accepts the legacy ITensors.jl keyword names +# (`ortho` / `cutoff` / `maxdim`) and maps them to the +# MatrixAlgebraKit.jl / TensorAlgebra.jl names (`orth` / `rtol` / `maxrank`). +function translate_factorize_kwargs(; + # MatrixAlgebraKit.jl/TensorAlgebra.jl kwargs. + orth = nothing, + rtol = nothing, + maxrank = nothing, + # ITensors.jl kwargs. + ortho = nothing, + cutoff = nothing, + maxdim = nothing, + kwargs... + ) + orth = Symbol(@something orth ortho :left) + rtol = @something rtol cutoff Some(nothing) + maxrank = @something maxrank maxdim Some(nothing) + trunc = (; rtol, maxrank) + return filter_namedtuple(!isnothing, (; orth, trunc, kwargs...)) +end + +# Other factorizations pass their keywords through unchanged; only `factorize` +# does the ITensors.jl keyword translation. +translate_factorization_kwargs(f, kwargs) = kwargs +function translate_factorization_kwargs(::typeof(TA.factorize), kwargs) + return translate_factorize_kwargs(; kwargs...) +end + for f in [ :factorize, :left_orth, :left_polar, :lq, :orth, :polar, :qr, :right_orth, :right_polar, @@ -145,9 +178,12 @@ for f in [ f_nameddims = Symbol(f, "_nameddims") @eval begin function TA.$f( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... + ) + return $f_nameddims( + a, dimnames_codomain, dimnames_domain; + translate_factorization_kwargs(TA.$f, kwargs)... ) - return $f_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end function $f_nameddims( a::AbstractArray, dimnames_codomain, dimnames_domain; kwargs... @@ -164,8 +200,10 @@ for f in [ y = nameddims(y_denamed, dimnames_y) return x, y end - function TA.$f(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) - return $f_nameddims(a, dimnames_codomain; kwargs...) + function TA.$f(a::AbstractITensor, dimnames_codomain; kwargs...) + return $f_nameddims( + a, dimnames_codomain; translate_factorization_kwargs(TA.$f, kwargs)... + ) end function $f_nameddims(a::AbstractArray, dimnames_codomain; kwargs...) codomain = name.(dimnames_codomain) @@ -176,13 +214,13 @@ for f in [ end # Overload LinearAlgebra functions where relevant. -function LA.qr(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.qr(a::AbstractITensor, args...; kwargs...) return TA.qr(a, args...; kwargs...) end -function LA.lq(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.lq(a::AbstractITensor, args...; kwargs...) return TA.lq(a, args...; kwargs...) end -function LA.factorize(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.factorize(a::AbstractITensor, args...; kwargs...) return TA.factorize(a, args...; kwargs...) end @@ -191,7 +229,7 @@ end # function TA.svd( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return svd_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -213,10 +251,10 @@ function svd_nameddims( v = nameddims(v_denamed, dimnames_v) return u, s, v end -function TA.svd(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function TA.svd(a::AbstractITensor, dimnames_codomain; kwargs...) return svd_nameddims(a, dimnames_codomain; kwargs...) end -function svd_nameddims(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function svd_nameddims(a::AbstractITensor, dimnames_codomain; kwargs...) return TA.svd( a, dimnames_codomain, @@ -224,12 +262,12 @@ function svd_nameddims(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) kwargs... ) end -function LA.svd(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.svd(a::AbstractITensor, args...; kwargs...) return TA.svd(a, args...; kwargs...) end function TA.svdvals( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return svdvals_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -245,7 +283,7 @@ function svdvals_nameddims( ) end -function TA.svdvals(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function TA.svdvals(a::AbstractITensor, dimnames_codomain; kwargs...) return svdvals_nameddims(a, dimnames_codomain; kwargs...) end function svdvals_nameddims(a::AbstractArray, dimnames_codomain; kwargs...) @@ -254,12 +292,12 @@ function svdvals_nameddims(a::AbstractArray, dimnames_codomain; kwargs...) return TA.svdvals(a, codomain, domain; kwargs...) end -function LA.svdvals(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.svdvals(a::AbstractITensor, args...; kwargs...) return TA.svdvals(a, args...; kwargs...) end function TA.eigen( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return eigen_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -279,12 +317,12 @@ function eigen_nameddims( return d, v end -function LA.eigen(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.eigen(a::AbstractITensor, args...; kwargs...) return TA.eigen(a, args...; kwargs...) end function TA.eigvals( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return eigvals_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -296,12 +334,12 @@ function eigvals_nameddims( return TA.eigvals(denamed(a), dimnames(a), codomain, domain; kwargs...) end -function LA.eigvals(a::AbstractNamedDimsArray, args...; kwargs...) +function LA.eigvals(a::AbstractITensor, args...; kwargs...) return TA.eigvals(a, args...; kwargs...) end function TA.left_null( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return left_null_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -316,17 +354,17 @@ function left_null_nameddims( return nameddims(n_denamed, dimnames_n) end -function TA.left_null(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function TA.left_null(a::AbstractITensor, dimnames_codomain; kwargs...) return left_null_nameddims(a, dimnames_codomain; kwargs...) end -function left_null_nameddims(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function left_null_nameddims(a::AbstractITensor, dimnames_codomain; kwargs...) codomain = name.(dimnames_codomain) domain = dimnames_setdiff(dimnames(a), codomain) return TA.left_null(a, codomain, domain; kwargs...) end function TA.right_null( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return right_null_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -341,7 +379,7 @@ function right_null_nameddims( return nameddims(n_denamed, dimnames_n) end -function TA.right_null(a::AbstractNamedDimsArray, dimnames_codomain; kwargs...) +function TA.right_null(a::AbstractITensor, dimnames_codomain; kwargs...) return right_null_nameddims(a, dimnames_codomain; kwargs...) end function right_null_nameddims(a::AbstractArray, dimnames_codomain; kwargs...) @@ -351,7 +389,7 @@ function right_null_nameddims(a::AbstractArray, dimnames_codomain; kwargs...) end """ - TensorAlgebra.gram_eigh_full(a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...) -> x + TensorAlgebra.gram_eigh_full(a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs...) -> x Gram factorization of a Hermitian positive semi-definite named array `a`, returning `x` such that `a ≈ x * x_cod`, where `x_cod` is `conj(x)` with @@ -383,7 +421,7 @@ true ``` """ function TA.gram_eigh_full( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return gram_eigh_full_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end @@ -399,7 +437,7 @@ function gram_eigh_full_nameddims( end """ - TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs...) -> x, y + TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs...) -> x, y Like `TensorAlgebra.gram_eigh_full`, but additionally returns a named array `y` that is a left inverse of `x`: `y * x ≈ I` on the rank @@ -435,7 +473,7 @@ true ``` """ function TA.gram_eigh_full_with_pinv( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return gram_eigh_full_with_pinv_nameddims( a, dimnames_codomain, dimnames_domain; kwargs... @@ -456,7 +494,7 @@ function gram_eigh_full_with_pinv_nameddims( end """ - Base.one(a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain) -> Id + Base.one(a::AbstractITensor, dimnames_codomain, dimnames_domain) -> Id Return an identity-operator-shaped named array sharing `a`'s dimension names, codomain/domain partition, and element type. The fused codomain and domain sizes @@ -484,7 +522,7 @@ true ``` """ function Base.one( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain + a::AbstractITensor, dimnames_codomain, dimnames_domain ) return one_nameddims(a, dimnames_codomain, dimnames_domain) end @@ -509,7 +547,7 @@ for f in MATRIX_FUNCTIONS f_nameddims = Symbol(f, "_nameddims") @eval begin function Base.$f( - a::AbstractNamedDimsArray, dimnames_codomain, dimnames_domain; kwargs... + a::AbstractITensor, dimnames_codomain, dimnames_domain; kwargs... ) return $f_nameddims(a, dimnames_codomain, dimnames_domain; kwargs...) end diff --git a/test/test_basics.jl b/test/test_basics.jl index 1ab4b88..fb9d0b7 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -77,10 +77,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test denamed(a′) == x @test issetequal(inds(a′), (prime(i), prime(j))) - # For now, the `ITensor` constructor is strict and only accepts a collection of - # `IndexName` as dimnames. - @test_throws ArgumentError ITensor(randn(elt, 2, 2), Index.((2, 2))) - @test_throws ArgumentError ITensor(randn(elt, 2, 2), Index.((2, 3))) + # The number of dimnames must match the array's `ndims`, and the dimnames are + # passed as a single collection. @test_throws ArgumentError ITensor(randn(elt, 4), Index.((2, 2))) @test_throws MethodError ITensor(randn(elt, 2, 2), Index(2), Index(2)) @@ -113,7 +111,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a isa ITensor @test dimnametype(a) === IndexName @test dimnametype(typeof(a)) === IndexName - @test dimnametype(ITensor) === IndexName + @test dimnametype(ITensor{IndexName}) === IndexName + # Unparameterized `ITensor` does not fix its dimname flavor, like `eltype(Array)`. + @test dimnametype(ITensor) === Any end @testset "show" begin id = rand(UInt64) diff --git a/test/test_exports.jl b/test/test_exports.jl index af3577b..8ed8ef5 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -2,7 +2,7 @@ using ITensorBase: ITensorBase using Test: @test, @testset @testset "Test exports" begin exports = [ - :ITensorBase, :ITensor, :Index, :NamedDimsArray, :aligndims, :dimnametype, + :ITensorBase, :ITensor, :Index, :aligndims, :dimnametype, :named, :nameddims, :operator, :similar_operator, ] publics = [:to_inds, Symbol("@names")] diff --git a/test/test_mooncakeext.jl b/test/test_mooncakeext.jl index 101c283..fc51909 100644 --- a/test/test_mooncakeext.jl +++ b/test/test_mooncakeext.jl @@ -1,4 +1,4 @@ -using ITensorBase: AbstractNamedUnitRange, Name, NamedDimsArray, NamedUnitRange, +using ITensorBase: AbstractNamedUnitRange, ITensor, Name, NamedUnitRange, blockedperm_nameddims, combine_nameddimsconstructors, dimnames, dimnames_setdiff, inds, name, nameddimsconstructorof, randname, to_inds using LinearAlgebra: mul! @@ -28,7 +28,7 @@ using Test: @test, @testset rng, blockedperm_nameddims, a1, (i,), (j,); mode, is_primitive ) Mooncake.TestUtils.test_rule( - rng, combine_nameddimsconstructors, NamedDimsArray, NamedDimsArray; + rng, combine_nameddimsconstructors, ITensor, ITensor; mode, is_primitive ) Mooncake.TestUtils.test_rule(rng, dimnames, a1; mode, is_primitive) diff --git a/test/test_nameddims_basics.jl b/test/test_nameddims_basics.jl index 3c6d9e3..7b73e3b 100644 --- a/test/test_nameddims_basics.jl +++ b/test/test_nameddims_basics.jl @@ -1,9 +1,9 @@ using Combinatorics: Combinatorics -using ITensorBase: @names, AbstractNamedDimsArray, AbstractNamedDimsMatrix, LittleSet, Name, - NameMismatch, NamedDimsArray, NamedDimsCartesianIndex, NamedDimsCartesianIndices, - NamedDimsMatrix, aligndims, aligneddims, apply, dename, denamed, denamedtype, dim, - dimnames, dimnametype, dims, fusednames, inds, isnamed, mapinds, name, named, nameddims, - namedoneto, product, replacedimnames, replaceinds, setinds +using ITensorBase: @names, AbstractITensor, ITensor, LittleSet, Name, NameMismatch, + NamedDimsCartesianIndex, NamedDimsCartesianIndices, aligndims, aligneddims, apply, + dename, denamed, denamedtype, dim, dimnames, dimnametype, dims, fusednames, inds, + isnamed, mapinds, name, named, nameddims, namedoneto, product, replacedimnames, + replaceinds, setinds using LinearAlgebra: LinearAlgebra using Test: @test, @test_throws, @testset using VectorInterface: scalartype @@ -17,10 +17,10 @@ end a = randn(elt, 3, 4) @test !isnamed(a) na = nameddims(a, ("i", "j")) - @test na isa NamedDimsMatrix{elt, Matrix{elt}} - @test na isa AbstractNamedDimsMatrix{elt} - @test na isa NamedDimsArray{elt} - @test na isa AbstractNamedDimsArray{elt} + @test na isa ITensor{String} + @test na isa AbstractITensor{String} + @test eltype(na) === elt + @test ndims(na) == 2 @test_throws MethodError denamed(a) @test_throws MethodError dename(a, ("i", "j")) @test_throws MethodError denamed(a, ("i", "j")) @@ -44,7 +44,10 @@ end @test dim(na, "j") == 2 @test dims(na, ("j", "i")) == (2, 1) @test na[1, 1] == a[1, 1] - @test denamedtype(typeof(na)) === typeof(a) + # The parent array's concrete type is erased from the type but is still + # recoverable from an instance. + @test denamedtype(na) === typeof(a) + @test denamedtype(typeof(na)) === AbstractArray @test dimnametype(typeof(na)) === String @test dimnametype(na) === String @@ -68,8 +71,8 @@ end @test CartesianIndices(na) == CartesianIndices(a) @test collect(pairs(na)) == (CartesianIndices(a) .=> a) - @test_throws ArgumentError NamedDimsArray(randn(4), namedoneto.((2, 2), ("i", "j"))) - ## @test_throws ErrorException NamedDimsArray(randn(2, 2), namedoneto.((2, 3), ("i", "j"))) + @test_throws ArgumentError ITensor(randn(4), namedoneto.((2, 2), ("i", "j"))) + ## @test_throws ErrorException ITensor(randn(2, 2), namedoneto.((2, 3), ("i", "j"))) a = randn(elt, 3, 4) na = nameddims(a, ("i", "j")) @@ -150,13 +153,13 @@ end i = namedoneto(2, "i") a = randn(elt, 2) na = a[i] - @test na isa NamedDimsArray{elt} + @test na isa ITensor{String} @test dimnames(na) == ("i",) @test denamed(na) == a # slicing a = randn(elt, 3, 3) - na = NamedDimsArray(a, ("i", "j")) + na = ITensor(a, ("i", "j")) for na′ in (na[named(2:3, "i"), named(2:3, "j")], na["i" => 2:3, "j" => 2:3]) @test inds(na′) == (named(1:2, "i"), named(1:2, "j")) @test denamed(na′) == a[2:3, 2:3] @@ -165,7 +168,7 @@ end # view slicing a = randn(elt, 3, 3) - na = NamedDimsArray(a, ("i", "j")) + na = ITensor(a, ("i", "j")) for na′ in (@view(na[named(2:3, "i"), named(2:3, "j")]), @view(na["i" => 2:3, "j" => 2:3])) @test inds(na′) == (named(1:2, "i"), named(1:2, "j")) @@ -233,7 +236,7 @@ end @test Tuple(size(na)) == (named(3, "i"), named(4, "j")) @test length(na) == named(12, fusednames("i", "j")) @test Tuple(axes(na)) == (named(1:3, "i"), named(1:4, "j")) - @test randn(named.((3, 4), ("i", "j"))) isa NamedDimsArray + @test randn(named.((3, 4), ("i", "j"))) isa ITensor @test na["i" => 1, "j" => 2] == a[1, 2] @test na["j" => 2, "i" => 1] == a[1, 2] na["j" => 2, "i" => 1] = 12 @@ -366,26 +369,26 @@ end value = rand(elt) for na in (zeros(elt, i, j), zeros(elt, (i, j))) @test eltype(na) ≡ elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test iszero(na) end for na in (fill(value, i, j), fill(value, (i, j))) @test eltype(na) ≡ elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test all(isequal(value), na) end for na in (rand(elt, i, j), rand(elt, (i, j))) @test eltype(na) ≡ elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test !iszero(na) @test all(x -> real(x) > 0, na) end for na in (randn(elt, i, j), randn(elt, (i, j))) @test eltype(na) ≡ elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test !iszero(na) end @@ -395,20 +398,20 @@ end default_elt = Float64 for na in (zeros(i, j), zeros((i, j))) @test eltype(na) ≡ default_elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test iszero(na) end for na in (rand(i, j), rand((i, j))) @test eltype(na) ≡ default_elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test !iszero(na) @test all(x -> real(x) > 0, na) end for na in (randn(i, j), randn((i, j))) @test eltype(na) ≡ default_elt - @test na isa NamedDimsArray + @test na isa ITensor @test inds(na) == Base.oneto.((i, j)) @test !iszero(na) end @@ -447,13 +450,13 @@ end @test values(sp) == (1.0, 2.0, 3.0) end @testset "show" begin - a = NamedDimsArray([1 2; 3 4], ("i", "j")) + a = ITensor([1 2; 3 4], ("i", "j")) @test sprint(show, "text/plain", a) == "named(Base.OneTo(2), \"i\")×named(Base.OneTo(2), \"j\") " * - "$NamedDimsArray{Int64, 2, Matrix{Int64}, Tuple{String, String}}:\n" * + "$ITensor{String}:\n" * "2×2 Matrix{Int64}:\n 1 2\n 3 4" - a = NamedDimsArray([1 2; 3 4], ("i", "j")) + a = ITensor([1 2; 3 4], ("i", "j")) @test sprint(show, a) == "[1 2; 3 4][named(Base.OneTo(2), \"i\"), named(Base.OneTo(2), \"j\")]" end diff --git a/test/test_operator.jl b/test/test_operator.jl index 729940d..6bad548 100644 --- a/test/test_operator.jl +++ b/test/test_operator.jl @@ -1,6 +1,6 @@ -using ITensorBase: ITensorBase as NDA, NamedDimsArray, NamedDimsOperator, apply, - codomainnames, dename, denamed, dimnames, domainnames, nameddims, namedoneto, operator, - product, replacedimnames, similar_operator, state +using ITensorBase: ITensorBase as NDA, ITensor, ITensorOperator, apply, codomainnames, + dename, denamed, dimnames, domainnames, nameddims, namedoneto, operator, product, + replacedimnames, similar_operator, state using LinearAlgebra: I, norm using Random: Random using StableRNGs: StableRNG @@ -9,27 +9,27 @@ using Test: @test, @testset @testset "operator" begin o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j")) - @test o isa NamedDimsOperator{Float64} + @test o isa ITensorOperator{String} @test eltype(o) ≡ Float64 @test issetequal(NDA.codomainnames(o), ("i'", "j'")) @test issetequal(NDA.domainnames(o), ("i", "j")) o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j")) õ = similar(o) - @test õ isa NamedDimsOperator{Float64} + @test õ isa ITensorOperator{String} @test eltype(õ) ≡ Float64 @test issetequal(NDA.codomainnames(õ), ("i'", "j'")) @test issetequal(NDA.domainnames(õ), ("i", "j")) o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j")) õ = similar(o, Float32) - @test õ isa NamedDimsOperator{Float32} + @test õ isa ITensorOperator{String} @test eltype(õ) ≡ Float32 @test issetequal(NDA.codomainnames(õ), ("i'", "j'")) @test issetequal(NDA.domainnames(õ), ("i", "j")) o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j")) - @test o isa NamedDimsOperator + @test o isa ITensorOperator o² = product(o, o) @test issetequal(dimnames(o²), ("i'", "j'", "i", "j")) õ = replacedimnames( @@ -39,25 +39,25 @@ using Test: @test, @testset @test state(o²) ≈ o²′ o = operator(randn(2, 2, 2, 2), ("i'", "j'"), ("i", "j")) - v = NamedDimsArray(randn(2, 2), ("i", "j")) + v = ITensor(randn(2, 2), ("i", "j")) ov = apply(o, v) @test issetequal(dimnames(ov), ("i", "j")) @test ov ≈ replacedimnames(o * v, "i'" => "i", "j'" => "j") end -@testset "one(::AbstractNamedDimsOperator)" begin +@testset "one(::AbstractITensorOperator)" begin # Identity-operator construction: matricized form is the identity matrix. i, j, k, l = namedoneto.((2, 3, 2, 3), ("i", "j", "k", "l")) op = operator(randn(i, j, k, l), ("i", "j"), ("k", "l")) Id = one(op) - @test Id isa NamedDimsOperator{Float64} + @test Id isa ITensorOperator{String} @test codomainnames(Id) == codomainnames(op) @test domainnames(Id) == domainnames(op) Id_mat = matricize(state(Id), (i, j) => "row", (k, l) => "col") @test dename(Id_mat, ("row", "col")) ≈ I(6) end -@testset "one(::AbstractNamedDimsArray, codomain, domain)" begin +@testset "one(::AbstractITensor, codomain, domain)" begin # Trivial codomain/domain layout. i, j, k, l = namedoneto.((2, 3, 2, 3), ("i", "j", "k", "l")) a = randn(i, j, k, l) @@ -77,13 +77,13 @@ end @testset "similar_operator" begin # Five-arg canonical: explicit element type, axes, codomain, domain names. op = similar_operator(randn(3, 3), Float32, (Base.OneTo(3),), ("i'",), ("i",)) - @test op isa NamedDimsOperator{Float32} + @test op isa ITensorOperator{String} @test issetequal(codomainnames(op), ("i'",)) @test issetequal(domainnames(op), ("i",)) # Codomain names default to fresh `randname` outputs. op = similar_operator(randn(3, 3), Float64, (Base.OneTo(3),), ("i",)) - @test op isa NamedDimsOperator{Float64} + @test op isa ITensorOperator{String} @test issetequal(domainnames(op), ("i",)) @test only(codomainnames(op)) != "i" @@ -98,7 +98,7 @@ end @test eltype(op) === ComplexF32 end -@testset "randn!(::AbstractNamedDimsOperator) / rand!" begin +@testset "randn!(::AbstractITensorOperator) / rand!" begin op = operator(zeros(3, 3), ("i'",), ("i",)) rng = StableRNG(123) Random.randn!(rng, op) @@ -121,8 +121,8 @@ end nms = ("i'", "i") for r in (o + o, o - o, -o, 2 * o, o * 2, 2 .* o, o .* 2) - @test r isa NamedDimsArray - @test !(r isa NamedDimsOperator) + @test r isa ITensor + @test !(r isa ITensorOperator) end @test dename(o + o, nms) ≈ 2 .* dename(s, nms) @@ -134,7 +134,7 @@ end @test dename(o .* 2, nms) ≈ 2 .* dename(s, nms) end -@testset "gram_eigh_full on AbstractNamedDimsOperator" begin +@testset "gram_eigh_full on AbstractITensorOperator" begin n = 5 B = randn(n, n) A = B * B' # Hermitian PSD