[WIP] Expose GoogleNoiseProperties at top level (#5209)

Adds `GoogleNoiseProperties` and related objects to the `__init__` files and adds JSON serialization.

This is a follow-up to #5188 to make these classes more convenient to use.

cirq-google/cirq_google/__init__.py

@@ -57,6 +57,8 @@
 from cirq_google.devices import (
     Bristlecone,
     Foxtail,
+    GoogleNoiseProperties,
+    NoiseModelFromGoogleNoiseProperties,
     SerializableDevice,
     Sycamore,
     Sycamore23,
@@ -78,6 +80,7 @@
     get_engine_calibration,
     get_engine_device,
     get_engine_sampler,
+    noise_properties_from_calibration,
 )
 
 from cirq_google.line import (

cirq-google/cirq_google/devices/__init__.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cirq_google.devices.google_noise_properties import (
+    GoogleNoiseProperties,
+    NoiseModelFromGoogleNoiseProperties,
+)
+
 from cirq_google.devices.known_devices import Bristlecone, Foxtail, Sycamore, Sycamore23
 
 from cirq_google.devices.serializable_device import SerializableDevice

cirq-google/cirq_google/devices/google_noise_properties.py

@@ -25,22 +25,6 @@
 from cirq.transformers.heuristic_decompositions import gate_tabulation_math_utils
 
 
-SINGLE_QUBIT_GATES: Set[Type['cirq.Gate']] = {
-    cirq.ZPowGate,
-    cirq.PhasedXZGate,
-    cirq.MeasurementGate,
-    cirq.ResetChannel,
-}
-SYMMETRIC_TWO_QUBIT_GATES: Set[Type['cirq.Gate']] = {
-    cirq_google.SycamoreGate,
-    cirq.FSimGate,
-    cirq.PhasedFSimGate,
-    cirq.ISwapPowGate,
-    cirq.CZPowGate,
-}
-ASYMMETRIC_TWO_QUBIT_GATES: Set[Type['cirq.Gate']] = set()
-
-
 T = TypeVar('T')
 V = TypeVar('V')
 
@@ -94,6 +78,26 @@ def __post_init__(self):
         # validate two qubit gate errors.
         self._validate_symmetric_errors('fsim_errors')
 
+    def __eq__(self, other):
+        if not isinstance(other, GoogleNoiseProperties):
+            return NotImplemented
+        if set(self.readout_errors) != set(other.readout_errors):
+            return False
+        return all(
+            [
+                self.gate_times_ns == other.gate_times_ns,
+                self.t1_ns == other.t1_ns,
+                self.tphi_ns == other.tphi_ns,
+                all(
+                    np.allclose(self.readout_errors[q], other.readout_errors[q])
+                    for q in self.readout_errors
+                ),
+                self.gate_pauli_errors == other.gate_pauli_errors,
+                self.validate == other.validate,
+                self.fsim_errors == other.fsim_errors,
+            ]
+        )
+
     def with_params(
         self,
         *,
@@ -181,15 +185,21 @@ def with_params(
 
     @classmethod
     def single_qubit_gates(cls) -> Set[type]:
-        return SINGLE_QUBIT_GATES
+        return {cirq.ZPowGate, cirq.PhasedXZGate, cirq.MeasurementGate, cirq.ResetChannel}
 
     @classmethod
     def symmetric_two_qubit_gates(cls) -> Set[type]:
-        return SYMMETRIC_TWO_QUBIT_GATES
+        return {
+            cirq_google.SycamoreGate,
+            cirq.FSimGate,
+            cirq.PhasedFSimGate,
+            cirq.ISwapPowGate,
+            cirq.CZPowGate,
+        }
 
     @classmethod
     def asymmetric_two_qubit_gates(cls) -> Set[type]:
-        return ASYMMETRIC_TWO_QUBIT_GATES
+        return set()
 
     @_compat.cached_property
     def _depolarizing_error(self) -> Dict[noise_utils.OpIdentifier, float]:
@@ -234,6 +244,62 @@ def build_noise_models(self) -> List['cirq.NoiseModel']:
 
         return noise_models
 
+    def __repr__(self):
+        gate_times_repr = ', '.join(
+            f'{key.__module__}.{key.__qualname__}: {val}' for key, val in self.gate_times_ns.items()
+        )
+        args = [
+            f'gate_times_ns={{{gate_times_repr}}}',
+            f't1_ns={self.t1_ns!r}',
+            f'tphi_ns={self.tphi_ns!r}',
+            f'readout_errors={_compat.proper_repr(self.readout_errors)}',
+            f'gate_pauli_errors={self.gate_pauli_errors!r}',
+            f'fsim_errors={self.fsim_errors!r}',
+            f'validate={self.validate!r}',
+        ]
+        arglines = ",\n".join(args)
+        return f'cirq_google.GoogleNoiseProperties({arglines})'
+
+    def _json_dict_(self):
+        storage_gate_times = {
+            cirq.json_cirq_type(key): val for key, val in self.gate_times_ns.items()
+        }
+        return {
+            # JSON requires mappings to have keys of basic types.
+            # Pairs must be sorted to ensure consistent serialization.
+            'gate_times_ns': tuple(storage_gate_times.items()),
+            't1_ns': tuple(self.t1_ns.items()),
+            'tphi_ns': tuple(self.tphi_ns.items()),
+            'readout_errors': tuple((k, v.tolist()) for k, v in self.readout_errors.items()),
+            'gate_pauli_errors': tuple(self.gate_pauli_errors.items()),
+            'fsim_errors': tuple(self.fsim_errors.items()),
+            'validate': self.validate,
+        }
+
+    @classmethod
+    def _from_json_dict_(
+        cls,
+        gate_times_ns,
+        t1_ns,
+        tphi_ns,
+        readout_errors,
+        gate_pauli_errors,
+        fsim_errors,
+        validate,
+        **kwargs,
+    ):
+        gate_type_times = {cirq.cirq_type_from_json(gate): val for gate, val in gate_times_ns}
+        # Known false positive: https://github.com/PyCQA/pylint/issues/5857
+        return GoogleNoiseProperties(  # pylint: disable=unexpected-keyword-arg
+            gate_times_ns=gate_type_times,
+            t1_ns=dict(t1_ns),
+            tphi_ns=dict(tphi_ns),
+            readout_errors={k: np.array(v) for k, v in readout_errors},
+            gate_pauli_errors=dict(gate_pauli_errors),
+            fsim_errors=dict(fsim_errors),
+            validate=validate,
+        )
+
 
 class NoiseModelFromGoogleNoiseProperties(devices.NoiseModelFromNoiseProperties):
     """A noise model defined from noise properties of a Google device."""

cirq-google/cirq_google/devices/google_noise_properties_test.py

@@ -18,10 +18,6 @@
 import pytest
 import cirq, cirq_google
 
-from cirq_google.devices.google_noise_properties import (
-    SYMMETRIC_TWO_QUBIT_GATES,
-    SINGLE_QUBIT_GATES,
-)
 from cirq.devices.noise_utils import OpIdentifier, PHYSICAL_GATE_TAG
 
 from cirq_google.devices.google_noise_properties import (
@@ -62,21 +58,41 @@ def sample_noise_properties(
         tphi_ns={q: 2e5 for q in system_qubits},
         readout_errors={q: np.array([SINGLE_QUBIT_ERROR, TWO_QUBIT_ERROR]) for q in system_qubits},
         gate_pauli_errors={
-            **{OpIdentifier(g, q): 0.001 for g in SINGLE_QUBIT_GATES for q in system_qubits},
+            **{
+                OpIdentifier(g, q): 0.001
+                for g in GoogleNoiseProperties.single_qubit_gates()
+                for q in system_qubits
+            },
             **{
                 OpIdentifier(g, q0, q1): 0.01
-                for g in SYMMETRIC_TWO_QUBIT_GATES
+                for g in GoogleNoiseProperties.symmetric_two_qubit_gates()
                 for q0, q1 in qubit_pairs
             },
         },
         fsim_errors={
             OpIdentifier(g, q0, q1): cirq.PhasedFSimGate(0.01, 0.03, 0.04, 0.05, 0.02)
-            for g in SYMMETRIC_TWO_QUBIT_GATES
+            for g in GoogleNoiseProperties.symmetric_two_qubit_gates()
             for q0, q1 in qubit_pairs
         },
     )
 
 
+def test_consistent_repr():
+    q0, q1 = cirq.LineQubit.range(2)
+    test_props = sample_noise_properties([q0, q1], [(q0, q1), (q1, q0)])
+    cirq.testing.assert_equivalent_repr(
+        test_props, setup_code="import cirq, cirq_google\nimport numpy as np"
+    )
+
+
+def test_equals():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    test_props = sample_noise_properties([q0, q1], [(q0, q1), (q1, q0)])
+    assert test_props != "mismatched_type"
+    test_props_v2 = test_props.with_params(readout_errors={q2: [0.01, 0.02]})
+    assert test_props != test_props_v2
+
+
 def test_zphase_gates():
     q0 = cirq.LineQubit(0)
     props = sample_noise_properties([q0], [])
@@ -106,6 +122,7 @@ def test_with_params_fill():
         gate_pauli_errors=expected_vals['gate_pauli_errors'],
         fsim_errors=expected_vals['fsim_errors'],
     )
+    assert props_v2 != props
     for key in props.gate_times_ns:
         assert key in props_v2.gate_times_ns
         assert props_v2.gate_times_ns[key] == expected_vals['gate_times_ns']
@@ -150,6 +167,7 @@ def test_with_params_target():
         gate_pauli_errors=expected_vals['gate_pauli_errors'],
         fsim_errors=expected_vals['fsim_errors'],
     )
+    assert props_v2 != props
     for field_name, expected in expected_vals.items():
         target_dict = getattr(props_v2, field_name)
         for key, val in expected.items():
@@ -171,6 +189,7 @@ def test_with_params_opid_with_gate():
         gate_pauli_errors={cirq.PhasedXZGate: expected_vals['gate_pauli_errors']},
         fsim_errors={cirq.CZPowGate: expected_vals['fsim_errors']},
     )
+    assert props_v2 != props
     gpe_op_id_0 = cirq.OpIdentifier(cirq.PhasedXZGate, q0)
     gpe_op_id_1 = cirq.OpIdentifier(cirq.PhasedXZGate, q1)
     assert props_v2.gate_pauli_errors[gpe_op_id_0] == expected_vals['gate_pauli_errors']

cirq-google/cirq_google/engine/__init__.py

@@ -42,6 +42,8 @@
 
 from cirq_google.engine.calibration_layer import CalibrationLayer
 from cirq_google.engine.calibration_result import CalibrationResult
+from cirq_google.engine.calibration_to_noise_properties import noise_properties_from_calibration
+
 from cirq_google.engine.engine import (
     Engine,
     get_engine,

cirq-google/cirq_google/engine/calibration_to_noise_properties.py

@@ -68,13 +68,13 @@ def noise_properties_from_calibration(
     """Translates between `cirq_google.Calibration` and NoiseProperties.
 
     The NoiseProperties object can then be used as input to the
-    `cirq.devices.noise_propertiesNoiseModelFromNoiseProperties` class to
-    create a `cirq.NoiseModel` that can be used with a simulator.
+    `cirq_google.NoiseModelFromGoogleNoiseProperties` class to create a
+    `cirq.NoiseModel` that can be used with a simulator.
 
-    To manually override noise properties, call `override` on the output:
+    To manually override noise properties, call `with_params` on the output:
 
-        # Set all gate durations to 37ns.
-        >>> noise_properties_from_calibration(cal).override(gate_times_ns=37)
+        >>> noise_props = noise_properties_from_calibration(cal).with_params(gate_times_ns=37)
+        # noise_props with all gate durations set to 37ns.
 
     See `cirq_google.GoogleNoiseProperties` for details.
 
@@ -96,6 +96,7 @@ def noise_properties_from_calibration(
         ops.FSimGate: 32.0,
         ops.ISwapPowGate: 32.0,
         ops.CZPowGate: 32.0,
+        cg_ops.SycamoreGate: 12.0,
         # ops.WaitGate is a special case.
     }
 
@@ -123,7 +124,7 @@ def noise_properties_from_calibration(
     gate_pauli_errors = {
         noise_utils.OpIdentifier(gate, q): pauli_err
         for q, pauli_err in rb_pauli_errors.items()
-        for gate in google_noise_properties.SINGLE_QUBIT_GATES
+        for gate in google_noise_properties.GoogleNoiseProperties.single_qubit_gates()
     }
 
     # 3b. Extract Pauli error for two-qubit gates.

cirq-google/cirq_google/engine/calibration_to_noise_properties_test.py

@@ -14,7 +14,6 @@
 
 import cirq, cirq_google
 from cirq.devices.noise_utils import OpIdentifier
-from cirq_google.engine.calibration_to_noise_properties import noise_properties_from_calibration
 
 from google.protobuf.text_format import Merge
 import numpy as np
@@ -211,7 +210,7 @@ def test_noise_properties_from_calibration():
 
     # Create NoiseProperties object from Calibration
     calibration = cirq_google.Calibration(_CALIBRATION_DATA)
-    prop = noise_properties_from_calibration(calibration)
+    prop = cirq_google.noise_properties_from_calibration(calibration)
 
     for i, q in enumerate(qubits):
         assert np.isclose(
@@ -337,4 +336,4 @@ def test_incomplete_calibration():
     # Create NoiseProperties object from Calibration
     calibration = cirq_google.Calibration(_CALIBRATION_DATA)
     with pytest.raises(ValueError, match='Keys specified for T1 and Tphi are not identical.'):
-        _ = noise_properties_from_calibration(calibration)
+        _ = cirq_google.noise_properties_from_calibration(calibration)

cirq-google/cirq_google/json_resolver_cache.py

@@ -30,6 +30,7 @@ def _class_resolver_dictionary() -> Dict[str, ObjectFactory]:
         'CalibrationLayer': cirq_google.CalibrationLayer,
         'CalibrationResult': cirq_google.CalibrationResult,
         'CouplerPulse': cirq_google.experimental.CouplerPulse,
+        'GoogleNoiseProperties': cirq_google.GoogleNoiseProperties,
         'SycamoreGate': cirq_google.SycamoreGate,
         'GateTabulation': cirq_google.GateTabulation,
         'PhysicalZTag': cirq_google.PhysicalZTag,