Correct formatting and add docs

qermit/noise_model/noise_model.py

@@ -12,7 +12,8 @@
 from numpy.random import Generator
 from enum import Enum
 from itertools import product
-from scipy.linalg import fractional_matrix_power
+from scipy.linalg import fractional_matrix_power  # type: ignore
+from numpy.typing import NDArray
 
 
 Direction = Enum('Direction', ['forward', 'backward'])
@@ -53,7 +54,7 @@ def __init__(
                     f"Probabilities sum to {sum(distribution.values())}"
                     + " but should be less than or equal to 1."
                 )
-        
+
         if distribution == {}:
             pass
         else:
@@ -65,33 +66,47 @@ def __init__(
         self.rng = rng
 
     @property
-    def identity_error_rate(self):
+    def identity_error_rate(self) -> float:
+        """Rate at which no error occurs.
+
+        :return: Rate at which no error occurs.
+            Calculated as 1 minus the total error rate of
+            error in this distribution.
+        :rtype: float
+        """
         return 1 - sum(self.distribution.values())
 
-    def to_ptm(self) -> Tuple[np.array, Dict[Tuple[Pauli], int]]:
-
+    def to_ptm(self) -> Tuple[NDArray, Dict[Tuple[Pauli, ...], int]]:
+        """Convert error distribution to Pauli Transfer Matrix (PTM) form.
+
+        :return: PTM of error distribution and Pauli index dictionary.
+            The Pauli index dictionary maps Pauli errors to their
+            index in the PTM
+        :rtype: Tuple[NDArray, Dict[Tuple[Pauli, ...], int]]
+        """
+
         ptm = np.zeros((4**self.n_qubits, 4**self.n_qubits))
         pauli_index = {
-            pauli:index
+            pauli: index
             for index, pauli
             in enumerate(product({Pauli.I, Pauli.X, Pauli.Y, Pauli.Z}, repeat=self.n_qubits))
         }
-        
+
         for pauli_tuple, index in pauli_index.items():
-                    
+
             pauli = QermitPauli.from_pauli_iterable(
                 pauli_iterable=pauli_tuple,
                 qubit_list=[Qubit(i) for i in range(self.n_qubits)]
             )
-            
+
             ptm[index][index] += self.identity_error_rate
-                        
+
             for error, error_rate in self.distribution.items():
                 error_pauli = QermitPauli.from_pauli_iterable(
                     pauli_iterable=error,
                     qubit_list=[Qubit(i) for i in range(self.n_qubits)]
                 )
-                
+
                 ptm[index][index] += error_rate * QermitPauli.commute_coeff(pauli_one=pauli, pauli_two=error_pauli)
 
         identity = tuple(Pauli.I for _ in range(self.n_qubits))
@@ -101,18 +116,29 @@ def to_ptm(self) -> Tuple[np.array, Dict[Tuple[Pauli], int]]:
                 + "This is a fault in Qermit. "
                 + "Please report this as an issue."
             )
-        
+
         if not self == ErrorDistribution.from_ptm(ptm=ptm, pauli_index=pauli_index):
             raise Exception(
                 "From PTM does not match to PTM. "
-                + "This is a bug. "
-                + "Please report to developers. "
+                + "This is a fault in Qermit. "
+                + "Please report this as an issue."
             )
-                
+
         return ptm, pauli_index
-    
+
     @classmethod
-    def from_ptm(cls, ptm, pauli_index):
+    def from_ptm(cls, ptm: NDArray, pauli_index: Dict[Tuple[Pauli, ...], int]) -> ErrorDistribution:
+        """Convert a Pauli Transfer Matrix (PTM) to an error distribution.
+
+        :param ptm: Pauli Transfer Matrix to convert. Should be a 4^n by 4^n matrix
+            where n is the number of qubits.
+        :type ptm: NDArray
+        :param pauli_index: A dictionary mapping Pauli errors to
+            their index in the PTM.
+        :type pauli_index: Dict[Tuple[Pauli, ...], int]
+        :return: The converted error distribution.
+        :rtype: ErrorDistribution
+        """
 
         assert ptm.ndim == 2
         assert ptm.shape[0] == ptm.shape[1]
@@ -143,9 +169,10 @@ def from_ptm(cls, ptm, pauli_index):
         }
         return cls(distribution=distribution)
 
-
     @property
     def n_qubits(self) -> int:
+        """The number of qubits this error distribution acts on.
+        """
         return len(list(self.distribution.keys())[0])
 
     def __eq__(self, other: object) -> bool:
@@ -301,7 +328,7 @@ def plot(self):
 
         return fig
 
-    def scale(self, scaling_factor:float) -> ErrorDistribution:
+    def scale(self, scaling_factor: float) -> ErrorDistribution:
 
         ptm, pauli_index = self.to_ptm()
         scaled_ptm = fractional_matrix_power(ptm, scaling_factor)

qermit/noise_model/qermit_pauli.py

@@ -632,12 +632,12 @@ def qubit_pauli_string(self) -> Tuple[QubitPauliString, complex]:
         )
 
         return qubit_pauli_string, operator_phase
-    
+
     @classmethod
     def from_pauli_iterable(cls, pauli_iterable: Iterable[Pauli], qubit_list: List[Qubit]) -> QermitPauli:
         return cls(
             Z_list=[int(pauli in (Pauli.Z, Pauli.Y)) for pauli in pauli_iterable],
             X_list=[int(pauli in (Pauli.X, Pauli.Y)) for pauli in pauli_iterable],
             qubit_list=qubit_list,
-            phase=sum(int(pauli==Pauli.Y) for pauli in pauli_iterable) % 4,
+            phase=sum(int(pauli == Pauli.Y) for pauli in pauli_iterable) % 4,
         )

