From b184463a04cc4bec0eea8d29b606be65194d1e97 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Wed, 10 Jul 2024 17:28:10 -0700
Subject: [PATCH 01/16] Add noise amplification transformer

---
 cirq-core/cirq/transformers/noise_adding.py   | 62 +++++++++++++++++++
 .../cirq/transformers/noise_adding_test.py    | 12 ++++
 2 files changed, 74 insertions(+)
 create mode 100644 cirq-core/cirq/transformers/noise_adding.py
 create mode 100644 cirq-core/cirq/transformers/noise_adding_test.py

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
new file mode 100644
index 00000000000..573c2e38def
--- /dev/null
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -0,0 +1,62 @@
+from collections.abc import Mapping
+
+import cirq
+import numpy as np
+
+
+def add_depolarizing_noise_to_two_qubit_gates(
+    circuit: cirq.Circuit,
+    p: float | Mapping[tuple[cirq.Qid, cirq.Qid], float],
+    target_gate: cirq.Gate = cirq.CZ,
+    rng: np.random.Generator | None = None,
+) -> cirq.Circuit:
+    """Add local depolarizing noise after two-qubit gates in a specified circuit. More specifically,
+    with probability p, append a random non-identity two-qubit Pauli operator after each specified
+    two-qubit gate.
+
+    Args:
+        circuit: The circuit to add noise to.
+        p: The probability with which to add noise.
+        target_gate: Add depolarizing nose after this type of gate
+        rng: The pseudorandom number generator to use.
+
+    Returns:
+        The transformed circuit.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    # add random Pauli gates with probability p after each of the specified gate
+    assert target_gate.num_qubits() == 2, "`target_gate` must be a two-qubit gate."
+    paulis = [cirq.I, cirq.X, cirq.Y, cirq.Z]
+    new_moments = []
+    for moment in circuit:
+        new_moments.append(moment)
+        if _gate_in_moment(target_gate, moment):
+            # add a new moment with the Paulis
+            target_pairs = {
+                tuple(sorted(op.qubits)) for op in moment.operations if op.gate == target_gate
+            }
+            added_moment_ops = []
+            for pair in target_pairs:
+                if isinstance(p, float):
+                    p_i = p
+                elif isinstance(p, Mapping):
+                    pair_sorted_tuple = (pair[0], pair[1])
+                    p_i = p[pair_sorted_tuple]
+                else:
+                    raise TypeError(
+                        "p must either be a float or a mapping from sorted qubit pairs to floats"
+                    )
+                apply = rng.choice([True, False], p=[p_i, 1 - p_i])
+                if apply:
+                    choices = [
+                        (pauli_a(pair[0]), pauli_b(pair[1]))
+                        for pauli_a in paulis
+                        for pauli_b in paulis
+                    ][1:]
+                    pauli_to_apply = rng.choice(np.array(choices, dtype=object))
+                    added_moment_ops.append(pauli_to_apply)
+            if len(added_moment_ops) > 0:
+                new_moments.append(cirq.Moment(*added_moment_ops))
+    return cirq.Circuit.from_moments(*new_moments)
diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
new file mode 100644
index 00000000000..7a859e2aaf7
--- /dev/null
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -0,0 +1,12 @@
+import cirq
+import cirq.transformers.noise_adding as na
+
+
+def test_noise_adding():
+    qubits = cirq.LineQubit.range(4)
+    circuit = cirq.Circuit(cirq.CZ(*qubits[:2]), cirq.CZ(*qubits[2:])) * 10
+    transformed_circuit_p0 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.0)
+    assert transformed_circuit_p0 == circuit
+    transformed_circuit_p1 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 1.0)
+    assert len(transformed_circuit_p1) == 20
+    transformed_circuit_p0_03 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.03)

From 813704e8fb2b99b55e4e5a2c91b0f325ce37bda3 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Wed, 10 Jul 2024 17:42:13 -0700
Subject: [PATCH 02/16] add _gate_in_moment

---
 cirq-core/cirq/transformers/noise_adding.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 573c2e38def..2f2b0b0284e 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -4,6 +4,16 @@
 import numpy as np
 
 
+def _gate_in_moment(gate: cirq.Gate, moment: cirq.Moment) -> bool:
+    """Check whether `gate` is in `moment`."""
+    target_gate_in_moment = False
+    for op in moment.operations:
+        if op.gate == gate:
+            target_gate_in_moment = True
+            break
+    return target_gate_in_moment
+
+
 def add_depolarizing_noise_to_two_qubit_gates(
     circuit: cirq.Circuit,
     p: float | Mapping[tuple[cirq.Qid, cirq.Qid], float],

From d7547da31fc9a72aa829ff13ea302cdd9ae8a2d0 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Wed, 10 Jul 2024 17:50:28 -0700
Subject: [PATCH 03/16] add copyright notice

---
 cirq-core/cirq/transformers/noise_adding.py      | 14 ++++++++++++++
 cirq-core/cirq/transformers/noise_adding_test.py | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 2f2b0b0284e..8ac63bbb364 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -1,3 +1,17 @@
+# 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 collections.abc import Mapping
 
 import cirq
diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
index 7a859e2aaf7..803dd5f7cfa 100644
--- a/cirq-core/cirq/transformers/noise_adding_test.py
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -1,3 +1,17 @@
+# 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
 import cirq.transformers.noise_adding as na
 

From 57382e5ddf4c2fd327431b812c1996663a7677c9 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Wed, 10 Jul 2024 18:06:39 -0700
Subject: [PATCH 04/16] add raises documentation, remove import cirq, fix
 coverage

---
 cirq-core/cirq/transformers/noise_adding.py   | 22 +++++++++++--------
 .../cirq/transformers/noise_adding_test.py    | 11 +++++++---
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 8ac63bbb364..332bf48f87c 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -14,11 +14,11 @@
 
 from collections.abc import Mapping
 
-import cirq
+from cirq import ops, circuits
 import numpy as np
 
 
-def _gate_in_moment(gate: cirq.Gate, moment: cirq.Moment) -> bool:
+def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
     """Check whether `gate` is in `moment`."""
     target_gate_in_moment = False
     for op in moment.operations:
@@ -28,12 +28,13 @@ def _gate_in_moment(gate: cirq.Gate, moment: cirq.Moment) -> bool:
     return target_gate_in_moment
 
 
+@transformer_api.transformer
 def add_depolarizing_noise_to_two_qubit_gates(
-    circuit: cirq.Circuit,
-    p: float | Mapping[tuple[cirq.Qid, cirq.Qid], float],
-    target_gate: cirq.Gate = cirq.CZ,
+    circuit: circuits.Circuit,
+    p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
+    target_gate: ops.Gate = ops.CZ,
     rng: np.random.Generator | None = None,
-) -> cirq.Circuit:
+) -> circuits.Circuit:
     """Add local depolarizing noise after two-qubit gates in a specified circuit. More specifically,
     with probability p, append a random non-identity two-qubit Pauli operator after each specified
     two-qubit gate.
