Introduce gauge compilation

cirq-core/cirq/transformers/__init__.py

@@ -119,3 +119,14 @@
     unroll_circuit_op_greedy_earliest,
     unroll_circuit_op_greedy_frontier,
 )
+
+
+from cirq.transformers.gauge_compiling import (
+    Gauge,
+    GaugeSelector,
+    GaugeTransformer,
+    ConstantGauge,
+    SqrtCZGaugeTransormer,
+    CZGaugeTransformer,
+    SpinInversionGaugeTransformer,
+)

cirq-core/cirq/transformers/gauge_compiling/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+
+
+from cirq.transformers.gauge_compiling.gauge_compiling import (
+    Gauge,
+    GaugeSelector,
+    GaugeTransformer,
+    ConstantGauge,
+)
+from cirq.transformers.gauge_compiling.sqrt_cz_gauge import SqrtCZGaugeTransormer
+from cirq.transformers.gauge_compiling.spin_inversion_gauge import SpinInversionGaugeTransformer
+from cirq.transformers.gauge_compiling.cz_gauge import CZGaugeTransformer

cirq-core/cirq/transformers/gauge_compiling/cz_gauge.py

@@ -0,0 +1,46 @@
+# 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.
+
+"""A Gauge Transformer for the CZ gate."""
+
+from cirq.transformers.gauge_compiling.gauge_compiling import (
+    GaugeTransformer,
+    GaugeSelector,
+    ConstantGauge,
+)
+from cirq.ops.common_gates import CZ
+from cirq import ops
+
+CZGaugeSelector = GaugeSelector(
+    gauges=[
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.I, pre_q1=ops.I, post_q0=ops.I, post_q1=ops.I),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.I, pre_q1=ops.X, post_q0=ops.Z, post_q1=ops.X),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.I, pre_q1=ops.Y, post_q0=ops.Z, post_q1=ops.Y),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.I, pre_q1=ops.Z, post_q0=ops.I, post_q1=ops.Z),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.X, pre_q1=ops.I, post_q0=ops.X, post_q1=ops.Z),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.X, pre_q1=ops.X, post_q0=ops.Y, post_q1=ops.Y),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.X, pre_q1=ops.Y, post_q0=ops.Y, post_q1=ops.X),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.X, pre_q1=ops.Z, post_q0=ops.X, post_q1=ops.I),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Y, pre_q1=ops.I, post_q0=ops.Y, post_q1=ops.Z),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Y, pre_q1=ops.X, post_q0=ops.X, post_q1=ops.Y),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Y, pre_q1=ops.Y, post_q0=ops.X, post_q1=ops.X),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Y, pre_q1=ops.Z, post_q0=ops.Y, post_q1=ops.I),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Z, pre_q1=ops.I, post_q0=ops.Z, post_q1=ops.I),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Z, pre_q1=ops.X, post_q0=ops.I, post_q1=ops.X),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Z, pre_q1=ops.Y, post_q0=ops.I, post_q1=ops.Y),
+        ConstantGauge(two_qubit_gate=CZ, pre_q0=ops.Z, pre_q1=ops.Z, post_q0=ops.Z, post_q1=ops.Z),
+    ]
+)
+
+CZGaugeTransformer = GaugeTransformer(target=CZ, gauge_selector=CZGaugeSelector)

cirq-core/cirq/transformers/gauge_compiling/cz_gauge_test.py

@@ -0,0 +1,23 @@
+# 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.transformers.gauge_compiling import CZGaugeTransformer
+from cirq.transformers.gauge_compiling.gauge_compiling_test_utils import GaugeTester
+
+
+class TestCZGauge(GaugeTester):
+    two_qubit_gate = cirq.CZ
+    gauge_transformer = CZGaugeTransformer

cirq-core/cirq/transformers/gauge_compiling/gauge_compiling.py

