Add option to do per-operation compensation (#3812)

* Add option to do per-operation compensation

* Add claryfing comment

* Missing semicolon

* Better comments

cirq/google/__init__.py

@@ -33,6 +33,7 @@
     SQRT_ISWAP_PARAMETERS,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     make_zeta_chi_gamma_compensation_for_moments,
+    merge_matching_results,
     prepare_floquet_characterization_for_moments,
     prepare_floquet_characterization_for_moment,
     run_calibrations,

cirq/google/calibration/__init__.py

@@ -27,6 +27,7 @@
     SQRT_ISWAP_PARAMETERS,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+    merge_matching_results,
 )
 
 from cirq.google.calibration.workflow import (

cirq/google/calibration/phased_fsim.py

@@ -15,6 +15,7 @@
     Any,
     Callable,
     Dict,
+    Iterable,
     List,
     MutableMapping,
     Optional,
@@ -286,6 +287,50 @@ def _json_dict_(self) -> Dict[str, Any]:
         }
 
 
+def merge_matching_results(
+    results: Iterable[PhasedFSimCalibrationResult],
+) -> Optional[PhasedFSimCalibrationResult]:
+    """Merges a collection of results into a single result.
+
+    Args:
+        results: List of results to merge. They must be compatible with each other: all gate and
+            options fields must be equal and every characterized pair must be present only in one of
+            the characterizations.
+
+    Returns:
+        New PhasedFSimCalibrationResult that contains all the parameters from every result in
+        results or None when the results list is empty.
+    """
+    all_parameters: Dict[Tuple[Qid, Qid], PhasedFSimCharacterization] = {}
+    common_gate = None
+    common_options = None
+    for result in results:
+        if common_gate is None:
+            common_gate = result.gate
+        elif common_gate != result.gate:
+            raise ValueError(
+                f'Only matching results can be merged, got gates {common_gate} and {result.gate}'
+            )
+
+        if common_options is None:
+            common_options = result.options
+        elif common_options != result.options:
+            raise ValueError(
+                f'Only matching results can be merged, got options {common_options} and '
+                f'{result.options}'
+            )
+
+        if not all_parameters.keys().isdisjoint(result.parameters):
+            raise ValueError(f'Only results with disjoint parameters sets can be merged')
+
+        all_parameters.update(result.parameters)
+
+    if common_gate is None or common_options is None:
+        return None
+
+    return PhasedFSimCalibrationResult(all_parameters, common_gate, common_options)
+
+
 # We have to relax a mypy constraint, see https://github.com/python/mypy/issues/5374
 @dataclasses.dataclass(frozen=True)  # type: ignore
 class PhasedFSimCalibrationRequest(abc.ABC):

cirq/google/calibration/phased_fsim_test.py

@@ -23,6 +23,7 @@
     PhasedFSimCharacterization,
     PhasedFSimCalibrationResult,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+    merge_matching_results,
     try_convert_sqrt_iswap_to_fsim,
 )
 
@@ -286,6 +287,106 @@ def test_get_parameters():
     assert result.get_parameters(q_00, q_03) is None
 
 
