Add JSON serialization support for NoiseModel classes

cirq-core/cirq/contrib/__init__.py

@@ -24,3 +24,4 @@
 from cirq.contrib.qcircuit import circuit_to_latex_using_qcircuit as circuit_to_latex_using_qcircuit
 from cirq.contrib import json  # noqa: F401
 from cirq.contrib.circuitdag import CircuitDag as CircuitDag, Unique as Unique
+from cirq.contrib import noise_models as noise_models

cirq-core/cirq/contrib/noise_models/noise_models.py

@@ -23,6 +23,7 @@
     import cirq
 
 
+@value.value_equality()
 class DepolarizingNoiseModel(devices.NoiseModel):
     """Applies depolarizing noise to each qubit individually at the end of
     every moment.
@@ -42,6 +43,16 @@ def __init__(self, depol_prob: float, prepend: bool = False):
         self.qubit_noise_gate = ops.DepolarizingChannel(depol_prob)
         self._prepend = prepend
 
+    def _value_equality_values_(self):
+        return self.qubit_noise_gate, self._prepend
+
+    def __repr__(self) -> str:
+        p = self.qubit_noise_gate.p
+        return (
+            f'cirq.contrib.noise_models.DepolarizingNoiseModel('
+            f'{p!r}, prepend={self._prepend!r})'
+        )
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if validate_all_measurements(moment) or self.is_virtual_moment(moment):  # pragma: no cover
             return moment
@@ -54,7 +65,18 @@ def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         ]
         return output[::-1] if self._prepend else output
 
+    def _json_dict_(self) -> dict[str, object]:
+        return {'qubit_noise_gate': self.qubit_noise_gate, 'prepend': self._prepend}
+
+    @classmethod
+    def _from_json_dict_(cls, qubit_noise_gate, prepend, **kwargs):
+        obj = cls.__new__(cls)
+        obj.qubit_noise_gate = qubit_noise_gate
+        obj._prepend = prepend
+        return obj
 
+
+@value.value_equality()
 class ReadoutNoiseModel(devices.NoiseModel):
     """NoiseModel with probabilistic bit flips preceding measurement.
 
@@ -78,6 +100,13 @@ def __init__(self, bitflip_prob: float, prepend: bool = True):
         self.readout_noise_gate = ops.BitFlipChannel(bitflip_prob)
         self._prepend = prepend
 
+    def _value_equality_values_(self):
+        return self.readout_noise_gate, self._prepend
+
+    def __repr__(self) -> str:
+        p = self.readout_noise_gate.p
+        return f'cirq.contrib.noise_models.ReadoutNoiseModel(' f'{p!r}, prepend={self._prepend!r})'
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if self.is_virtual_moment(moment):
             return moment
@@ -91,7 +120,18 @@ def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
             return output if self._prepend else output[::-1]
         return moment
 
+    def _json_dict_(self) -> dict[str, object]:
+        return {'readout_noise_gate': self.readout_noise_gate, 'prepend': self._prepend}
+
+    @classmethod
+    def _from_json_dict_(cls, readout_noise_gate, prepend, **kwargs):
+        obj = cls.__new__(cls)
+        obj.readout_noise_gate = readout_noise_gate
+        obj._prepend = prepend
+        return obj
+
 
+@value.value_equality()
 class DampedReadoutNoiseModel(devices.NoiseModel):
     """NoiseModel with T1 decay preceding measurement.
 
@@ -115,6 +155,16 @@ def __init__(self, decay_prob: float, prepend: bool = True):
         self.readout_decay_gate = ops.AmplitudeDampingChannel(decay_prob)
         self._prepend = prepend
 
+    def _value_equality_values_(self):
+        return self.readout_decay_gate, self._prepend
+
+    def __repr__(self) -> str:
+        p = self.readout_decay_gate.gamma
+        return (
+            f'cirq.contrib.noise_models.DampedReadoutNoiseModel('
+            f'{p!r}, prepend={self._prepend!r})'
+        )
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if self.is_virtual_moment(moment):
             return moment
@@ -128,7 +178,18 @@ def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
             return output if self._prepend else output[::-1]
         return moment
 
+    def _json_dict_(self) -> dict[str, object]:
+        return {'readout_decay_gate': self.readout_decay_gate, 'prepend': self._prepend}
 
+    @classmethod
+    def _from_json_dict_(cls, readout_decay_gate, prepend, **kwargs):
+        obj = cls.__new__(cls)
+        obj.readout_decay_gate = readout_decay_gate
+        obj._prepend = prepend
+        return obj
+
+
+@value.value_equality()
 class DepolarizingWithReadoutNoiseModel(devices.NoiseModel):
     """DepolarizingNoiseModel with probabilistic bit flips preceding
     measurement.
