Convenience methods for modifying GoogleNoiseProperties (#5188)

* Add override methods.

* Mypy passes, weak typing

* Resolved mypy conundrum

* with_params

cirq-google/cirq_google/devices/google_noise_properties.py

@@ -15,8 +15,8 @@
 
 """Class for representing noise on a Google device."""
 
-from dataclasses import dataclass, field
-from typing import Dict, List, Set, Type
+import dataclasses
+from typing import Any, Dict, List, Sequence, Set, Type, TypeVar, Union
 import numpy as np
 
 import cirq, cirq_google
@@ -41,39 +41,144 @@
 ASYMMETRIC_TWO_QUBIT_GATES: Set[Type['cirq.Gate']] = set()
 
 
-@dataclass
+T = TypeVar('T')
+V = TypeVar('V')
+
+
+def _with_values(original: Dict[T, V], val: Union[V, Dict[T, V]]) -> Dict[T, V]:
+    """Returns a copy of `original` using values from `val`.
+
+    If val is a single value, all keys are mapped to that value. If val is a
+    dict, the union of original and val is returned, using values from val for
+    any conflicting keys.
+    """
+    if isinstance(val, dict):
+        return {**original, **val}
+    return {k: val for k in original}
+
+
+@dataclasses.dataclass
 class GoogleNoiseProperties(devices.SuperconductingQubitsNoiseProperties):
     """Noise-defining properties for a Google device.
 
     Inherited args:
-        gate_times_ns: Dict[type, float] of gate types to their duration on
-            quantum hardware. Used with t(1|phi)_ns to specify thermal noise.
-        t1_ns: Dict[cirq.Qid, float] of qubits to their T_1 time, in ns.
-        tphi_ns: Dict[cirq.Qid, float] of qubits to their T_phi time, in ns.
-        readout_errors: Dict[cirq.Qid, np.ndarray] of qubits to their readout
+        gate_times_ns: Dict[Type[`cirq.Gate`], float] of gate types to their
+            duration on quantum hardware. Used with t(1|phi)_ns to specify
+            thermal noise.
+        t1_ns: Dict[`cirq.Qid`, float] of qubits to their T_1 time, in ns.
+        tphi_ns: Dict[`cirq.Qid`, float] of qubits to their T_phi time, in ns.
+        readout_errors: Dict[`cirq.Qid`, np.ndarray] of qubits to their readout
             errors in matrix form: [P(read |1> from |0>), P(read |0> from |1>)].
             Used to prepend amplitude damping errors to measurements.
-        gate_pauli_errors: dict of noise_utils.OpIdentifiers (a gate and the qubits it
-            targets) to the Pauli error for that operation. Used to construct
-            depolarizing error. Keys in this dict must have defined qubits.
+        gate_pauli_errors: dict of `noise_utils.OpIdentifiers` (a gate and the
+            qubits it targets) to the Pauli error for that operation. Used to
+            construct depolarizing error. Keys in this dict must have defined
+            qubits.
         validate: If True, verifies that t1 and tphi qubits sets match, and
             that all symmetric two-qubit gates have errors which are
             symmetric over the qubits they affect. Defaults to True.
 
     Additional args:
-        fsim_errors: Dict[noise_utils.OpIdentifier, cirq.PhasedFSimGate] of gate types
-            (potentially on specific qubits) to the PhasedFSim fix-up operation
-            for that gate. Defaults to no-op for all gates.
+        fsim_errors: Dict[`noise_utils.OpIdentifier`, `cirq.PhasedFSimGate`] of
+            gate types (potentially on specific qubits) to the PhasedFSim
+            fix-up operation for that gate. Defaults to no-op for all gates.
     """
 
-    fsim_errors: Dict[noise_utils.OpIdentifier, cirq.PhasedFSimGate] = field(default_factory=dict)
+    fsim_errors: Dict[noise_utils.OpIdentifier, cirq.PhasedFSimGate] = dataclasses.field(
+        default_factory=dict
+    )
 
     def __post_init__(self):
         super().__post_init__()
 
         # validate two qubit gate errors.
         self._validate_symmetric_errors('fsim_errors')
 
+    def with_params(
+        self,
+        *,
+        gate_times_ns: Union[None, float, Dict[Type['cirq.Gate'], float]] = None,
+        t1_ns: Union[None, float, Dict['cirq.Qid', float]] = None,
+        tphi_ns: Union[None, float, Dict['cirq.Qid', float]] = None,
+        readout_errors: Union[None, Sequence[float], Dict['cirq.Qid', Sequence[float]]] = None,
+        gate_pauli_errors: Union[
+            None, float, Dict[Union[Type['cirq.Gate'], noise_utils.OpIdentifier], float],
+        ] = None,
+        fsim_errors: Union[
+            None,
+            'cirq.PhasedFSimGate',
+            Dict[Union[Type['cirq.Gate'], noise_utils.OpIdentifier], 'cirq.PhasedFSimGate'],
+        ] = None,
+    ):
+        """Returns a copy of this object with the given params overridden.
+
+        This method supports partial replacement: each arg can accept a single
+        value (which will replace all existing values) or a mapping (which will
+        replace matching entries in the old object). Otherwise, all fields are
+        the same as those used in the constructor.
+
+        Args:
+        gate_times_ns: float or Dict[Type[`cirq.Gate`], float].
+        t1_ns: float or Dict[`cirq.Qid`, float].
+        tphi_ns: float or Dict[`cirq.Qid`, float].
+        readout_errors: Sequence or Dict[`cirq.Qid`, Sequence]. Converted to
+            np.ndarray if not provided in that format.
+        gate_pauli_errors: float or Dict[`cirq.OpIdentifier`, float].
+            Dict key can also be Type[`cirq.Gate`]; this will apply the given
+            error to all placements of that gate that appear in the original
+            object.
+        fsim_errors: `cirq.PhasedFSimGate` or Dict[`cirq.OpIdentifier`,
+            `cirq.PhasedFSimGate`] Dict key can also be Type[`cirq.Gate`]; this
+            will apply the given error to all placements of that gate that
+            appear in the original object.
+
+        """
+        replace_args: Dict[str, Dict[Any, Any]] = {}
+        if gate_times_ns is not None:
+            replace_args['gate_times_ns'] = _with_values(self.gate_times_ns, gate_times_ns)
+        if t1_ns is not None:
+            replace_args['t1_ns'] = _with_values(self.t1_ns, t1_ns)
+        if tphi_ns is not None:
+            replace_args['tphi_ns'] = _with_values(self.tphi_ns, tphi_ns)
+        if readout_errors is not None:
+            if isinstance(readout_errors, dict):
+                readout_errors = {k: np.array(v) for k, v in readout_errors.items()}
+            else:
+                readout_errors = np.array(readout_errors)
+            replace_args['readout_errors'] = _with_values(self.readout_errors, readout_errors)
+        if gate_pauli_errors is not None:
+            if isinstance(gate_pauli_errors, dict):
+                combined_pauli_errors: Dict[
+                    Union[Type['cirq.Gate'], noise_utils.OpIdentifier], float
+                ] = {}
+                for op_id in self.gate_pauli_errors:
+                    if op_id in gate_pauli_errors:
+                        combined_pauli_errors[op_id] = gate_pauli_errors[op_id]
+                    elif op_id.gate_type in gate_pauli_errors:
+                        combined_pauli_errors[op_id] = gate_pauli_errors[op_id.gate_type]
+                gate_pauli_errors = combined_pauli_errors
+            replace_args['gate_pauli_errors'] = _with_values(
+                self.gate_pauli_errors, gate_pauli_errors
+            )
+        if fsim_errors is not None:
+            if isinstance(fsim_errors, dict):
+                combined_fsim_errors: Dict[
+                    Union[Type['cirq.Gate'], noise_utils.OpIdentifier], 'cirq.PhasedFSimGate'
+                ] = {}
+                for op_id in self.fsim_errors:
+                    op_id_swapped = noise_utils.OpIdentifier(op_id.gate_type, *op_id.qubits[::-1])
+                    if op_id in fsim_errors:
+                        combined_fsim_errors[op_id] = fsim_errors[op_id]
+                        combined_fsim_errors[op_id_swapped] = fsim_errors[op_id]
+                    elif op_id_swapped in fsim_errors:
+                        combined_fsim_errors[op_id] = fsim_errors[op_id_swapped]
+                        combined_fsim_errors[op_id_swapped] = fsim_errors[op_id_swapped]
+                    elif op_id.gate_type in fsim_errors:
+                        combined_fsim_errors[op_id] = fsim_errors[op_id.gate_type]
+                fsim_errors = combined_fsim_errors
+            replace_args['fsim_errors'] = _with_values(self.fsim_errors, fsim_errors)
+        return dataclasses.replace(self, **replace_args)
+
     @classmethod
     def single_qubit_gates(cls) -> Set[type]:
         return SINGLE_QUBIT_GATES

cirq-google/cirq_google/devices/google_noise_properties_test.py

@@ -86,6 +86,102 @@ def test_zphase_gates():
     assert noisy_circuit == circuit
 
 
+def test_with_params_fill():
+    # Test single-value with_params.
+    q0, q1 = cirq.LineQubit.range(2)
+    props = sample_noise_properties([q0, q1], [(q0, q1), (q1, q0)])
+    expected_vals = {
+        'gate_times_ns': 91,
+        't1_ns': 92,
+        'tphi_ns': 93,
+        'readout_errors': [0.094, 0.095],
+        'gate_pauli_errors': 0.096,
+        'fsim_errors': cirq.PhasedFSimGate(0.0971, 0.0972, 0.0973, 0.0974, 0.0975),
+    }
+    props_v2 = props.with_params(
+        gate_times_ns=expected_vals['gate_times_ns'],
+        t1_ns=expected_vals['t1_ns'],
+        tphi_ns=expected_vals['tphi_ns'],
+        readout_errors=expected_vals['readout_errors'],
+        gate_pauli_errors=expected_vals['gate_pauli_errors'],
+        fsim_errors=expected_vals['fsim_errors'],
+    )
+    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']
+    for key in props.t1_ns:
+        assert key in props_v2.t1_ns
+        assert props_v2.t1_ns[key] == expected_vals['t1_ns']
+    for key in props.tphi_ns:
+        assert key in props_v2.tphi_ns
+        assert props_v2.tphi_ns[key] == expected_vals['tphi_ns']
+    for key in props.readout_errors:
+        assert key in props_v2.readout_errors
+        assert np.allclose(props_v2.readout_errors[key], expected_vals['readout_errors'])
+    for key in props.gate_pauli_errors:
+        assert key in props_v2.gate_pauli_errors
+        assert props_v2.gate_pauli_errors[key] == expected_vals['gate_pauli_errors']
+    for key in props.fsim_errors:
+        assert key in props_v2.fsim_errors
+        assert props_v2.fsim_errors[key] == expected_vals['fsim_errors']
+
+
+def test_with_params_target():
+    # Test targeted-value with_params.
+    q0, q1 = cirq.LineQubit.range(2)
+    props = sample_noise_properties([q0, q1], [(q0, q1), (q1, q0)])
+    expected_vals = {
+        'gate_times_ns': {cirq.ZPowGate: 91},
+        't1_ns': {q0: 92},
+        'tphi_ns': {q1: 93},
+        'readout_errors': {q0: [0.094, 0.095]},
+        'gate_pauli_errors': {cirq.OpIdentifier(cirq.PhasedXZGate, q1): 0.096},
+        'fsim_errors': {
+            cirq.OpIdentifier(cirq.CZPowGate, q0, q1): cirq.PhasedFSimGate(
+                0.0971, 0.0972, 0.0973, 0.0974, 0.0975
+            )
+        },
+    }
+    props_v2 = props.with_params(
+        gate_times_ns=expected_vals['gate_times_ns'],
+        t1_ns=expected_vals['t1_ns'],
+        tphi_ns=expected_vals['tphi_ns'],
+        readout_errors=expected_vals['readout_errors'],
+        gate_pauli_errors=expected_vals['gate_pauli_errors'],
+        fsim_errors=expected_vals['fsim_errors'],
+    )
+    for field_name, expected in expected_vals.items():
+        target_dict = getattr(props_v2, field_name)
+        for key, val in expected.items():
+            if isinstance(target_dict[key], np.ndarray):
+                assert np.allclose(target_dict[key], val)
+            else:
+                assert target_dict[key] == val
+
+
+def test_with_params_opid_with_gate():
+    # Test gate-based opid with_params.
+    q0, q1 = cirq.LineQubit.range(2)
+    props = sample_noise_properties([q0, q1], [(q0, q1), (q1, q0)])
+    expected_vals = {
+        'gate_pauli_errors': 0.096,
+        'fsim_errors': cirq.PhasedFSimGate(0.0971, 0.0972, 0.0973, 0.0974, 0.0975),
+    }
+    props_v2 = props.with_params(
+        gate_pauli_errors={cirq.PhasedXZGate: expected_vals['gate_pauli_errors']},
+        fsim_errors={cirq.CZPowGate: expected_vals['fsim_errors']},
+    )
+    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']
+    assert props_v2.gate_pauli_errors[gpe_op_id_1] == expected_vals['gate_pauli_errors']
+
+    fsim_op_id_0 = cirq.OpIdentifier(cirq.CZPowGate, q0, q1)
+    fsim_op_id_1 = cirq.OpIdentifier(cirq.CZPowGate, q1, q0)
+    assert props_v2.fsim_errors[fsim_op_id_0] == expected_vals['fsim_errors']
+    assert props_v2.fsim_errors[fsim_op_id_1] == expected_vals['fsim_errors']
+
+
 @pytest.mark.parametrize(
     'op',
     [