Support serialization of GateFamilies (quantumlib#4715)

Exactly what it says: this PR uses the new type-serialization behavior to allow GateFamilies to be serialized.

Additional work is still required for Gateset serialization - that behavior is not included in this PR.

cirq/json_resolver_cache.py

@@ -51,7 +51,9 @@ def _parallel_gate_op(gate, qubits):
     import sympy
 
     return {
+        'AnyIntegerPowerGateFamily': cirq.AnyIntegerPowerGateFamily,
         'AmplitudeDampingChannel': cirq.AmplitudeDampingChannel,
+        'AnyUnitaryGateFamily': cirq.AnyUnitaryGateFamily,
         'AsymmetricDepolarizingChannel': cirq.AsymmetricDepolarizingChannel,
         'BitFlipChannel': cirq.BitFlipChannel,
         'BitstringAccumulator': cirq.work.BitstringAccumulator,
@@ -81,6 +83,7 @@ def _parallel_gate_op(gate, qubits):
         'MutableDensePauliString': cirq.MutableDensePauliString,
         'MutablePauliString': cirq.MutablePauliString,
         'ObservableMeasuredResult': cirq.work.ObservableMeasuredResult,
+        'GateFamily': cirq.GateFamily,
         'GateOperation': cirq.GateOperation,
         'GeneralizedAmplitudeDampingChannel': cirq.GeneralizedAmplitudeDampingChannel,
         'GlobalPhaseOperation': cirq.GlobalPhaseOperation,
@@ -115,6 +118,7 @@ def _parallel_gate_op(gate, qubits):
         '_PauliZ': cirq.ops.pauli_gates._PauliZ,
         'ParamResolver': cirq.ParamResolver,
         'ParallelGate': cirq.ParallelGate,
+        'ParallelGateFamily': cirq.ParallelGateFamily,
         'PauliMeasurementGate': cirq.PauliMeasurementGate,
         'PauliString': cirq.PauliString,
         'PhaseDampingChannel': cirq.PhaseDampingChannel,

cirq/ops/common_gate_families.py

@@ -53,6 +53,13 @@ def __repr__(self) -> str:
     def _value_equality_values_(self) -> Any:
         return self._num_qubits
 
+    def _json_dict_(self):
+        return {'num_qubits': self._num_qubits}
+
+    @classmethod
+    def _from_json_dict_(cls, num_qubits, **kwargs):
+        return cls(num_qubits)
+
 
 class AnyIntegerPowerGateFamily(gateset.GateFamily):
     """GateFamily which accepts instances of a given `cirq.EigenGate`, raised to integer power."""
@@ -87,6 +94,15 @@ def __repr__(self) -> str:
     def _value_equality_values_(self) -> Any:
         return self.gate
 
+    def _json_dict_(self):
+        return {'gate': self._gate_json()}
+
+    @classmethod
+    def _from_json_dict_(cls, gate, **kwargs):
+        if isinstance(gate, str):
+            gate = protocols.cirq_type_from_json(gate)
+        return cls(gate)
+
 
 class ParallelGateFamily(gateset.GateFamily):
     """GateFamily which accepts instances of `cirq.ParallelGate` and it's sub_gate.
@@ -175,3 +191,19 @@ def __repr__(self) -> str:
     def _value_equality_values_(self) -> Any:
         # `isinstance` is used to ensure the a gate type and gate instance is not compared.
         return super()._value_equality_values_() + (self._max_parallel_allowed,)
+
+    def _json_dict_(self):
+        return {
+            'gate': self._gate_json(),
+            'name': self.name,
+            'description': self.description,
+            'max_parallel_allowed': self._max_parallel_allowed,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, gate, name, description, max_parallel_allowed, **kwargs):
+        if isinstance(gate, str):
+            gate = protocols.cirq_type_from_json(gate)
+        return cls(
+            gate, name=name, description=description, max_parallel_allowed=max_parallel_allowed
+        )

cirq/ops/gateset.py

@@ -97,6 +97,9 @@ def __init__(
     def _gate_str(self, gettr: Callable[[Any], str] = str) -> str:
         return _gate_str(self.gate, gettr)
 
+    def _gate_json(self) -> Union[raw_types.Gate, str]:
+        return self.gate if not isinstance(self.gate, type) else protocols.json_cirq_type(self.gate)
+
     def _default_name(self) -> str:
         family_type = 'Instance' if isinstance(self.gate, raw_types.Gate) else 'Type'
         return f'{family_type} GateFamily: {self._gate_str()}'
@@ -167,6 +170,22 @@ def _value_equality_values_(self) -> Any:
             self._ignore_global_phase,
         )
 
+    def _json_dict_(self):
+        return {
+            'gate': self._gate_json(),
+            'name': self.name,
+            'description': self.description,
+            'ignore_global_phase': self._ignore_global_phase,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, gate, name, description, ignore_global_phase, **kwargs):
+        if isinstance(gate, str):
+            gate = protocols.cirq_type_from_json(gate)
+        return cls(
+            gate, name=name, description=description, ignore_global_phase=ignore_global_phase
+        )
+
 
 @value.value_equality()
 class Gateset:

cirq/ops/gateset_test.py

@@ -97,6 +97,14 @@ def test_gate_family_repr_and_str(gate, name, description):
     assert g.description in str(g)
 
 
+@pytest.mark.parametrize('gate', [cirq.X, cirq.XPowGate(), cirq.XPowGate])
+@pytest.mark.parametrize('name, description', [(None, None), ('custom_name', 'custom_description')])
+def test_gate_family_json(gate, name, description):
+    g = cirq.GateFamily(gate, name=name, description=description)
+    g_json = cirq.to_json(g)
+    assert cirq.read_json(json_text=g_json) == g
+
+
 def test_gate_family_eq():
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(cirq.GateFamily(CustomX))

cirq/protocols/json_test_data/AnyIntegerPowerGateFamily.json

@@ -0,0 +1,4 @@
+{
+  "cirq_type": "AnyIntegerPowerGateFamily",
+  "gate": "XPowGate"
+}

cirq/protocols/json_test_data/AnyIntegerPowerGateFamily.repr

@@ -0,0 +1 @@
+cirq.AnyIntegerPowerGateFamily(cirq.ops.common_gates.XPowGate)

cirq/protocols/json_test_data/AnyUnitaryGateFamily.json

@@ -0,0 +1,4 @@
+{
+  "cirq_type": "AnyUnitaryGateFamily",
+  "num_qubits": 2
+}

cirq/protocols/json_test_data/AnyUnitaryGateFamily.repr

@@ -0,0 +1 @@
+cirq.AnyUnitaryGateFamily(num_qubits = 2)

cirq/protocols/json_test_data/GateFamily.json

@@ -0,0 +1,20 @@
+[
+  {
+    "cirq_type": "GateFamily",
+    "gate": "XPowGate",
+    "name": "Type GateFamily: cirq.ops.common_gates.XPowGate",
+    "description": "Accepts `cirq.Gate` instances `g` s.t. `isinstance(g, cirq.ops.common_gates.XPowGate)`",
+    "ignore_global_phase": true
+  },
+  {
+    "cirq_type": "GateFamily",
+    "gate": {
+      "cirq_type": "_PauliX",
+      "exponent": 1.0,
+      "global_shift": 0.0
+    },
+    "name": "XFamily",
+    "description": "Just the X gate.",
+    "ignore_global_phase": false
+  }
+]

cirq/protocols/json_test_data/GateFamily.repr

@@ -0,0 +1,4 @@
+[
+    cirq.GateFamily(gate=cirq.ops.common_gates.XPowGate, ignore_global_phase=True),
+    cirq.GateFamily(gate=cirq.X, name="XFamily", description="Just the X gate.", ignore_global_phase=False)
+]

cirq/protocols/json_test_data/ParallelGateFamily.json

@@ -0,0 +1,20 @@
+[
+  {
+    "cirq_type": "ParallelGateFamily",
+    "gate": "XPowGate",
+    "name": "INF Parallel Type GateFamily: cirq.ops.common_gates.XPowGate",
+    "description": "Accepts\n1. `cirq.Gate` instances `g` s.t. `isinstance(g, cirq.ops.common_gates.XPowGate)` OR\n2. `cirq.ParallelGate` instance `g` s.t. `g.sub_gate` satisfies 1. and `cirq.num_qubits(g) <= INF` OR\n3. `cirq.Operation` instance `op` s.t. `op.gate` satisfies 1. or 2.",
+    "max_parallel_allowed": null
+  },
+  {
+    "cirq_type": "ParallelGateFamily",
+    "gate": {
+      "cirq_type": "_PauliX",
+      "exponent": 1.0,
+      "global_shift": 0.0
+    },
+    "name": "ParallelXFamily",
+    "description": "Up to 4 parallel X gates",
+    "max_parallel_allowed": 4
+  }
+]

cirq/protocols/json_test_data/ParallelGateFamily.repr

@@ -0,0 +1,9 @@
+[
+    cirq.ParallelGateFamily(gate=cirq.ops.common_gates.XPowGate, max_parallel_allowed=None),
+    cirq.ParallelGateFamily(
+        gate=cirq.X,
+        name="ParallelXFamily",
+        description=r'''Up to 4 parallel X gates''',
+        max_parallel_allowed=4
+    )
+]

cirq/protocols/json_test_data/spec.py

@@ -25,8 +25,6 @@
     resolver_cache=_class_resolver_dictionary(),
     not_yet_serializable=[
         'Alignment',
-        'AnyIntegerPowerGateFamily',
-        'AnyUnitaryGateFamily',
         'AxisAngleDecomposition',
         'CircuitDag',
         'CircuitDiagramInfo',
@@ -39,7 +37,6 @@
         'DensityMatrixStepResult',
         'DensityMatrixTrialResult',
         'ExpressionMap',
-        'GateFamily',
         'Gateset',
         'InsertStrategy',
         'IonDevice',
@@ -50,7 +47,6 @@
         'ListSweep',
         'DiagonalGate',
         'NeutralAtomDevice',
-        'ParallelGateFamily',
         'PauliInteractionGate',
         'PauliStringPhasor',
         'PauliSum',