+def test_merge_matching_results():
+    q_00, q_01, q_02, q_03 = [cirq.GridQubit(0, index) for index in range(4)]
+    gate = cirq.FSimGate(theta=np.pi / 4, phi=0.0)
+    options = WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION
+    parameters_1 = {
+        (q_00, q_01): PhasedFSimCharacterization(
+            theta=0.1, zeta=0.2, chi=None, gamma=None, phi=0.3
+        ),
+    }
+    parameters_2 = {
+        (q_02, q_03): PhasedFSimCharacterization(
+            theta=0.4, zeta=0.5, chi=None, gamma=None, phi=0.6
+        ),
+    }
+
+    results = [
+        PhasedFSimCalibrationResult(
+            parameters=parameters_1,
+            gate=gate,
+            options=options,
+        ),
+        PhasedFSimCalibrationResult(
+            parameters=parameters_2,
+            gate=gate,
+            options=options,
+        ),
+    ]
+
+    assert merge_matching_results(results) == PhasedFSimCalibrationResult(
+        parameters={**parameters_1, **parameters_2},
+        gate=gate,
+        options=options,
+    )
+
+
+def test_merge_matching_results_when_empty_none():
+    assert merge_matching_results([]) is None
+
+
+def test_merge_matching_results_when_incompatible_fails():
+    q_00, q_01, q_02, q_03 = [cirq.GridQubit(0, index) for index in range(4)]
+    gate = cirq.FSimGate(theta=np.pi / 4, phi=0.0)
+    options = WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION
+    parameters_1 = {
+        (q_00, q_01): PhasedFSimCharacterization(
+            theta=0.1, zeta=0.2, chi=None, gamma=None, phi=0.3
+        ),
+    }
+    parameters_2 = {
+        (q_02, q_03): PhasedFSimCharacterization(
+            theta=0.4, zeta=0.5, chi=None, gamma=None, phi=0.6
+        ),
+    }
+
+    with pytest.raises(ValueError):
+        results = [
+            PhasedFSimCalibrationResult(
+                parameters=parameters_1,
+                gate=gate,
+                options=options,
+            ),
+            PhasedFSimCalibrationResult(
+                parameters=parameters_1,
+                gate=gate,
+                options=options,
+            ),
+        ]
+        assert merge_matching_results(results)
+
+    with pytest.raises(ValueError):
+        results = [
+            PhasedFSimCalibrationResult(
+                parameters=parameters_1,
+                gate=gate,
+                options=options,
+            ),
+            PhasedFSimCalibrationResult(
+                parameters=parameters_2,
+                gate=cirq.CZ,
+                options=options,
+            ),
+        ]
+        assert merge_matching_results(results)
+
+    with pytest.raises(ValueError):
+        results = [
+            PhasedFSimCalibrationResult(
+                parameters=parameters_1,
+                gate=gate,
+                options=options,
+            ),
+            PhasedFSimCalibrationResult(
+                parameters=parameters_2,
+                gate=gate,
+                options=ALL_ANGLES_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+            ),
+        ]
+        assert merge_matching_results(results)
+
+
 @pytest.mark.parametrize('phase_exponent', np.linspace(0, 1, 5))
 def test_phase_calibrated_fsim_gate_as_characterized_phased_fsim_gate(phase_exponent: float):
     a, b = cirq.LineQubit.range(2)

cirq/google/calibration/workflow.py

@@ -40,6 +40,7 @@
     PhasedFSimCharacterization,
     WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
     THETA_ZETA_GAMMA_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+    merge_matching_results,
     try_convert_sqrt_iswap_to_fsim,
 )
 from cirq.google.engine import Engine
@@ -60,7 +61,7 @@ class CircuitWithCalibration:
     """
 
     circuit: Circuit
-    moment_to_calibration: List[Optional[int]]
+    moment_to_calibration: Sequence[Optional[int]]
 
 
 def prepare_floquet_characterization_for_moment(
@@ -522,7 +523,7 @@ def make_zeta_chi_gamma_compensation_for_moments(
         [Gate], Optional[PhaseCalibratedFSimGate]
     ] = try_convert_sqrt_iswap_to_fsim,
 ) -> CircuitWithCalibration:
-    """Compensates circuit against errors in zeta, chi and gamma angles.
+    """Compensates circuit moments against errors in zeta, chi and gamma angles.
 
     This method creates a new circuit with a single-qubit Z gates added in a such way so that
     zeta, chi and gamma angles discovered by characterizations are cancelled-out and set to 0.
@@ -546,6 +547,74 @@ def make_zeta_chi_gamma_compensation_for_moments(
         The moment to calibration mapping is updated for the new circuit so that successive
         calibrations could be applied.
     """