@@ -148,12 +209,34 @@ def __init__(self, depol_prob: float, bitflip_prob: float):
         self.qubit_noise_gate = ops.DepolarizingChannel(depol_prob)
         self.readout_noise_gate = ops.BitFlipChannel(bitflip_prob)
 
+    def _value_equality_values_(self):
+        return self.qubit_noise_gate, self.readout_noise_gate
+
+    def __repr__(self) -> str:
+        p = self.qubit_noise_gate.p
+        b = self.readout_noise_gate.p
+        return 'cirq.contrib.noise_models.DepolarizingWithReadoutNoiseModel(' f'{p!r}, {b!r})'
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if validate_all_measurements(moment):
             return [circuits.Moment(self.readout_noise_gate(q) for q in system_qubits), moment]
         return [moment, circuits.Moment(self.qubit_noise_gate(q) for q in system_qubits)]
 
+    def _json_dict_(self) -> dict[str, object]:
+        return {
+            'qubit_noise_gate': self.qubit_noise_gate,
+            'readout_noise_gate': self.readout_noise_gate,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, qubit_noise_gate, readout_noise_gate, **kwargs):
+        obj = cls.__new__(cls)
+        obj.qubit_noise_gate = qubit_noise_gate
+        obj.readout_noise_gate = readout_noise_gate
+        return obj
 
+
+@value.value_equality()
 class DepolarizingWithDampedReadoutNoiseModel(devices.NoiseModel):
     """DepolarizingWithReadoutNoiseModel with T1 decay preceding
     measurement.
@@ -178,6 +261,18 @@ def __init__(self, depol_prob: float, bitflip_prob: float, decay_prob: float):
         self.readout_noise_gate = ops.BitFlipChannel(bitflip_prob)
         self.readout_decay_gate = ops.AmplitudeDampingChannel(decay_prob)
 
+    def _value_equality_values_(self):
+        return (self.qubit_noise_gate, self.readout_noise_gate, self.readout_decay_gate)
+
+    def __repr__(self) -> str:
+        p = self.qubit_noise_gate.p
+        b = self.readout_noise_gate.p
+        d = self.readout_decay_gate.gamma
+        return (
+            'cirq.contrib.noise_models.DepolarizingWithDampedReadoutNoiseModel('
+            f'{p!r}, {b!r}, {d!r})'
+        )
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if validate_all_measurements(moment):
             return [
@@ -187,3 +282,18 @@ def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
             ]
         else:
             return [moment, circuits.Moment(self.qubit_noise_gate(q) for q in system_qubits)]
+
+    def _json_dict_(self) -> dict[str, object]:
+        return {
+            'qubit_noise_gate': self.qubit_noise_gate,
+            'readout_noise_gate': self.readout_noise_gate,
+            'readout_decay_gate': self.readout_decay_gate,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, qubit_noise_gate, readout_noise_gate, readout_decay_gate, **kwargs):
+        obj = cls.__new__(cls)
+        obj.qubit_noise_gate = qubit_noise_gate
+        obj.readout_noise_gate = readout_noise_gate
+        obj.readout_decay_gate = readout_decay_gate
+        return obj

cirq-core/cirq/contrib/noise_models/noise_models_test.py

@@ -269,3 +269,19 @@ def test_aggregate_decay_noise_after_moment() -> None:
         ]
     )
     assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
+
+
+def test_noise_model_repr_and_json():
+    models = [
+        ccn.DepolarizingNoiseModel(0.1),
+        ccn.ReadoutNoiseModel(0.2),
+        ccn.DampedReadoutNoiseModel(0.3),
+        ccn.DepolarizingWithReadoutNoiseModel(0.1, 0.2),
+        ccn.DepolarizingWithDampedReadoutNoiseModel(0.1, 0.2, 0.3),
+    ]
+    for m in models:
+        cirq.testing.assert_equivalent_repr(
+            m, setup_code="import cirq\nimport cirq.contrib.noise_models as ccn"
+        )
+        restored = type(m)._from_json_dict_(**m._json_dict_())
+        assert restored == m

cirq-core/cirq/devices/noise_properties.py

@@ -24,7 +24,7 @@
 import abc
 from typing import Iterable, Sequence, TYPE_CHECKING
 
-from cirq import _import, devices, ops, protocols
+from cirq import _import, devices, ops, protocols, value
 from cirq.devices.noise_utils import PHYSICAL_GATE_TAG
 
 circuits = _import.LazyLoader("circuits", globals(), "cirq.circuits.circuit")
@@ -41,6 +41,7 @@ def build_noise_models(self) -> list[cirq.NoiseModel]:
         """Construct all NoiseModels associated with this NoiseProperties."""
 
 
+@value.value_equality
 class NoiseModelFromNoiseProperties(devices.NoiseModel):
     def __init__(self, noise_properties: NoiseProperties) -> None:
         """Creates a Noise Model from a NoiseProperties object that can be used with a Simulator.
