Add map-shaped array construction and TensorMap broadcasting

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.16.1"
+version = "0.16.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -10,6 +10,7 @@ projects = ["benchmark", "dev", "docs", "examples", "test"]
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
@@ -29,6 +30,7 @@ EllipsisNotation = "1.8"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
+Random = "1.10"
 Strided = "2.6"
 StridedViews = "0.5"
 TensorKit = "0.17"

ext/TensorAlgebraTensorKitExt.jl

@@ -1,8 +1,9 @@
 module TensorAlgebraTensorKitExt
 
+using Random: AbstractRNG
 using TensorAlgebra: TensorAlgebra
-using TensorKit: TensorKit, AbstractTensorMap, ProductSpace, numind, permute, space,
-    spacetype, zerovector!, ←
+using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, numind,
+    permute, space, spacetype, zerovector!, ←
 using TensorOperations: TensorOperations as TO
 
 # ============================  AbstractArray-vocabulary bridge  ============================
@@ -26,6 +27,44 @@ function TensorAlgebra.similar_map(
     return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...))
 end
 
+# ===============================  zeros_map / randn_map / rand_map  ========================
+# A `TensorMap` keeps its codomain and domain as separate `ProductSpace`s rather than a single
+# flattened axis, so build the `codomain ← domain` space directly instead of the dense
+# flatten-and-dualize fallback. As with `similar_map`, the axes arrive codomain-facing
+# (un-dualized), which is TensorKit's own codomain/domain convention. The elementary space type
+# `S` is passed to `_map_homspace` explicitly so the two dispatch entries per constructor can
+# read it from whichever of the codomain/domain is non-empty and share one builder; an empty
+# axis tuple gives the unit space `ProductSpace{S}()`.
+function _map_homspace(::Type{S}, codomain_axes, domain_axes) where {S <: ElementarySpace}
+    return ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...)
+end
+function TensorAlgebra.zeros_map(
+        ::Type{T}, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+    ) where {T, S <: ElementarySpace}
+    return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes))
+end
+function TensorAlgebra.zeros_map(
+        ::Type{T}, codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+    ) where {T, S <: ElementarySpace}
+    return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes))
+end
+for (f, g) in ((:randn_map, :randn), (:rand_map, :rand))
+    @eval begin
+        function TensorAlgebra.$f(
+                rng::AbstractRNG, ::Type{T},
+                codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+            ) where {T, S <: ElementarySpace}
+            return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes))
+        end
+        function TensorAlgebra.$f(
+                rng::AbstractRNG, ::Type{T},
+                codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+            ) where {T, S <: ElementarySpace}
+            return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes))
+        end
+    end
+end
+
 # ================================  bipermutedimsopadd!  =====================================
 # `dest = β * dest + α * permutedims(op.(src), (perm_codomain, perm_domain))`. Delegate to
 # TensorKit's TensorOperations interface: `tensoradd!` realizes the permutation, the `op === conj`
