Reimplement Product on top of #223

ext/TenetAdaptExt.jl

@@ -8,5 +8,6 @@ Adapt.adapt_structure(to, x::TensorNetwork) = TensorNetwork(adapt.(Ref(to), tens
 
 Adapt.adapt_structure(to, x::Quantum) = Quantum(adapt(to, TensorNetwork(x)), x.sites)
 Adapt.adapt_structure(to, x::Ansatz) = Ansatz(adapt(to, Quantum(x)), Tenet.lattice(x))
+Adapt.adapt_structure(to, x::Product) = Product(adapt(to, Ansatz(x)))
 
 end

ext/TenetChainRulesCoreExt/frules.jl

@@ -16,6 +16,9 @@ function ChainRulesCore.frule((_, ẋ), ::Type{Ansatz}, x::Quantum, lattice)
     return Ansatz(x, lattice), Tangent{Ansatz}(; tn=ẋ, lattice=NoTangent())
 end
 
+# `AbstractAnsatz`-subtype constructors
+ChainRulesCore.frule((_, ẋ), ::Type{Product}, x::Ansatz) = Product(x), Tangent{Product}(; tn=ẋ)
+
 # `Base.conj` methods
 ChainRulesCore.frule((_, Δ), ::typeof(Base.conj), tn::Tensor) = conj(tn), conj(Δ)
 

ext/TenetChainRulesCoreExt/rrules.jl

@@ -20,6 +20,11 @@ Ansatz_pullback(ȳ) = (NoTangent(), ȳ.tn, NoTangent())
 Ansatz_pullback(ȳ::AbstractThunk) = Ansatz_pullback(unthunk(ȳ))
 ChainRulesCore.rrule(::Type{Ansatz}, x::Quantum, lattice) = Ansatz(x, lattice), Ansatz_pullback
 
+# `AbstractAnsatz`-subtype constructors
+Product_pullback(ȳ) = (NoTangent(), ȳ.tn)
+Product_pullback(ȳ::AbstractThunk) = Product_pullback(unthunk(ȳ))
+ChainRulesCore.rrule(::Type{Product}, x::Ansatz) = Product(x), Product_pullback
+
 # `Base.conj` methods
 conj_pullback(Δ::Tensor) = (NoTangent(), conj(Δ))
 conj_pullback(Δ::Tangent{Tensor}) = (NoTangent(), conj(Δ))

ext/TenetReactantExt.jl

@@ -36,6 +36,12 @@ function Reactant.make_tracer(seen, prev::Ansatz, path::Tuple, mode::Reactant.Tr
     return Ansatz(tracetn, copy(Tenet.lattice(prev)))
 end
 
+# TODO try rely on generic fallback for ansatzes
+function Reactant.make_tracer(seen, prev::Tenet.Product, path::Tuple, mode::Reactant.TraceMode; kwargs...)
+    tracetn = Reactant.make_tracer(seen, Ansatz(prev), Reactant.append_path(path, :tn), mode; kwargs...)
+    return Tenet.Product(tracetn)
+end
+
 function Reactant.create_result(@nospecialize(tocopy::Tensor), @nospecialize(path), result_stores)
     data = Reactant.create_result(parent(tocopy), Reactant.append_path(path, :data), result_stores)
     return :($Tensor($data, $(inds(tocopy))))
@@ -58,6 +64,12 @@ function Reactant.create_result(tocopy::Ansatz, @nospecialize(path), result_stor
     return :($Ansatz($tn, $(copy(Tenet.lattice(tocopy)))))
 end
 
+# TODO try rely on generic fallback for ansatzes
+function Reactant.create_result(tocopy::Tenet.Product, @nospecialize(path), result_stores)
+    tn = Reactant.create_result(Ansatz(tocopy), Reactant.append_path(path, :tn), result_stores)
+    return :($(Tenet.Product)($tn))
+end
+
 # TODO try rely on generic fallback for ansatzes
 # function Reactant.create_result(tocopy::Tenet.Product, @nospecialize(path), result_stores)
 #     tn = Reactant.create_result(Ansatz(tocopy), Reactant.append_path(path, :tn), result_stores)

src/Product.jl

@@ -0,0 +1,95 @@
+using LinearAlgebra
+using Graphs
+
+"""
+    Product <: AbstractAnsatz
+
+An [`Ansatz`](@ref) represented as a tensor product.
+
+# Constructors
+
+If you pass an `Abstract{<:AbstractVector}` to the constructor, it will create a [`State`](@ref).
+If you pass an `Abstract{<:AbstractMatrix}` to the constructor, it will create an [`Operator`](@ref).
+"""
+struct Product <: AbstractAnsatz
+    tn::Ansatz
+end
+
+Ansatz(tn::Product) = tn.tn
+
+Base.copy(x::Product) = Product(copy(Ansatz(x)))
+Base.similar(x::Product) = Product(similar(Ansatz(x)))
+Base.zero(x::Product) = Product(zero(Ansatz(x)))
+
+function Product(arrays::AbstractArray{<:AbstractVector})
+    n = length(arrays)
+    gen = IndexCounter()
+    symbols = map(arrays) do _
+        nextindex!(gen)
+    end
+    _tensors = map(eachindex(arrays)) do i
+        Tensor(arrays[i], [symbols[i]])
+    end
+
+    sitemap = Dict(Site(i) => symbols[i] for i in eachindex(arrays))
+    qtn = Quantum(TensorNetwork(_tensors), sitemap)
+    graph = Graph(n)
+    mapping = BijectiveIdDict{Site,Int}(Pair{Site,Int}[site => i for (i, site) in enumerate(lanes(qtn))])
+    lattice = Lattice(mapping, graph)
+    ansatz = Ansatz(qtn, lattice)
+    return Product(ansatz)
+end
+
+function Product(arrays::AbstractArray{<:AbstractMatrix})
+    n = length(arrays)
+    gen = IndexCounter()
+    symbols = map(arrays) do _
+        (nextindex!(gen), nextindex!(gen))
+    end
+    _tensors = map(eachindex(arrays)) do i
+        Tensor(arrays[i], [symbols[i][1], symbols[i][2]])
+    end
+
+    sitemap = merge!(
+        Dict(Site(i; dual=true) => symbols[i][1] for i in eachindex(arrays)),
+        Dict(Site(i) => symbols[i][2] for i in eachindex(arrays)),
+    )
+    qtn = Quantum(TensorNetwork(_tensors), sitemap)
+    graph = Graph(n)
+    mapping = BijectiveIdDict{Site,Int}(Pair{Site,Int}[site => i for (i, site) in enumerate(lanes(qtn))])
+    lattice = Lattice(mapping, graph)
+    ansatz = Ansatz(qtn, lattice)
+    return Product(ansatz)
+end
+
+LinearAlgebra.norm(tn::Product, p::Real=2) = LinearAlgebra.norm(socket(tn), tn, p)
+function LinearAlgebra.norm(::Union{State,Operator}, tn::Product, p::Real)
+    return mapreduce(*, tensors(tn)) do tensor
+        norm(tensor, p)
+    end^(1//p)
+end
+
+LinearAlgebra.opnorm(tn::Product, p::Real=2) = LinearAlgebra.opnorm(socket(tn), tn, p)
+function LinearAlgebra.opnorm(::Operator, tn::Product, p::Real)
+    return mapreduce(*, tensors(tn)) do tensor
+        opnorm(parent(tensor), p)
+    end^(1//p)
+end
+
+LinearAlgebra.normalize!(tn::Product, p::Real=2) = LinearAlgebra.normalize!(socket(tn), tn, p)
+function LinearAlgebra.normalize!(::Union{State,Operator}, tn::Product, p::Real)
+    for tensor in tensors(tn)
+        normalize!(tensor, p)
+    end
+    return tn
+end
+
+overlap(a::Product, b::Product) = overlap(socket(a), a, socket(b), b)
+
+function overlap(::State, a::Product, ::State, b::Product)
+    @assert issetequal(sites(a), sites(b)) "Ansatzes must have the same sites"
+
+    mapreduce(*, zip(tensors(a), tensors(b))) do (ta, tb)
+        dot(parent(ta), conj(parent(tb)))
+    end
+end

src/Tenet.jl

@@ -34,6 +34,9 @@ export form
 export canonize_site, canonize_site!, canonize, canonize!, mixed_canonize, mixed_canonize!, truncate!
 export evolve!, expect, overlap
 
+include("Product.jl")
+export Product
+
 # reexports from EinExprs
 export einexpr, inds
 

test/Product_test.jl

@@ -1,4 +1,4 @@
-@testset_skip "Product ansatz" begin
+@testset "Product ansatz" begin
     using LinearAlgebra
 
     # TODO test `Product` with `Scalar` socket
@@ -14,9 +14,7 @@
     end
     @test adjoint(qtn) isa Product
     @test socket(adjoint(qtn)) == State(; dual=true)
-
-    # conversion to `Quantum`
-    @test Quantum(qtn) isa Quantum
+    @test Ansatz(qtn) isa Ansatz
 
     qtn = Product([rand(2, 2) for _ in 1:3])
     @test socket(qtn) == Operator()
@@ -30,4 +28,5 @@
     end
     @test adjoint(qtn) isa Product
     @test socket(adjoint(qtn)) == Operator()
+    @test Ansatz(qtn) isa Ansatz
 end

test/integration/ChainRules_test.jl

@@ -199,4 +199,10 @@
         test_frule(Ansatz, tn, lattice)
         test_rrule(Ansatz, tn, lattice)
     end
+
+    @testset "Product" begin
+        tn = Product([ones(2), ones(2), ones(2)])
+        test_frule(Product, Ansatz(tn))
+        test_rrule(Product, Ansatz(tn))
+    end
 end