tests/noise_model_test.py

@@ -21,7 +21,7 @@
 def test_to_ptm() -> None:
 
     # A simple test with an only X noise model on one qubit
-    error_distribution = ErrorDistribution(distribution={(Pauli.X,):0.1})
+    error_distribution = ErrorDistribution(distribution={(Pauli.X,): 0.1})
     ptm, pauli_index = error_distribution.to_ptm()
 
     assert ptm[pauli_index[(Pauli.I, )]][pauli_index[(Pauli.I, )]] == 1
@@ -32,8 +32,8 @@ def test_to_ptm() -> None:
     # A slightly more complicated example with some verified entries.
     error_distribution = ErrorDistribution(
         distribution={
-            (Pauli.X, Pauli.Z):0.08,
-            (Pauli.Y, Pauli.Z):0.02,
+            (Pauli.X, Pauli.Z): 0.08,
+            (Pauli.Y, Pauli.Z): 0.02,
         }
     )
 
@@ -44,13 +44,14 @@ def test_to_ptm() -> None:
     assert abs(ptm[pauli_index[(Pauli.X, Pauli.Z)]][pauli_index[(Pauli.X, Pauli.Z)]] - 0.96) < 10**(-6)
     assert abs(ptm[pauli_index[(Pauli.X, Pauli.X)]][pauli_index[(Pauli.X, Pauli.X)]] - 0.84) < 10**(-6)
 
+
 def test_from_ptm() -> None:
 
     # Test that the error distribution to and from ptm is the same as the initial
-    distribution={
-        (Pauli.X, Pauli.X):0.1,
-        (Pauli.Y, Pauli.Z):0.2,
-        (Pauli.Z, Pauli.X):0.3,
+    distribution = {
+        (Pauli.X, Pauli.X): 0.1,
+        (Pauli.Y, Pauli.Z): 0.2,
+        (Pauli.Z, Pauli.X): 0.3,
     }
 
     error_distribution = ErrorDistribution(
@@ -82,7 +83,8 @@ def test_qermit_pauli_from_iterable() -> None:
         pauli_iterable=qubit_pauli_string.map.values(),
         qubit_list=list(qubit_pauli_string.map.keys())
     )
-    pauli.qubit_pauli_string == (qubit_pauli_string, 1+0j)
+    pauli.qubit_pauli_string == (qubit_pauli_string, 1 + 0j)
+
 
 def test_qermit_pauli_commute_coeff() -> None:
 
@@ -103,7 +105,7 @@ def test_qermit_pauli_commute_coeff() -> None:
     ]
 
     for verified in verified_list:
-        
+
         n_qubits = len(verified[0][0][0])
 
         pauli_one = QermitPauli(Z_list=verified[0][0][0], X_list=verified[0][0][1], qubit_list=[Qubit(i) for i in range(n_qubits)])

tests/zne_test.py

@@ -844,18 +844,19 @@ def test_end_to_end_noise_scaled_mitex():
     assert abs(qubit_pauli_operator_list[1]._dict[qps_noisless_one] - 1) < 0.1
     assert abs(qubit_pauli_operator_list[1]._dict[qps_noisless_zero]) < 0.1
 
+
 @pytest.mark.high_compute
 def test_end_to_end_noise_aware_zne_mitex_starting_from_ptm() -> None:
 
     # Here we are creating the PTM for a noise model acting
     # XI with rate 0.1
-    ptm = np.diag([1,1,1,1,1,1,1,1,0.8,0.8,0.8,0.8,0.8,0.8,0.8,0.8])
+    ptm = np.diag([1, 1, 1, 1, 1, 1, 1, 1, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8])
     pauli_index = {pauli: index for index, pauli in enumerate(product([Pauli.I, Pauli.X, Pauli.Y, Pauli.Z], repeat=2))}
     error_distribution = ErrorDistribution.from_ptm(ptm, pauli_index)
 
     noise_model = NoiseModel(
         noise_model={
-            OpType.CZ:error_distribution
+            OpType.CZ: error_distribution
         }
     )
     transpiler = PauliErrorTranspile(noise_model=noise_model)
@@ -872,21 +873,19 @@ def test_end_to_end_noise_aware_zne_mitex_starting_from_ptm() -> None:
         n_noisy_circuit_samples=1000,
     )
 
-    circuit_noisy = Circuit(2).CZ(0,1)
+    circuit_noisy = Circuit(2).CZ(0, 1)
 
-    qps_noisy_noisy = QubitPauliString(map={Qubit(0):Pauli.Z, Qubit(1):Pauli.Z})
+    qps_noisy_noisy = QubitPauliString(map={Qubit(0): Pauli.Z, Qubit(1): Pauli.Z})
 
     qubit_pauli_operator_noisy = QubitPauliOperator(
-        dictionary = {
-            qps_noisy_noisy:1,
-        }
+        dictionary={qps_noisy_noisy: 1}
     )
 
     observable_tracker_noisy = ObservableTracker(
         qubit_pauli_operator=qubit_pauli_operator_noisy
     )
 
-    shots=10000
+    shots = 10000
     ansatz_circuit_noisy = AnsatzCircuit(
         Circuit=circuit_noisy,
         Shots=shots,