@@ -84,4 +123,28 @@ function TensorAlgebra.default_contract_algorithm(
     return TensorAlgebra.ContractAlgorithm(TO.DefaultBackend())
 end
 
+# ==================================  linear-combination broadcast  =========================
+# A `TensorMap` is not an `AbstractArray`, so it needs a `BroadcastStyle` to broadcast lazily
+# (otherwise Base tries to `collect` it). A linear combination flattens (via `tryflattenlinear`)
+# to a `LinearBroadcasted` that materializes through `add!`/`bipermutedimsopadd!` above; the
+# `copyto!` here is not piracy because `LinearBroadcasted` is TensorAlgebra-owned. Element-wise
+# (nonlinear) broadcast is not a meaningful operation on a symmetric tensor, so it errors rather
+# than dense-converting.
+struct TensorMapStyle <: Base.Broadcast.BroadcastStyle end
+Base.Broadcast.BroadcastStyle(::Type{<:AbstractTensorMap}) = TensorMapStyle()
+Base.Broadcast.BroadcastStyle(s::TensorMapStyle, ::TensorMapStyle) = s
+Base.Broadcast.BroadcastStyle(s::TensorMapStyle, ::Base.Broadcast.BroadcastStyle) = s
+Base.Broadcast.broadcastable(a::AbstractTensorMap) = a
+
+function Base.copyto!(dest::AbstractTensorMap, src::TensorAlgebra.LinearBroadcasted)
+    return TensorAlgebra.add!(dest, src, true, false)
+end
+
+function Base.copy(::Base.Broadcast.Broadcasted{TensorMapStyle})
+    return error(
+        "element-wise broadcast is not supported for a `TensorMap`; only linear combinations \
+        such as `a .+ b` and `2 .* a` are supported"
+    )
+end
+
 end

src/similar_map.jl

@@ -1,3 +1,5 @@
+using Random: Random, AbstractRNG
+
 """
     similar_map(prototype, [T,] codomain_axes, domain_axes) -> M
 
@@ -28,3 +30,76 @@ end
 function similar_map(prototype, codomain_axes, domain_axes)
     return similar_map(prototype, eltype(prototype), codomain_axes, domain_axes)
 end
+
+"""
+    zeros([T,] axes) -> A
+    randn([rng,] [T,] axes) -> A
+    rand([rng,] [T,] axes) -> A
+
+Axis-friendly counterparts of `Base.zeros`/`Base.randn`/`Base.rand`, taking the axes
+as a single tuple. `Base.zeros` already accepts axes, but `Base.randn`/`Base.rand`
+accept only integer dims, so these fill that gap for dense `Base.OneTo` axes and
+otherwise forward to `Base` (so a graded-axis backend that extends `Base.randn`/`rand`
+on its axis type is picked up). These are the flat (non-map) companions of
+[`zeros_map`](@ref).
+"""
+function zeros end
+function randn end
+function rand end
+@doc (@doc zeros) randn
+@doc (@doc zeros) rand
+
+zeros(::Type{T}, axes::Tuple) where {T} = Base.zeros(T, axes)
+for (f, g) in ((:randn, :randn), (:rand, :rand))
+    @eval begin
+        $f(rng::AbstractRNG, ::Type{T}, axes::Tuple) where {T} = Base.$g(rng, T, axes)
+        function $f(
+                rng::AbstractRNG, ::Type{T}, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}}
+            ) where {T}
+            return Base.$g(rng, T, map(length, axes))
+        end
+    end
+end
+
+"""
+    zeros_map([T,] codomain_axes, domain_axes) -> M
+    randn_map([rng,] [T,] codomain_axes, domain_axes) -> M
+    rand_map([rng,] [T,] codomain_axes, domain_axes) -> M
+
+Construct an array shaped as a linear map from `domain_axes` to `codomain_axes`,
+filled with zeros (`zeros_map`), normally-distributed values (`randn_map`), or
+uniformly-distributed values (`rand_map`), with element type `T` (defaulting to
+`Float64`). These are the value-filling companions of [`similar_map`](@ref): the
+domain axes are given un-dualized (codomain facing) and stored dual, so the default
+flattens to the axis-friendly [`zeros`](@ref)/[`randn`](@ref)/[`rand`](@ref) over
+`(codomain_axes..., conj.(domain_axes)...)` (`conj` dualizes a graded axis and is a
+no-op on a dense one). Backends with map-shaped storage (e.g. a `TensorMap`) overload
+these to build the codomain/domain directly.
+"""
+function zeros_map end
+function randn_map end
+function rand_map end
+@doc (@doc zeros_map) randn_map
+@doc (@doc zeros_map) rand_map
+
+zeros_map(codomain_axes, domain_axes) = zeros_map(Float64, codomain_axes, domain_axes)
+function zeros_map(::Type{T}, codomain_axes, domain_axes) where {T}
+    return zeros(T, (codomain_axes..., conj.(domain_axes)...))
+end
+
+for f in (:randn_map, :rand_map)
+    g = Symbol(chopsuffix(String(f), "_map"))
+    @eval begin
+        $f(codomain_axes, domain_axes) =
+            $f(Random.default_rng(), codomain_axes, domain_axes)
+        function $f(rng::AbstractRNG, codomain_axes, domain_axes)
+            return $f(rng, Float64, codomain_axes, domain_axes)
+        end
+        function $f(::Type{T}, codomain_axes, domain_axes) where {T}
+            return $f(Random.default_rng(), T, codomain_axes, domain_axes)
+        end
+        function $f(rng::AbstractRNG, ::Type{T}, codomain_axes, domain_axes) where {T}
+            return $g(rng, T, (codomain_axes..., conj.(domain_axes)...))
+        end
+    end
+end

test/test_similar_map.jl

@@ -1,4 +1,5 @@
-using TensorAlgebra: similar_map
+using Random: default_rng
+using TensorAlgebra: TensorAlgebra, rand_map, randn_map, similar_map, zeros_map
 using Test: @test, @testset
 
 @testset "similar_map ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64)
@@ -21,3 +22,44 @@ using Test: @test, @testset
     @test eltype(O3) === ComplexF32
     @test size(O3) == (2, 3)
 end
+
+# The flat axis-friendly constructors fill Base's gap: `randn`/`rand` reject `Base.OneTo`.
+@testset "flat construction ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    ax = (Base.OneTo(2), Base.OneTo(3))
+    z = TensorAlgebra.zeros(T, ax)
+    @test z isa Matrix{T}
+    @test size(z) == (2, 3)
+    @test iszero(z)
+    for f in (TensorAlgebra.randn, TensorAlgebra.rand)
+        a = f(default_rng(), T, ax)
+        @test a isa Matrix{T}
+        @test size(a) == (2, 3)
+    end
+end
+
+# The dense map constructors flatten `(codomain_axes..., conj.(domain_axes)...)`; `conj` is a
+# no-op on a dense axis, so the shape is the concatenation of codomain and domain lengths.
+@testset "map construction ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    cod = (Base.OneTo(2), Base.OneTo(3))
+    dom = (Base.OneTo(4),)
+
+    z = zeros_map(T, cod, dom)
+    @test z isa Array{T, 3}
+    @test size(z) == (2, 3, 4)
+    @test iszero(z)
+
+    for f in (randn_map, rand_map)
+        # Fully specified.
+        a = f(default_rng(), T, cod, dom)
+        @test a isa Array{T, 3}
+        @test size(a) == (2, 3, 4)
+        # Element type defaults to `Float64`.
+        b = f(cod, dom)
+        @test b isa Array{Float64, 3}
+        @test size(b) == (2, 3, 4)
+    end
+
+    # An empty domain (all-codomain map) is how a plain tensor is constructed.
+    zc = zeros_map(T, cod, ())
+    @test size(zc) == (2, 3)
+end

test/test_tensorkitext.jl

@@ -1,6 +1,9 @@
+using Base.Broadcast: broadcasted
 using StableRNGs: StableRNG
-using TensorAlgebra: contract, matricize, similar_map, unmatricize
-using TensorKit: @tensor, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗
+using TensorAlgebra: TensorAlgebra, contract, matricize, rand_map, randn_map, similar_map,
+    tryflattenlinear, unmatricize, zeros_map
+using TensorKit:
+    @tensor, AbstractTensorMap, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗
 using Test: @test, @test_throws, @testset
 
 # A shared bond contracts when it sits in one operand's domain and the other's codomain, i.e.
@@ -99,4 +102,49 @@ using Test: @test, @test_throws, @testset
         sm_domain = similar_map(t_domain, elt, (), (A1, A2))
         @test space(sm_domain) == space(t_domain)
     end
+
+    # The map constructors build a `TensorMap` from the codomain/domain spaces directly rather
+    # than flattening, mirroring the `similar_map` space convention (domain given un-dualized).
+    @testset "map construction on spaces" begin
+        A1 = Rep[U₁](0 => 2, 1 => 1)
+        A2 = Rep[U₁](0 => 1, 1 => 1)
+        B = Rep[U₁](0 => 1, -1 => 2)
+
+        for f in (randn_map, rand_map)
+            t = f(rng, elt, (A1, A2), (B,))
+            @test t isa AbstractTensorMap
+            @test space(t) == ((A1 ⊗ A2) ← B)
+        end
+        z = zeros_map(elt, (A1, A2), (B,))
+        @test z isa AbstractTensorMap
+        @test space(z) == ((A1 ⊗ A2) ← B)
+        @test iszero(z)
+
+        # An empty domain gives an all-codomain `TensorMap`, the plain-tensor case.
+        tc = randn_map(rng, elt, (A1, A2), ())
+        @test space(tc) == ((A1 ⊗ A2) ← one(A1))
+
+        # An empty codomain is the mirror case: the space type comes from the domain.
+        td = randn_map(rng, elt, (), (A1, A2))
+        @test space(td) == (one(A1) ← (A1 ⊗ A2))
+        zd = zeros_map(elt, (), (A1, B))
+        @test space(zd) == (one(A1) ← (A1 ⊗ B))
+    end
+
+    # A linear combination of `TensorMap`s flattens to a `LinearBroadcasted` that materializes
+    # into a `TensorMap` destination via `copyto!`; a nonlinear broadcast has no linear form.
+    @testset "linear-combination broadcast" begin
+        V = Rep[SU₂](0 => 1, 1 // 2 => 2)
+        a = randn(rng, elt, V ← V)
+        b = randn(rng, elt, V ← V)
+
+        lb = tryflattenlinear(broadcasted(+, a, broadcasted(*, 2, b)))
+        dest = similar(a)
+        copyto!(dest, lb)
+        @test dest ≈ a + 2 * b
+
+        # A nonlinear (element-wise) broadcast is not expressible as a `LinearBroadcasted`.
+        @test isnothing(tryflattenlinear(broadcasted(*, a, b)))
+        @test_throws ErrorException copy(broadcasted(*, a, b))
+    end
 end