@@ -54,6 +55,12 @@ def __init__(self, noise_properties: NoiseProperties) -> None:
         self._noise_properties = noise_properties
         self.noise_models = self._noise_properties.build_noise_models()
 
+    def _value_equality_values_(self):
+        return self._noise_properties
+
+    def __repr__(self) -> str:
+        return "cirq.devices.NoiseModelFromNoiseProperties(" f"{self._noise_properties!r})"
+
     def is_virtual(self, op: cirq.Operation) -> bool:
         """Returns True if an operation is virtual.
 
@@ -126,3 +133,10 @@ def noisy_moments(
                 combined_measure_ops.append(multi_measurements[key])
             final_moments.append(circuits.Moment(combined_measure_ops))
         return final_moments
+
+    def _json_dict_(self) -> dict[str, object]:
+        return {'noise_properties': self._noise_properties}
+
+    @classmethod
+    def _from_json_dict_(cls, noise_properties, **kwargs):
+        return cls(noise_properties)

cirq-core/cirq/devices/noise_properties_test.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import cirq
+import cirq.testing
 from cirq.devices.insertion_noise_model import InsertionNoiseModel
 from cirq.devices.noise_properties import NoiseModelFromNoiseProperties, NoiseProperties
 from cirq.devices.noise_utils import OpIdentifier, PHYSICAL_GATE_TAG
@@ -54,3 +55,23 @@ def test_sample_model() -> None:
         cirq.Moment(cirq.H(q0), cirq.H(q1)),
     )
     assert noisy_circuit == expected_circuit
+
+
+def test_noise_model_from_noise_properties_repr_and_json():
+    q0 = cirq.LineQubit(0)
+    props = SampleNoiseProperties([q0], [])
+    model = NoiseModelFromNoiseProperties(props)
+    assert 'NoiseModelFromNoiseProperties' in repr(model)
+    restored = NoiseModelFromNoiseProperties._from_json_dict_(**model._json_dict_())
+    assert restored._noise_properties is props
+
+
+def test_noise_model_from_noise_properties_equality() -> None:
+    q0 = cirq.LineQubit(0)
+    props1 = SampleNoiseProperties([q0], [])
+    props2 = SampleNoiseProperties([q0], [])
+    eq = cirq.testing.EqualsTester()
+    eq.make_equality_group(
+        lambda: NoiseModelFromNoiseProperties(props1), lambda: NoiseModelFromNoiseProperties(props1)
+    )
+    eq.add_equality_group(NoiseModelFromNoiseProperties(props2))

cirq-core/cirq/devices/thermal_noise_model.py

@@ -21,7 +21,8 @@
 import numpy as np
 import sympy
 
-from cirq import devices, ops, protocols, qis
+import cirq
+from cirq import devices, ops, protocols, qis, value
 from cirq._import import LazyLoader
 from cirq.devices.noise_utils import PHYSICAL_GATE_TAG
 
@@ -160,7 +161,8 @@ def _validate_rates(qubits: set[cirq.Qid], rates: dict[cirq.Qid, np.ndarray]) ->
             )
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(eq=False)
+@value.value_equality
 class ThermalNoiseModel(devices.NoiseModel):
     """NoiseModel representing simulated thermalization of a qubit.
 
