Implement serialization logic for ProjectorSums. (#623)

* Serialize Projector Sums

* nits

* Use Boolean to encode the basis state

* Address review comments

Author: tonybruguier

tensorflow_quantum/core/BUILD

@@ -12,6 +12,7 @@ py_library(
     deps = [
         "//tensorflow_quantum/core/ops",
         "//tensorflow_quantum/core/proto:pauli_sum_py_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_py_proto",
         "//tensorflow_quantum/core/serialize",
     ],
 )

tensorflow_quantum/core/ops/BUILD

@@ -82,6 +82,7 @@ cc_binary(
         ":parse_context",
         # cirq cc proto
         # pauli sum cc proto
+        # projector sum cc proto
         ":tfq_simulate_utils",
         "//tensorflow_quantum/core/src:adj_util",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
@@ -204,6 +205,7 @@ cc_binary(
         # cirq cc proto
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
         "//tensorflow_quantum/core/src:program_resolution",
         "//tensorflow_quantum/core/src:util_qsim",
@@ -286,6 +288,7 @@ cc_library(
         ":tfq_simulate_utils",
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "//tensorflow_quantum/core/src:program_resolution",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
@@ -344,6 +347,7 @@ cc_binary(
         # cirq cc proto
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
         "//tensorflow_quantum/core/src:util_qsim",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -374,6 +378,7 @@ py_library(
     deps = [
         ":load_module",
         # pauli sum cc proto
+        # projector sum cc proto
         # tensorflow framework for wrappers
     ],
 )
@@ -464,6 +469,7 @@ py_library(
         ":batch_util",
         "//tensorflow_quantum/core/proto:program_py_proto",
         "//tensorflow_quantum/core/proto:pauli_sum_py_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_py_proto",
         "//tensorflow_quantum/core/serialize:serializer",
     ],
 )

tensorflow_quantum/core/proto/BUILD

@@ -38,3 +38,20 @@ cc_proto_library(
         ":pauli_sum_proto",
     ],
 )
+
+py_proto_library(
+    name = "projector_sum_py_proto",
+    srcs = ["projector_sum.proto"],
+)
+
+proto_library(
+    name = "projector_sum_proto",
+    srcs = ["projector_sum.proto"],
+)
+
+cc_proto_library(
+    name = "projector_sum_cc_proto",
+    deps = [
+        ":projector_sum_proto",
+    ],
+)

tensorflow_quantum/core/proto/projector_sum.proto

@@ -0,0 +1,23 @@
+syntax = "proto3";
+
+package tfq.proto;
+
+// Store the sum of simpler terms.
+message ProjectorSum {
+  repeated ProjectorTerm terms = 1;
+}
+
+// Store a term which is a coefficient and the qubits of the projection.
+message ProjectorTerm {
+  float coefficient_real = 1;
+  float coefficient_imag = 2;
+  repeated ProjectorDictEntry projector_dict = 3;
+}
+
+// Store a single projection.
+message ProjectorDictEntry {
+  string qubit_id = 1;
+  // False means |0> and true means |1>.
+  bool basis_state = 2;
+}
+

tensorflow_quantum/core/serialize/BUILD

@@ -82,6 +82,7 @@ py_library(
         ":serializable_gate_set",
         "//tensorflow_quantum/core/proto:pauli_sum_py_proto",
         "//tensorflow_quantum/core/proto:program_py_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_py_proto",
     ],
 )
 
@@ -93,5 +94,6 @@ py_test(
         ":serializer",
         # cirq proto
         "//tensorflow_quantum/core/proto:pauli_sum_py_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_py_proto",
     ],
 )

tensorflow_quantum/core/serialize/serializer.py

@@ -24,6 +24,7 @@
     serializable_gate_set
 from tensorflow_quantum.core.proto import pauli_sum_pb2
 from tensorflow_quantum.core.proto import program_pb2
+from tensorflow_quantum.core.proto import projector_sum_pb2
 
 # Needed to allow autograph to crawl AST without erroring.
 _CONSTANT_TRUE = lambda x: True
@@ -971,3 +972,68 @@ def _process_pauli_type(char):
     if char == 'Y':
         return cirq.Y
     raise ValueError("Invalid pauli type.")
+
+
+def serialize_projectorsum(projectorsum):
+    """Constructs a projector_sum proto from `cirq.ProjectorSum`.
+
+    Args:
+        projectorsum: A `cirq.ProjectorSum` or `cirq.ProjectorString` object.
+
+    Returns:
+        A projector_sum proto object.
+    """
+    if isinstance(projectorsum, cirq.ProjectorString):
+        projectorsum = cirq.ProjectorSum.from_pauli_strings(projectorsum)
+
+    if not isinstance(projectorsum, cirq.ProjectorSum):
+        raise TypeError("serialize requires a cirq.ProjectorSum object."
+                        " Given: " + str(type(projectorsum)))
+
+    if any(not isinstance(qubit, (cirq.GridQubit, cirq.LineQubit))
+           for qubit in projectorsum.qubits):
+        raise ValueError("Attempted to serialize a paulisum that doesn't use "
+                         "only cirq.GridQubit or cirq.LineQubit.")
+
+    projectorsum_proto = projector_sum_pb2.ProjectorSum()
+    for term in projectorsum:
+        projectorterm_proto = projector_sum_pb2.ProjectorTerm()
+
+        projectorterm_proto.coefficient_real = term.coefficient.real
+        projectorterm_proto.coefficient_imag = term.coefficient.imag
+        for qubit, basis_state in sorted(
+                term.projector_dict.items()):  # sort to keep qubits ordered
+            projectorterm_proto.projector_dict.add(
+                qubit_id=op_serializer.qubit_to_proto(qubit),
+                basis_state=basis_state)
+
+        projectorsum_proto.terms.extend([projectorterm_proto])
+
+    return projectorsum_proto
+
+
+def deserialize_projectorsum(proto):
+    """Constructs a `cirq.ProjectorSum` from projector_sum proto.
+
+    Args:
+        proto: A projector_sum proto object.
+
+    Returns:
+        A `cirq.ProjectorSum` object.
+    """
+    if not isinstance(proto, projector_sum_pb2.ProjectorSum):
+        raise TypeError("deserialize requires a projector_sum_pb2 object."
+                        " Given: " + str(type(proto)))
+
+    res = cirq.ProjectorSum()
+    for term_proto in proto.terms:
+        coef = float(_round(term_proto.coefficient_real)) + \
+            1.0j * float(_round(term_proto.coefficient_imag))
+        projector_dict = {}
+        for projector_dict_entry in term_proto.projector_dict:
+            qubit = op_deserializer.qubit_from_proto(
+                projector_dict_entry.qubit_id)
+            projector_dict[qubit] = 1 if projector_dict_entry.basis_state else 0
+        res += cirq.ProjectorString(projector_dict, coef)
+
+    return res

tensorflow_quantum/core/serialize/serializer_test.py

@@ -22,6 +22,7 @@
 from absl.testing import parameterized
 from tensorflow_quantum.core.proto import pauli_sum_pb2
 from tensorflow_quantum.core.proto import program_pb2
+from tensorflow_quantum.core.proto import projector_sum_pb2
 from tensorflow_quantum.core.serialize import serializer
 
 
@@ -438,6 +439,37 @@ def _get_valid_pauli_proto_pairs(qubit_type='grid'):
     return pairs
 
 
+def _get_valid_projector_proto_pairs(qubit_type='grid'):
+    """Generate valid projectorsum proto pairs."""
+    q0 = cirq.GridQubit(0, 0)
+    q1 = cirq.GridQubit(1, 0)
+    q0_str = '0_0'
+    q1_str = '1_0'
+
+    if qubit_type == 'line':
+        q0 = cirq.LineQubit(0)
+        q1 = cirq.LineQubit(1)
+        q0_str = '0'
+        q1_str = '1'
+
+    pairs = [
+        (cirq.ProjectorSum.from_projector_strings(
+            cirq.ProjectorString(projector_dict={q0: 0})),
+         _build_projector_proto([1.0], [[0]], [[q0_str]])),
+        (cirq.ProjectorSum.from_projector_strings(
+            cirq.ProjectorString(projector_dict={q0: 0}, coefficient=0.125j)),
+         _build_projector_proto([0.125j], [[0]], [[q0_str]])),
+        (cirq.ProjectorSum.from_projector_strings([
+            cirq.ProjectorString(projector_dict={
+                q0: 0,
+                q1: 1
+            }),
+        ]), _build_projector_proto([1.0], [[0, 1]], [[q0_str, q1_str]])),
+    ]
+
+    return pairs
+
+
 def _get_noise_proto_pairs(qubit_type='grid'):
     q0 = cirq.GridQubit(0, 0)
     q0_str = '0_0'
@@ -506,6 +538,26 @@ def _build_pauli_proto(coefs, ops, qubit_ids):
     return a
 
 
+def _build_projector_proto(coefs, basis_states, qubit_ids):
+    """Construct projector_sum proto explicitly."""
+    terms = []
+    for i in range(len(coefs)):
+        term = projector_sum_pb2.ProjectorTerm()
+        term.coefficient_real = coefs[i].real
+        term.coefficient_imag = coefs[i].imag
+        for j in range(len(qubit_ids[i])):
+            if basis_states[i][j] == 0:
+                term.projector_dict.add(qubit_id=qubit_ids[i][j])
+            else:
+                term.projector_dict.add(qubit_id=qubit_ids[i][j],
+                                        basis_state=True)
+        terms.append(term)
+
+    a = projector_sum_pb2.ProjectorSum()
+    a.terms.extend(terms)
+    return a
+
+
 class SerializerTest(tf.test.TestCase, parameterized.TestCase):
     """Tests basic serializer functionality"""
 
@@ -715,6 +767,65 @@ def test_serialize_deserialize_paulisum_consistency(self, sum_proto_pair):
                 serializer.serialize_paulisum(sum_proto_pair[0])),
             sum_proto_pair[0])
 
