Repository: google/qhbm-library (language: python)

Issue:
Expectation of a modular Hamiltonian against a QHBM
The logarithm of a QHBM can be treated as a Hamiltonian.  Call a Hamiltonian represented in this way a "modular Hamiltonian".  We should be able to take the expectation value of a modular Hamiltonian against a QHBM.


Discussion / hints:
It might also be a good idea to have a proper representation for a sum of modular Hamiltonians, perhaps in the same vein as `cirq.PauliSum`, though this might be something you are already considering in the design of QVR.

Fix patch:
diff --git a/qhbmlib/hamiltonian_infer.py b/qhbmlib/hamiltonian_infer.py
index 1cc7342e..ce83e229 100644
--- a/qhbmlib/hamiltonian_infer.py
+++ b/qhbmlib/hamiltonian_infer.py
@@ -20,12 +20,47 @@
 
 from qhbmlib import circuit_infer
 from qhbmlib import energy_infer
+from qhbmlib import energy_model
 from qhbmlib import hamiltonian_model
 from qhbmlib import utils
 
 
 class QHBM(tf.keras.layers.Layer):
-  """Methods for inference on Hamiltonian objects."""
+  r"""Methods for inference involving normalized exponentials of Hamiltonians.
+
+  We also call the normalized exponential of a Hamiltonian a "thermal state".
+  Here we formalize some aspects of thermal states, which will be used later
+  to explain particular methods of this class.
+
+  # TODO(#119): add reference to updated QHBM paper.
+
+  Each method takes as input some modular Hamiltonian
+  $$K_{\theta\phi} = U_\phi K_\theta U_\phi^\dagger.$$
+  The [thermal state][1] corresponding to the model is
+  $$ \rho_T = Z^{-1} e^{-\beta K_{\theta\phi}}.$$
+  For QHBMs, we assume $\beta = 1$, effectively absorbing it into the definition
+  of the modular Hamiltonian.  Then $\rho_T$ can be expanded as
+  $$\rho_T = \sum_x p_\theta(x)U_\phi\ket{x}\bra{x}U_\phi^\dagger,$$
+  where the probability is given by
+  $$p_\theta(x) = \tr[\exp(-K_\theta)]\bra{x}\exp(-K_\theta)\ket{x}$$
+  for $x\in\{1, \ldots, \dim(K_{\theta\phi})\} = \mathcal{X}$. Note that each
+  $U_\phi\ket{x}$ is an eigenvector of both $\rho_T$ and $K_{\theta\phi}$.
+
+  Corresponding to this density operator is an [ensemble of quantum states][2].
+  Using the terms above, we define the particular ensemble
+  $$\mathcal{E} = \{p_\theta(x), U_\phi\ket{x}\}_{x\in\mathcal{X}},$$
+  also known as the [canonical ensemble][2] corresponding to $\rho_T$.
+  Each method of this class implicitly samples from this ensemble, then
+  post-processes to perform a particular inference task.
+
+  #### References
+  [1]: Nielsen, Michael A. and Chuang, Isaac L. (2010).
+       Quantum Computation and Quantum Information.
+       Cambridge University Press.
+  [2]: Wilde, Mark M. (2017).
+       Quantum Information Theory (second edition).
+       Cambridge University Press.
+  """
 
   def __init__(self,
                e_inference: energy_infer.EnergyInference,
@@ -55,23 +90,7 @@ def q_inference(self):
   def circuits(self, model: hamiltonian_model.Hamiltonian, num_samples: int):
     r"""Draws thermally distributed eigenstates from the model Hamiltonian.
 
-    We define some terms that will be used to explain the algorithm.
-    `model` is some modular Hamiltonian
-    $$K_{\theta\phi} = U_\phi K_\theta U_\phi^\dagger.$$
-    The [thermal state][1] corresponding to the model is
-    $$ \rho_T = Z^{-1} e^{-\beta K_{\theta\phi}}.$$
-    For QHBMs, we fix $\beta = 1$, effectively absorbing it into the definition
-    of the modular Hamiltonian.  Then $\rho_T$ can be expanded as
-    $$\rho_T = \sum_x p_\theta(x)U_\phi\ket{x}\bra{x}U_\phi^\dagger,$$
-    where the probability is given by
-    $$p_\theta(x) = \tr[\exp(-K_\theta)]\bra{x}\exp(-K_\theta)\ket{x}$$
-    for $x\in\{1, \ldots, \dim(K_{\theta\phi})\} = \mathcal{X}$. Note that each
-    $U_\phi\ket{x}$ is an eigenvector of both $\rho_T$ and $K_{\theta\phi}$.
-    Corresponding to this operator is the [ensemble of quantum states][2]
-    $$\mathcal{E} = \{p_\theta(x), U_\phi\ket{x}\}_{x\in\mathcal{X}}.$$
-    This function returns pure state samples from the ensemble.
-
-    We now explain the algorithm.  First, construct $X$ to be a classical
+    Here we explain the algorithm.  First, construct $X$ to be a classical
     random variable with probability distribution $p_\theta(x)$ set by
     `model.energy`.  Then, draw $n = $`num\_samples` bitstrings,
     $S=\{x_1, \ldots, x_n\}$, from $X$.  For each unique $x_i\in S$, set
@@ -79,14 +98,6 @@ def circuits(self, model: hamiltonian_model.Hamiltonian, num_samples: int):
     $U_\phi$ is set by `model.circuit`.  Finally, set `counts[i]` equal to the
     number of times $x_i$ occurs in $S$.
 
-    #### References
-    [1]: Nielsen, Michael A. and Chuang, Isaac L. (2010).
-         Quantum Computation and Quantum Information.
-         Cambridge University Press.
-    [2]: Wilde, Mark M. (2017).
-         Quantum Information Theory (second edition).
-         Cambridge University Press.
-
     Args:
       model: The modular Hamiltonian whose normalized exponential is the
         density operator governing the ensemble of states from which to sample.
@@ -103,3 +114,45 @@ def circuits(self, model: hamiltonian_model.Hamiltonian, num_samples: int):
     bitstrings, counts = utils.unique_bitstrings_with_counts(samples)
     states = model.circuit(bitstrings)
     return states, counts
+
+  def expectation(self, model: hamiltonian_model.Hamiltonian,
+                  ops: Union[tf.Tensor,
+                             hamiltonian_model.Hamiltonian], num_samples: int):
+    """Estimates observable expectation values against the density operator.
+
+    TODO(#119): add expectation and derivative equations and discussions
+                from updated paper.
+
+    Implicitly sample `num_samples` pure states from the canonical ensemble
+    corresponding to the thermal state defined by `model`.  For each such state
+    |psi>, estimate the expectation value <psi|op_j|psi> for each `ops[j]`.
+    Then, average these expectation values over the sampled states.
+
+    Args:
+      model: The modular Hamiltonian whose normalized exponential is the
+        density operator against which expectation values will be estimated.
+      ops: The observables to measure.  If `tf.Tensor`, strings with shape
+        [n_ops], result of calling `tfq.convert_to_tensor` on a list of
+        cirq.PauliSum, `[op1, op2, ...]`.  Otherwise, a Hamiltonian.
+      num_samples: Number of draws from the EBM associated with `model` to
+        average over.
+
+    Returns:
+      `tf.Tensor` with shape [n_ops] whose entries are are the sample averaged
+      expectation values of each entry in `ops`.
+    """
+    self.e_inference.infer(model.energy)
+    samples = self.e_inference.sample(num_samples)
+    bitstrings, counts = utils.unique_bitstrings_with_counts(samples)
+    if isinstance(ops, tf.Tensor):
+      return self.q_inference.expectation(
+          model.circuit, bitstrings, counts, ops, reduce=True)
+    elif isinstance(ops.energy, energy_model.PauliMixin):
+      u_dagger_u = model.circuit + ops.circuit_dagger
+      expectation_shards = self.q_inference.expectation(
+          u_dagger_u, bitstrings, counts, ops.operator_shards, reduce=True)
+      return tf.expand_dims(
+          ops.energy.operator_expectation(expectation_shards), 0)
+    else:
+      raise NotImplementedError(
+          "General `BitstringEnergy` models not yet supported.")


Test patch:
diff --git a/tests/hamiltonian_infer_test.py b/tests/hamiltonian_infer_test.py
index 48fa4973..d4ddc32e 100644
--- a/tests/hamiltonian_infer_test.py
+++ b/tests/hamiltonian_infer_test.py
@@ -15,6 +15,7 @@
 """Tests for the hamiltonian_infer module."""
 
 import absl
+from absl.testing import parameterized
 import random
 import string
 
@@ -25,13 +26,16 @@
 
 from qhbmlib import circuit_infer
 from qhbmlib import circuit_model
+from qhbmlib import circuit_model_utils
 from qhbmlib import energy_infer
 from qhbmlib import energy_model
 from qhbmlib import hamiltonian_model
 from qhbmlib import hamiltonian_infer
+from qhbmlib import utils
+from tests import test_util
 
 
-class QHBMTest(tf.test.TestCase):
+class QHBMTest(parameterized.TestCase, tf.test.TestCase):
   """Tests the QHBM class."""
 
   def setUp(self):
@@ -115,6 +119,122 @@ def circuit_func(model, num_samples):
       self.assertAllClose(
           actual_counts[0], actual_counts[1], atol=num_samples / 1000)
 
+  def test_expectation_cirq(self):
+    """Compares library expectation values to those from Cirq."""
+    # observable
+    num_bits = 4
+    num_ops = 3
+    qubits = cirq.GridQubit.rect(1, num_bits)
+    raw_ops = [test_util.get_random_pauli_sum(qubits) for _ in range(num_ops)]
+    ops = tfq.convert_to_tensor(raw_ops)
+
+    # unitary
+    n_moments = 10
+    act_fraction = 0.9
+    raw_circuit = cirq.testing.random_circuit(qubits, n_moments, act_fraction)
+
+    # hamiltonian model and inference
+    seed = tf.constant([5, 6], dtype=tf.int32)
+    energy = energy_model.BernoulliEnergy(list(range(num_bits)))
+    energy.build([None, num_bits])
+    circuit = circuit_model.DirectQuantumCircuit(raw_circuit)
+    circuit.build([])
+    actual_hamiltonian = hamiltonian_model.Hamiltonian(energy, circuit)
+    e_infer = energy_infer.BernoulliEnergyInference(seed=seed)
+    q_infer = circuit_infer.QuantumInference()
+    actual_h_infer = hamiltonian_infer.QHBM(e_infer, q_infer)
+
+    # sample bitstrings
+    num_samples = 1e6
+    e_infer.infer(energy)
+    samples = e_infer.sample(num_samples)
+    bitstrings, counts = utils.unique_bitstrings_with_counts(samples)
+    bit_list = bitstrings.numpy().tolist()
+
+    # bitstring injectors
+    bitstring_circuit = circuit_model_utils.bit_circuit(qubits)
+    bitstring_symbols = sorted(tfq.util.get_circuit_symbols(bitstring_circuit))
+    bitstring_resolvers = [
+        dict(zip(bitstring_symbols, bstr)) for bstr in bit_list
+    ]
+
+    # calculate expected values
+    total_circuit = bitstring_circuit + raw_circuit
+    raw_expectation_list = [[
+        cirq.Simulator().simulate_expectation_values(total_circuit, o,
+                                                     r)[0].real for o in raw_ops
+    ] for r in bitstring_resolvers]
+    expected_expectations = utils.weighted_average(counts, raw_expectation_list)
+
+    actual_expectations = actual_h_infer.expectation(actual_hamiltonian, ops,
+                                                     num_samples)
+    self.assertAllClose(actual_expectations, expected_expectations)
+
+  @parameterized.parameters({
+      "energy_class": energy_class,
+      "energy_args": energy_args,
+  } for energy_class, energy_args in zip(
+      [energy_model.BernoulliEnergy, energy_model.KOBE], [[], [2]]))
+  def test_expectation_modular_hamiltonian(self, energy_class, energy_args):
+    """Confirm expectation of modular Hamiltonians works."""
+    # set up the modular Hamiltonian to measure
+    num_bits = 3
+    n_moments = 5
+    act_fraction = 1.0
+    qubits = cirq.GridQubit.rect(1, num_bits)
+    energy_h = energy_class(*([list(range(num_bits))] + energy_args))
+    energy_h.build([None, num_bits])
+    raw_circuit_h = cirq.testing.random_circuit(qubits, n_moments, act_fraction)
+    circuit_h = circuit_model.DirectQuantumCircuit(raw_circuit_h)
+    circuit_h.build([])
+    hamiltonian_measure = hamiltonian_model.Hamiltonian(energy_h, circuit_h)
+    raw_shards = tfq.from_tensor(hamiltonian_measure.operator_shards)
+
+    # hamiltonian model and inference
+    seed = tf.constant([5, 6], dtype=tf.int32)
+    model_energy = energy_model.BernoulliEnergy(list(range(num_bits)))
+    model_energy.build([None, num_bits])
+    model_raw_circuit = cirq.testing.random_circuit(qubits, n_moments,
+                                                    act_fraction)
+    model_circuit = circuit_model.DirectQuantumCircuit(model_raw_circuit)
+    model_circuit.build([])
+    model_hamiltonian = hamiltonian_model.Hamiltonian(model_energy,
+                                                      model_circuit)
+    e_infer = energy_infer.BernoulliEnergyInference(seed=seed)
+    q_infer = circuit_infer.QuantumInference()
+    model_h_infer = hamiltonian_infer.QHBM(e_infer, q_infer)
+
+    # sample bitstrings
+    num_samples = 1e6
+    e_infer.infer(model_energy)
+    samples = e_infer.sample(num_samples)
+    bitstrings, counts = utils.unique_bitstrings_with_counts(samples)
+    bit_list = bitstrings.numpy().tolist()
+
+    # bitstring injectors
+    bitstring_circuit = circuit_model_utils.bit_circuit(qubits)
+    bitstring_symbols = sorted(tfq.util.get_circuit_symbols(bitstring_circuit))
+    bitstring_resolvers = [
+        dict(zip(bitstring_symbols, bstr)) for bstr in bit_list
+    ]
+
+    # calculate expected values
+    total_circuit = bitstring_circuit + model_raw_circuit + raw_circuit_h**-1
+    raw_expectation_list = [[
+        hamiltonian_measure.energy.operator_expectation([
+            cirq.Simulator().simulate_expectation_values(total_circuit, o,
+                                                         r)[0].real
+            for o in raw_shards
+        ])
+    ]
+                            for r in bitstring_resolvers]
+    expected_expectations = utils.weighted_average(counts, raw_expectation_list)
+
+    actual_expectations = model_h_infer.expectation(model_hamiltonian,
+                                                    hamiltonian_measure,
+                                                    num_samples)
+    self.assertAllClose(actual_expectations, expected_expectations)
+
 
 if __name__ == "__main__":
   absl.logging.info("Running hamiltonian_infer_test.py ...")