@@ -46,13 +47,16 @@ def add_depolarizing_noise_to_two_qubit_gates(
 
     Returns:
         The transformed circuit.
+
+    Raises:
+        TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to floats.
     """
     if rng is None:
         rng = np.random.default_rng()
 
     # add random Pauli gates with probability p after each of the specified gate
     assert target_gate.num_qubits() == 2, "`target_gate` must be a two-qubit gate."
-    paulis = [cirq.I, cirq.X, cirq.Y, cirq.Z]
+    paulis = [ops.I, ops.X, ops.Y, ops.Z]
     new_moments = []
     for moment in circuit:
         new_moments.append(moment)
@@ -82,5 +86,5 @@ def add_depolarizing_noise_to_two_qubit_gates(
                     pauli_to_apply = rng.choice(np.array(choices, dtype=object))
                     added_moment_ops.append(pauli_to_apply)
             if len(added_moment_ops) > 0:
-                new_moments.append(cirq.Moment(*added_moment_ops))
-    return cirq.Circuit.from_moments(*new_moments)
+                new_moments.append(circuits.Moment(*added_moment_ops))
+    return circuits.Circuit.from_moments(*new_moments)
diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
index 803dd5f7cfa..b07a28d919f 100644
--- a/cirq-core/cirq/transformers/noise_adding_test.py
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -12,15 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import cirq
+from cirq import ops, circuits, devices
 import cirq.transformers.noise_adding as na
 
 
 def test_noise_adding():
-    qubits = cirq.LineQubit.range(4)
-    circuit = cirq.Circuit(cirq.CZ(*qubits[:2]), cirq.CZ(*qubits[2:])) * 10
+    qubits = devices.LineQubit.range(4)
+    circuit = circuits.Circuit(ops.CZ(*qubits[:2]), ops.CZ(*qubits[2:])) * 10
     transformed_circuit_p0 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.0)
     assert transformed_circuit_p0 == circuit
     transformed_circuit_p1 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 1.0)
     assert len(transformed_circuit_p1) == 20
     transformed_circuit_p0_03 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.03)
+    assert 10 <= len(transformed_circuit_p0_03) <= 20
+    transformed_circuit_p_dict = na.add_depolarizing_noise_to_two_qubit_gates(
+        circuit, {tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
+    )
+    assert len(transformed_circuit_p_dict) == 20

From 69355c54cd87f375c1d24544c48e93340c94dfe4 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Wed, 10 Jul 2024 18:12:36 -0700
Subject: [PATCH 05/16] transformer api

---
 cirq-core/cirq/transformers/noise_adding.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 332bf48f87c..e8c52aa8aa6 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -15,6 +15,7 @@
 from collections.abc import Mapping
 
 from cirq import ops, circuits
+from cirq.transformers import transformer_api
 import numpy as np
 
 
@@ -34,6 +35,7 @@ def add_depolarizing_noise_to_two_qubit_gates(
     p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
     target_gate: ops.Gate = ops.CZ,
     rng: np.random.Generator | None = None,
+    context: transformer_api.TransformerContext | None = None,
 ) -> circuits.Circuit:
     """Add local depolarizing noise after two-qubit gates in a specified circuit. More specifically,
     with probability p, append a random non-identity two-qubit Pauli operator after each specified
@@ -44,6 +46,7 @@ def add_depolarizing_noise_to_two_qubit_gates(
         p: The probability with which to add noise.
         target_gate: Add depolarizing nose after this type of gate
         rng: The pseudorandom number generator to use.
+        context: Not used; to satisfy transformer API.
 
     Returns:
         The transformed circuit.

From 6409766efb423b8f625cc884e24f6145dceaad80 Mon Sep 17 00:00:00 2001
From: eliottrosenberg <61400172+eliottrosenberg@users.noreply.github.com>
Date: Thu, 11 Jul 2024 12:50:07 -0400
Subject: [PATCH 06/16] * after `circuit` input

Co-authored-by: Noureldin <noureldinyosri@google.com>
---
 cirq-core/cirq/transformers/noise_adding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index e8c52aa8aa6..fda698491dd 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -32,7 +32,7 @@ def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
 @transformer_api.transformer
 def add_depolarizing_noise_to_two_qubit_gates(
     circuit: circuits.Circuit,
-    p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
+    *, p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
     target_gate: ops.Gate = ops.CZ,
     rng: np.random.Generator | None = None,
     context: transformer_api.TransformerContext | None = None,

From bccffec1d6fd364e012f03b1ae20b434a368cdd6 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 09:51:38 -0700
Subject: [PATCH 07/16] format

---
 cirq-core/cirq/transformers/noise_adding.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index fda698491dd..131cec16d7e 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -32,7 +32,8 @@ def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
 @transformer_api.transformer
 def add_depolarizing_noise_to_two_qubit_gates(
     circuit: circuits.Circuit,
-    *, p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
+    *,
+    p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
     target_gate: ops.Gate = ops.CZ,
     rng: np.random.Generator | None = None,
     context: transformer_api.TransformerContext | None = None,

From 4078f86dc45bbe7d56f4fce04ea75673cefb175b Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 09:55:47 -0700
Subject: [PATCH 08/16] update test

---
 cirq-core/cirq/transformers/noise_adding_test.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
index b07a28d919f..fcaa7e383ba 100644
--- a/cirq-core/cirq/transformers/noise_adding_test.py
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -19,13 +19,13 @@
 def test_noise_adding():
     qubits = devices.LineQubit.range(4)
     circuit = circuits.Circuit(ops.CZ(*qubits[:2]), ops.CZ(*qubits[2:])) * 10
-    transformed_circuit_p0 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.0)
+    transformed_circuit_p0 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=0.0)
     assert transformed_circuit_p0 == circuit
-    transformed_circuit_p1 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 1.0)
+    transformed_circuit_p1 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=1.0)
     assert len(transformed_circuit_p1) == 20
