From 496a97fb847f1efdf4d862a0a6a595056d951757 Mon Sep 17 00:00:00 2001 From: Miles Stoudenmire Date: Tue, 25 Aug 2026 13:20:25 +0200 Subject: [PATCH] Improve implementation and tests of expect to ensure return value is a dictionary when no vertices passed and a vector when a vector of vertices is passed. --- src/expect.jl | 15 +++-- test/test_expect.jl | 132 ++++++++++++++++++++++++++++++-------------- 2 files changed, 101 insertions(+), 46 deletions(-) diff --git a/src/expect.jl b/src/expect.jl index 2ad9e774c..65faa38ec 100644 --- a/src/expect.jl +++ b/src/expect.jl @@ -56,25 +56,29 @@ contraction. See also: [`expect(ψ, op::String)`](@ref). """ function expect(ψ::AbstractITensorNetwork, op::Op; alg = default_expect_alg(), kwargs...) - return expect(Algorithm(alg), ψ, [op]; kwargs...) + return only(expect(Algorithm(alg), ψ, [op]; kwargs...)) end """ - expect(ψ::AbstractITensorNetwork, op::String, vertices; alg="bp", kwargs...) -> Dictionary + expect(ψ::AbstractITensorNetwork, op::String, vertices; alg="bp", kwargs...) Compute local expectation values ⟨ψ|op_v|ψ⟩ / ⟨ψ|ψ⟩ for the operator named `op` at each vertex in `vertices`. +The result is a container of the same kind as `vertices`: passing a `Vector` of vertices +returns a `Vector` of values in the same order, while passing a dictionary of vertex indices +(such as `vertices(ψ)`) returns a dictionary keyed by vertex. + See [`expect(ψ, op::String)`](@ref) for full documentation. """ function expect( ψ::AbstractITensorNetwork, op::String, vertices; alg = default_expect_alg(), kwargs... ) - return expect(Algorithm(alg), ψ, [Op(op, vertex) for vertex in vertices]; kwargs...) + return expect(Algorithm(alg), ψ, map(v -> Op(op, v), vertices); kwargs...) end """ - expect(ψ::AbstractITensorNetwork, op::String; alg="bp", kwargs...) -> Dictionary + expect(ψ::AbstractITensorNetwork, op::String; alg="bp", kwargs...) -> AbstractDictionary Compute local expectation values ⟨ψ|op_v|ψ⟩ / ⟨ψ|ψ⟩ for the operator named `op` at every vertex of `ψ`. @@ -94,7 +98,8 @@ vertex of `ψ`. # Returns -A `Dictionary` mapping each vertex of `ψ` to its expectation value. +A dictionary mapping each vertex of `ψ` to its expectation value, so results +are looked up by vertex: `expect(ψ, "Sz")[v]`. See also: [`expect(ψ, op::String, vertices)`](@ref), [`expect(operator, state::AbstractTreeTensorNetwork)`](@ref). diff --git a/test/test_expect.jl b/test/test_expect.jl index 98040d989..73b9046af 100644 --- a/test/test_expect.jl +++ b/test/test_expect.jl @@ -1,6 +1,7 @@ +using Dictionaries: AbstractDictionary using Graphs: SimpleGraph, uniform_tree -using ITensorNetworks: - BeliefPropagationCache, ITensorNetwork, expect, original_state_vertex, siteinds +using ITensorNetworks: ITensorNetwork, expect, siteinds +using ITensors: Op using NamedGraphs.NamedGraphGenerators: named_grid using NamedGraphs: NamedGraph, vertices using SplitApplyCombine: group @@ -8,44 +9,93 @@ using StableRNGs: StableRNG using TensorOperations: TensorOperations using Test: @test, @testset include("utils.jl") + +# Whole-network `expect` returns a dictionary, so compare the two algorithms vertex by +# vertex. +agree(a, b) = collect(keys(a)) == collect(keys(b)) && all(v -> a[v] ≈ b[v], keys(a)) + @testset "Test Expect" begin - #Test on a tree - L, χ = 4, 2 - g = NamedGraph(SimpleGraph(uniform_tree(L))) - s = siteinds("S=1/2", g) - rng = StableRNG(1234) - ψ = random_tensornetwork(rng, s; link_space = χ) - sz_bp = expect(ψ, "Sz"; alg = "bp") - sz_exact = expect(ψ, "Sz"; alg = "exact") - @test sz_bp ≈ sz_exact - - #Test on a grid, group by column to make BP exact - L, χ = 2, 2 - g = named_grid((L, L)) - s = siteinds("S=1/2", g) - rng = StableRNG(1234) - ψ = random_tensornetwork(rng, s; link_space = χ) - quadratic_form_vertices = reduce( - vcat, [[(v, "ket"), (v, "bra"), (v, "operator")] for v in vertices(ψ)] - ) - cache_construction_kwargs = (; - partitioned_vertices = group(v -> first(first(v)), quadratic_form_vertices), - ) - sz_bp = expect( - ψ, "Sz"; alg = "bp", cache_construction_kwargs, - cache_update_kwargs = (; maxiter = 20) - ) - sz_exact = expect(ψ, "Sz"; alg = "exact") - @test sz_bp ≈ sz_exact - - #Test with Quantum Numbers, product state so BP should be exact - L, χ = 2, 2 - g = named_grid((L, L)) - s = siteinds("S=1/2", g; conserve_qns = true) - - ψ = productstate(v -> isodd(sum(v)) ? "↑" : "↓", s) - - sz_bp = expect(ψ, "Sz"; alg = "bp", cache_update_kwargs = (; maxiter = 20)) - sz_exact = expect(ψ, "Sz"; alg = "exact") - @test sz_bp ≈ sz_exact + @testset "Whole-network result is indexed by vertex" begin + # Product state: BP is exact and the expected values are known analytically. + g = named_grid((2, 2)) + s = siteinds("S=1/2", g) + ψ = productstate(v -> isodd(sum(v)) ? "↑" : "↓", s) + + sz = expect(ψ, "Sz"; alg = "bp", cache_update_kwargs = (; maxiter = 20)) + @test sz isa AbstractDictionary + @test collect(keys(sz)) == collect(vertices(ψ)) + # Index the result directly with a vertex of `ψ`. + @test sz[(1, 1)] ≈ -0.5 + @test sz[(1, 2)] ≈ +0.5 + @test sz[(2, 1)] ≈ +0.5 + @test sz[(2, 2)] ≈ -0.5 + @test all(v -> sz[v] ≈ (isodd(sum(v)) ? +0.5 : -0.5), vertices(ψ)) + end + + @testset "Vertex subsets and single operators" begin + L, χ = 4, 2 + g = NamedGraph(SimpleGraph(uniform_tree(L))) + s = siteinds("S=1/2", g) + rng = StableRNG(1234) + ψ = random_tensornetwork(rng, s; link_space = χ) + + sz = expect(ψ, "Sz"; alg = "exact") + # Reversed relative to `vertices(ψ)`, to pin down the ordering of the result. + vs = reverse(collect(vertices(ψ))) + + # Passing an explicit `vertices` argument mirrors its container type: a `Vector` + # of vertices in, a `Vector` of values in the same order out. + sz_vec = expect(ψ, "Sz", vs; alg = "exact") + @test sz_vec isa Vector + @test sz_vec ≈ [sz[v] for v in vs] + @test expect(ψ, "Sz", vs; alg = "bp") ≈ sz_vec + + # A subset works the same way. + @test expect(ψ, "Sz", vs[1:2]; alg = "exact") ≈ sz_vec[1:2] + + # A single `Op` returns a bare number, not a collection. + for v in vs + sz_v = expect(ψ, Op("Sz", v); alg = "exact") + @test sz_v isa Number + @test sz_v ≈ sz[v] + end + end + + @testset "Tree: BP is exact" begin + L, χ = 4, 2 + g = NamedGraph(SimpleGraph(uniform_tree(L))) + s = siteinds("S=1/2", g) + rng = StableRNG(1234) + ψ = random_tensornetwork(rng, s; link_space = χ) + @test agree(expect(ψ, "Sz"; alg = "bp"), expect(ψ, "Sz"; alg = "exact")) + end + + @testset "Grid: BP grouped by column is exact" begin + L, χ = 2, 2 + g = named_grid((L, L)) + s = siteinds("S=1/2", g) + rng = StableRNG(1234) + ψ = random_tensornetwork(rng, s; link_space = χ) + quadratic_form_vertices = reduce( + vcat, [[(v, "ket"), (v, "bra"), (v, "operator")] for v in vertices(ψ)] + ) + cache_construction_kwargs = (; + partitioned_vertices = group(v -> first(first(v)), quadratic_form_vertices), + ) + sz_bp = expect( + ψ, "Sz"; alg = "bp", cache_construction_kwargs, + cache_update_kwargs = (; maxiter = 20) + ) + @test agree(sz_bp, expect(ψ, "Sz"; alg = "exact")) + end + + @testset "Quantum numbers" begin + # Product state, so BP should be exact. + L = 2 + g = named_grid((L, L)) + s = siteinds("S=1/2", g; conserve_qns = true) + ψ = productstate(v -> isodd(sum(v)) ? "↑" : "↓", s) + sz_bp = expect(ψ, "Sz"; alg = "bp", cache_update_kwargs = (; maxiter = 20)) + @test agree(sz_bp, expect(ψ, "Sz"; alg = "exact")) + end end