+    return _make_zeta_chi_gamma_compensation(
+        circuit_with_calibration, characterizations, gates_translator, permit_mixed_moments=False
+    )
+
+
+def make_zeta_chi_gamma_compensation_for_operations(
+    circuit: Circuit,
+    characterizations: List[PhasedFSimCalibrationResult],
+    gates_translator: Callable[
+        [Gate], Optional[PhaseCalibratedFSimGate]
+    ] = try_convert_sqrt_iswap_to_fsim,
+    permit_mixed_moments: bool = False,
+) -> Circuit:
+    """Compensates circuit operations against errors in zeta, chi and gamma angles.
+
+    This method creates a new circuit with a single-qubit Z gates added in a such way so that
+    zeta, chi and gamma angles discovered by characterizations are cancelled-out and set to 0.
+
+    Contrary to make_zeta_chi_gamma_compensation_for_moments this method does not match
+    characterizations to the moment structure of the circuits and thus is less accurate because
+    some errors caused by cross-talks are not mitigated.
+
+    The major advantage of this method over make_zeta_chi_gamma_compensation_for_moments is that it
+    can work with arbitrary set of characterizations that cover all the interactions of the circuit
+    (up to assumptions of merge_matching_results method). In particular, for grid-like devices the
+    number of characterizations is bounded by four, where in the case of
+    make_zeta_chi_gamma_compensation_for_moments the number of characterizations is bounded by
+    number of moments in a circuit.
+
+    This function preserves a moment structure of the circuit. All single qubit gates appear on new
+    moments in the final circuit.
+
+    Args:
+        circuit: Circuit to calibrate.
+        characterizations: List of characterization results (likely returned from run_calibrations).
+            All the characterizations must be compatible in sense of merge_matching_results, they
+            will be merged together.
+        gates_translator: Function that translates a gate to a supported FSimGate which will undergo
+            characterization. Defaults to sqrt_iswap_gates_translator.
+        permit_mixed_moments: Whether to allow mixing single-qubit and two-qubit gates in a single
+            moment.
+
+    Returns:
+        Calibrated circuit with a single-qubit Z gates added which compensates for the true gates
+        imperfections.
+    """
+
+    characterization = merge_matching_results(characterizations)
+    moment_to_calibration = [0] * len(circuit)
+    calibrated = _make_zeta_chi_gamma_compensation(
+        CircuitWithCalibration(circuit, moment_to_calibration),
+        [characterization] if characterization is not None else [],
+        gates_translator,
+        permit_mixed_moments=permit_mixed_moments,
+    )
+    return calibrated.circuit
+
+
+def _make_zeta_chi_gamma_compensation(
+    circuit_with_calibration: CircuitWithCalibration,
+    characterizations: List[PhasedFSimCalibrationResult],
+    gates_translator: Callable[[Gate], Optional[PhaseCalibratedFSimGate]],
+    permit_mixed_moments: bool,
+) -> CircuitWithCalibration:
+
+    if permit_mixed_moments:
+        raise NotImplementedError('Mixed moments compensation ist supported yet')
+
     if len(circuit_with_calibration.circuit) != len(circuit_with_calibration.moment_to_calibration):
         raise ValueError('Moment allocations does not match circuit length')
 

cirq/google/calibration/workflow_test.py

@@ -359,6 +359,62 @@ def test_make_floquet_request_for_operations_when_multiple_gates_fails() -> None
         )
 
 
+def test_make_zeta_chi_gamma_compensation_for_operations():
+    a, b, c, d = cirq.LineQubit.range(4)
+    parameters_ab = cirq.google.PhasedFSimCharacterization(zeta=0.5, chi=0.4, gamma=0.3)
+    parameters_bc = cirq.google.PhasedFSimCharacterization(zeta=-0.5, chi=-0.4, gamma=-0.3)
+    parameters_cd = cirq.google.PhasedFSimCharacterization(zeta=0.2, chi=0.3, gamma=0.4)
+
+    parameters_dict = {(a, b): parameters_ab, (b, c): parameters_bc, (c, d): parameters_cd}
+
+    engine_simulator = cirq.google.PhasedFSimEngineSimulator.create_from_dictionary_sqrt_iswap(
+        parameters={
+            pair: parameters.merge_with(SQRT_ISWAP_PARAMETERS)
+            for pair, parameters in parameters_dict.items()
+        }
+    )
+
+    circuit = cirq.Circuit(
+        [
+            [cirq.X(a), cirq.Y(c)],
+            [SQRT_ISWAP_GATE.on(a, b), SQRT_ISWAP_GATE.on(c, d)],
+            [SQRT_ISWAP_GATE.on(b, c)],
+        ]
+    )
+
+    options = cirq.google.FloquetPhasedFSimCalibrationOptions(
+        characterize_theta=False,
+        characterize_zeta=True,
+        characterize_chi=True,
+        characterize_gamma=True,
+        characterize_phi=False,
+    )
+
+    characterizations = [
+        PhasedFSimCalibrationResult(
+            parameters={pair: parameters}, gate=SQRT_ISWAP_GATE, options=options
+        )
+        for pair, parameters in parameters_dict.items()
+    ]
+
+    calibrated_circuit = workflow.make_zeta_chi_gamma_compensation_for_operations(
+        circuit,
+        characterizations,
+    )
+
+    assert cirq.allclose_up_to_global_phase(
+        engine_simulator.final_state_vector(calibrated_circuit),
+        cirq.final_state_vector(circuit),
+    )
+
+
+def test_make_zeta_chi_gamma_compensation_for_operations_permit_mixed_moments():
+    with pytest.raises(NotImplementedError):
+        workflow.make_zeta_chi_gamma_compensation_for_operations(
+            cirq.Circuit(), [], permit_mixed_moments=True
+        )
+
+
 def test_run_characterization():
     q_00, q_01, q_02, q_03 = [cirq.GridQubit(0, index) for index in range(4)]
     gate = cirq.FSimGate(theta=np.pi / 4, phi=0.0)