Fix formatting of fidelity test.

tests/atoms_test.py

@@ -2,19 +2,21 @@
 import aisim as ais
 import numpy as np
 
+
 def generate_random_atoms(n_atom, n_int):
-        random_phase_space_vectors = np.random.rand(n_atom, 6)
-        random_kets = np.random.rand(n_atom, n_int)
-        # normalize random ket
-        norm = np.einsum('ij,ij->i', random_kets, random_kets)
-        norm_random_kets = np.einsum('ij,i->ij', random_kets, 1/norm**(1/2))
-        # apply random phase shifts
-        random_phases = 2*np.pi*np.random.rand(n_atom, n_int)
-        norm_random_kets = np.exp(1j*random_phases) * norm_random_kets
-        norm_random_kets = np.reshape(norm_random_kets, (n_atom, n_int, 1))
-        random_atoms = ais.AtomicEnsemble(
-            random_phase_space_vectors, norm_random_kets)
-        return random_atoms
+    random_phase_space_vectors = np.random.rand(n_atom, 6)
+    random_kets = np.random.rand(n_atom, n_int)
+    # normalize random ket
+    norm = np.einsum('ij,ij->i', random_kets, random_kets)
+    norm_random_kets = np.einsum('ij,i->ij', random_kets, 1/norm**(1/2))
+    # apply random phase shifts
+    random_phases = 2*np.pi*np.random.rand(n_atom, n_int)
+    norm_random_kets = np.exp(1j*random_phases) * norm_random_kets
+    norm_random_kets = np.reshape(norm_random_kets, (n_atom, n_int, 1))
+    random_atoms = ais.AtomicEnsemble(
+        random_phase_space_vectors, norm_random_kets)
+    return random_atoms
+
 
 def test_atomic_ensemble_methods():
     # This tests the AtomicEnsemble class' methods for consistency. It creates
@@ -50,7 +52,7 @@ def atomic_ensemble_test_function(atoms):
             state_occupation += atoms.state_occupation(m)
         np.testing.assert_almost_equal(state_occupation, 1)
         # Test properties of the fidelity function
-        ## fidelity of density matrix with itself is 1
+        # fidelity of density matrix with itself is 1
         np.testing.assert_almost_equal(1, atoms.fidelity(atoms.density_matrix))
 
     # Test AtomicEnsemble from very general randomly generated states
@@ -91,18 +93,20 @@ def test_make_grid():
     assert atomic_ensemble[:, 2].max() == 3.0
     np.testing.assert_almost_equal(atomic_ensemble[:, 2].mean(), 0.0)
 
+
 def test_fidelity():
-    for n_atom in [1,3,1000]:
-        for n_int in [2,3,10,50]:
+    for n_atom in [1, 3, 1000]:
+        for n_int in [2, 3, 10, 50]:
             # Generate random density matrices
-            rhoA = generate_random_atoms(n_atom, n_int).density_matrix
-            rhoB = generate_random_atoms(n_atom, n_int).density_matrix
+            rho_a = generate_random_atoms(n_atom, n_int).density_matrix
+            rho_b = generate_random_atoms(n_atom, n_int).density_matrix
             # Test that return is float
-            assert isinstance(ais.atoms._fidelity(rhoA, rhoB), float)
+            assert isinstance(ais.atoms._fidelity(rho_a, rho_b), float)
             # Test symmetry F(rhoA,rhoB) == F(rhoB,rhoA)
-            np.testing.assert_almost_equal(ais.atoms._fidelity(rhoA, rhoB),
-                ais.atoms._fidelity(rhoB, rhoA))
+            np.testing.assert_almost_equal(ais.atoms._fidelity(rho_a, rho_b),
+                                           ais.atoms._fidelity(rho_b, rho_a))
             # Test if F(rhoA, rhoB) is equal or less than 1
-            assert np.all(ais.atoms._fidelity(rhoA, rhoB) <= 1)
+            assert np.all(ais.atoms._fidelity(rho_a, rho_b) <= 1)
             # Test if F(rhoA, rhoA) is equal to 1
-            np.testing.assert_array_almost_equal(ais.atoms._fidelity(rhoA, rhoA), 1)
+            np.testing.assert_array_almost_equal(
+                ais.atoms._fidelity(rho_a, rho_a), 1)