From f3ddfef5ad8bb6620aa76ca51886098112596b6e Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Mon, 31 Jan 2022 14:30:36 -0500
Subject: [PATCH 1/7] feature: adding TwoQubitPauliChannel

---
 src/braket/circuits/noise.py                  |  84 ++++++++++-
 src/braket/circuits/noises.py                 | 137 +++++++++++++++++-
 test/unit_tests/braket/circuits/test_noise.py |  35 +++++
 .../unit_tests/braket/circuits/test_noises.py |  51 +++++++
 4 files changed, 296 insertions(+), 11 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index e5c0dab4b..68d4ae3e7 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import Any, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence
 
 from braket.circuits.quantum_operator import QuantumOperator
 from braket.circuits.qubit_set import QubitSet
@@ -121,7 +121,7 @@ def __init__(
     def probability(self) -> float:
         """
         Returns:
-            probability (float): The probability that parameterizes the noise channel.
+            probability (float): The probability that parametrizes the noise channel.
         """
         return self._probability
 
@@ -163,7 +163,7 @@ def __init__(
     def probability(self) -> float:
         """
         Returns:
-            probability (float): The probability that parameterizes the noise channel.
+            probability (float): The probability that parametrizes the noise channel.
         """
         return self._probability
 
@@ -205,7 +205,7 @@ def __init__(
     def probability(self) -> float:
         """
         Returns:
-            probability (float): The probability that parameterizes the noise channel.
+            probability (float): The probability that parametrizes the noise channel.
         """
         return self._probability
 
@@ -213,10 +213,76 @@ def __repr__(self):
         return f"{self.name}('probability': {self.probability}, 'qubit_count': {self.qubit_count})"
 
 
-class PauliNoise(Noise):
+class MultiQubitPauliNoise(Noise):
     """
-    Class `PauliNoise` represents the general Pauli noise channel on N qubits
-    parameterized by three probabilities.
+    Class `MultiQubitPauliNoise` represents a general multi-qubit Pauli channel,
+    parameterized by up to 4**N - 1 probabilities.
+    """
+
+    def __init__(
+        self,
+        probabilities: Dict[str, float],
+        qubit_count: Optional[int],
+        ascii_symbols: Sequence[str],
+    ):
+        """[summary]
+
+        Args:
+            probabilities (Dict[str, float]): A dictionary with Pauli string as the keys,
+            and the probabilities as values, i.e. {"XX": 0.1. "IZ": 0.2}.
+            qubit_count (Optional[int]): The number of qubits the Pauli noise acts on.
+            ascii_symbols (Sequence[str]): ASCII string symbols for the noise. These are used when
+                printing a diagram of a circuit. The length must be the same as `qubit_count`, and
+                index ordering is expected to correlate with the target ordering on the instruction.
+
+        Raises:
+            ValueError: If the `qubit_count` is less than 1, `ascii_symbols` are `None`, or
+                `ascii_symbols` length != `qubit_count`. Also if `probabilities` are not `float`s,
+                any `probabilities` > 1, or `probabilities` < 0, or if the sum of all
+                probabilities is > 1,
+                or if "II" is specified as a Pauli string.
+                Also if any Pauli string contains invalid strings.
+                Also if the length of probabilities is greater than 4**qubit_count.
+            TypeError: If the type of the dictionary keys are not strings.
+                If the probabilities are not floats.
+        """
+
+        super().__init__(qubit_count=qubit_count, ascii_symbols=ascii_symbols)
+        self.probabilities = probabilities
+
+        allowed_substrings = {"I", "X", "Y", "Z"}
+
+        if "I" * self.qubit_count in probabilities:
+            raise ValueError("Leakage error is not supported.")
+
+        if sum(probabilities.values()) > 1.0 or sum(probabilities.values()) < 0.0:
+            raise ValueError("Total probability must be a real number in the interval [0,1]")
+
+        if len(probabilities) >= 4 ** self.qubit_count:
+            raise ValueError(
+                "Too many probabilities provided. Only 4 ** self.qubit_count - 1 are allowed."
+            )
+
+        for pauli_string, prob in probabilities.items():
+            if not isinstance(pauli_string, str):
+                raise TypeError("Keys must be a string type")
+            if len(pauli_string) != self.qubit_count:
+                raise ValueError("Length of each Pauli strings must be equal to number of qubits.")
+            if not isinstance(prob, float):
+                raise TypeError("Keys must be a float type")
+            if not set(pauli_string) <= allowed_substrings:
+                raise ValueError("Strings must be Pauli strings consisting of only [I, X, Y, Z]")
+            if prob < 0.0 or prob > 1.0:
+                raise ValueError("Individual values must be a real number in the interval [0,1]")
+
+    def __repr__(self):
+        return f"{self.name}('probabilities' : {self.probabilities}, 'qubit_count': {self.qubit_count})"  # noqa
+
+
+class PauliNoise(MultiQubitPauliNoise):
+    """
+    Class `PauliNoise` represents the a single-qubit Pauli noise channel
+    acting on one qubit. It is parameterized by three probabilities.
     """
 
     def __init__(
@@ -242,7 +308,9 @@ def __init__(
                 is not `float`, `probX` or `probY` or `probZ` > 1.0, or
                 `probX` or `probY` or `probZ` < 0.0, or `probX`+`probY`+`probZ` > 1
         """
-        super().__init__(qubit_count=qubit_count, ascii_symbols=ascii_symbols)
+        super().__init__(
+            {"X": probX, "Y": probY, "Z": probZ}, qubit_count=1, ascii_symbols=ascii_symbols
+        )
 
         if not isinstance(probX, float):
             raise TypeError("probX must be float type")
diff --git a/src/braket/circuits/noises.py b/src/braket/circuits/noises.py
index 1f98e0fde..3da5dcef8 100644
--- a/src/braket/circuits/noises.py
+++ b/src/braket/circuits/noises.py
@@ -11,7 +11,8 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import Iterable
+import itertools
+from typing import Dict, Iterable
 
 import numpy as np
 
@@ -21,6 +22,7 @@
 from braket.circuits.noise import (
     DampingNoise,
     GeneralizedAmplitudeDampingNoise,
+    MultiQubitPauliNoise,
     Noise,
     PauliNoise,
     SingleProbabilisticNoise,
@@ -789,7 +791,7 @@ class Kraus(Noise):
 
     Args:
         matrices (Iterable[np.array]): A list of matrices that define a noise
-            channel. These matrices need to satisify the requirement of CPTP map.
+            channel. These matrices need to satisfy the requirement of CPTP map.
         display_name (str): Name to be used for an instance of this general noise
             channel for circuit diagrams. Defaults to `KR`.
 
@@ -797,7 +799,7 @@ class Kraus(Noise):
         ValueError: If any matrix in `matrices` is not a two-dimensional square
             matrix,
             or has a dimension length which is not a positive exponent of 2,
-            or the `matrices` do not satisify CPTP condition.
+            or the `matrices` do not satisfy CPTP condition.
     """
 
     def __init__(self, matrices: Iterable[np.ndarray], display_name: str = "KR"):
@@ -869,3 +871,132 @@ def kraus(
 
 
 Noise.register_noise(Kraus)
+
+
+class TwoQubitPauliChannel(MultiQubitPauliNoise):
+    """Two-Qubit Pauli noise channel which transforms a
+        density matrix :math:`\\rho` according to:
+
+    .. math::
+        \\rho \\Rightarrow (1-p) \\rho +
+            p_{IX} IX \\rho IX^{\\dagger} +
+            p_{IY} IY \\rho IY^{\\dagger} +
+            p_{IZ} IZ \\rho IZ^{\\dagger} +
+            p_{XI} XI \\rho XI^{\\dagger} +
+            p_{XX} XX \\rho XX^{\\dagger} +
+            p_{XY} XY \\rho XY^{\\dagger} +
+            p_{XZ} XZ \\rho XZ^{\\dagger} +
+            p_{YI} YI \\rho YI^{\\dagger} +
+            p_{YX} YX \\rho YX^{\\dagger} +
+            p_{YY} YY \\rho YY^{\\dagger} +
+            p_{YZ} YZ \\rho YZ^{\\dagger} +
+            p_{ZI} ZI \\rho ZI^{\\dagger} +
+            p_{ZX} ZX \\rho ZX^{\\dagger} +
+            p_{ZY} ZY \\rho ZY^{\\dagger} +
+            p_{ZZ} ZZ \\rho ZZ^{\\dagger})
+    where
+
+    .. math::
+        I = \\left(
+                \\begin{matrix}
+                    1 & 0 \\\\
+                    0 & 1
+                \\end{matrix}
+            \\right)
+
+        X = \\left(
+                \\begin{matrix}
+                    0 & 1 \\\\
+                    1 & 0
+                \\end{matrix}
+            \\right)
+
+        Y = \\left(
+                \\begin{matrix}
+                    0 & -i \\\\
+                    i &  0
+                \\end{matrix}
+            \\right)
+
+        Z = \\left(
+                \\begin{matrix}
+                    1 & 0 \\\\
+                    0 & -1
+                \\end{matrix}
+            \\right)
+
+        p = \\text{sum of all probabilities}
+
+    This noise channel is shown as `PC_2({"pauli_string": probability})` in circuit diagrams.
+    """
+
+    _paulis = {
+        "I": np.array([[1.0, 0.0], [0.0, 1.0]], dtype=complex),
+        "X": np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex),
+        "Y": np.array([[0.0, -1.0j], [1.0j, 0.0]], dtype=complex),
+        "Z": np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex),
+    }
+    _tensor_products_strings = itertools.product(_paulis.keys(), repeat=2)
+    _names_list = ["".join(x) for x in _tensor_products_strings]
+
+    def __init__(self, probabilities: Dict[str, float]):
+        super().__init__(
+            probabilities=probabilities,
+            qubit_count=None,
+            ascii_symbols=[
+                f"PC2({probabilities})",
+                f"PC2({probabilities})",
+            ],
+        )
+
+        total_prob = sum(self.probabilities.values())
+
+        K_list = [np.sqrt(1 - total_prob) * np.identity(4)]  # "II" element
+        for pstring in self._names_list[1:]:  # ignore "II"
+            if pstring in self.probabilities:
+                mat = np.sqrt(self.probabilities[pstring]) * np.kron(
+                    self._paulis[pstring[0]], self._paulis[pstring[1]]
+                )
+                K_list.append(mat)
+            else:
+                K_list.append(np.zeros((4, 4)))
+        self._matrix = K_list
+
+    def to_ir(self, target: QubitSet):
+        return ir.MultiQubitPauliChannel.construct(
+            targets=[target[0], target[1]], probabilities=self.probabilities
+        )
+
+    def to_matrix(self) -> Iterable[np.ndarray]:
+        return self._matrix
+
+    @staticmethod
+    def fixed_qubit_count() -> int:
+        return 2
+
+    @staticmethod
+    @circuit.subroutine(register=True)
+    def two_qubit_pauli_channel(
+        target1: QubitInput, target2: QubitInput, probabilities: Dict[str, float]
+    ) -> Iterable[Instruction]:
+        """Registers this function into the circuit class.
+
+        Args:
+            target (Qubit, int, or iterable of Qubit / int): Target qubits
+            probability (float): Probability of two-qubit Pauli channel.
+
+        Returns:
+            Iterable[Instruction]: `Iterable` of Depolarizing instructions.
+
+        Examples:
+            >>> circ = Circuit().two_qubit_pauli_channel(0, 1, {"XX": 0.1})
+        """
+        return [
+            Instruction(
+                Noise.TwoQubitPauliChannel(probabilities=probabilities),
+                target=[target1, target2],
+            )
+        ]
+
+
+Noise.register_noise(TwoQubitPauliChannel)
diff --git a/test/unit_tests/braket/circuits/test_noise.py b/test/unit_tests/braket/circuits/test_noise.py
index edab29ac5..dcb5576e8 100644
--- a/test/unit_tests/braket/circuits/test_noise.py
+++ b/test/unit_tests/braket/circuits/test_noise.py
@@ -17,6 +17,7 @@
 from braket.circuits.noise import (
     DampingNoise,
     GeneralizedAmplitudeDampingNoise,
+    MultiQubitPauliNoise,
     Noise,
     PauliNoise,
     SingleProbabilisticNoise,
@@ -348,3 +349,37 @@ def __init__(self):
 
     Noise.register_noise(_FooNoise)
     assert Noise._FooNoise().name == _FooNoise().name
+
+
+@pytest.mark.parametrize("probs", [{"X": 0.1}, {"XX": 0.1, "YY": 0.1}])
+def test_valid_data_pauli_nqubit(probs):
+
+    qubit_count = len(list(probs.keys())[0])
+    ascii_symbols = [["foo"] for _ in range(qubit_count)]
+    MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
+
+
+@pytest.mark.xfail(raises=ValueError)
+@pytest.mark.parametrize(
+    "probs",
+    [
+        {"X": -0.1},
+        {"XY": 1.1},
+        {"TX": 0.1},
+        {"X": 0.5, "Y": 0.6},
+        {"X": 0.1, "YY": 0.2},
+        {"II": 0.9, "XX": 0.1},
+    ],
+)
+def test_invalid_values_pauli_nqubit(probs):
+    qubit_count = len(list(probs.keys())[0])
+    ascii_symbols = [["foo"] for _ in range(qubit_count)]
+    MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
+
+
+@pytest.mark.xfail(raises=TypeError)
+@pytest.mark.parametrize("probs", [{0.1: 0.1}, {"X": "F"}])
+def test_invalid_types_pauli_nqubit(probs):
+    qubit_count = len(list(probs.keys())[0])
+    ascii_symbols = [["foo"] for _ in range(qubit_count)]
+    MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
diff --git a/test/unit_tests/braket/circuits/test_noises.py b/test/unit_tests/braket/circuits/test_noises.py
index 0de9f0e90..cb22068b9 100644
--- a/test/unit_tests/braket/circuits/test_noises.py
+++ b/test/unit_tests/braket/circuits/test_noises.py
@@ -21,6 +21,7 @@
     DampingSingleProbability,
     DoubleControl,
     DoubleTarget,
+    MultiProbability,
     MultiTarget,
     SingleControl,
     SingleProbability,
@@ -70,6 +71,13 @@
         [DoubleTarget, SingleProbability_34],
         {},
     ),
+    (
+        Noise.TwoQubitPauliChannel,
+        "multi_qubit_pauli_channel",
+        ir.MultiQubitPauliChannel,
+        [DoubleTarget, MultiProbability],
+        {},
+    ),
     (
         Noise.PauliChannel,
         "pauli_channel",
@@ -178,6 +186,14 @@ def two_dimensional_matrix_list_valid_input(**kwargs):
     }
 
 
+def multi_probability_valid_input(**kwargs):
+    return {"probabilities": {"XX": 0.1}}
+
+
+def multi_probability_invalid_input(**kwargs):
+    return {"probabilities": {"XX": 1.1}}
+
+
 valid_ir_switcher = {
     "SingleTarget": single_target_valid_input,
     "DoubleTarget": double_target_valid_ir_input,
@@ -187,6 +203,7 @@ def two_dimensional_matrix_list_valid_input(**kwargs):
     "DampingProbability": damping_probability_valid_input,
     "DampingSingleProbability": damping_single_probability_valid_input,
     "TripleProbability": triple_probability_valid_input,
+    "MultiProbability": multi_probability_valid_input,
     "SingleControl": single_control_valid_input,
     "DoubleControl": double_control_valid_ir_input,
     "MultiTarget": multi_target_valid_input,
@@ -245,6 +262,7 @@ def create_valid_target_input(irsubclasses):
                 DampingProbability,
                 TripleProbability,
                 TwoDimensionalMatrixList,
+                MultiProbability,
             ]
         ):
             pass
@@ -268,6 +286,8 @@ def create_valid_noise_class_input(irsubclasses, **kwargs):
         input.update(damping_probability_valid_input())
     if TripleProbability in irsubclasses:
         input.update(triple_probability_valid_input())
+    if MultiProbability in irsubclasses:
+        input.update(multi_probability_valid_input())
     if TwoDimensionalMatrixList in irsubclasses:
         input.update(two_dimensional_matrix_list_valid_input(**kwargs))
     return input
@@ -301,6 +321,7 @@ def calculate_qubit_count(irsubclasses):
                 DampingSingleProbability,
                 DampingProbability,
                 TripleProbability,
+                MultiProbability,
                 TwoDimensionalMatrixList,
             ]
         ):
@@ -358,6 +379,8 @@ def test_noise_subroutine(testclass, subroutine_name, irclass, irsubclasses, kwa
             subroutine_input.update(damping_probability_valid_input())
         if TripleProbability in irsubclasses:
             subroutine_input.update(triple_probability_valid_input())
+        if MultiProbability in irsubclasses:
+            subroutine_input.update(multi_probability_valid_input())
 
         circuit1 = subroutine(**subroutine_input)
         circuit2 = Circuit(instruction_list)
@@ -394,3 +417,31 @@ def test_kraus_matrix_target_size_mismatch():
     Circuit().kraus(
         matrices=[np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]])], targets=[0]
     )
