From e1abe27b05e1b8471819940dbfbc4a1b7e52d08e Mon Sep 17 00:00:00 2001
From: b-kloss <bkloss@flatironinstitute.org>
Date: Fri, 26 Jan 2024 14:20:34 -0500
Subject: [PATCH] Efficient implementation of `expect` for `AbstractTTN` (#129)

---
 .../abstracttreetensornetwork.jl              | 24 ++++++++++++++
 test/test_treetensornetworks/test_expect.jl   | 32 +++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 test/test_treetensornetworks/test_expect.jl

diff --git a/src/treetensornetworks/abstracttreetensornetwork.jl b/src/treetensornetworks/abstracttreetensornetwork.jl
index 72642687..289449ab 100644
--- a/src/treetensornetworks/abstracttreetensornetwork.jl
+++ b/src/treetensornetworks/abstracttreetensornetwork.jl
@@ -420,3 +420,27 @@ function inner(
   end
   return O[]
 end
+
+function expect(
+  operator::String,
+  state::AbstractTTN;
+  vertices=vertices(state),
+  # ToDo: verify that this is a sane default
+  root_vertex=default_root_vertex(siteinds(state)),
+)
+  # ToDo: for performance it may be beneficial to also implement expect! and change the orthogonality center in place
+  # assuming that the next algorithmic step can make use of the orthogonality center being moved to a different vertex
+  # ToDo: Verify that this is indeed the correct order for performance
+  sites = siteinds(state)
+  ordered_vertices = reverse(post_order_dfs_vertices(sites, root_vertex))
+  res = Dictionary(vertices, undef)
+  for v in ordered_vertices
+    !(v in vertices) && continue
+    state = orthogonalize(state, v)
+    @assert isone(length(sites[v]))
+    #ToDo: Add compatibility with more than a single index per vertex
+    op_v = op(operator, only(sites[v]))
+    res[v] = (dag(state[v]) * apply(op_v, state[v]))[]
+  end
+  return mapreduce(typeof, promote_type, res).(res)
+end
diff --git a/test/test_treetensornetworks/test_expect.jl b/test/test_treetensornetworks/test_expect.jl
new file mode 100644
index 00000000..37dc9f7c
--- /dev/null
+++ b/test/test_treetensornetworks/test_expect.jl
@@ -0,0 +1,32 @@
+using ITensors
+using ITensorNetworks
+using Test
+
+@testset "MPS expect comparison with ITensors" begin
+  N = 25
+  s = siteinds("S=1/2", N)
+  a = random_mps(s; internal_inds_space=100)
+  b = MPS([a[v] for v in vertices(a)])
+  res_a = expect("Sz", a)
+  res_b = expect(b, "Sz")
+  res_a = [res_a[v] for v in vertices(a)]
+  @test res_a ≈ res_b rtol = 1e-6
+end
+
+@testset "TTN expect" begin
+  tooth_lengths = fill(5, 6)
+  c = named_comb_tree(tooth_lengths)
+  s = siteinds("S=1/2", c)
+  d = Dict()
+  magnetization = Dict()
+  for (i, v) in enumerate(vertices(s))
+    d[v] = isodd(i) ? "Up" : "Dn"
+    magnetization[v] = isodd(i) ? 0.5 : -0.5
+  end
+  states = v -> d[v]
+  state = TTN(s, states)
+  res = expect("Sz", state)
+  @test all([isapprox(res[v], magnetization[v]; atol=1e-8) for v in vertices(s)])
+end
+
+nothing