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

Issue:
Density matrix tools
Bring back fidelity and density matrix tools for new style QHBMs.


Discussion / hints:

Fix patch:
diff --git a/qhbmlib/circuit_infer_utils.py b/qhbmlib/circuit_infer_utils.py
new file mode 100644
index 00000000..33dc45da
--- /dev/null
+++ b/qhbmlib/circuit_infer_utils.py
@@ -0,0 +1,32 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for metrics on QuantumCircuit."""
+
+import tensorflow as tf
+import tensorflow_quantum as tfq
+
+from qhbmlib import circuit_model
+
+
+def unitary(circuit: circuit_model.QuantumCircuit):
+  """Returns the unitary matrix corresponding to the given circuit.
+
+  Args:
+    circuit: Quantum circuit whose unitary matrix is to be calculated.
+  """
+  return tfq.layers.Unitary()(
+      circuit.pqc,
+      symbol_names=circuit.symbol_names,
+      symbol_values=tf.expand_dims(circuit.symbol_values, 0)).to_tensor()[0]
diff --git a/qhbmlib/energy_infer_utils.py b/qhbmlib/energy_infer_utils.py
new file mode 100644
index 00000000..b0991172
--- /dev/null
+++ b/qhbmlib/energy_infer_utils.py
@@ -0,0 +1,35 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for metrics on BitstringEnergy."""
+
+import itertools
+
+import tensorflow as tf
+
+from qhbmlib import energy_model
+
+
+def probabilities(energy: energy_model.BitstringEnergy):
+  """Returns the probabilities of the EBM.
+
+  Args:
+    energy: The energy function defining the EBM.
+  """
+  all_bitstrings = tf.constant(
+      list(itertools.product([0, 1], repeat=energy.num_bits)), dtype=tf.int8)
+  all_energies = energy(all_bitstrings)
+  energy_exp = tf.math.exp(-all_energies)
+  partition = tf.math.reduce_sum(energy_exp)
+  return energy_exp / partition
diff --git a/qhbmlib/hamiltonian_infer_utils.py b/qhbmlib/hamiltonian_infer_utils.py
new file mode 100644
index 00000000..ad1e4dec
--- /dev/null
+++ b/qhbmlib/hamiltonian_infer_utils.py
@@ -0,0 +1,117 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for metrics on Hamiltonian."""
+
+import tensorflow as tf
+
+from qhbmlib import circuit_infer_utils
+from qhbmlib import energy_infer_utils
+from qhbmlib import hamiltonian_model
+
+
+def density_matrix(model: hamiltonian_model.Hamiltonian):
+  r"""Returns the thermal state corresponding to a modular Hamiltonian.
+
+  Given a modular Hamiltonian $K_{\theta\phi} = U_\phi K_\theta U_\phi^\dagger$,
+  the corresponding thermal state is
+
+  \begin{align*}
+    \rho &= (\text{tr}[e^{-K_{\theta\phi}}])^{-1} e^{-K_{\theta\phi}}
+    \\&= Z_\theta^{-1} U_\phi e^{-K_\theta} U_\phi^\dagger
+    \\&= U_\phi P_\theta U_\phi^\dagger
+  \end{align*}
+  where we defined the diagonal matrix $P_\theta$ as
+  $$
+    \langle x|P_\theta|y\rangle = \begin{cases}
+                           p_\theta(x), & \text{if}\ x = y \\
+                           0, & \text{otherwise}
+                    \end{cases}
+  $$
+  Continuing, using the definition of matrix multiplication, we have
+  \begin{align*}
+    \rho &= U_\phi \sum_{i,k,j} |i\rangle\langle i|P_\theta|k\rangle
+            \langle k| U_\phi^\dagger|j\rangle\langle j|
+    \\&= U_\phi \sum_{k,j} p_\theta(k)|k\rangle
+            \langle k| U_\phi^\dagger|j\rangle\langle j|
+    \\&= \sum_{i,k,j} p_\theta(k)|i\rangle\langle i|U_\phi|k\rangle
+            \langle k| U_\phi^\dagger|j\rangle\langle j|
+  \end{align*}
+
+  Args:
+    model: Modular Hamiltonian whose corresponding thermal state is the density
+      matrix to be calculated.
+  """
+  probabilities = tf.cast(
+      energy_infer_utils.probabilities(model.energy), tf.complex64)
+  unitary_matrix = circuit_infer_utils.unitary(model.circuit)
+  return tf.einsum("k,ik,kj->ij", probabilities, unitary_matrix,
+                   tf.linalg.adjoint(unitary_matrix))
+
+
+def fidelity(model: hamiltonian_model.Hamiltonian, sigma: tf.Tensor):
+  r"""Calculate the fidelity between a QHBM and a density matrix.
+
+  Definition of the fidelity between two quantum states $\rho$ and $\sigma$ is
+  $$
+  F(\rho, \sigma) = \left(\text{tr}\sqrt{\sqrt{\rho}\sigma\sqrt{\rho}}\right)^2.
+  $$
+  When the first argument is a QHBM, we can write
+  $$
+    F(\rho, \sigma) = \left(\text{tr}\sqrt{
+        U_\phi\sqrt{K_\theta}U_\phi^\dagger
+        \sigma U_\phi\sqrt{K_\theta}U_\phi^\dagger}\right)^2.
+  $$
+  By the definition of matrix functions, we can pull the unitaries outside
+  the square root, and use cyclicity of trace to remove them.  Then we have
+  $$
+    F(\rho, \sigma) = \left(\text{tr}\sqrt{
+        \sqrt{K_\theta}U_\phi^\dagger\sigma U_\phi\sqrt{K_\theta}}\right)^2.
+  $$
+  Let $\omega = \sqrt{K_\theta}U_\phi^\dagger\sigma U_\phi\sqrt{K_\theta}$,
+  and let $WD W^\dagger$ be a unitary diagonalization of $\omega$.  Note that
+  $U_\phi^\dagger\sigma U_\phi$ is Hermitian since it is a unitary
+  conjugation of a Hermitian matrix.  Also, for Hermitian matrices A and B,
+  we have
+  $$
+    (ABA)^\dagger = (BA)^\dagger A^\dagger
+                  = A^\dagger B^\dagger A^\dagger
+                  = ABA.
+  $$
+  Therefore $ABA$ is also Hermitian.  Thus $\omega$ is Hermitian, which
+  allows the use of faster eigenvalue finding routines.  Then we have
+  $$
+    F(\rho, \sigma) = \left(\text{tr}\sqrt{D}\right)^2
+                    = \left(\sum_i\sqrt{D_{ii}}\right)^2.
+  $$
+
+  Args:
+    model: Modular Hamiltonian whose corresponding thermal state is to be
+      compared to `sigma`, as the left density matrix in fidelity.
+    sigma: 2-D `tf.Tensor` of dtype `complex64` representing the right
+      density matrix in the fidelity calculation.
+
+  Returns:
+    A scalar `tf.Tensor` which is the fidelity between the density matrix
+      represented by this QHBM and `sigma`.
+  """
+  k_theta = tf.cast(
+      energy_infer_utils.probabilities(model.energy), tf.complex64)
+  u_phi = circuit_infer_utils.unitary(model.circuit)
+  u_phi_dagger = tf.linalg.adjoint(u_phi)
+  sqrt_k_theta = tf.sqrt(k_theta)
+  omega = tf.einsum("a,ab,bc,cd,d->ad", sqrt_k_theta, u_phi_dagger, sigma,
+                    u_phi, sqrt_k_theta)
+  d_omega = tf.linalg.eigvalsh(omega)
+  return tf.math.reduce_sum(tf.math.sqrt(d_omega))**2


Test patch:
diff --git a/tests/circuit_infer_utils_test.py b/tests/circuit_infer_utils_test.py
new file mode 100644
index 00000000..64fe54bd
--- /dev/null
+++ b/tests/circuit_infer_utils_test.py
@@ -0,0 +1,74 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the circuit_infer_utils module."""
+
+from absl import logging
+import random
+import string
+
+import cirq
+import tensorflow as tf
+import tensorflow_quantum as tfq
+from tensorflow_quantum.python import util as tfq_util
+
+from qhbmlib import circuit_infer_utils
+from qhbmlib import circuit_model
+
+
+class UnitaryTest(tf.test.TestCase):
+  """Tests the unitary function."""
+
+  def setUp(self):
+    """Initializes test objects."""
+    super().setUp()
+
+    self.num_bits = 4
+    self.tf_random_seed = 7
+    self.close_rtol = 1e-4
+
+  def test_unitary(self):
+    """Confirms unitary is correct for a random circuit."""
+    # TODO(#171) this random circuit construction happens in a few places.
+    qubits = cirq.GridQubit.rect(1, self.num_bits)
+    batch_size = 1
+    n_moments = 10
+    act_fraction = 0.9
+    num_symbols = 2
+    symbols = set()
+    for _ in range(num_symbols):
+      symbols.add("".join(random.sample(string.ascii_letters, 10)))
+    symbols = sorted(list(symbols))
+    raw_circuits, _ = tfq_util.random_symbol_circuit_resolver_batch(
+        qubits, symbols, batch_size, n_moments=n_moments, p=act_fraction)
+    raw_circuit = raw_circuits[0]
+    random_values = tf.random.uniform([len(symbols)], -1, 1, tf.float32,
+                                      self.tf_random_seed).numpy().tolist()
+    resolver = dict(zip(symbols, random_values))
+
+    circuit = circuit_model.QuantumCircuit(
+        tfq.convert_to_tensor([raw_circuit]), qubits, tf.constant(symbols),
+        [tf.Variable([resolver[s] for s in symbols])], [[]])
+
+    resolved_circuit = cirq.protocols.resolve_parameters(raw_circuit, resolver)
+    expected_unitary = resolved_circuit.unitary()
+
+    unitary_wrapper = tf.function(circuit_infer_utils.unitary)
+    actual_unitary = unitary_wrapper(circuit)
+    self.assertAllClose(actual_unitary, expected_unitary, rtol=self.close_rtol)
+
+
+if __name__ == "__main__":
+  logging.info("Running circuit_infer_utils_test.py ...")
+  tf.test.main()
diff --git a/tests/energy_infer_utils_test.py b/tests/energy_infer_utils_test.py
new file mode 100644
index 00000000..c60b0127
--- /dev/null
+++ b/tests/energy_infer_utils_test.py
@@ -0,0 +1,60 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the energy_infer_utils module."""
+
+import random
+
+import tensorflow as tf
+
+from qhbmlib import energy_infer
+from qhbmlib import energy_infer_utils
+from qhbmlib import energy_model
+from qhbmlib import utils
+
+
+class ProbabilitiesTest(tf.test.TestCase):
+  """Tests the probabilities function."""
+
+  def test_probabilities(self):
+    """Confirms probabilities are correct for an MLP."""
+    num_bits = 5
+    num_layers = 3
+    units = random.sample(range(1, 100), num_layers)
+    activations = random.sample([
+        "elu", "exponential", "gelu", "hard_sigmoid", "linear", "relu", "selu",
+        "sigmoid", "softmax", "softplus", "softsign", "swish", "tanh"
+    ], num_layers)
+    expected_layer_list = []
+    for i in range(num_layers):
+      expected_layer_list.append(
+          tf.keras.layers.Dense(units[i], activation=activations[i]))
+    expected_layer_list.append(tf.keras.layers.Dense(1))
+    expected_layer_list.append(utils.Squeeze(-1))
+    energy = energy_model.BitstringEnergy(
+        list(range(num_bits)), expected_layer_list)
+
+    num_expectation_samples = 1  # Required but unused
+    infer = energy_infer.AnalyticEnergyInference(energy,
+                                                 num_expectation_samples)
+    expected_probabilities = infer.distribution.probs_parameter()
+
+    probabilities_wrapped = tf.function(energy_infer_utils.probabilities)
+    actual_probabilities = probabilities_wrapped(energy)
+    self.assertAllClose(actual_probabilities, expected_probabilities)
+
+
+if __name__ == "__main__":
+  print("Running energy_infer_utils_test.py ...")
+  tf.test.main()
diff --git a/tests/hamiltonian_infer_utils_test.py b/tests/hamiltonian_infer_utils_test.py
new file mode 100644
index 00000000..c9f978b6
--- /dev/null
+++ b/tests/hamiltonian_infer_utils_test.py
@@ -0,0 +1,116 @@
+# Copyright 2021 The QHBM Library Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the hamiltonian_infer_utils module."""
+
+import cirq
+import tensorflow as tf
+
+from qhbmlib import circuit_model
+from qhbmlib import energy_model
+from qhbmlib import hamiltonian_model
+from qhbmlib import hamiltonian_infer_utils
+from tests import test_util
+
+
+class DensityMatrixTest(tf.test.TestCase):
+  """Tests the density_matrix function."""
+
+  @test_util.eager_mode_toggle
+  def test_density_matrix(self):
+    """Confirms the density matrix represented by the QHBM is correct."""
+
+    # Check density matrix of Bell state.
+    num_bits = 2
+    energy = energy_model.BernoulliEnergy(list(range(num_bits)))
+    energy.build([None, num_bits])
+    energy.set_weights([tf.constant([-10.0, -10.0])])  # pin at |00>
+
+    qubits = cirq.GridQubit.rect(1, num_bits)
+    test_u = cirq.Circuit([cirq.H(qubits[0]), cirq.CNOT(qubits[0], qubits[1])])
+    circuit = circuit_model.DirectQuantumCircuit(test_u)
+    circuit.build([])
+    model = hamiltonian_model.Hamiltonian(energy, circuit)
+
+    expected_dm = tf.constant(
+        [[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]],
+        tf.complex64,
+    )
+
+    density_matrix_wrapper = tf.function(hamiltonian_infer_utils.density_matrix)
+    actual_dm = density_matrix_wrapper(model)
+    self.assertAllClose(actual_dm, expected_dm)
+
+
+class FidelityTest(tf.test.TestCase):
+  """Tests the fidelity function."""
+
+  def setUp(self):
+    """Initializes test objects."""
+    super().setUp()
+    self.close_rtol = 1e-4
+
+  @test_util.eager_mode_toggle
+  def test_fidelity_self(self):
+    """Confirms the fidelity of a model with itself is 1."""
+    num_bits = 4
+    num_samples = 1  # required but not used
+    qubits = cirq.GridQubit.rect(1, num_bits)
+    num_layers = 3
+    model, _ = test_util.get_random_hamiltonian_and_inference(
+        qubits, num_layers, "test_fidelity", num_samples)
+
+    density_matrix_wrapper = tf.function(hamiltonian_infer_utils.density_matrix)
+    model_dm = density_matrix_wrapper(model)
+    fidelity_wrapper = tf.function(hamiltonian_infer_utils.fidelity)
+    actual_fidelity = fidelity_wrapper(model, model_dm)
+
+    expected_fidelity = 1.0
+    self.assertAllClose(
+        actual_fidelity, expected_fidelity, rtol=self.close_rtol)
+
+  def test_fidelity_random(self):
+    """Confirms correct fidelity against slower direct formula."""
+
+    def direct_fidelity(rho, sigma):
+      """Direct matrix to matrix fidelity function."""
+      sqrt_rho = tf.linalg.sqrtm(rho)
+      intermediate = tf.linalg.sqrtm(sqrt_rho @ sigma @ sqrt_rho)
+      return tf.linalg.trace(intermediate)**2
+
+    num_rerolls = 5
+    for _ in range(num_rerolls):
+      num_qubits = 4
+      sigma, _ = test_util.generate_mixed_random_density_operator(
+          num_qubits, 2**num_qubits)
+      sigma = tf.cast(sigma, tf.complex64)
+
+      qubits = cirq.GridQubit.rect(num_qubits, 1)
+      num_layers = 3
+      identifier = "fidelity_test"
+      num_samples = 1  # required but unused
+      h, _ = test_util.get_random_hamiltonian_and_inference(
+          qubits, num_layers, identifier, num_samples)
+      h_dm = hamiltonian_infer_utils.density_matrix(h)
+
+      expected_fidelity = direct_fidelity(h_dm, sigma)
+      fidelity_wrapper = tf.function(hamiltonian_infer_utils.fidelity)
+      actual_fidelity = fidelity_wrapper(h, sigma)
+      self.assertAllClose(
+          actual_fidelity, expected_fidelity, rtol=self.close_rtol)
+
+
+if __name__ == "__main__":
+  print("Running hamiltonian_infer_utils_test.py ...")
+  tf.test.main()
diff --git a/tests/vqt_loss_test.py b/tests/vqt_loss_test.py
index b6bace60..9f8e4fe3 100644
--- a/tests/vqt_loss_test.py
+++ b/tests/vqt_loss_test.py
@@ -36,7 +36,7 @@ class VQTTest(tf.test.TestCase):
   def setUp(self):
     """Initializes test objects."""
     super().setUp()
-    self.num_qubits_list = [1, 2, 3]
+    self.num_qubits_list = [1, 2]
     self.tf_random_seed = 5
     self.tf_random_seed_alt = 6
     self.tfp_seed = tf.constant([7, 8], tf.int32)