+
+
+@pytest.mark.xfail(raises=ValueError)
+@pytest.mark.parametrize(
+    "probs",
+    [
+        {"X": -0.1},
+        {"XY": 1.1},
+        {"TX": 0.1},
+        {"X": 0.5, "Y": 0.6},
+        {"X": 0.1, "YY": 0.2},
+        {"II": 0.9, "XX": 0.1},
+    ],
+)
+def test_invalid_values_pauli_channel_two_qubit(probs):
+    Noise.TwoQubitPauliChannel(probs)
+
+
+@pytest.mark.parametrize(
+    "probs",
+    [
+        {"XY": 0.1},
+        {"XX": 0.1, "ZZ": 0.2},
+    ],
+)
+def test_valid_values_pauli_channel_two_qubit(probs):
+    noise = Noise.TwoQubitPauliChannel(probs)
+    assert len(noise.to_matrix()) == 16

From fc52c47de3affb7b762d05dab1fe3e183e003b68 Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Mon, 31 Jan 2022 15:36:57 -0500
Subject: [PATCH 2/7] fix: tox formatting errors

---
 src/braket/circuits/noise.py                        | 2 +-
 src/braket/circuits/observables.py                  | 2 +-
 src/braket/circuits/quantum_operator_helpers.py     | 2 +-
 src/braket/circuits/unitary_calculation.py          | 4 ++--
 src/braket/tasks/gate_model_quantum_task_result.py  | 2 +-
 test/integ_tests/test_tensor_network_simulator.py   | 4 ++--
 test/unit_tests/braket/circuits/test_observables.py | 2 +-
 7 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index 68d4ae3e7..58a4f0531 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -258,7 +258,7 @@ def __init__(
         if sum(probabilities.values()) > 1.0 or sum(probabilities.values()) < 0.0:
             raise ValueError("Total probability must be a real number in the interval [0,1]")
 
-        if len(probabilities) >= 4 ** self.qubit_count:
+        if len(probabilities) >= 4**self.qubit_count:
             raise ValueError(
                 "Too many probabilities provided. Only 4 ** self.qubit_count - 1 are allowed."
             )
diff --git a/src/braket/circuits/observables.py b/src/braket/circuits/observables.py
index 0fc0c9e4a..62b523f00 100644
--- a/src/braket/circuits/observables.py
+++ b/src/braket/circuits/observables.py
@@ -230,7 +230,7 @@ def eigenvalues(self):
     def eigenvalue(self, index: int) -> float:
         if index in self._eigenvalue_indices:
             return self._eigenvalue_indices[index]
-        dimension = 2 ** self.qubit_count
+        dimension = 2**self.qubit_count
         if index >= dimension:
             raise ValueError(
                 f"Index {index} requested but observable has at most {dimension} eigenvalues"
diff --git a/src/braket/circuits/quantum_operator_helpers.py b/src/braket/circuits/quantum_operator_helpers.py
index 5a922100f..b61a3184b 100644
--- a/src/braket/circuits/quantum_operator_helpers.py
+++ b/src/braket/circuits/quantum_operator_helpers.py
@@ -35,7 +35,7 @@ def verify_quantum_operator_matrix_dimensions(matrix: np.array) -> None:
     matrix = np.array(matrix, dtype=complex)
     qubit_count = int(np.log2(matrix.shape[0]))
 
-    if 2 ** qubit_count != matrix.shape[0] or qubit_count < 1:
+    if 2**qubit_count != matrix.shape[0] or qubit_count < 1:
         raise ValueError(f"`matrix` dimension {matrix.shape[0]} is not a positive power of 2")
 
 
diff --git a/src/braket/circuits/unitary_calculation.py b/src/braket/circuits/unitary_calculation.py
index d532e17b2..18ddba6bc 100644
--- a/src/braket/circuits/unitary_calculation.py
+++ b/src/braket/circuits/unitary_calculation.py
@@ -60,7 +60,7 @@ def calculate_unitary(qubit_count: int, instructions: Iterable[Instruction]) ->
         TypeError: If `instructions` is not composed only of `Gate` instances,
             i.e. a circuit with `Noise` operators will raise this error.
     """
-    unitary = np.eye(2 ** qubit_count, dtype=complex)
+    unitary = np.eye(2**qubit_count, dtype=complex)
     un_tensor = np.reshape(unitary, qubit_count * [2, 2])
 
     for instr in instructions:
@@ -83,4 +83,4 @@ def calculate_unitary(qubit_count: int, instructions: Iterable[Instruction]) ->
             casting="no",
         )
 
-    return np.reshape(un_tensor, 2 * [2 ** qubit_count])
+    return np.reshape(un_tensor, 2 * [2**qubit_count])
diff --git a/src/braket/tasks/gate_model_quantum_task_result.py b/src/braket/tasks/gate_model_quantum_task_result.py
index 3e4a5714d..b001f01af 100644
--- a/src/braket/tasks/gate_model_quantum_task_result.py
+++ b/src/braket/tasks/gate_model_quantum_task_result.py
@@ -424,7 +424,7 @@ def _probability_from_measurements(
 
         # count the basis state occurrences, and construct the probability vector
         basis_states, counts = np.unique(indices, return_counts=True)
-        probabilities = np.zeros([2 ** num_measured_qubits], dtype=np.float64)
+        probabilities = np.zeros([2**num_measured_qubits], dtype=np.float64)
         probabilities[basis_states] = counts / shots
         return probabilities
 
diff --git a/test/integ_tests/test_tensor_network_simulator.py b/test/integ_tests/test_tensor_network_simulator.py
index e96d0cbcf..e14f8f730 100644
--- a/test/integ_tests/test_tensor_network_simulator.py
+++ b/test/integ_tests/test_tensor_network_simulator.py
@@ -64,7 +64,7 @@ def _qft(circuit, num_qubits):
     for i in range(num_qubits):
         circuit.h(i)
         for j in range(1, num_qubits - i):
-            circuit.cphaseshift(i + j, i, math.pi / (2 ** j))
+            circuit.cphaseshift(i + j, i, math.pi / (2**j))
 
     for qubit in range(math.floor(num_qubits / 2)):
         circuit.swap(qubit, num_qubits - qubit - 1)
@@ -78,7 +78,7 @@ def _inverse_qft(circuit, num_qubits):
 
     for i in reversed(range(num_qubits)):
         for j in reversed(range(1, num_qubits - i)):
-            circuit.cphaseshift(i + j, i, -math.pi / (2 ** j))
+            circuit.cphaseshift(i + j, i, -math.pi / (2**j))
         circuit.h(i)
 
     return circuit
diff --git a/test/unit_tests/braket/circuits/test_observables.py b/test/unit_tests/braket/circuits/test_observables.py
index 2792efbbe..fbe72eb8d 100644
--- a/test/unit_tests/braket/circuits/test_observables.py
+++ b/test/unit_tests/braket/circuits/test_observables.py
@@ -294,6 +294,6 @@ def test_observable_from_ir_tensor_product_value_error():
 def compare_eigenvalues(observable, expected):
     assert np.allclose(observable.eigenvalues, expected)
     assert np.allclose(
-        np.array([observable.eigenvalue(i) for i in range(2 ** observable.qubit_count)]),
+        np.array([observable.eigenvalue(i) for i in range(2**observable.qubit_count)]),
         expected,
     )

From f7f9b0981cc7a3dc5e6a7ca60b507881870d3a9c Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Tue, 1 Feb 2022 10:39:51 -0500
Subject: [PATCH 3/7] fix: fix unit test

---
 test/unit_tests/braket/circuits/test_noises.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/unit_tests/braket/circuits/test_noises.py b/test/unit_tests/braket/circuits/test_noises.py
index cb22068b9..37e420aba 100644
--- a/test/unit_tests/braket/circuits/test_noises.py
+++ b/test/unit_tests/braket/circuits/test_noises.py
@@ -73,7 +73,7 @@
     ),
     (
         Noise.TwoQubitPauliChannel,
-        "multi_qubit_pauli_channel",
+        "two_qubit_pauli_channel",
         ir.MultiQubitPauliChannel,
         [DoubleTarget, MultiProbability],
         {},

From 2bbb95134cc65600fd2ee5dcb4a50ffeb4028bb6 Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Tue, 1 Feb 2022 11:34:36 -0500
Subject: [PATCH 4/7] fix: fix test coverage

---
 src/braket/circuits/noise.py                  | 25 +++++----
 test/unit_tests/braket/circuits/test_noise.py | 56 ++++++++++---------
 2 files changed, 44 insertions(+), 37 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index 58a4f0531..a8fc76e0b 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -219,6 +219,8 @@ class MultiQubitPauliNoise(Noise):
     parameterized by up to 4**N - 1 probabilities.
     """
 
+    _allowed_substrings = {"I", "X", "Y", "Z"}
+
     def __init__(
         self,
         probabilities: Dict[str, float],
@@ -250,19 +252,14 @@ def __init__(
         super().__init__(qubit_count=qubit_count, ascii_symbols=ascii_symbols)
         self.probabilities = probabilities
 
-        allowed_substrings = {"I", "X", "Y", "Z"}
-
-        if "I" * self.qubit_count in probabilities:
-            raise ValueError("Leakage error is not supported.")
+        if not probabilities:
+            raise ValueError("Pauli dictionary must not be empty.")
 
-        if sum(probabilities.values()) > 1.0 or sum(probabilities.values()) < 0.0:
-            raise ValueError("Total probability must be a real number in the interval [0,1]")
-
-        if len(probabilities) >= 4**self.qubit_count:
+        identity = self.qubit_count * "I"
+        if identity in probabilities:
             raise ValueError(
-                "Too many probabilities provided. Only 4 ** self.qubit_count - 1 are allowed."
+                f"{identity} is not allowed as a key. Please enter only non-identity Pauli strings."
             )
-
         for pauli_string, prob in probabilities.items():
             if not isinstance(pauli_string, str):
                 raise TypeError("Keys must be a string type")
@@ -270,11 +267,17 @@ def __init__(
                 raise ValueError("Length of each Pauli strings must be equal to number of qubits.")
             if not isinstance(prob, float):
                 raise TypeError("Keys must be a float type")
-            if not set(pauli_string) <= allowed_substrings:
+            if not set(pauli_string) <= self._allowed_substrings:
                 raise ValueError("Strings must be Pauli strings consisting of only [I, X, Y, Z]")
             if prob < 0.0 or prob > 1.0:
                 raise ValueError("Individual values must be a real number in the interval [0,1]")
 
+        total_prob = sum(probabilities.values())
+        if total_prob > 1.0 or total_prob < 0.0:
+            raise ValueError(
+                f"Total probability must be a real number in the interval [0, 1]. Total probability was {total_prob}."  # noqa: E501
+            )
+
     def __repr__(self):
         return f"{self.name}('probabilities' : {self.probabilities}, 'qubit_count': {self.qubit_count})"  # noqa
 
diff --git a/test/unit_tests/braket/circuits/test_noise.py b/test/unit_tests/braket/circuits/test_noise.py
index dcb5576e8..b39346fca 100644
--- a/test/unit_tests/braket/circuits/test_noise.py
+++ b/test/unit_tests/braket/circuits/test_noise.py
@@ -351,35 +351,39 @@ def __init__(self):
     assert Noise._FooNoise().name == _FooNoise().name
 
 
-@pytest.mark.parametrize("probs", [{"X": 0.1}, {"XX": 0.1, "YY": 0.1}])
-def test_valid_data_pauli_nqubit(probs):
-
-    qubit_count = len(list(probs.keys())[0])
-    ascii_symbols = [["foo"] for _ in range(qubit_count)]
+@pytest.mark.parametrize(
+    "probs, qubit_count, ascii_symbols", [({"X": 0.1}, 1, ["PC"]), ({"XX": 0.1}, 2, ["PC2", "PC2"])]
+)
+def test_multi_qubit_noise(probs, qubit_count, ascii_symbols):
     MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
 
 
 @pytest.mark.xfail(raises=ValueError)
-@pytest.mark.parametrize(
-    "probs",
-    [
-        {"X": -0.1},
-        {"XY": 1.1},
-        {"TX": 0.1},
-        {"X": 0.5, "Y": 0.6},
-        {"X": 0.1, "YY": 0.2},
-        {"II": 0.9, "XX": 0.1},
-    ],
-)
-def test_invalid_values_pauli_nqubit(probs):
-    qubit_count = len(list(probs.keys())[0])
-    ascii_symbols = [["foo"] for _ in range(qubit_count)]
-    MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
+class TestMultiQubitNoise:
+    qubit_count = 1
+    ascii_symbols = ["PC2"]
 
+    def test_non_empty(self):
+        MultiQubitPauliNoise({}, self.qubit_count, self.ascii_symbols)
 
-@pytest.mark.xfail(raises=TypeError)
-@pytest.mark.parametrize("probs", [{0.1: 0.1}, {"X": "F"}])
-def test_invalid_types_pauli_nqubit(probs):
-    qubit_count = len(list(probs.keys())[0])
-    ascii_symbols = [["foo"] for _ in range(qubit_count)]
-    MultiQubitPauliNoise(probs, qubit_count, ascii_symbols)
+    def test_non_identity(self):
+        MultiQubitPauliNoise({"I": 0.1}, self.qubit_count, self.ascii_symbols)
+
+    def test_non_equal_length_paulis(self):
+        MultiQubitPauliNoise({"X": 0.1, "XY": 0.1}, 1, self.ascii_symbols)
+        MultiQubitPauliNoise({"X": 0.1, "Y": 0.1}, 2, ["PC2", "PC2"])
+
+    def test_prob_over_one(self):
+        MultiQubitPauliNoise({"X": 0.9, "Y": 0.9}, 1, self.ascii_symbols)
+        MultiQubitPauliNoise({"XX": 0.9, "YY": 0.9}, 1, self.ascii_symbols)
+
+    def test_prob_under_one(self):
+        MultiQubitPauliNoise({"X": -0.6, "Y": -0.9}, 1, self.ascii_symbols)
+        MultiQubitPauliNoise({"XX": -0.9, "YY": -0.9}, 2, ["PC2", "PC2"])
+
+    def test_non_pauli_string(self):
+        MultiQubitPauliNoise({"T": 0.1}, 1, self.ascii_symbols)
+
+    def test_individual_probs(self):
+        MultiQubitPauliNoise({"X": -0.1}, 1, self.ascii_symbols)
+        MultiQubitPauliNoise({"X": 1.1}, 1, self.ascii_symbols)

From 1b76f3e27bfc797ba226bae3804dc0a47138f8b9 Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Tue, 1 Feb 2022 14:33:30 -0500
Subject: [PATCH 5/7] fix: unit tests

---
 src/braket/circuits/noise.py                  |   6 +-
 src/braket/circuits/noises.py                 | 258 +++++++++---------
 test/unit_tests/braket/circuits/test_noise.py |   8 +
 3 files changed, 139 insertions(+), 133 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index a8fc76e0b..2e6582bd0 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -282,7 +282,7 @@ def __repr__(self):
         return f"{self.name}('probabilities' : {self.probabilities}, 'qubit_count': {self.qubit_count})"  # noqa
 
 
-class PauliNoise(MultiQubitPauliNoise):
+class PauliNoise(Noise):
     """
     Class `PauliNoise` represents the a single-qubit Pauli noise channel
     acting on one qubit. It is parameterized by three probabilities.
@@ -311,9 +311,7 @@ def __init__(
                 is not `float`, `probX` or `probY` or `probZ` > 1.0, or
                 `probX` or `probY` or `probZ` < 0.0, or `probX`+`probY`+`probZ` > 1
         """
-        super().__init__(
-            {"X": probX, "Y": probY, "Z": probZ}, qubit_count=1, ascii_symbols=ascii_symbols
-        )
+        super().__init__(qubit_count=qubit_count, ascii_symbols=ascii_symbols)
 
         if not isinstance(probX, float):
             raise TypeError("probX must be float type")
diff --git a/src/braket/circuits/noises.py b/src/braket/circuits/noises.py
index 3da5dcef8..7d79fbaa2 100644
--- a/src/braket/circuits/noises.py
+++ b/src/braket/circuits/noises.py
@@ -554,6 +554,135 @@ def two_qubit_dephasing(
 Noise.register_noise(TwoQubitDephasing)
 
 
+class TwoQubitPauliChannel(MultiQubitPauliNoise):
+    """Two-Qubit Pauli noise channel which transforms a
+        density matrix :math:`\\rho` according to:
+
+    .. math::
+        \\rho \\Rightarrow (1-p) \\rho +
+            p_{IX} IX \\rho IX^{\\dagger} +
+            p_{IY} IY \\rho IY^{\\dagger} +
+            p_{IZ} IZ \\rho IZ^{\\dagger} +
+            p_{XI} XI \\rho XI^{\\dagger} +
+            p_{XX} XX \\rho XX^{\\dagger} +
+            p_{XY} XY \\rho XY^{\\dagger} +
+            p_{XZ} XZ \\rho XZ^{\\dagger} +
+            p_{YI} YI \\rho YI^{\\dagger} +
+            p_{YX} YX \\rho YX^{\\dagger} +
+            p_{YY} YY \\rho YY^{\\dagger} +
+            p_{YZ} YZ \\rho YZ^{\\dagger} +
+            p_{ZI} ZI \\rho ZI^{\\dagger} +
+            p_{ZX} ZX \\rho ZX^{\\dagger} +
+            p_{ZY} ZY \\rho ZY^{\\dagger} +
+            p_{ZZ} ZZ \\rho ZZ^{\\dagger})
+    where
+
+    .. math::
+        I = \\left(
+                \\begin{matrix}
+                    1 & 0 \\\\
+                    0 & 1
+                \\end{matrix}
+            \\right)
+
+        X = \\left(
+                \\begin{matrix}
+                    0 & 1 \\\\
+                    1 & 0
+                \\end{matrix}
+            \\right)
+
+        Y = \\left(
+                \\begin{matrix}
+                    0 & -i \\\\
+                    i &  0
+                \\end{matrix}
+            \\right)
+
+        Z = \\left(
+                \\begin{matrix}
+                    1 & 0 \\\\
+                    0 & -1
+                \\end{matrix}
+            \\right)
+
+        p = \\text{sum of all probabilities}
+
+    This noise channel is shown as `PC_2({"pauli_string": probability})` in circuit diagrams.
+    """
+
+    _paulis = {
+        "I": np.array([[1.0, 0.0], [0.0, 1.0]], dtype=complex),
+        "X": np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex),
+        "Y": np.array([[0.0, -1.0j], [1.0j, 0.0]], dtype=complex),
+        "Z": np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex),
+    }
+    _tensor_products_strings = itertools.product(_paulis.keys(), repeat=2)
+    _names_list = ["".join(x) for x in _tensor_products_strings]
+
+    def __init__(self, probabilities: Dict[str, float]):
+        super().__init__(
+            probabilities=probabilities,
+            qubit_count=None,
+            ascii_symbols=[
+                f"PC2({probabilities})",
+                f"PC2({probabilities})",
+            ],
+        )
+
+        total_prob = sum(self.probabilities.values())
+
+        K_list = [np.sqrt(1 - total_prob) * np.identity(4)]  # "II" element
+        for pstring in self._names_list[1:]:  # ignore "II"
+            if pstring in self.probabilities:
+                mat = np.sqrt(self.probabilities[pstring]) * np.kron(
+                    self._paulis[pstring[0]], self._paulis[pstring[1]]
+                )
+                K_list.append(mat)
+            else:
+                K_list.append(np.zeros((4, 4)))
+        self._matrix = K_list
+
+    def to_ir(self, target: QubitSet):
+        return ir.MultiQubitPauliChannel.construct(
+            targets=[target[0], target[1]], probabilities=self.probabilities
+        )
+
+    def to_matrix(self) -> Iterable[np.ndarray]:
+        return self._matrix
+
+    @staticmethod
+    def fixed_qubit_count() -> int:
+        return 2
+
+    @staticmethod
+    @circuit.subroutine(register=True)
+    def two_qubit_pauli_channel(
+        target1: QubitInput, target2: QubitInput, probabilities: Dict[str, float]
+    ) -> Iterable[Instruction]:
+        """Registers this function into the circuit class.
+
+        Args:
+            target (Qubit, int, or iterable of Qubit / int): Target qubits
+            probability (float): Probability of two-qubit Pauli channel.
+
+        Returns:
+            Iterable[Instruction]: `Iterable` of Depolarizing instructions.
+
+        Examples:
+            >>> circ = Circuit().two_qubit_pauli_channel(0, 1, {"XX": 0.1})
+        """
+        return [
+            Instruction(
+                Noise.TwoQubitPauliChannel(probabilities=probabilities),
+                target=[target1, target2],
+            )
+        ]
+
+
+Noise.register_noise(TwoQubitPauliChannel)
+
+
 class AmplitudeDamping(DampingNoise):
     """AmplitudeDamping noise channel which transforms a density matrix :math:`\\rho` according to:
 
@@ -871,132 +1000,3 @@ def kraus(
 
 
 Noise.register_noise(Kraus)
-
-
-class TwoQubitPauliChannel(MultiQubitPauliNoise):
-    """Two-Qubit Pauli noise channel which transforms a
-        density matrix :math:`\\rho` according to:
-
-    .. math::
-        \\rho \\Rightarrow (1-p) \\rho +
-            p_{IX} IX \\rho IX^{\\dagger} +
-            p_{IY} IY \\rho IY^{\\dagger} +
-            p_{IZ} IZ \\rho IZ^{\\dagger} +
-            p_{XI} XI \\rho XI^{\\dagger} +
-            p_{XX} XX \\rho XX^{\\dagger} +
-            p_{XY} XY \\rho XY^{\\dagger} +
-            p_{XZ} XZ \\rho XZ^{\\dagger} +
-            p_{YI} YI \\rho YI^{\\dagger} +
-            p_{YX} YX \\rho YX^{\\dagger} +
-            p_{YY} YY \\rho YY^{\\dagger} +
-            p_{YZ} YZ \\rho YZ^{\\dagger} +
-            p_{ZI} ZI \\rho ZI^{\\dagger} +
-            p_{ZX} ZX \\rho ZX^{\\dagger} +
-            p_{ZY} ZY \\rho ZY^{\\dagger} +
-            p_{ZZ} ZZ \\rho ZZ^{\\dagger})
-    where
-
-    .. math::
-        I = \\left(
-                \\begin{matrix}
-                    1 & 0 \\\\
-                    0 & 1
-                \\end{matrix}
-            \\right)
-
-        X = \\left(
-                \\begin{matrix}
-                    0 & 1 \\\\
-                    1 & 0
-                \\end{matrix}
-            \\right)
-
-        Y = \\left(
-                \\begin{matrix}
-                    0 & -i \\\\
-                    i &  0
-                \\end{matrix}
-            \\right)
-
-        Z = \\left(
-                \\begin{matrix}
-                    1 & 0 \\\\
-                    0 & -1
-                \\end{matrix}
-            \\right)
-
-        p = \\text{sum of all probabilities}
-
-    This noise channel is shown as `PC_2({"pauli_string": probability})` in circuit diagrams.
-    """
-
-    _paulis = {
-        "I": np.array([[1.0, 0.0], [0.0, 1.0]], dtype=complex),
-        "X": np.array([[0.0, 1.0], [1.0, 0.0]], dtype=complex),
-        "Y": np.array([[0.0, -1.0j], [1.0j, 0.0]], dtype=complex),
-        "Z": np.array([[1.0, 0.0], [0.0, -1.0]], dtype=complex),
-    }
-    _tensor_products_strings = itertools.product(_paulis.keys(), repeat=2)
-    _names_list = ["".join(x) for x in _tensor_products_strings]
-
-    def __init__(self, probabilities: Dict[str, float]):
-        super().__init__(
-            probabilities=probabilities,
-            qubit_count=None,
-            ascii_symbols=[
-                f"PC2({probabilities})",
-                f"PC2({probabilities})",
-            ],
-        )
-
-        total_prob = sum(self.probabilities.values())
-
-        K_list = [np.sqrt(1 - total_prob) * np.identity(4)]  # "II" element
-        for pstring in self._names_list[1:]:  # ignore "II"
-            if pstring in self.probabilities:
-                mat = np.sqrt(self.probabilities[pstring]) * np.kron(
-                    self._paulis[pstring[0]], self._paulis[pstring[1]]
-                )
-                K_list.append(mat)
-            else:
-                K_list.append(np.zeros((4, 4)))
-        self._matrix = K_list
-
-    def to_ir(self, target: QubitSet):
-        return ir.MultiQubitPauliChannel.construct(
-            targets=[target[0], target[1]], probabilities=self.probabilities
-        )
-
-    def to_matrix(self) -> Iterable[np.ndarray]:
-        return self._matrix
-
-    @staticmethod
-    def fixed_qubit_count() -> int:
-        return 2
-
-    @staticmethod
-    @circuit.subroutine(register=True)
-    def two_qubit_pauli_channel(
-        target1: QubitInput, target2: QubitInput, probabilities: Dict[str, float]
-    ) -> Iterable[Instruction]:
-        """Registers this function into the circuit class.
-
-        Args:
-            target (Qubit, int, or iterable of Qubit / int): Target qubits
-            probability (float): Probability of two-qubit Pauli channel.
-
-        Returns:
-            Iterable[Instruction]: `Iterable` of Depolarizing instructions.
-
-        Examples:
-            >>> circ = Circuit().two_qubit_pauli_channel(0, 1, {"XX": 0.1})
-        """
-        return [
-            Instruction(
-                Noise.TwoQubitPauliChannel(probabilities=probabilities),
-                target=[target1, target2],
-            )
-        ]
-
-
-Noise.register_noise(TwoQubitPauliChannel)
diff --git a/test/unit_tests/braket/circuits/test_noise.py b/test/unit_tests/braket/circuits/test_noise.py
index b39346fca..0c367bbc5 100644
--- a/test/unit_tests/braket/circuits/test_noise.py
+++ b/test/unit_tests/braket/circuits/test_noise.py
@@ -387,3 +387,11 @@ def test_non_pauli_string(self):
     def test_individual_probs(self):
         MultiQubitPauliNoise({"X": -0.1}, 1, self.ascii_symbols)
         MultiQubitPauliNoise({"X": 1.1}, 1, self.ascii_symbols)
+
+    @pytest.mark.xfail(raises=TypeError)
+    def test_keys_strings(self):
+        MultiQubitPauliNoise({1: 1.1}, 1, self.ascii_symbols)
+
+    @pytest.mark.xfail(raises=TypeError)
+    def test_values_floats(self):
+        MultiQubitPauliNoise({"X": "str"}, 1, self.ascii_symbols)

From c31ed45655a74cf4dc5ad63fc7de4499c0b123ba Mon Sep 17 00:00:00 2001
From: mbeach-aws <85963088+mbeach-aws@users.noreply.github.com>
Date: Tue, 1 Feb 2022 15:34:20 -0500
Subject: [PATCH 6/7] fix: apply suggestions from code review

Co-authored-by: Cody Wang <speller26@gmail.com>
---
 src/braket/circuits/noise.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index 2e6582bd0..565c324b1 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -262,11 +262,11 @@ def __init__(
             )
         for pauli_string, prob in probabilities.items():
             if not isinstance(pauli_string, str):
-                raise TypeError("Keys must be a string type")
+                raise TypeError("Keys must be of string type")
             if len(pauli_string) != self.qubit_count:
-                raise ValueError("Length of each Pauli strings must be equal to number of qubits.")
+                raise ValueError("Length of each Pauli string must be equal to number of qubits.")
             if not isinstance(prob, float):
-                raise TypeError("Keys must be a float type")
+                raise TypeError("Values must be of float type")
             if not set(pauli_string) <= self._allowed_substrings:
                 raise ValueError("Strings must be Pauli strings consisting of only [I, X, Y, Z]")
             if prob < 0.0 or prob > 1.0:

From 37bfd075358fe25745210b2fb4433013989f507a Mon Sep 17 00:00:00 2001
From: mbeach-aws <drbeach@amazon.com>
Date: Tue, 1 Feb 2022 15:52:06 -0500
Subject: [PATCH 7/7] fix: better error message

---
 src/braket/circuits/noise.py | 37 +++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/braket/circuits/noise.py b/src/braket/circuits/noise.py
index 565c324b1..303080aa7 100644
--- a/src/braket/circuits/noise.py
+++ b/src/braket/circuits/noise.py
@@ -260,22 +260,45 @@ def __init__(
             raise ValueError(
                 f"{identity} is not allowed as a key. Please enter only non-identity Pauli strings."
             )
+
         for pauli_string, prob in probabilities.items():
             if not isinstance(pauli_string, str):
-                raise TypeError("Keys must be of string type")
+                raise TypeError(f"Type of {pauli_string} was not a string.")
             if len(pauli_string) != self.qubit_count:
-                raise ValueError("Length of each Pauli string must be equal to number of qubits.")
+                raise ValueError(
+                    (
+                        "Length of each Pauli string must be equal to number of qubits. "
+                        f"{pauli_string} had length {len(pauli_string)} instead of length {self.qubit_count}."  # noqa
+                    )
+                )
             if not isinstance(prob, float):
-                raise TypeError("Values must be of float type")
+                raise TypeError(
+                    (
+                        "Probabilities must be a float type. "
+                        f"The probability for {pauli_string} was of type {type(prob)}."
+                    )
+                )
             if not set(pauli_string) <= self._allowed_substrings:
-                raise ValueError("Strings must be Pauli strings consisting of only [I, X, Y, Z]")
+                raise ValueError(
+                    (
+                        "Strings must be Pauli strings consisting of only [I, X, Y, Z]. "
+                        f"Received {pauli_string}."
+                    )
+                )
             if prob < 0.0 or prob > 1.0:
-                raise ValueError("Individual values must be a real number in the interval [0,1]")
-
+                raise ValueError(
+                    (
+                        "Individual probabilities must be real numbers in the interval [0, 1]. "
+                        f"Probability for {pauli_string} was {prob}."
+                    )
+                )
         total_prob = sum(probabilities.values())
         if total_prob > 1.0 or total_prob < 0.0:
             raise ValueError(
-                f"Total probability must be a real number in the interval [0, 1]. Total probability was {total_prob}."  # noqa: E501
+                (
+                    "Total probability must be a real number in the interval [0, 1]. "
+                    f"Total probability was {total_prob}."
+                )
             )
 
     def __repr__(self):