+    @parameterized.parameters([{'inp': v} for v in ['wrong', 1.0, None, []]])
+    def test_serialize_projectorsum_wrong_type(self, inp):
+        """Attempt to serialize invalid object types."""
+        with self.assertRaises(TypeError):
+            serializer.serialize_projectorsum(inp)
+
+    @parameterized.parameters([{'inp': v} for v in ['wrong', 1.0, None, []]])
+    def test_deserialize_projectorsum_wrong_type(self, inp):
+        """Attempt to deserialize invalid object types."""
+        with self.assertRaises(TypeError):
+            serializer.deserialize_projectorsum(inp)
+
+    def test_serialize_projectorsum_invalid(self):
+        """Ensure we don't support anything but GridQubits."""
+        q0 = cirq.NamedQubit('wont work')
+        a = cirq.ProjectorSum.from_projector_strings(
+            cirq.ProjectorString(projector_dict={q0: 0}))
+        with self.assertRaises(ValueError):
+            serializer.serialize_projectorsum(a)
+
+    @parameterized.parameters(
+        [{
+            'sum_proto_pair': v
+        } for v in _get_valid_projector_proto_pairs(qubit_type='grid') +
+         _get_valid_projector_proto_pairs(qubit_type='line')])
+    def test_serialize_projectorsum_simple(self, sum_proto_pair):
+        """Ensure serialization is correct."""
+        self.assertProtoEquals(
+            sum_proto_pair[1],
+            serializer.serialize_projectorsum(sum_proto_pair[0]))
+
+    @parameterized.parameters(
+        [{
+            'sum_proto_pair': v
+        } for v in _get_valid_projector_proto_pairs(qubit_type='grid') +
+         _get_valid_projector_proto_pairs(qubit_type='line')])
+    def test_deserialize_projectorsum_simple(self, sum_proto_pair):
+        """Ensure deserialization is correct."""
+        self.assertEqual(serializer.deserialize_projectorsum(sum_proto_pair[1]),
+                         sum_proto_pair[0])
+
+    @parameterized.parameters(
+        [{
+            'sum_proto_pair': v
+        } for v in _get_valid_projector_proto_pairs(qubit_type='grid') +
+         _get_valid_projector_proto_pairs(qubit_type='line')])
+    def test_serialize_deserialize_projectorsum_consistency(
+            self, sum_proto_pair):
+        """Serialize and deserialize and ensure nothing changed."""
+        self.assertEqual(
+            serializer.serialize_projectorsum(
+                serializer.deserialize_projectorsum(sum_proto_pair[1])),
+            sum_proto_pair[1])
+
+        self.assertEqual(
+            serializer.deserialize_projectorsum(
+                serializer.serialize_projectorsum(sum_proto_pair[0])),
+            sum_proto_pair[0])
+
     @parameterized.parameters([
         {
             'gate': cirq.rx(3.0 * sympy.Symbol('alpha'))

tensorflow_quantum/core/src/BUILD

@@ -50,6 +50,7 @@ cc_library(
     deps = [
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "@com_google_absl//absl/container:flat_hash_map",
         "@local_config_tf//:libtensorflow_framework",
         "@local_config_tf//:tf_header_lib",
@@ -73,6 +74,7 @@ cc_test(
         ":circuit_parser_qsim",
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
@@ -91,6 +93,7 @@ cc_library(
     deps = [
         ":circuit_parser_qsim",
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "@com_google_absl//absl/container:inlined_vector",  # unclear why needed.
         "@local_config_tf//:libtensorflow_framework",
         "@local_config_tf//:tf_header_lib",
@@ -120,6 +123,7 @@ cc_library(
     deps = [
         "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
         "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/proto:projector_sum_cc_proto",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/strings",