Add serialization support for InsertionNoiseModel (#6282)

quantumlib · Sep 14, 2023 · 432d57a · 432d57a
1 parent cf005c2
commit 432d57a
Show file tree

Hide file tree

Showing 5 changed files with 124 additions and 1 deletion.
diff --git a/cirq-core/cirq/devices/insertion_noise_model.py b/cirq-core/cirq/devices/insertion_noise_model.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import dataclasses
-from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
 
 from cirq import devices
 from cirq.devices import noise_utils
@@ -74,3 +74,23 @@ def noisy_moment(
         if self.prepend:
             return [*noise_steps.moments, moment]
         return [moment, *noise_steps.moments]
+
+    def __repr__(self) -> str:
+        return (
+            f'cirq.devices.InsertionNoiseModel(ops_added={self.ops_added},'
+            + f' prepend={self.prepend},'
+            + f' require_physical_tag={self.require_physical_tag})'
+        )
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        return {
+            'ops_added': list(self.ops_added.items()),
+            'prepend': self.prepend,
+            'require_physical_tag': self.require_physical_tag,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, ops_added, prepend, require_physical_tag, **kwargs):
+        return cls(
+            ops_added=dict(ops_added), prepend=prepend, require_physical_tag=require_physical_tag
+        )
diff --git a/cirq-core/cirq/devices/insertion_noise_model_test.py b/cirq-core/cirq/devices/insertion_noise_model_test.py
@@ -47,6 +47,8 @@ def test_insertion_noise():
     moment_3 = cirq.Moment(cirq.Z(q0), cirq.X(q1))
     assert model.noisy_moment(moment_3, system_qubits=[q0, q1]) == [moment_3]
 
+    cirq.testing.assert_equivalent_repr(model)
+
 
 def test_colliding_noise_qubits():
     # Check that noise affecting other qubits doesn't cause issues.
@@ -61,6 +63,8 @@ def test_colliding_noise_qubits():
         cirq.Moment(cirq.CNOT(q1, q2)),
     ]
 
+    cirq.testing.assert_equivalent_repr(model)
+
 
 def test_prepend():
     q0, q1 = cirq.LineQubit.range(2)
@@ -106,3 +110,5 @@ def test_supertype_matching():
 
     moment_1 = cirq.Moment(cirq.Y(q0))
     assert model.noisy_moment(moment_1, system_qubits=[q0]) == [moment_1, cirq.Moment(cirq.T(q0))]
+
+    cirq.testing.assert_equivalent_repr(model)
diff --git a/cirq-core/cirq/json_resolver_cache.py b/cirq-core/cirq/json_resolver_cache.py
@@ -47,6 +47,7 @@ def _class_resolver_dictionary() -> Dict[str, ObjectFactory]:
     import pandas as pd
     import numpy as np
     from cirq.devices.noise_model import _NoNoiseModel
+    from cirq.devices import InsertionNoiseModel
     from cirq.experiments import GridInteractionLayer
     from cirq.experiments.grid_parallel_two_qubit_xeb import GridParallelXEBMetadata
 
@@ -147,6 +148,7 @@ def _symmetricalqidpair(qids):
         'ISwapPowGate': cirq.ISwapPowGate,
         'IdentityGate': cirq.IdentityGate,
         'InitObsSetting': cirq.work.InitObsSetting,
+        'InsertionNoiseModel': InsertionNoiseModel,
         'KeyCondition': cirq.KeyCondition,
         'KrausChannel': cirq.KrausChannel,
         'LinearDict': cirq.LinearDict,

diff --git a/cirq-core/cirq/protocols/json_test_data/InsertionNoiseModel.json b/cirq-core/cirq/protocols/json_test_data/InsertionNoiseModel.json
@@ -0,0 +1,91 @@
+[
+    {
+        "cirq_type": "InsertionNoiseModel",
+        "ops_added": [
+            [
+            {
+                "cirq_type": "OpIdentifier",
+                "gate_type": "XPowGate",
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 0
+                }
+                ]
+            },
+            {
+                "cirq_type": "GateOperation",
+                "gate": {
+                "cirq_type": "BitFlipChannel",
+                "p": 0.2
+                },
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 0
+                }
+                ]
+            }
+            ]
+        ],
+        "prepend": false,
+        "require_physical_tag": false
+    },
+    {
+        "cirq_type": "InsertionNoiseModel",
+        "ops_added": [
+            [
+            {
+                "cirq_type": "OpIdentifier",
+                "gate_type": "XPowGate",
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 0
+                }
+                ]
+            },
+            {
+                "cirq_type": "GateOperation",
+                "gate": {
+                "cirq_type": "BitFlipChannel",
+                "p": 0.2
+                },
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 0
+                }
+                ]
+            }
+            ],
+            [
+            {
+                "cirq_type": "OpIdentifier",
+                "gate_type": "HPowGate",
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 1
+                }
+                ]
+            },
+            {
+                "cirq_type": "GateOperation",
+                "gate": {
+                "cirq_type": "BitFlipChannel",
+                "p": 0.1
+                },
+                "qubits": [
+                {
+                    "cirq_type": "LineQubit",
+                    "x": 1
+                }
+                ]
+            }
+            ]
+        ],
+        "prepend": false,
+        "require_physical_tag": false
+    }
+]
diff --git a/cirq-core/cirq/protocols/json_test_data/InsertionNoiseModel.repr b/cirq-core/cirq/protocols/json_test_data/InsertionNoiseModel.repr
@@ -0,0 +1,4 @@
+[
+cirq.devices.InsertionNoiseModel(ops_added={cirq.devices.noise_utils.OpIdentifier(cirq.ops.common_gates.XPowGate, cirq.LineQubit(0)): cirq.bit_flip(p=0.2).on(cirq.LineQubit(0))}, prepend=False, require_physical_tag=False),
+cirq.devices.InsertionNoiseModel(ops_added={cirq.devices.noise_utils.OpIdentifier(cirq.ops.common_gates.XPowGate, cirq.LineQubit(0)): cirq.bit_flip(p=0.2).on(cirq.LineQubit(0)), cirq.devices.noise_utils.OpIdentifier(cirq.ops.common_gates.HPowGate, cirq.LineQubit(1)): cirq.bit_flip(p=0.1).on(cirq.LineQubit(1))}, prepend=False, require_physical_tag=False)
+]