@@ -0,0 +1,182 @@
+# 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.
+
+"""Creates the abstraction for gauge compiling as a cirq transformer."""
+
+from typing import Callable, Tuple, Optional, Sequence, Union, List, TYPE_CHECKING
+import abc
+import itertools
+import functools
+from dataclasses import dataclass
+
+import numpy as np
+
+from cirq.transformers import transformer_api
+from cirq import ops, circuits
+
+if TYPE_CHECKING:
+    import cirq
+
+
+_SINGLE_QUBIT_GATES_T = Optional[Union[ops.Gate, Sequence[ops.Gate]]]
+
+
+class Gauge(abc.ABC):
+    """A gauge replaces a two qubit gate with an equivalent subcircuit.
+    0: pre_q0───────two_qubit_gate───────post_q0
+                        |
+    1: pre_q1───────two_qubit_gate───────post_q1
+
+    The Gauge class in general represents a family of closely related gauges
+    (e.g. random z-rotations); Use `sample` method to get a specific gauge.
+    """
+
+    def weight(self) -> float:
+        """Returns the relative frequency for selecting this gauge."""
+        return 1.0
+
+    @abc.abstractmethod
+    def sample(self, gate: ops.Gate, prng: np.random.Generator) -> "ConstantGauge":
+        """Returns a ConstantGauge sampled from a family of gauges.
+
+        Args:
+            gate: The two qubit gate to replace.
+            prng: A numpy random number generator.
+
+        Returns:
+            A ConstantGauge.
+        """
+
+
+@dataclass(frozen=True)
+class ConstantGauge(Gauge):
+    """A gauge that replaces a two qubit gate with a constant gauge."""
+
+    two_qubit_gate: ops.Gate
+    pre_q0: _SINGLE_QUBIT_GATES_T = None
+    pre_q1: _SINGLE_QUBIT_GATES_T = None
+    post_q0: _SINGLE_QUBIT_GATES_T = None
+    post_q1: _SINGLE_QUBIT_GATES_T = None
+
+    def sample(self, gate: ops.Gate, prng: np.random.Generator) -> "ConstantGauge":
+        return self
+
+    @functools.cached_property
+    def pre(self) -> Tuple[Sequence[ops.Gate], Sequence[ops.Gate]]:
+        """A tuple (ops to apply to q0, ops to apply to q1)."""
+        return _as_sequence(self.pre_q0), _as_sequence(self.pre_q1)
+
+    @functools.cached_property
+    def post(self) -> Tuple[Sequence[ops.Gate], Sequence[ops.Gate]]:
+        """A tuple (ops to apply to q0, ops to apply to q1)."""
+        return _as_sequence(self.post_q0), _as_sequence(self.post_q1)
+
+
+@dataclass(frozen=True)
+class GaugeSelector:
+    """Samples a gauge from a list of gauges."""
+
+    gauges: Sequence[Gauge]
+
+    @functools.cached_property
+    def _weights(self) -> np.ndarray:
+        weights = np.array([g.weight() for g in self.gauges])
+        return weights / np.sum(weights)
+
+    def __call__(self, prng: np.random.Generator) -> "Gauge":
+        """Randomly selects a gauge with probability proportional to its weight."""
+        return self.gauges[prng.choice(len(self.gauges), p=self._weights)]
+
+
+@transformer_api.transformer
+class GaugeTransformer:
+    def __init__(
+        self,
+        # target can be either a specific gate, gatefamily or gateset
+        # which allows matching parametric gates.
+        target: Union[ops.Gate, ops.Gateset, ops.GateFamily],
+        gauge_selector: Callable[[np.random.Generator], Gauge],
+    ) -> None:
+        """Constructs a GaugeTransformer.
+
+        Args:
+            target: Target two-qubit gate, a gate-family or a gate-set of two-qubit gates.
+            gauge_selector: A callable that takes a numpy random number generator
+                as an argument and returns a Gauge.
+        """
+        if isinstance(target, ops.Gate):
+            self.target: Union[ops.GateFamily, ops.Gateset] = ops.GateFamily(target)
+        else:
+            self.target = target  # pragma: no cover
+        self.gauge_selector = gauge_selector
+
+    def __call__(
+        self,
+        circuit: circuits.AbstractCircuit,
+        *,
+        context: Optional["cirq.TransformerContext"] = None,
+        prng: Optional[np.random.Generator] = None,
+    ) -> circuits.AbstractCircuit:
+        rng = np.random.default_rng() if prng is None else prng
+        if context is None:
+            context = transformer_api.TransformerContext(deep=False)
+        if context.deep:
+            raise ValueError('GaugeTransformer cannot be used with deep=True')
+        new_moments = []
+        left: List[List[ops.Operation]] = []
+        right: List[List[ops.Operation]] = []
+        for moment in circuit:
+            left.clear()
+            right.clear()
+            center: List[ops.Operation] = []
+            for op in moment:
+                if isinstance(op, ops.TaggedOperation) and set(op.tags).intersection(
+                    context.tags_to_ignore
+                ):
+                    center.append(op)
+                    continue
+                if op.gate is not None and len(op.qubits) == 2 and op in self.target:
+                    gauge = self.gauge_selector(rng).sample(op.gate, rng)
+                    q0, q1 = op.qubits
+                    left.extend([g(q) for g in gs] for q, gs in zip(op.qubits, gauge.pre))
+                    center.append(gauge.two_qubit_gate(q0, q1))
+                    right.extend([g(q) for g in gs] for q, gs in zip(op.qubits, gauge.post))
+                else:
+                    center.append(op)
+            if left:
+                new_moments.extend(_build_moments(left))
+            new_moments.append(center)
+            if right:
+                new_moments.extend(_build_moments(right))
+        return circuits.Circuit.from_moments(*new_moments)
+
+
+def _as_sequence(gate: _SINGLE_QUBIT_GATES_T) -> Sequence[ops.Gate]:
+    if gate is None:
+        return []
+    if isinstance(gate, ops.Gate):
+        return [gate]
+    return gate
+
+
+def _build_moments(operation_by_qubits: List[List[ops.Operation]]) -> List[List[ops.Operation]]:
+    """Builds moments from a list of operations grouped by qubits.
+
+    Returns a list of moments from a list whose ith element is a list of operations applied
+    to qubit i.
+    """
+    moments = []
+    for moment in itertools.zip_longest(*operation_by_qubits):
+        moments.append([op for op in moment if op is not None])
+    return moments

cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test.py

@@ -0,0 +1,30 @@
+# 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 pytest
+import cirq
+from cirq.transformers.gauge_compiling import GaugeTransformer, CZGaugeTransformer
+
+
+def test_deep_transformation_not_supported():
+
+    with pytest.raises(ValueError):
+        _ = GaugeTransformer(target=cirq.CZ, gauge_selector=lambda _: None)(
+            cirq.Circuit(), context=cirq.TransformerContext(deep=True)
+        )
+
+
+def test_ignore_tags():
+    c = cirq.Circuit(cirq.CZ(*cirq.LineQubit.range(2)).with_tags('foo'))
+    assert c == CZGaugeTransformer(c, context=cirq.TransformerContext(tags_to_ignore={"foo"}))

cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test_utils.py

@@ -0,0 +1,60 @@
+# 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.
+
+
+from unittest.mock import patch
+import pytest
+
+import numpy as np
+
+import cirq
+from cirq.transformers.gauge_compiling import GaugeTransformer, GaugeSelector
+
+
+class GaugeTester:
+
+    two_qubit_gate: cirq.Gate
+    gauge_transformer: GaugeTransformer
+
+    @pytest.mark.parametrize('generation_seed', [*range(5)])
+    @pytest.mark.parametrize('transformation_seed', [*range(5)])
+    def test_gauge_transformer(self, generation_seed, transformation_seed):
+        c = cirq.testing.random_circuit(
+            qubits=3,
+            n_moments=3,
+            op_density=1,
+            gate_domain={self.two_qubit_gate: 2, cirq.X: 1, cirq.Y: 1, cirq.H: 1, cirq.Z: 1},
+            random_state=generation_seed,
+        )
+        nc = self.gauge_transformer(c, prng=np.random.default_rng(transformation_seed))
+        cirq.testing.assert_circuits_have_same_unitary_given_final_permutation(
+            nc, c, qubit_map={q: q for q in c.all_qubits()}
+        )
+
+    @patch('numpy.random.Generator', autospec=True)
+    def test_all_gauges(self, generator_mock):
+        assert isinstance(
+            self.gauge_transformer.gauge_selector, GaugeSelector
+        ), 'When using a custom selector, please override this method to properly test all gauges'
+        c = cirq.Circuit(self.two_qubit_gate(cirq.LineQubit(0), cirq.LineQubit(1)))
+        prng = generator_mock()
+        for idx, gauge in enumerate(self.gauge_transformer.gauge_selector.gauges):
+            prng.choice.return_value = idx
+            nc = self.gauge_transformer(c, prng=prng)
+            _ = (
+                cirq.testing.assert_circuits_have_same_unitary_given_final_permutation(
+                    nc, c, qubit_map={q: q for q in c.all_qubits()}
+                ),
+                f'{gauge=}',
+            )