mypy & coverage

Author: NoureldinYosri

cirq-core/cirq/transformers/pauli_insertion.py

@@ -23,24 +23,7 @@
 from cirq import circuits, ops
 from cirq.transformers import transformer_api
 
-_PAULIS = [ops.I, ops.X, ops.Y, ops.Z]
-
-
-def _is_target(
-    op: ops.Operation,
-    target: ops.Gate | ops.GateFamily | ops.Gateset | type[ops.Gate | ops.Operation],
-):
-    if inspect.isclass(target):
-        if issubclass(target, ops.Operation):
-            return isinstance(op, target)
-        if not hasattr(op, 'gate'):
-            return False
-        return isinstance(op.gate, target)
-    if isinstance(target, ops.Gate):
-        if not hasattr(op, 'gate') or op.gate is None:
-            return False
-        return op.gate == target
-    return op in target
+_PAULIS: tuple[ops.Gate] = (ops.I, ops.X, ops.Y, ops.Z)  # type: ignore[has-type]
 
 
 @transformer_api.transformer
@@ -54,13 +37,13 @@ class PauliInsertionTransformer:
 
     def __init__(
         self,
-        target: ops.Gate | ops.GateFamily | ops.Gateset | type[ops.Gate | ops.Operation],
+        target: ops.Gate | ops.GateFamily | ops.Gateset | type[ops.Gate],
         probabilities: np.ndarray | None = None,
     ):
         """Makes a pauli insertion transformer that samples 2Q paulis with the given probabilities.
 
         Args:
-            target: The target gate, gatefamily, gateset, or type (e.g. PauliSumExponential).
+            target: The target gate, gatefamily, gateset, or type (e.g. ZZPowGAte).
             probabilities: Optional ndarray representing the probabilities of sampling 2Q paulis.
                 The order of the paulis is IXYZ. If None, assume uniform distribution.
         Returns:
@@ -72,7 +55,13 @@ def __init__(
         assert probabilities.shape == (4, 4)
         assert np.isclose(probabilities.sum(), 1)
 
-        self.target = target
+        if inspect.isclass(target):
+            self.target = ops.GateFamily(target)
+        elif isinstance(target, ops.Gate):
+            self.target = ops.Gateset(target)
+        else:
+            assert isinstance(target, (ops.Gateset, ops.GateFamily))
+            self.target = target
         self._flat_probs = probabilities.reshape(-1)
 
     def __call__(
@@ -106,7 +95,7 @@ def __call__(
             for op in moment:
                 if any(tag in tags_to_ignore for tag in op.tags):
                     continue
-                if not _is_target(op, self.target):
+                if op not in self.target:
                     continue
                 pair = np.unravel_index(rng.choice(16, p=self._flat_probs), (4, 4))
                 for pauli_index, q in zip(pair, op.qubits):

cirq-core/cirq/transformers/pauli_insertion_test.py

@@ -29,12 +29,17 @@ def _random_probs(n: int, seed: int | None = None):
 
 
 @pytest.mark.parametrize('probs', _random_probs(3, 0))
-def test_pauli_insertion_with_probabilities(probs):
+@pytest.mark.parametrize(
+    'target',
+    [cirq.ZZPowGate, cirq.ZZ**0.324, cirq.Gateset(cirq.ZZ**0.324), cirq.GateFamily(cirq.ZZ**0.324)],
+)
+def test_pauli_insertion_with_probabilities(probs, target):
     c = cirq.Circuit(cirq.ZZ(*cirq.LineQubit.range(2)) ** 0.324)
-    transformer = cirq.transformers.PauliInsertionTransformer(cirq.ZZPowGate, probs)
+    transformer = cirq.transformers.PauliInsertionTransformer(target, probs)
     count = np.zeros((4, 4))
+    rng = np.random.default_rng(0)
     for _ in range(100):
-        nc = transformer(c)
+        nc = transformer(c, rng_or_seed=rng)
         assert len(nc) == 2
         u, v = nc[0]
         i = _PAULIS.index(u.gate)
@@ -49,12 +54,36 @@ def test_pauli_insertion_with_probabilities_doesnot_create_moment(probs):
     c = cirq.Circuit.from_moments([], [cirq.ZZ(*cirq.LineQubit.range(2)) ** 0.324])
     transformer = cirq.transformers.PauliInsertionTransformer(cirq.ZZPowGate, probs)
     count = np.zeros((4, 4))
+    rng = np.random.default_rng(0)
     for _ in range(100):
-        nc = transformer(c)
+        nc = transformer(c, rng_or_seed=rng)
         assert len(nc) == 2
         u, v = nc[0]
         i = _PAULIS.index(u.gate)
         j = _PAULIS.index(v.gate)
         count[i, j] += 1
     count = count / count.sum()
     np.testing.assert_allclose(count, probs, atol=0.1)
+
+
+def test_invalid_context_raises():
+    c = cirq.Circuit(cirq.ZZ(*cirq.LineQubit.range(2)) ** 0.324)
+    transformer = cirq.transformers.PauliInsertionTransformer(cirq.ZZPowGate)
+    with pytest.raises(ValueError):
+        _ = transformer(c, context=cirq.TransformerContext(deep=True))
+
+
+def test_transformer_ignores_tagged_ops():
+    op = cirq.ZZ(*cirq.LineQubit.range(2)) ** 0.324
+    c = cirq.Circuit(op.with_tags('ignore'))
+    transformer = cirq.transformers.PauliInsertionTransformer(cirq.ZZPowGate)
+
+    assert transformer(c, context=cirq.TransformerContext(tags_to_ignore=('ignore',))) == c
+
+
+def test_transformer_ignores_tagged_moments():
+    op = cirq.ZZ(*cirq.LineQubit.range(2)) ** 0.324
+    c = cirq.Circuit(cirq.Moment(op).with_tags('ignore'))
+    transformer = cirq.transformers.PauliInsertionTransformer(cirq.ZZPowGate)
+
+    assert transformer(c, context=cirq.TransformerContext(tags_to_ignore=('ignore',))) == c