Named-operator constructors, projection verbs, and balanced gram_eigh_full

Project.toml

@@ -1,6 +1,6 @@
 name = "NamedDimsArrays"
 uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
-version = "0.15.6"
+version = "0.15.7"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -45,7 +45,7 @@ Mooncake = "0.4.202, 0.5"
 OrderedCollections = "1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.8, 0.9"
+TensorAlgebra = "0.9.5"
 TupleTools = "1.6"
 TypeParameterAccessors = "0.4"
 VectorInterface = "0.5, 0.6"

src/NamedDimsArrays.jl

@@ -1,6 +1,7 @@
 module NamedDimsArrays
 
-export NamedDimsArray, aligndims, dimnametype, named, nameddims, operator
+export NamedDimsArray, aligndims, dimnametype, named, nameddims, operator,
+    similar_operator
 using Compat: @compat
 @compat public to_inds
 @compat public @names

src/abstractnameddimsarray.jl

@@ -1,4 +1,6 @@
-import FunctionImplementations as FI
+using FunctionImplementations: FunctionImplementations as FI
+using LinearAlgebra: LinearAlgebra
+using Random: Random
 using TypeParameterAccessors: unspecify_type_parameters
 
 # Some of the interface is inspired by:
@@ -52,7 +54,13 @@ dimnametype(type::Type{<:AbstractNamedDimsArray}) = throw(MethodError(dimnametyp
 # Unwrapping the names (`NamedDimsArrays.jl` interface).
 # TODO: Use `IsNamed` trait?
 denamed(a::AbstractNamedDimsArray) = throw(MethodError(denamed, a))
-denamed(a::AbstractNamedDimsArray, inds) = denamed(aligneddims(a, inds))
+function denamed(a::AbstractNamedDimsArray, inds)
+    # Skip the lazy `PermutedDimsArray` wrap when the requested order already
+    # matches `a`'s; compare via `Tuple` because `LittleSet ==` is
+    # set-equality and would mask a permutation.
+    Tuple(name.(inds)) == Tuple(dimnames(a)) && return denamed(a)
+    return denamed(aligneddims(a, inds))
+end
 dename(a::AbstractNamedDimsArray, inds) = denamed(aligndims(a, inds))
 
 # Output the named axes/indices of the named dims array.
@@ -161,6 +169,29 @@ 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.
@@ -766,6 +797,11 @@ for dimtype in [:AbstractNamedInteger, :AbstractNamedUnitRange]
     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.
@@ -784,6 +820,10 @@ function Base.mapreduce(f, op, a::AbstractNamedDimsArray; kwargs...)
     return mapreduce(f, op, 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`.

src/abstractnamedunitrange.jl

@@ -36,6 +36,11 @@ function Base.hash(r::AbstractNamedUnitRange, h::UInt)
     return hash(name(r), h)
 end
 
+# Forward `conj` to the underlying range so graded axes flip their sector
+# arrows. The `Base.conj(::AbstractArray{<:Real}) = x` fallback would
+# otherwise short-circuit before the inner range is touched.
+Base.conj(r::AbstractNamedUnitRange) = named(conj(denamed(r)), name(r))
+
 # Unit range functionality.
 Base.first(r::AbstractNamedUnitRange) = named(first(denamed(r)), name(r))
 Base.last(r::AbstractNamedUnitRange) = named(last(denamed(r)), name(r))

src/broadcast.jl

@@ -1,8 +1,8 @@
-import TensorAlgebra as TA
 using ..NamedDimsArrays: AbstractNamedDimsArray, AbstractNamedUnitRange, LittleSet,
     NamedDimsArrays, 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
 

src/linearalgebra.jl

@@ -1,4 +1,4 @@
-import LinearAlgebra as LA
+using LinearAlgebra: LinearAlgebra as LA
 
 # We overload `LinearAlgebra.norm` because the LinearAlgebra.jl AbstractArray definition
 # uses scalar indexing:

src/nameddimsoperator.jl

@@ -1,4 +1,5 @@
 using OrderedCollections: OrderedDict
+using Random: Random
 
 # Named dimension operator minimal interface.
 
@@ -81,13 +82,16 @@ end
 function inverse(b::Bijection)
     return Bijection(b.codomain_to_domain, b.domain_to_codomain)
 end
-# Both accessors iterate `codomain_to_domain` so that successive calls return
-# values in lock-step positional order (codomain[i] paired with domain[i]).
+# Both accessors return the `keys(::OrderedDict)` of the dict that has the
+# requested side as its key type, so the result is a `Base.KeySet` that
+# compares correctly with `==`. The two dicts are constructed from the same
+# pairs in the constructor, so `codomain(b)[i]` and `domain(b)[i]` remain in
+# lock-step positional order.
 function codomain(b::Bijection)
     return keys(b.codomain_to_domain)
 end
 function domain(b::Bijection)
-    return values(b.codomain_to_domain)
+    return keys(b.domain_to_codomain)
 end
 Base.iterate(b::Bijection) = iterate(b.domain_to_codomain)
 Base.iterate(b::Bijection, state) = iterate(b.domain_to_codomain, state)
@@ -136,9 +140,10 @@ const _gram_eigh_full_operator_docstring = """
     TensorAlgebra.gram_eigh_full(a::AbstractNamedDimsOperator; kwargs...) -> x
 
 Gram factorization of a Hermitian positive semi-definite named operator
-`a`, returning `x` such that `conj(x) * x_dom ≈ state(a)`, where `x_dom`
-is `x` with its codomain dimension names replaced by the corresponding
-domain names of `a`. The codomain and domain partition is taken from
+`a`, returning `x` such that `x * x_cod ≈ state(a)`, where `x_cod` is
+`conj(x)` with its domain dimension names replaced by the corresponding
+codomain names of `a`. `x` carries `a`'s domain dimension names and a
+fresh trailing rank name. The codomain and domain partition is taken from
 `codomainnames(a)` and `domainnames(a)`.
 
 `kwargs` are forwarded to `TensorAlgebra.gram_eigh_full` on the
@@ -159,7 +164,7 @@ julia> a = operator(conj(b) * replacedimnames(b, "i" => "j", "k" => "l"), ("i",
 
 julia> x = gram_eigh_full(a);
 
-julia> conj(x) * replacedimnames(x, "i" => "j", "k" => "l") ≈ state(a)
+julia> replacedimnames(x, "j" => "i", "l" => "k") * conj(x) ≈ state(a)
 true
 ```
 """
@@ -168,9 +173,10 @@ const _gram_eigh_full_with_pinv_operator_docstring = """
     TensorAlgebra.gram_eigh_full_with_pinv(a::AbstractNamedDimsOperator; kwargs...) -> x, y
 
 Like `TensorAlgebra.gram_eigh_full`, but additionally returns a
-named array `y` such that `x * y` projects onto the rank subspace
-(equal to the identity when `a` is full rank). The codomain and domain
-partition is taken from `codomainnames(a)` and `domainnames(a)`.
+named array `y` that is a left inverse of `x`: `y * x ≈ I` on the
+rank subspace (equal to the identity when `a` is full rank). The
+codomain and domain partition is taken from `codomainnames(a)` and
+`domainnames(a)`.
 
 # Examples
 
@@ -189,10 +195,10 @@ julia> a = operator(conj(b) * replacedimnames(b, "i" => "j", "k" => "l"), ("i",
 
 julia> x, y = gram_eigh_full_with_pinv(a);
 
-julia> rname = only(setdiff(dimnames(x), ("i", "k")));
+julia> rname = only(setdiff(dimnames(x), ("j", "l")));
 
-julia> reshape(dename(x, (rname, "i", "k")), :, 4) *
-       reshape(dename(y, ("i", "k", rname)), 4, :) ≈ I
+julia> reshape(dename(y, (rname, "j", "l")), :, 4) *
+       reshape(dename(x, ("j", "l", rname)), 4, :) ≈ I
 true
 ```
 """
@@ -206,6 +212,105 @@ for f in (:gram_eigh_full, :gram_eigh_full_with_pinv)
     end
 end
 
+"""
+    Base.one(op::AbstractNamedDimsOperator) -> 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.
+
+The identity acts as the multiplicative identity for `NamedDimsArrays.apply`: it
+contracts on the domain names and renames the resulting codomain names back to
+the domain names, leaving the input unchanged.
+
+# Examples
+
+```jldoctest
+julia> using NamedDimsArrays: apply, namedoneto, operator
+
+julia> i, j, k, l = namedoneto.((2, 3, 2, 3), ("i", "j", "k", "l"));
+
+julia> op = operator(randn(i, j, k, l), ("i", "j"), ("k", "l"));
+
+julia> Id = one(op);
+
+julia> v = randn(k, l);
+
+julia> apply(Id, v) ≈ v
+true
+```
+"""
+function Base.one(op::AbstractNamedDimsOperator)
+    co, dom = codomainnames(op), domainnames(op)
+    return operator(one(state(op), co, dom), co, dom)
+end
+
+# === similar_operator ===
+#
+# Allocate an operator with the user-supplied axes as the domain (input). The
+# codomain (output) shares the domain direction and either takes
+# explicitly-supplied names or fresh `randname` outputs. The 5-arg form is
+# canonical, the others fill in defaults. The bra/ket flip on the storage side
+# is handled inside `TA.similar_map`.
+
+"""
+    similar_operator(prototype, [T,] unnamed_domain_axes, [codomain_names,] domain_names) -> op
+    similar_operator(prototype, [T,] named_domain_axes) -> op
+
+Allocate an operator-shaped named array with undefined data, with the
+user-supplied side as the domain (input) and a matching codomain (output).
+Element type defaults to `eltype(prototype)`. Codomain names default to fresh
+`randname`-generated names. The first form takes unnamed (raw) axes and
+explicit names, the second takes already-named axes and reuses their names as
+the domain. Storage layout (including the bra/ket flip on the domain side for
+graded axes) is delegated to `TensorAlgebra.similar_map`.
+"""
+function similar_operator(
+        prototype, ::Type{T}, unnamed_domain_axes, codomain_names, domain_names
+    ) where {T}
+    codomain_axes = named.(unnamed_domain_axes, codomain_names)
+    domain_axes = named.(unnamed_domain_axes, domain_names)
+    raw = TA.similar_map(prototype, T, codomain_axes, domain_axes)
+    return operator(raw, codomain_names, domain_names)
+end
+function similar_operator(
+        prototype, ::Type{T}, unnamed_domain_axes, domain_names
+    ) where {T}
+    codomain_names = randname.(domain_names)
+    return similar_operator(
+        prototype, T, unnamed_domain_axes, codomain_names, domain_names
+    )
+end
+function similar_operator(prototype, ::Type{T}, named_domain_axes) where {T}
+    return similar_operator(
+        prototype, T, denamed.(named_domain_axes), name.(named_domain_axes)
+    )
+end
+function similar_operator(prototype, unnamed_domain_axes, codomain_names, domain_names)
+    return similar_operator(
+        prototype, eltype(prototype), unnamed_domain_axes, codomain_names, domain_names
+    )
+end
+function similar_operator(prototype, unnamed_domain_axes, domain_names)
+    return similar_operator(prototype, eltype(prototype), unnamed_domain_axes, domain_names)
+end
+function similar_operator(prototype, named_domain_axes)
+    return similar_operator(prototype, eltype(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
+# method.
+
+function Random.randn!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator)
+    Random.randn!(rng, state(op))
+    return op
+end
+
+function Random.rand!(rng::Random.AbstractRNG, op::AbstractNamedDimsOperator)
+    Random.rand!(rng, state(op))
+    return op
+end
+
 struct NamedDimsOperator{T, N, P <: AbstractNamedDimsArray{T, N}, D, C} <:
     AbstractNamedDimsOperator{T, N}
     parent::P