Impl dynamical decoupling transfomer.

Author: babacry

cirq-core/cirq/__init__.py

@@ -336,6 +336,8 @@
 
 from cirq.transformers import (
     AbstractInitialMapper,
+    add_dynamical_decoupling,
+    DynamicalDecouplingModel,
     align_left,
     align_right,
     CompilationTargetGateset,

cirq-core/cirq/json_resolver_cache.py

@@ -102,6 +102,7 @@ def _symmetricalqidpair(qids):
     import sympy
 
     return {
+        'DynamicalDecouplingModel': cirq.DynamicalDecouplingModel,
         'AmplitudeDampingChannel': cirq.AmplitudeDampingChannel,
         'AnyIntegerPowerGateFamily': cirq.AnyIntegerPowerGateFamily,
         'AnyUnitaryGateFamily': cirq.AnyUnitaryGateFamily,

cirq-core/cirq/protocols/json_test_data/DynamicalDecouplingModel.json

@@ -0,0 +1,21 @@
+[
+    {
+        "cirq_type": "DynamicalDecouplingModel",
+        "schema": "XX_PAIR"
+    },
+    {
+        "cirq_type": "DynamicalDecouplingModel",
+        "base_dd_sequence": [
+            {
+                "cirq_type": "XPowGate",
+                "exponent": 1.0,
+                "global_shift": 0.0
+            },
+            {
+                "cirq_type": "XPowGate",
+                "exponent": 1.0,
+                "global_shift": 0.0
+            }
+        ]
+    }
+]

cirq-core/cirq/protocols/json_test_data/DynamicalDecouplingModel.repr

@@ -0,0 +1,2 @@
+[cirq.DynamicalDecouplingModel.from_schema("XX_PAIR"),
+cirq.DynamicalDecouplingModel(base_dd_sequence=[cirq.XPowGate(), cirq.XPowGate()])]

cirq-core/cirq/transformers/__init__.py

@@ -78,6 +78,11 @@
 
 from cirq.transformers.drop_negligible_operations import drop_negligible_operations
 
+from cirq.transformers.dynamical_decoupling import (
+    add_dynamical_decoupling,
+    DynamicalDecouplingModel,
+)
+
 from cirq.transformers.eject_z import eject_z
 
 from cirq.transformers.measurement_transformers import (

cirq-core/cirq/transformers/dynamical_decoupling.py

@@ -0,0 +1,146 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformer pass that adds dynamical decoupling moments to a circuit."""
+
+import enum
+from functools import reduce
+from typing import Any, Dict, Optional, Tuple
+
+import cirq
+from cirq import value
+import numpy as np
+
+
+@enum.unique
+class _DynamicalDecouplingSchema(enum.Enum):
+    """Supported schemes of dynamical decoupling."""
+
+    XX_PAIR = 'XX_PAIR'
+    YY_PAIR = 'YY_PAIR'
+
+
+def _repeat_sequence(base_sequence: list['cirq.Gate'], num_idle_moments: int):
+    repeat_times = num_idle_moments // len(base_sequence)
+    return base_sequence * repeat_times
+
+
+def _generate_dd_sequence_from_schema(
+    schema: _DynamicalDecouplingSchema, num_idle_moments: int = 2
+) -> list['cirq.Gate']:
+    match schema:
+        case _DynamicalDecouplingSchema.XX_PAIR:
+            return _repeat_sequence([cirq.XPowGate(), cirq.XPowGate()], num_idle_moments)
+        case _DynamicalDecouplingSchema.YY_PAIR:
+            return _repeat_sequence([cirq.YPowGate(), cirq.YPowGate()], num_idle_moments)
+
+
+def _validate_dd_sequence(dd_sequence: list['cirq.Gate']) -> None:
+    if len(dd_sequence) < 2:
+        raise ValueError('Invalid dynamical decoupling sequence. Expect more than one gates.')
+    matrices = [cirq.unitary(gate) for gate in dd_sequence]
+    product = reduce(np.matmul, matrices)
+    if not np.array_equal(product, np.eye(2)):
+        raise ValueError(
+            "Invalid dynamical decoupling sequence, sequence product doesn't equal" ' identity.'
+        )
+
+
+@value.value_equality
+class DynamicalDecouplingModel:
+    """Dynamical decoupling model that generates dynamical decoupling gate sequences."""
+
+    def __init__(
+        self,
+        schema: Optional[_DynamicalDecouplingSchema] = None,
+        base_dd_sequence: Optional[list['cirq.Gate']] = None,
+    ):
+        if not schema and not base_dd_sequence:
+            raise ValueError(
+                'Specify either schema or base_dd_sequence to construct a valid'
+                ' DynamicalDecouplingModel.'
+            )
+        self.schema = schema
+        self.base_dd_sequence = base_dd_sequence
+        if base_dd_sequence:
+            _validate_dd_sequence(base_dd_sequence)
+
+    def generate_dd_sequence(self, num_idle_moments: int = 2) -> list['cirq.Gate']:
+        """Returns the longest possible dynamical decoupling sequence."""
+        if num_idle_moments <= 0:
+            return []
+        if self.schema:
+            return _generate_dd_sequence_from_schema(self.schema, num_idle_moments)
+        if self.base_dd_sequence:
+            return _repeat_sequence(self.base_dd_sequence, num_idle_moments)
+        return []
+
+    @classmethod
+    def from_schema(cls, schema: str):
+        """Create dynamical decoupling model according to a given schema."""
+        if not schema in _DynamicalDecouplingSchema.__members__:
+            raise ValueError("Invalid schema name.")
+        return cls(schema=_DynamicalDecouplingSchema[schema])
+
+    @classmethod
+    def from_base_dd_sequence(cls, base_dd_sequence: list['cirq.Gate']):
+        """Create dynamical decoupling model according to a base sequence."""
+        return cls(base_dd_sequence=base_dd_sequence)
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        d: Dict[str, Any] = {}
+        if self.schema:
+            d['schema'] = self.schema.name
+        if self.base_dd_sequence:
+            d['base_dd_sequence'] = self.base_dd_sequence
+        return d
+
+    @classmethod
+    def _from_json_dict_(cls, schema=None, base_dd_sequence=None, **kwargs):
+        if schema:
+            return cls(schema=_DynamicalDecouplingSchema[schema])
+        if base_dd_sequence:
+            return cls(base_dd_sequence=base_dd_sequence)
+
+    def _value_equality_values_(self) -> Any:
+        return self.schema, self.base_dd_sequence
+
+
+def add_dynamical_decoupling(
+    circuit: 'cirq.AbstractCircuit', dd_model: DynamicalDecouplingModel
+) -> 'cirq.Circuit':
+    """Add dynamical decoupling gates in a given circuit.
+
+    Args:
+          circuit: Input circuit to transform.
+          dd_model: Dynamical decoupling model that defines the schema to generate
+            dynamical decoupling sequences.
+
+    Return:
+          A circuit with dynamical decoupling operations.
+    """
+    last_busy_moment_by_qubits: Dict['cirq.Qid', int] = {q: 0 for q in circuit.all_qubits()}
+    insert_into: list[Tuple[int, 'cirq.OP_TREE']] = []
+    for moment_id, moment in enumerate(circuit):
+        for q in moment.qubits:
+            insert_gates = dd_model.generate_dd_sequence(
+                num_idle_moments=moment_id - last_busy_moment_by_qubits[q] - 1
+            )
+            for idx, gate in enumerate(insert_gates):
+                insert_into.append((last_busy_moment_by_qubits[q] + idx + 1, gate.on(q)))
+            last_busy_moment_by_qubits[q] = moment_id
+
+    circuit.batch_insert_into(insert_into)
+
+    return circuit

cirq-core/cirq/transformers/dynamical_decoupling_test.py

@@ -0,0 +1,121 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import cirq
+from cirq import DynamicalDecouplingModel, add_dynamical_decoupling
+import pytest
+
+
+def assert_dd(
+    input_circuit: cirq.Circuit, expected_circuit: cirq.Circuit, dd_model: DynamicalDecouplingModel
+):
+    updated_circuit = add_dynamical_decoupling(input_circuit, dd_model=dd_model)
+    cirq.testing.assert_same_circuits(updated_circuit, expected_circuit)
+
+
+def test_insert_provided_schema():
+    a = cirq.NamedQubit("a")
+    b = cirq.NamedQubit("b")
+    c = cirq.NamedQubit("c")
+
+    # No insertion as there is no room for a dd sequence.
+    assert_dd(
+        input_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)), cirq.Moment(cirq.CNOT(a, b)), cirq.Moment(cirq.H(b))
+        ),
+        expected_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)), cirq.Moment(cirq.CNOT(a, b)), cirq.Moment(cirq.H(b))
+        ),
+        dd_model=DynamicalDecouplingModel.from_schema("XX_PAIR"),
+    )
+
+    # Insert one XX_PAIR dynamical decoupling sequence in idle moments.
+    assert_dd(
+        input_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        expected_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.X(a)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.X(a)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        dd_model=DynamicalDecouplingModel.from_schema("XX_PAIR"),
+    )
+
+    # Insert one XX_PAIR dynamical decoupling sequence in idle moments.
+    assert_dd(
+        input_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        expected_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.Y(a)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.Y(a)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        dd_model=DynamicalDecouplingModel.from_schema("YY_PAIR"),
+    )
+
+
+def test_insert_by_customized_dd_sequence():
+    a = cirq.NamedQubit("a")
+    b = cirq.NamedQubit("b")
+    c = cirq.NamedQubit("c")
+
+    assert_dd(
+        input_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.CNOT(b, c)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        expected_circuit=cirq.Circuit(
+            cirq.Moment(cirq.H(a)),
+            cirq.Moment(cirq.CNOT(a, b)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.X(a)),
+            cirq.Moment(cirq.CNOT(b, c), cirq.X(a)),
+            cirq.Moment(cirq.measure_each(a, b, c)),
+        ),
+        dd_model=DynamicalDecouplingModel.from_base_dd_sequence([cirq.XPowGate(), cirq.XPowGate()]),
+    )
+
+
+def test_dd_model_constructor():
+    # Succeed
+    DynamicalDecouplingModel.from_schema("XX_PAIR")
+    DynamicalDecouplingModel.from_schema("YY_PAIR")
+    DynamicalDecouplingModel.from_base_dd_sequence(
+        [cirq.XPowGate(), cirq.XPowGate(), cirq.YPowGate(), cirq.YPowGate()]
+    )
+    # Fail
+    with pytest.raises(ValueError, match="Specify either schema or base_dd_sequence"):
+        DynamicalDecouplingModel()
+    with pytest.raises(ValueError, match="Invalid schema name."):
+        DynamicalDecouplingModel.from_schema("unimplemented_schema")
+    with pytest.raises(ValueError, match="Invalid dynamical decoupling sequence. Expect more than one gates."):
+        DynamicalDecouplingModel.from_base_dd_sequence([cirq.XPowGate()])
+    with pytest.raises(ValueError, match="Invalid dynamical decoupling sequence"):
+        DynamicalDecouplingModel.from_base_dd_sequence([cirq.XPowGate(), cirq.YPowGate()])