Improve expect implementation and tests

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).

test/test_expect.jl

@@ -1,51 +1,101 @@
+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
 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