JSON Protocol (#1880)

Protocol to support serializing cirq objects

 - `cirq.to_json`
 - `cirq.read_json`

The former looks for a magic method `_json_dict_(self)` which returns a dictionary keyed by strings and whose values are either basic python types or are objects that support the SupportsJSON protocol. They dictionary must have an entry with key `"cirq_type"`, to be used during loading. The latter looks for a magic classmethod `_from_json_dict_(cls, **kwargs)` or falls back to calling the constructor.

`_json_dict_` is required for objects to support serialization, so `cirq.protocols.to_json_dict` is provided as convenience to implementers. You must only provide a list of attribute names which will be included in the dictionary.

During deserialization, the `"cirq_type"` value is passed to an ordered list of functions. Each function either *resolves* the string name into a python type or the next function is called. A default mapping from class name is provided as a dictionary. 

There is a test fixture that tests any object exposed at the top `cirq`-level import by roundtripping to JSON. There is a blacklist of classes for which it doesn't make sense to serialize, and another "xfail" list for objects that haven't been upgraded to support serialization yet, but there's no reason not for them to.

This pull request includes support for Qubits, Gates, Operations, Moments, Circuits, PauliString, and a couple of other things. I recommend reviewing and merging this PR which extends support for these high-priority objects and then working on the rest over time. Any new objects should implement the serialization protocol unless there's good reason not to. 

This pull request makes some small changes to objects, mostly adding `@property`s to provide a consistent expectation that arguments passed to the constructor are available under the same name as attributes/properties.

Author: mpharrigan

cirq/__init__.py

@@ -335,6 +335,7 @@
     qasm,
     QasmArgs,
     qid_shape,
+    read_json,
     resolve_parameters,
     SupportsApplyChannel,
     SupportsConsistentApplyUnitary,
@@ -353,6 +354,8 @@
     SupportsQasmWithArgsAndQubits,
     SupportsTraceDistanceBound,
     SupportsUnitary,
+    to_json,
+    to_json_dict,
     trace_distance_bound,
     unitary,
     validate_mixture,

cirq/circuits/circuit.py

@@ -1578,6 +1578,13 @@ def save_qasm(self,
         """
         self._to_qasm_output(header, precision, qubit_order).save(file_path)
 
+    @property
+    def moments(self):
+        return self._moments
+
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['moments', 'device'])
+
 
 def _resolve_operations(
         operations: Iterable[ops.Operation],

cirq/circuits/circuit_test.py

@@ -3180,3 +3180,20 @@ def test_deprecated_to_unitary_matrix():
 def test_deprecated_apply_unitary_effect_to_state():
     np.testing.assert_allclose(cirq.Circuit().apply_unitary_effect_to_state(),
                                cirq.Circuit().final_wavefunction())
+
+
+def test_moments_property():
+    q = cirq.NamedQubit('q')
+    c = cirq.Circuit.from_ops(cirq.X(q), cirq.Y(q))
+    assert c.moments[0] == cirq.Moment([cirq.X(q)])
+    assert c.moments[1] == cirq.Moment([cirq.Y(q)])
+
+
+def test_json_dict():
+    q0, q1 = cirq.LineQubit.range(2)
+    c = cirq.Circuit.from_ops(cirq.CNOT(q0, q1))
+    assert c._json_dict_() == {
+        'cirq_type': 'Circuit',
+        'moments': [cirq.Moment([cirq.CNOT(q0, q1)])],
+        'device': cirq.UNCONSTRAINED_DEVICE,
+    }

cirq/devices/grid_qubit.py

@@ -15,7 +15,7 @@
 
 from typing import Dict, List, Tuple
 
-from cirq import ops
+from cirq import ops, protocols
 
 
 class GridQubit(ops.Qid):
@@ -139,6 +139,9 @@ def __repr__(self):
     def __str__(self):
         return '({}, {})'.format(self.row, self.col)
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['row', 'col'])
+
     def __add__(self, other: Tuple[int, int]) -> 'GridQubit':
         if not (isinstance(other, tuple) and len(other) == 2 and
                 all(isinstance(x, int) for x in other)):

cirq/devices/grid_qubit_test.py

@@ -167,3 +167,13 @@ def test_from_proto_bad_dict():
         cirq.GridQubit.from_proto_dict({})
     with pytest.raises(ValueError):
         cirq.GridQubit.from_proto_dict({'nothing': 1})
+
+
+def test_to_json():
+    q = cirq.GridQubit(5, 6)
+    d = q._json_dict_()
+    assert d == {
+        'cirq_type': 'GridQubit',
+        'row': 5,
+        'col': 6,
+    }

cirq/devices/line_qubit.py

@@ -15,7 +15,7 @@
 import functools
 from typing import List
 
-from cirq import ops
+from cirq import ops, protocols
 
 
 @functools.total_ordering
@@ -85,3 +85,6 @@ def __rsub__(self, other: int) -> 'LineQubit':
 
     def __neg__(self) -> 'LineQubit':
         return LineQubit(-self.x)
+
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['x'])

cirq/devices/line_qubit_test.py

@@ -107,3 +107,10 @@ def test_addition_subtraction_type_error():
 
 def test_neg():
     assert -cirq.LineQubit(1) == cirq.LineQubit(-1)
+
+
+def test_json_dict():
+    assert cirq.LineQubit(5)._json_dict_() == {
+        'cirq_type': 'LineQubit',
+        'x': 5,
+    }

cirq/devices/unconstrained_device.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cirq import value
+from cirq import value, protocols
 from cirq.devices import device
 
 
+@value.value_equality()
 class _UnconstrainedDevice(device.Device):
     """A device that allows everything, infinitely fast."""
 
@@ -37,5 +38,11 @@ def validate_schedule(self, schedule):
     def __repr__(self):
         return 'cirq.UNCONSTRAINED_DEVICE'
 
+    def _value_equality_values_(self):
+        return ()
+
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, [])
+
 
 UNCONSTRAINED_DEVICE: device.Device = _UnconstrainedDevice()

cirq/ops/common_gates.py

@@ -552,6 +552,20 @@ def __repr__(self):
     def _value_equality_values_(self):
         return self.num_qubits(), self.key, self.invert_mask
 
+    def _json_dict_(self):
+        return {
+            'cirq_type': self.__class__.__name__,
+            'num_qubits': self.num_qubits(),
+            'key': self.key,
+            'invert_mask': self.invert_mask
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, num_qubits, key, invert_mask, **kwargs):
+        return cls(num_qubits=num_qubits,
+                   key=key,
+                   invert_mask=tuple(invert_mask))
+
 
 def _default_measurement_key(qubits: Iterable[raw_types.Qid]) -> str:
     return ','.join(str(q) for q in qubits)
@@ -697,6 +711,12 @@ def _qasm_(self, args: protocols.QasmArgs,
     def _value_equality_values_(self):
         return self.num_qubits(),
 
+    def _json_dict_(self):
+        return {
+            'cirq_type': self.__class__.__name__,
+            'num_qubits': self.num_qubits(),
+        }
+
 
 class HPowGate(eigen_gate.EigenGate, gate_features.SingleQubitGate):
     """A Gate that performs a rotation around the X+Z axis of the Bloch sphere.

cirq/ops/eigen_gate.py

@@ -111,6 +111,10 @@ def __init__(
     def exponent(self) -> value.TParamVal:
         return self._exponent
 
+    @property
+    def global_shift(self) -> float:
+        return self._global_shift
+
     # virtual method
     def _with_exponent(self: TSelf, exponent: value.TParamVal) -> TSelf:
         """Return the same kind of gate, but with a different exponent.
@@ -319,6 +323,9 @@ def _resolve_parameters_(self: TSelf, param_resolver) -> TSelf:
         return self._with_exponent(
                 exponent=param_resolver.value_of(self._exponent))
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['exponent', 'global_shift'])
+
 
 def _lcm(vals: Iterable[int]) -> int:
     t = 1

cirq/ops/fsim_gate.py

@@ -154,6 +154,9 @@ def __repr__(self):
         return 'cirq.FSimGate(theta={}, phi={})'.format(proper_repr(self.theta),
                                                         proper_repr(self.phi))
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['theta', 'phi'])
+
 
 def _format_rads(args: 'cirq.CircuitDiagramInfoArgs', radians: float) -> str:
     if cirq.is_parameterized(radians):

cirq/ops/fsim_gate_test.py

@@ -234,3 +234,11 @@ def test_fsim_iswap_cphase(theta, phi):
     iswap_cphase = cirq.Circuit.from_ops((iswap.on(q0, q1), cphase.on(q0, q1)))
     fsim = cirq.FSimGate(theta=theta, phi=phi)
     assert np.allclose(cirq.unitary(iswap_cphase), cirq.unitary(fsim))
+
+
+def test_fsim_json_dict():
+    assert cirq.FSimGate(theta=0.123, phi=0.456)._json_dict_() == {
+        'cirq_type': 'FSimGate',
+        'theta': 0.123,
+        'phi': 0.456,
+    }

cirq/ops/gate_operation.py

@@ -71,6 +71,9 @@ def __str__(self):
         return '{}({})'.format(self.gate,
                                ', '.join(str(e) for e in self.qubits))
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['gate', 'qubits'])
+
     def _group_interchangeable_qubits(self) -> Tuple[
             Union[raw_types.Qid,
                   Tuple[int, FrozenSet[raw_types.Qid]]],

cirq/ops/global_phase_op.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 
-from cirq import value
+from cirq import value, protocols
 from cirq.ops import raw_types
 
 
@@ -62,3 +62,6 @@ def __str__(self):
 
     def __repr__(self):
         return 'cirq.GlobalPhaseOperation({!r})'.format(self.coefficient)
+
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['coefficient'])

cirq/ops/global_phase_op_test.py

@@ -182,3 +182,10 @@ def test_diagram():
         ]), """
 global phase:   -0.5π
         """)
+
+
+def test_global_phase_op_json_dict():
+    assert cirq.GlobalPhaseOperation(-1j)._json_dict_() == {
+        'cirq_type': 'GlobalPhaseOperation',
+        'coefficient': -1j,
+    }

cirq/ops/moment.py

@@ -16,6 +16,7 @@
 
 from typing import Any, Callable, Iterable, Sequence, TypeVar, Union
 
+from cirq import protocols
 from cirq.protocols import approx_eq
 from cirq.ops import raw_types
 
@@ -147,6 +148,9 @@ def transform_qubits(self: TSelf_Moment,
         return self.__class__(op.transform_qubits(func)
                 for op in self.operations)
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['operations'])
+
 
 def _list_repr_with_indented_item_lines(items: Sequence[Any]) -> str:
     block = '\n'.join([repr(op) + ',' for op in items])

cirq/ops/moment_test.py

@@ -203,7 +203,7 @@ def test_qubits():
     a = cirq.NamedQubit('a')
     b = cirq.NamedQubit('b')
 
-    assert Moment([cirq.X(a), cirq.X(b)]).qubits == {a , b}
+    assert Moment([cirq.X(a), cirq.X(b)]).qubits == {a, b}
     assert Moment([cirq.X(a)]).qubits == {a}
     assert Moment([cirq.CZ(a, b)]).qubits == {a, b}
 
@@ -225,3 +225,13 @@ def test_bool():
     assert not Moment()
     a = cirq.NamedQubit('a')
     assert Moment([cirq.X(a)])
+
+
+def test_json_dict():
+    a = cirq.NamedQubit('a')
+    b = cirq.NamedQubit('b')
+    mom = Moment([cirq.CZ(a, b)])
+    assert mom._json_dict_() == {
+        'cirq_type': 'Moment',
+        'operations': (cirq.CZ(a, b),)
+    }

cirq/ops/named_qubit.py

@@ -11,7 +11,7 @@
 # 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.
-
+from cirq import protocols
 from cirq.ops import raw_types
 
 
@@ -60,6 +60,9 @@ def range(*args, prefix: str):
         """
         return [NamedQubit(prefix + str(i)) for i in range(*args)]
 
+    def _json_dict_(self):
+        return protocols.to_json_dict(self, ['name'])
+
 
 def _pad_digits(text: str) -> str:
     """A str method with hacks to support better lexicographic ordering.

cirq/ops/pauli_gates.py

@@ -96,6 +96,11 @@ def __pow__(self: '_PauliX',
                 exponent: value.TParamVal) -> common_gates.XPowGate:
         return common_gates.XPowGate(exponent=exponent)
 
+    @classmethod
+    def _from_json_dict_(cls, exponent, global_shift, **kwargs):
+        assert global_shift == 0
+        return cls(exponent=exponent)
+
 
 class _PauliY(Pauli, common_gates.YPowGate):
 
@@ -107,6 +112,11 @@ def __pow__(self: '_PauliY',
                 exponent: value.TParamVal) -> common_gates.YPowGate:
         return common_gates.YPowGate(exponent=exponent)
 
+    @classmethod
+    def _from_json_dict_(cls, exponent, global_shift, **kwargs):
+        assert global_shift == 0
+        return cls(exponent=exponent)
+
 
 class _PauliZ(Pauli, common_gates.ZPowGate):
 
@@ -118,6 +128,11 @@ def __pow__(self: '_PauliZ',
                 exponent: value.TParamVal) -> common_gates.ZPowGate:
         return common_gates.ZPowGate(exponent=exponent)
 
+    @classmethod
+    def _from_json_dict_(cls, exponent, global_shift, **kwargs):
+        assert global_shift == 0
+        return cls(exponent=exponent)
+
 
 # The Pauli X gate.
 #
@@ -127,7 +142,6 @@ def __pow__(self: '_PauliZ',
 #    [1, 0]]
 X = _PauliX()
 
-
 # The Pauli Y gate.
 #
 # Matrix:
@@ -136,7 +150,6 @@ def __pow__(self: '_PauliZ',
 #      [i, 0]]
 Y = _PauliY()
 
-
 # The Pauli Z gate.
 #
 # Matrix:
@@ -145,5 +158,4 @@ def __pow__(self: '_PauliZ',
 #      [0, -1]]
 Z = _PauliZ()
 
-
 Pauli._XYZ = (X, Y, Z)