@@ -232,6 +234,26 @@ def __init__(
         self.skip_measurements: bool = skip_measurements
         self._prepend = prepend
 
+    def _value_equality_values_(self):
+        gate_times = tuple(sorted((k, v) for k, v in self.gate_durations_ns.items()))
+        rates = tuple(
+            sorted(
+                ((q, tuple(m.reshape(-1))) for q, m in self.rate_matrix_GHz.items()),
+                key=lambda x: repr(x[0]),
+            )
+        )
+        return gate_times, rates, self.require_physical_tag, self.skip_measurements, self._prepend
+
+    def __repr__(self) -> str:
+        return (
+            "cirq.devices.ThermalNoiseModel("
+            f"qubits={set(self.rate_matrix_GHz.keys())!r}, "
+            f"gate_durations_ns={self.gate_durations_ns!r}, "
+            f"heat_rate_GHz=None, cool_rate_GHz=None, dephase_rate_GHz=None, "
+            f"require_physical_tag={self.require_physical_tag!r}, "
+            f"skip_measurements={self.skip_measurements!r}, prepend={self._prepend!r})"
+        )
+
     def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]) -> cirq.OP_TREE:
         if not moment.operations:
             return [moment]
@@ -283,3 +305,31 @@ def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]) -
             return [moment]
         output = [moment, moment_module.Moment(noise_ops)]
         return output[::-1] if self._prepend else output
+
+    def _json_dict_(self) -> dict[str, object]:
+        gate_times = {cirq.json_cirq_type(k): v for k, v in self.gate_durations_ns.items()}
+        return {
+            'gate_durations_ns': tuple(gate_times.items()),
+            'rate_matrix_GHz': tuple((q, m.tolist()) for q, m in self.rate_matrix_GHz.items()),
+            'require_physical_tag': self.require_physical_tag,
+            'skip_measurements': self.skip_measurements,
+            'prepend': self._prepend,
+        }
+
+    @classmethod
+    def _from_json_dict_(
+        cls,
+        gate_durations_ns,
+        rate_matrix_GHz,
+        require_physical_tag,
+        skip_measurements,
+        prepend,
+        **kwargs,
+    ):
+        obj = cls.__new__(cls)
+        obj.gate_durations_ns = {cirq.cirq_type_from_json(k): v for k, v in gate_durations_ns}
+        obj.rate_matrix_GHz = {q: np.array(m) for q, m in rate_matrix_GHz}
+        obj.require_physical_tag = require_physical_tag
+        obj.skip_measurements = skip_measurements
+        obj._prepend = prepend
+        return obj

cirq-core/cirq/devices/thermal_noise_model_test.py

@@ -342,3 +342,19 @@ def test_noisy_moment_two_qubit():
             [9.87330937e-01, 0, 0, 9.95013725e-01],
         ],
     )
+
+
+def test_thermal_noise_model_repr_and_json():
+    q0 = cirq.LineQubit(0)
+    model = ThermalNoiseModel(
+        qubits={q0},
+        gate_durations_ns={cirq.ZPowGate: 5.0},
+        heat_rate_GHz={q0: 1e-5},
+        cool_rate_GHz={q0: 1e-4},
+        dephase_rate_GHz={q0: 2e-4},
+        require_physical_tag=False,
+        skip_measurements=False,
+    )
+    assert 'ThermalNoiseModel' in repr(model)
+    restored = ThermalNoiseModel._from_json_dict_(**model._json_dict_())
+    assert restored == model