Support for stimcirq gates and operations in cirq_google protos

cirq-core/cirq/ops/gate_operation_test.py

@@ -490,7 +490,9 @@ def all_subclasses(cls):
     gate_subclasses = {
         g
         for g in all_subclasses(cirq.Gate)
-        if "cirq." in g.__module__ and "contrib" not in g.__module__ and "test" not in g.__module__
+        if g.__module__.startswith("cirq.")
+        and "contrib" not in g.__module__
+        and "test" not in g.__module__
     }
 
     test_module_spec = cirq.testing.json.spec_for("cirq.protocols")

cirq-google/cirq_google/serialization/circuit_serializer.py

@@ -38,6 +38,9 @@
 # CircuitSerializer is the dedicated serializer for the v2.5 format.
 _SERIALIZER_NAME = 'v2_5'
 
+# Package name for stimcirq
+_STIMCIRQ_MODULE = "stimcirq"
+
 
 class CircuitSerializer(serializer.Serializer):
     """A class for serializing and deserializing programs and operations.
@@ -193,7 +196,6 @@ def _serialize_gate_op(
             ValueError: If the operation cannot be serialized.
         """
         gate = op.gate
-
         if isinstance(gate, InternalGate):
             arg_func_langs.internal_gate_arg_to_proto(gate, out=msg.internalgate)
         elif isinstance(gate, cirq.XPowGate):
@@ -260,6 +262,30 @@ def _serialize_gate_op(
             arg_func_langs.float_arg_to_proto(
                 gate.q1_detune_mhz, out=msg.couplerpulsegate.q1_detune_mhz
             )
+        elif getattr(op, "__module__", "").startswith(_STIMCIRQ_MODULE) or getattr(
+            gate, "__module__", ""
+        ).startswith(_STIMCIRQ_MODULE):
+            # Special handling for stimcirq objects, which can be both operations and gates.
+            stimcirq_obj = (
+                op if getattr(op, "__module__", "").startswith(_STIMCIRQ_MODULE) else gate
+            )
+            if stimcirq_obj is not None and hasattr(stimcirq_obj, '_json_dict_'):
+                # All stimcirq gates currently have _json_dict_defined
+                msg.internalgate.name = type(stimcirq_obj).__name__
+                msg.internalgate.module = _STIMCIRQ_MODULE
+                if isinstance(stimcirq_obj, cirq.Gate):
+                    msg.internalgate.num_qubits = stimcirq_obj.num_qubits()
+                else:
+                    msg.internalgate.num_qubits = len(stimcirq_obj.qubits)
+
+                # Store json_dict objects in gate_args
+                for k, v in stimcirq_obj._json_dict_().items():
+                    arg_func_langs.arg_to_proto(value=v, out=msg.internalgate.gate_args[k])
+            else:
+                # New stimcirq op without a json dict has been introduced
+                raise ValueError(
+                    f'Cannot serialize stimcirq {op!r}:{type(gate)}'
+                )  # pragma: no cover
         else:
             raise ValueError(f'Cannot serialize op {op!r} of type {type(gate)}')
 
@@ -670,7 +696,31 @@ def _deserialize_gate_op(
                 raise ValueError(f"dimensions {dimensions} for ResetChannel must be an integer!")
             op = cirq.ResetChannel(dimension=dimensions)(*qubits)
         elif which_gate_type == 'internalgate':
-            op = arg_func_langs.internal_gate_from_proto(operation_proto.internalgate)(*qubits)
+            parsed_as_stimcirq = False
+            msg = operation_proto.internalgate
+            if msg.module == _STIMCIRQ_MODULE:
+                # special handling for stimcirq
+                try:
+                    import stimcirq
+
+                    # Use JSON resolver to instantiate the object
+                    if msg.name in stimcirq.JSON_RESOLVERS_DICT:
+                        kwargs = {}
+                        for k, v in msg.gate_args.items():
+                            arg = arg_func_langs.arg_from_proto(v)
+                            if arg is not None:
+                                kwargs[k] = arg
+                        op = stimcirq.JSON_RESOLVERS_DICT[msg.name](**kwargs)
+                        if qubits:
+                            op = op(*qubits)
+                        parsed_as_stimcirq = True
+
+                except ModuleNotFoundError:  # pragma: no cover
+                    # fall back to creating internal gates if stimcirq not installed
+                    pass
+            if not parsed_as_stimcirq:
+                # all other internal gates
+                op = arg_func_langs.internal_gate_from_proto(msg)(*qubits)
         elif which_gate_type == 'couplerpulsegate':
             gate = CouplerPulse(
                 hold_time=cirq.Duration(

cirq-google/cirq_google/serialization/circuit_serializer_test.py

@@ -1043,3 +1043,22 @@ def test_reset_gate_with_improper_argument():
 
     with pytest.raises(ValueError, match="must be an integer"):
         serializer.deserialize(circuit_proto)
+
+
+def test_stimcirq_gates():
+    try:
+        import stimcirq
+    except ModuleNotFoundError:  # pragma: no cover
+        # Stimcirq not found, these are optional tests.
+        return
+    serializer = cg.CircuitSerializer()
+    q = cirq.q(1, 2)
+    q2 = cirq.q(2, 2)
+    c = cirq.Circuit(
+        cirq.Moment(stimcirq.CXSwapGate(inverted=True)(q, q2)),
+        cirq.Moment(cirq.measure(q, key="m")),
+        cirq.Moment(stimcirq.DetAnnotation(parity_keys=["m"])),
+    )
+    msg = serializer.serialize(c)
+    deserialized_circuit = serializer.deserialize(msg)
+    assert deserialized_circuit == c

dev_tools/conf/mypy.ini

@@ -21,7 +21,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # Non-Google
-[mypy-IPython.*,sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,ruamel.*,absl.*,tensorflow_docs.*,ipywidgets.*,cachetools.*]
+[mypy-IPython.*,sympy.*,matplotlib.*,proto.*,pandas.*,scipy.*,freezegun.*,mpl_toolkits.*,networkx.*,ply.*,astroid.*,pytest.*,_pytest.*,pylint.*,setuptools.*,qiskit.*,quimb.*,pylatex.*,filelock.*,sortedcontainers.*,tqdm.*,ruamel.*,absl.*,tensorflow_docs.*,ipywidgets.*,cachetools.*,stimcirq.*]
 follow_imports = silent
 ignore_missing_imports = true
 

dev_tools/requirements/deps/dev-tools.txt

@@ -12,3 +12,6 @@ asv
 
 # For verifying behavior of qasm output.
 qiskit-aer~=0.16.1
+
+# For testing stimcirq compatibility (cirq-google)
+stimcirq