-    transformed_circuit_p0_03 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, 0.03)
+    transformed_circuit_p0_03 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=0.03)
     assert 10 <= len(transformed_circuit_p0_03) <= 20
     transformed_circuit_p_dict = na.add_depolarizing_noise_to_two_qubit_gates(
-        circuit, {tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
+        circuit, p={tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
     )
     assert len(transformed_circuit_p_dict) == 20

From 1119ad941e166d49454f40c5abe168e2e4c2d5ff Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 10:45:11 -0700
Subject: [PATCH 09/16] Convert to a class

---
 cirq-core/cirq/transformers/noise_adding.py   | 123 ++++++++++--------
 .../cirq/transformers/noise_adding_test.py    |  12 +-
 2 files changed, 77 insertions(+), 58 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 131cec16d7e..8ddccd352da 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -30,65 +30,84 @@ def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
 
 
 @transformer_api.transformer
-def add_depolarizing_noise_to_two_qubit_gates(
-    circuit: circuits.Circuit,
-    *,
-    p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
-    target_gate: ops.Gate = ops.CZ,
-    rng: np.random.Generator | None = None,
-    context: transformer_api.TransformerContext | None = None,
-) -> circuits.Circuit:
+class DepolerizingNoiseTransformer:
     """Add local depolarizing noise after two-qubit gates in a specified circuit. More specifically,
     with probability p, append a random non-identity two-qubit Pauli operator after each specified
     two-qubit gate.
 
-    Args:
-        circuit: The circuit to add noise to.
+    Attrs:
         p: The probability with which to add noise.
         target_gate: Add depolarizing nose after this type of gate
         rng: The pseudorandom number generator to use.
-        context: Not used; to satisfy transformer API.
+    """
 
-    Returns:
-        The transformed circuit.
+    def __init__(
+        self,
+        p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
+        target_gate: ops.Gate = ops.CZ,
+        rng: np.random.Generator | None = None,
+    ):
+        if rng is None:
+            rng = np.random.default_rng()
+        self.p = p
+        self.target_gate = target_gate
+        self.rng = rng
 
-    Raises:
-        TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to floats.
-    """
-    if rng is None:
-        rng = np.random.default_rng()
+    def __call__(
+        self,
+        circuit: circuits.Circuit,
+        *,
+        context: transformer_api.TransformerContext | None = None,
+    ):
+        """
+        Apply the transformer to the given circuit.
+
+        Args:
+            circuit: The circuit to add noise to.
+            context: Not used; to satisfy transformer API.
+
+        Returns:
+            The transformed circuit.
+
+        Raises:
+            TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to floats.
+        """
+
+        p = self.p
+        rng = self.rng
+        target_gate = self.target_gate
 
-    # add random Pauli gates with probability p after each of the specified gate
-    assert target_gate.num_qubits() == 2, "`target_gate` must be a two-qubit gate."
-    paulis = [ops.I, ops.X, ops.Y, ops.Z]
-    new_moments = []
-    for moment in circuit:
-        new_moments.append(moment)
-        if _gate_in_moment(target_gate, moment):
-            # add a new moment with the Paulis
-            target_pairs = {
-                tuple(sorted(op.qubits)) for op in moment.operations if op.gate == target_gate
-            }
-            added_moment_ops = []
-            for pair in target_pairs:
-                if isinstance(p, float):
-                    p_i = p
-                elif isinstance(p, Mapping):
-                    pair_sorted_tuple = (pair[0], pair[1])
-                    p_i = p[pair_sorted_tuple]
-                else:
-                    raise TypeError(
-                        "p must either be a float or a mapping from sorted qubit pairs to floats"
-                    )
-                apply = rng.choice([True, False], p=[p_i, 1 - p_i])
-                if apply:
-                    choices = [
-                        (pauli_a(pair[0]), pauli_b(pair[1]))
-                        for pauli_a in paulis
-                        for pauli_b in paulis
-                    ][1:]
-                    pauli_to_apply = rng.choice(np.array(choices, dtype=object))
-                    added_moment_ops.append(pauli_to_apply)
-            if len(added_moment_ops) > 0:
-                new_moments.append(circuits.Moment(*added_moment_ops))
-    return circuits.Circuit.from_moments(*new_moments)
+        # add random Pauli gates with probability p after each of the specified gate
+        assert target_gate.num_qubits() == 2, "`target_gate` must be a two-qubit gate."
+        paulis = [ops.I, ops.X, ops.Y, ops.Z]
+        new_moments = []
+        for moment in circuit:
+            new_moments.append(moment)
+            if _gate_in_moment(target_gate, moment):
+                # add a new moment with the Paulis
+                target_pairs = {
+                    tuple(sorted(op.qubits)) for op in moment.operations if op.gate == target_gate
+                }
+                added_moment_ops = []
+                for pair in target_pairs:
+                    if isinstance(p, float):
+                        p_i = p
+                    elif isinstance(p, Mapping):
+                        pair_sorted_tuple = (pair[0], pair[1])
+                        p_i = p[pair_sorted_tuple]
+                    else:  # pragma: no cover
+                        raise TypeError(  # pragma: no cover
+                            "p must either be a float or a mapping from sorted qubit pairs to floats"  # pragma: no cover
+                        )  # pragma: no cover
+                    apply = rng.choice([True, False], p=[p_i, 1 - p_i])
+                    if apply:
+                        choices = [
+                            (pauli_a(pair[0]), pauli_b(pair[1]))
+                            for pauli_a in paulis
+                            for pauli_b in paulis
+                        ][1:]
+                        pauli_to_apply = rng.choice(np.array(choices, dtype=object))
+                        added_moment_ops.append(pauli_to_apply)
+                if len(added_moment_ops) > 0:
+                    new_moments.append(circuits.Moment(*added_moment_ops))
+        return circuits.Circuit.from_moments(*new_moments)
diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
index fcaa7e383ba..2277a12f2b0 100644
--- a/cirq-core/cirq/transformers/noise_adding_test.py
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -19,13 +19,13 @@
 def test_noise_adding():
     qubits = devices.LineQubit.range(4)
     circuit = circuits.Circuit(ops.CZ(*qubits[:2]), ops.CZ(*qubits[2:])) * 10
-    transformed_circuit_p0 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=0.0)
+    transformed_circuit_p0 = na.DepolerizingNoiseTransformer(0.0)(circuit)
     assert transformed_circuit_p0 == circuit
-    transformed_circuit_p1 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=1.0)
+    transformed_circuit_p1 = na.DepolerizingNoiseTransformer(1.0)(circuit)
     assert len(transformed_circuit_p1) == 20
-    transformed_circuit_p0_03 = na.add_depolarizing_noise_to_two_qubit_gates(circuit, p=0.03)
+    transformed_circuit_p0_03 = na.DepolerizingNoiseTransformer(0.03)(circuit)
     assert 10 <= len(transformed_circuit_p0_03) <= 20
-    transformed_circuit_p_dict = na.add_depolarizing_noise_to_two_qubit_gates(
-        circuit, p={tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
-    )
+    transformed_circuit_p_dict = na.DepolerizingNoiseTransformer(
+        {tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
+    )(circuit)
     assert len(transformed_circuit_p_dict) == 20

From 65a5e16779e3e4ef393dbe552817f715de53dfa1 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 12:49:05 -0700
Subject: [PATCH 10/16] types and lint

---
 cirq-core/cirq/transformers/noise_adding.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 8ddccd352da..98bbd4a624b 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -55,12 +55,11 @@ def __init__(
 
     def __call__(
         self,
-        circuit: circuits.Circuit,
+        circuit: circuits.AbstractCircuit,
         *,
         context: transformer_api.TransformerContext | None = None,
     ):
-        """
-        Apply the transformer to the given circuit.
+        """Apply the transformer to the given circuit.
 
         Args:
             circuit: The circuit to add noise to.
@@ -70,7 +69,8 @@ def __call__(
             The transformed circuit.
 
         Raises:
-            TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to floats.
+            TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to
+                       floats.
         """
 
         p = self.p
@@ -97,7 +97,8 @@ def __call__(
                         p_i = p[pair_sorted_tuple]
                     else:  # pragma: no cover
                         raise TypeError(  # pragma: no cover
-                            "p must either be a float or a mapping from sorted qubit pairs to floats"  # pragma: no cover
+                            "p must either be a float or a mapping from"  # pragma: no cover
+                            + "sorted qubit pairs to floats"  # pragma: no cover
                         )  # pragma: no cover
                     apply = rng.choice([True, False], p=[p_i, 1 - p_i])
                     if apply:

From 9472dfd6171b8f8dc6b65f8db60e3a334f0a5dd7 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 16:31:47 -0700
Subject: [PATCH 11/16] suggestions from Nour

---
 cirq-core/cirq/transformers/noise_adding.py   | 48 +++++++++----------
 .../cirq/transformers/noise_adding_test.py    | 30 ++++++++++--
 2 files changed, 49 insertions(+), 29 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 98bbd4a624b..17953538351 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -21,12 +21,7 @@
 
 def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
     """Check whether `gate` is in `moment`."""
-    target_gate_in_moment = False
-    for op in moment.operations:
-        if op.gate == gate:
-            target_gate_in_moment = True
-            break
-    return target_gate_in_moment
+    return any(op.gate == gate for op in moment)
 
 
 @transformer_api.transformer
@@ -38,24 +33,36 @@ class DepolerizingNoiseTransformer:
     Attrs:
         p: The probability with which to add noise.
         target_gate: Add depolarizing nose after this type of gate
-        rng: The pseudorandom number generator to use.
     """
 
     def __init__(
-        self,
-        p: float | Mapping[tuple[ops.Qid, ops.Qid], float],
-        target_gate: ops.Gate = ops.CZ,
-        rng: np.random.Generator | None = None,
+        self, p: float | Mapping[tuple[ops.Qid, ops.Qid], float], target_gate: ops.Gate = ops.CZ
     ):
-        if rng is None:
-            rng = np.random.default_rng()
+        """Initialize the depolarizing noise transformer with some depolarizing probability and
+        target gate.
+
+        Args:
+            p: The depolarizing probability, either a single float or a mapping from pairs of qubits
+               to floats.
+           target_gate: The gate after which to add depolarizing noise.
+
+        Raises:
+            TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to
+                       floats.
+        """
+
+        if not (isinstance(p, float) or isinstance(p, Mapping)):
+            raise TypeError(  # pragma: no cover
+                "p must either be a float or a mapping from"  # pragma: no cover
+                + "sorted qubit pairs to floats"  # pragma: no cover
+            )  # pragma: no cover
         self.p = p
         self.target_gate = target_gate
-        self.rng = rng
 
     def __call__(
         self,
         circuit: circuits.AbstractCircuit,
+        rng: np.random.Generator | None = None,
         *,
         context: transformer_api.TransformerContext | None = None,
     ):
@@ -67,14 +74,10 @@ def __call__(
 
         Returns:
             The transformed circuit.
-
-        Raises:
-            TypeError: If `p` is not either be a float or a mapping from sorted qubit pairs to
-                       floats.
         """
-
+        if rng is None:
+            rng = np.random.default_rng()
         p = self.p
-        rng = self.rng
         target_gate = self.target_gate
 
         # add random Pauli gates with probability p after each of the specified gate
@@ -95,11 +98,6 @@ def __call__(
                     elif isinstance(p, Mapping):
                         pair_sorted_tuple = (pair[0], pair[1])
                         p_i = p[pair_sorted_tuple]
-                    else:  # pragma: no cover
-                        raise TypeError(  # pragma: no cover
-                            "p must either be a float or a mapping from"  # pragma: no cover
-                            + "sorted qubit pairs to floats"  # pragma: no cover
-                        )  # pragma: no cover
                     apply = rng.choice([True, False], p=[p_i, 1 - p_i])
                     if apply:
                         choices = [
diff --git a/cirq-core/cirq/transformers/noise_adding_test.py b/cirq-core/cirq/transformers/noise_adding_test.py
index 2277a12f2b0..ff82ccbaedd 100644
--- a/cirq-core/cirq/transformers/noise_adding_test.py
+++ b/cirq-core/cirq/transformers/noise_adding_test.py
@@ -14,18 +14,40 @@
 
 from cirq import ops, circuits, devices
 import cirq.transformers.noise_adding as na
+import numpy as np
 
 
 def test_noise_adding():
     qubits = devices.LineQubit.range(4)
-    circuit = circuits.Circuit(ops.CZ(*qubits[:2]), ops.CZ(*qubits[2:])) * 10
+    one_layer = circuits.Circuit(ops.CZ(*qubits[:2]), ops.CZ(*qubits[2:]))
+    circuit = one_layer * 10
+
+    # test that p=0 does nothing
     transformed_circuit_p0 = na.DepolerizingNoiseTransformer(0.0)(circuit)
     assert transformed_circuit_p0 == circuit
+
+    # test that p=1 doubles the circuit depth
     transformed_circuit_p1 = na.DepolerizingNoiseTransformer(1.0)(circuit)
     assert len(transformed_circuit_p1) == 20
-    transformed_circuit_p0_03 = na.DepolerizingNoiseTransformer(0.03)(circuit)
-    assert 10 <= len(transformed_circuit_p0_03) <= 20
+
+    # test that we get a deterministic result when using a specific rng
+    rng = np.random.default_rng(0)
+    transformed_circuit_p0_03 = na.DepolerizingNoiseTransformer(0.03)(circuit, rng=rng)
+    expected_circuit = (
+        one_layer * 2
+        + circuits.Circuit(ops.I(qubits[2]), ops.Z(qubits[3]))
+        + one_layer * 4
+        + circuits.Circuit(ops.Z(qubits[0]), ops.X(qubits[1]))
+        + one_layer * 4
+        + circuits.Circuit(ops.I(qubits[2]), ops.X(qubits[3]))
+    )
+    assert transformed_circuit_p0_03 == expected_circuit
+
+    # test that supplying a dictionary for p works
     transformed_circuit_p_dict = na.DepolerizingNoiseTransformer(
         {tuple(qubits[:2]): 1.0, tuple(qubits[2:]): 0.0}
     )(circuit)
-    assert len(transformed_circuit_p_dict) == 20
+    assert len(transformed_circuit_p_dict) == 20  # depth should be doubled
+    assert transformed_circuit_p_dict[1::2].all_qubits() == frozenset(
+        qubits[:2]
+    )  # no single-qubit gates get added to qubits[2:]

From 174553d1d644706e9672ab0a6d2278a40583b899 Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 16:38:26 -0700
Subject: [PATCH 12/16] lint

---
 cirq-core/cirq/transformers/noise_adding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 17953538351..1e5e68c1d7c 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -51,7 +51,7 @@ def __init__(
                        floats.
         """
 
-        if not (isinstance(p, float) or isinstance(p, Mapping)):
+        if not isinstance(p, (Mapping, float)):
             raise TypeError(  # pragma: no cover
                 "p must either be a float or a mapping from"  # pragma: no cover
                 + "sorted qubit pairs to floats"  # pragma: no cover

From 8555599abba7fab80cb57922c84f289ef60c0e3b Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 16:56:11 -0700
Subject: [PATCH 13/16] suggestion from Nour

---
 cirq-core/cirq/transformers/noise_adding.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 1e5e68c1d7c..d80e4669a9f 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -57,6 +57,7 @@ def __init__(
                 + "sorted qubit pairs to floats"  # pragma: no cover
             )  # pragma: no cover
         self.p = p
+        self.p_func = lambda _: p if isinstance(p, (int, float)) else lambda pair: p.get(pair, 0.0)
         self.target_gate = target_gate
 
     def __call__(
@@ -77,7 +78,6 @@ def __call__(
         """
         if rng is None:
             rng = np.random.default_rng()
-        p = self.p
         target_gate = self.target_gate
 
         # add random Pauli gates with probability p after each of the specified gate
@@ -93,11 +93,8 @@ def __call__(
                 }
                 added_moment_ops = []
                 for pair in target_pairs:
-                    if isinstance(p, float):
-                        p_i = p
-                    elif isinstance(p, Mapping):
-                        pair_sorted_tuple = (pair[0], pair[1])
-                        p_i = p[pair_sorted_tuple]
+                    pair_sorted_tuple = (pair[0], pair[1])
+                    p_i = self.p_func(pair_sorted_tuple)
                     apply = rng.choice([True, False], p=[p_i, 1 - p_i])
                     if apply:
                         choices = [

From 71c10dd0b9f0e557fa4bc219df3c2206eb3f8bfd Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 17:06:20 -0700
Subject: [PATCH 14/16] add ()

---
 cirq-core/cirq/transformers/noise_adding.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index d80e4669a9f..ccc2c65be3b 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -57,7 +57,9 @@ def __init__(
                 + "sorted qubit pairs to floats"  # pragma: no cover
             )  # pragma: no cover
         self.p = p
-        self.p_func = lambda _: p if isinstance(p, (int, float)) else lambda pair: p.get(pair, 0.0)
+        self.p_func = (
+            (lambda _: p) if isinstance(p, (int, float)) else (lambda pair: p.get(pair, 0.0))
+        )
         self.target_gate = target_gate
 
     def __call__(

From a76f38205434f0ba3db9da8db20b54c416ebdb4a Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 17:10:02 -0700
Subject: [PATCH 15/16] types

---
 cirq-core/cirq/transformers/noise_adding.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index ccc2c65be3b..3eb5e7c28e8 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -17,6 +17,7 @@
 from cirq import ops, circuits
 from cirq.transformers import transformer_api
 import numpy as np
+from typing import cast
 
 
 def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool:
@@ -58,7 +59,9 @@ def __init__(
             )  # pragma: no cover
         self.p = p
         self.p_func = (
-            (lambda _: p) if isinstance(p, (int, float)) else (lambda pair: p.get(pair, 0.0))
+            (lambda _: p)
+            if isinstance(p, (int, float))
+            else (lambda pair: cast(Mapping, p).get(pair, 0.0))
         )
         self.target_gate = target_gate
 

From 6a34db2aee4ac793db410fbfc9f6d7242700c91a Mon Sep 17 00:00:00 2001
From: Eliott Rosenberg <eliottrosenbrg@google.com>
Date: Thu, 11 Jul 2024 17:17:35 -0700
Subject: [PATCH 16/16] lint

---
 cirq-core/cirq/transformers/noise_adding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cirq-core/cirq/transformers/noise_adding.py b/cirq-core/cirq/transformers/noise_adding.py
index 3eb5e7c28e8..28f0b3350ad 100644
--- a/cirq-core/cirq/transformers/noise_adding.py
+++ b/cirq-core/cirq/transformers/noise_adding.py
@@ -14,10 +14,10 @@
 
 from collections.abc import Mapping
 
+from typing import cast
 from cirq import ops, circuits
 from cirq.transformers import transformer_api
 import numpy as np
-from typing import cast
 
 
 def _gate_in_moment(gate: ops.Gate, moment: circuits.Moment) -> bool: