Clean up tests (#3713)

* clean up tests

e.g. remove unnecessary for loops when using pytest.approx()

* Fix typo in op_params

pymatgen/io/lmto.py

@@ -184,13 +184,13 @@ def from_str(cls, data: str, sigfigs: int = 8) -> LMTOCtrl:
 
         for cat in ["STRUC", "CLASS", "SITE"]:
             fields = struct_lines[cat].split("=")
-            for f, field in enumerate(fields):
+            for idx, field in enumerate(fields):
                 token = field.split()[-1]
                 if token == "ALAT":
-                    alat = round(float(fields[f + 1].split()[0]), sigfigs)
-                    structure_tokens["ALAT"] = alat
+                    a_lat = round(float(fields[idx + 1].split()[0]), sigfigs)
+                    structure_tokens["ALAT"] = a_lat
                 elif token == "ATOM":
-                    atom = fields[f + 1].split()[0]
+                    atom = fields[idx + 1].split()[0]
                     if not bool(re.match("E[0-9]*$", atom)):
                         if cat == "CLASS":
                             structure_tokens["CLASS"].append(atom)
@@ -200,9 +200,9 @@ def from_str(cls, data: str, sigfigs: int = 8) -> LMTOCtrl:
                         pass
                 elif token in ["PLAT", "POS"]:
                     try:
-                        arr = np.array([round(float(i), sigfigs) for i in fields[f + 1].split()])
+                        arr = np.array([round(float(i), sigfigs) for i in fields[idx + 1].split()])
                     except ValueError:
-                        arr = np.array([round(float(i), sigfigs) for i in fields[f + 1].split()[:-1]])
+                        arr = np.array([round(float(i), sigfigs) for i in fields[idx + 1].split()[:-1]])
                     if token == "PLAT":
                         structure_tokens["PLAT"] = arr.reshape([3, 3])
                     elif not bool(re.match("E[0-9]*$", atom)):
@@ -212,9 +212,9 @@ def from_str(cls, data: str, sigfigs: int = 8) -> LMTOCtrl:
                 else:
                     pass
         try:
-            spcgrp_index = struct_lines["SYMGRP"].index("SPCGRP")
-            spcgrp = struct_lines["SYMGRP"][spcgrp_index : spcgrp_index + 12]
-            structure_tokens["SPCGRP"] = spcgrp.split("=")[1].split()[0]
+            spc_grp_index = struct_lines["SYMGRP"].index("SPCGRP")
+            spc_grp = struct_lines["SYMGRP"][spc_grp_index : spc_grp_index + 12]
+            structure_tokens["SPCGRP"] = spc_grp.split("=")[1].split()[0]
         except ValueError:
             pass
 

pymatgen/io/nwchem.py

@@ -920,15 +920,15 @@ def isfloatstring(in_str):
             for _freq, mode in normal_frequencies:
                 mode[:] = zip(*[iter(mode)] * 3)
         if hessian:
-            n = len(hessian)
-            for i in range(n):
-                for j in range(i + 1, n):
-                    hessian[i].append(hessian[j][i])
+            len_hess = len(hessian)
+            for ii in range(len_hess):
+                for jj in range(ii + 1, len_hess):
+                    hessian[ii].append(hessian[jj][ii])
         if projected_hessian:
-            n = len(projected_hessian)
-            for i in range(n):
-                for j in range(i + 1, n):
-                    projected_hessian[i].append(projected_hessian[j][i])
+            len_hess = len(projected_hessian)
+            for ii in range(len_hess):
+                for jj in range(ii + 1, len_hess):
+                    projected_hessian[ii].append(projected_hessian[jj][ii])
 
         data.update(
             {

tests/analysis/elasticity/test_elastic.py

@@ -411,8 +411,8 @@ def test_get_strain_state_dict(self):
             strain_states.append(tuple(ss))
             vec = np.zeros((4, 6))
             rand_values = np.random.uniform(0.1, 1, 4)
-            for i in strain_ind:
-                vec[:, i] = rand_values
+            for idx in strain_ind:
+                vec[:, idx] = rand_values
             vecs[strain_ind] = vec
         all_strains = [Strain.from_voigt(v).zeroed() for vec in vecs.values() for v in vec]
         random.shuffle(all_strains)

tests/analysis/ferroelectricity/test_polarization.py

@@ -17,7 +17,7 @@
 
 TEST_DIR = f"{TEST_FILES_DIR}/vasp/fixtures/BTO_221_99_polarization"
 bto_folders = ["nonpolar_polarization"]
-bto_folders += [f"interpolation_{i}_polarization" for i in range(1, 9)][::-1]
+bto_folders += [f"interpolation_{idx}_polarization" for idx in range(8, 0, -1)]
 bto_folders += ["polar_polarization"]
 
 structures = [Structure.from_file(f"{TEST_DIR}/{folder}/POSCAR") for folder in bto_folders]

tests/analysis/test_graphs.py

@@ -772,8 +772,8 @@ def test_split(self):
         diff_spec_mg.add_edge(0, 3)
         diff_spec_mg.add_edge(0, 4)
 
-        for i in range(1, 5):
-            bond = (0, i)
+        for idx in range(1, 5):
+            bond = (0, idx)
 
             split_mgs = diff_spec_mg.split_molecule_subgraphs([bond])
             for split_mg in split_mgs:

tests/analysis/test_local_env.py

@@ -534,18 +534,18 @@ def setUp(self):
         )
 
     def test_site_is_of_motif_type(self):
-        for i in range(len(self.diamond)):
-            assert site_is_of_motif_type(self.diamond, i) == "tetrahedral"
-        for i in range(len(self.nacl)):
-            assert site_is_of_motif_type(self.nacl, i) == "octahedral"
-        for i in range(len(self.cscl)):
-            assert site_is_of_motif_type(self.cscl, i) == "bcc"
+        for idx in range(len(self.diamond)):
+            assert site_is_of_motif_type(self.diamond, idx) == "tetrahedral"
+        for idx in range(len(self.nacl)):
+            assert site_is_of_motif_type(self.nacl, idx) == "octahedral"
+        for idx in range(len(self.cscl)):
+            assert site_is_of_motif_type(self.cscl, idx) == "bcc"
         assert site_is_of_motif_type(self.square_pyramid, 0) == "square pyramidal"
-        for i in range(1, len(self.square_pyramid)):
-            assert site_is_of_motif_type(self.square_pyramid, i) == "unrecognized"
+        for idx in range(1, len(self.square_pyramid)):
+            assert site_is_of_motif_type(self.square_pyramid, idx) == "unrecognized"
         assert site_is_of_motif_type(self.trigonal_bipyramid, 0) == "trigonal bipyramidal"
-        for i in range(1, len(self.trigonal_bipyramid)):
-            assert site_is_of_motif_type(self.trigonal_bipyramid, i) == "unrecognized"
+        for idx in range(1, len(self.trigonal_bipyramid)):
+            assert site_is_of_motif_type(self.trigonal_bipyramid, idx) == "unrecognized"
 
     def test_get_neighbors_of_site_with_index(self):
         assert len(get_neighbors_of_site_with_index(self.diamond, 0)) == 4
@@ -1003,15 +1003,10 @@ def test_get_order_parameters(self):
             "tet_max",
             "sq_face_cap_trig_pris",
         ]
-        op_params = [None for i in range(len(op_types))]
+        op_params = [None] * len(op_types)
         op_params[1] = {"TA": 1, "IGW_TA": 1.0 / 0.0667}
         op_params[2] = {"TA": 45.0 / 180, "IGW_TA": 1.0 / 0.0667}
-        op_params[33] = {
-            "TA": 0.6081734479693927,
-            "IGW_TA": 18.33,
-            "fac_AA": 1.5,
-            "exp_cos_AA": 2,
-        }
+        op_params[33] = {"TA": 0.6081734479693927, "IGW_TA": 18.33, "fac_AA": 1.5, "exp_cos_AA": 2}
         ops_044 = LocalStructOrderParams(op_types, parameters=op_params, cutoff=0.44)
         ops_071 = LocalStructOrderParams(op_types, parameters=op_params, cutoff=0.71)
         ops_087 = LocalStructOrderParams(op_types, parameters=op_params, cutoff=0.87)

tests/core/test_interface.py

@@ -61,8 +61,8 @@ def test_in_plane_offset_setter(self):
         assert_allclose(interface.in_plane_offset, [0.2, 0.2])
 
         test_coords = np.array(init_coords)
-        for i in interface.film_indices:
-            test_coords[i] += [0.2, 0.2, 0]
+        for idx in interface.film_indices:
+            test_coords[idx] += [0.2, 0.2, 0]
         assert_allclose(np.mod(test_coords, 1.0), np.mod(interface.frac_coords, 1.0))
 
     def test_vacuum_over_film_setter(self):

tests/core/test_lattice.py

@@ -435,8 +435,8 @@ def test_get_distance_and_image(self):
 
     def test_get_distance_and_image_strict(self):
         for _ in range(10):
-            lengths = [np.random.randint(1, 100) for i in range(3)]
-            lattice = [np.random.rand(3) * lengths[i] for i in range(3)]
+            lengths = np.random.randint(1, 100, 3)
+            lattice = np.random.rand(3, 3) * lengths
             lattice = Lattice(lattice)
 
             f1 = np.random.rand(3)

tests/core/test_surface.py

@@ -749,9 +749,9 @@ def test_get_symmetrically_distinct_miller_indices(self):
         assert all(len(hkl) == 4 for hkl in indices)
 
         # Test to see if the output with max_index i is a subset of the output with max_index i+1
-        for i in range(1, 4):
-            assert set(get_symmetrically_distinct_miller_indices(self.trigBi, i)) <= set(
-                get_symmetrically_distinct_miller_indices(self.trigBi, i + 1)
+        for idx in range(1, 4):
+            assert set(get_symmetrically_distinct_miller_indices(self.trigBi, idx)) <= set(
+                get_symmetrically_distinct_miller_indices(self.trigBi, idx + 1)
             )
 
     def test_get_symmetrically_equivalent_miller_indices(self):

tests/core/test_trajectory.py

@@ -71,8 +71,8 @@ def _get_species_and_coords(self):
         return species, coords, 0, 1
 
     def test_single_index_slice(self):
-        assert all(self.traj[i] == self.structures[i] for i in range(0, len(self.structures), 19))
-        assert all(self.traj_mols[i] == self.molecules[i] for i in range(len(self.molecules)))
+        assert all(self.traj[idx] == self.structures[idx] for idx in range(0, len(self.structures), 19))
+        assert all(self.traj_mols[idx] == self.molecules[idx] for idx in range(len(self.molecules)))
 
     def test_slice(self):
         sliced_traj = self.traj[2:99:3]
@@ -87,7 +87,7 @@ def test_slice(self):
         sliced_traj_from_structs = Trajectory.from_structures(self.structures[:-4:2])
 
         if len(sliced_traj) == len(sliced_traj_from_structs):
-            assert all(sliced_traj[i] == sliced_traj_from_structs[i] for i in range(len(sliced_traj)))
+            assert all(sliced_traj[idx] == sliced_traj_from_structs[idx] for idx in range(len(sliced_traj)))
         else:
             raise AssertionError
 
@@ -421,11 +421,11 @@ def test_displacements(self):
         structures = [Structure.from_file(f"{VASP_IN_DIR}/POSCAR")]
         displacements = np.zeros((11, *np.shape(structures[-1].frac_coords)))
 
-        for i in range(10):
+        for idx in range(10):
             displacement = np.random.random_sample(np.shape(structures[-1].frac_coords)) / 20
             new_coords = displacement + structures[-1].frac_coords
             structures.append(Structure(structures[-1].lattice, structures[-1].species, new_coords))
-            displacements[i + 1, :, :] = displacement
+            displacements[idx + 1, :, :] = displacement
 
         traj = Trajectory.from_structures(structures, constant_lattice=True)
         traj.to_displacements()

tests/core/test_units.py

@@ -93,40 +93,37 @@ def test_memory(self):
 
     def test_unitized(self):
         @unitized("eV")
-        def f():
+        def func1():
             return [1, 2, 3]
 
-        assert str(f()[0]) == "1.0 eV"
-        assert isinstance(f(), list)
+        assert str(func1()[0]) == "1.0 eV"
+        assert isinstance(func1(), list)
 
         @unitized("eV")
-        def g():
+        def func2():
             return 2, 3, 4
 
-        assert str(g()[0]) == "2.0 eV"
-        assert isinstance(g(), tuple)
+        assert str(func2()[0]) == "2.0 eV"
+        assert isinstance(func2(), tuple)
 
         @unitized("pm")
-        def h():
-            dct = {}
-            for i in range(3):
-                dct[i] = i * 20
-            return dct
+        def func3():
+            return {idx: idx * 20 for idx in range(3)}
 
-        assert str(h()[1]) == "20.0 pm"
-        assert isinstance(h(), dict)
+        assert str(func3()[1]) == "20.0 pm"
+        assert isinstance(func3(), dict)
 
         @unitized("kg")
-        def i():
+        def func4():
             return FloatWithUnit(5, "g")
 
-        assert i() == FloatWithUnit(0.005, "kg")
+        assert func4() == FloatWithUnit(0.005, "kg")
 
         @unitized("kg")
-        def j():
+        def func5():
             return ArrayWithUnit([5, 10], "g")
 
-        j_out = j()
+        j_out = func5()
         assert j_out.unit == Unit("kg")
         assert j_out[0] == 0.005
         assert j_out[1] == 0.01

tests/electronic_structure/test_boltztrap.py

@@ -93,9 +93,8 @@ def test_get_seebeck_eff_mass(self):
         sbk_mass_avg_mu = self.bz.get_seebeck_eff_mass(output="average", doping_levels=False, temp=300)[3]
         sbk_mass_avg_dop = self.bz.get_seebeck_eff_mass(output="average", doping_levels=True, temp=300)["n"][2]
 
-        for i in range(3):
-            assert sbk_mass_tens_mu[i] == approx(ref2[i], abs=1e-1)
-            assert sbk_mass_tens_dop[i] == approx(ref[i], abs=1e-4)
+        assert sbk_mass_tens_mu == approx(ref2, abs=1e-1)
+        assert sbk_mass_tens_dop == approx(ref, abs=1e-4)
 
         assert sbk_mass_avg_mu == approx(4361.4744008038842, abs=1e-1)
         assert sbk_mass_avg_dop == approx(1.661553842105382, abs=1e-4)
@@ -109,52 +108,40 @@ def test_get_complexity_factor(self):
         sbk_mass_avg_mu = self.bz.get_complexity_factor(output="average", doping_levels=False, temp=300)[3]
         sbk_mass_avg_dop = self.bz.get_complexity_factor(output="average", doping_levels=True, temp=300)["n"][2]
 
-        for i in range(3):
-            assert sbk_mass_tens_mu[i] == approx(ref2[i], abs=1e-4)
-            assert sbk_mass_tens_dop[i] == approx(ref[i], abs=1e-4)
+        assert sbk_mass_tens_mu == approx(ref2, abs=1e-4)
+        assert sbk_mass_tens_dop == approx(ref, abs=1e-4)
 
         assert sbk_mass_avg_mu == approx(0.00628677029221, abs=1e-4)
         assert sbk_mass_avg_dop == approx(1.12322832119, abs=1e-4)
 
     def test_get_seebeck(self):
         ref = [-768.99078999999995, -724.43919999999991, -686.84682999999973]
-        for i in range(3):
-            assert self.bz.get_seebeck()["n"][800][3][i] == approx(ref[i])
+        assert self.bz.get_seebeck()["n"][800][3] == approx(ref)
         assert self.bz.get_seebeck(output="average")["p"][800][3] == approx(697.608936667)
         assert self.bz.get_seebeck(output="average", doping_levels=False)[500][520] == approx(1266.7056)
-        assert self.bz.get_seebeck(output="average", doping_levels=False)[300][65] == approx(
-            -36.2459389333
-        )  # TODO: this was originally "eigs"
+        assert self.bz.get_seebeck(output="average", doping_levels=False)[300][65] == approx(-36.2459389333)
 
     def test_get_conductivity(self):
         ref = [5.9043185000000022, 17.855599000000002, 26.462935000000002]
-        for i in range(3):
-            assert self.bz.get_conductivity()["p"][600][2][i] == approx(ref[i])
+        assert self.bz.get_conductivity()["p"][600][2] == approx(ref)
         assert self.bz.get_conductivity(output="average")["n"][700][1] == approx(1.58736609667)
         assert self.bz.get_conductivity(output="average", doping_levels=False)[300][457] == approx(2.87163566667)
-        assert self.bz.get_conductivity(
-            output="average",
-            doping_levels=False,
-            # TODO: this was originally "eigs"
-            relaxation_time=1e-15,
-        )[200][63] == approx(16573.0536667)
+        assert self.bz.get_conductivity(output="average", doping_levels=False, relaxation_time=1e-15)[200][
+            63
+        ] == approx(16573.0536667)
 
     def test_get_power_factor(self):
         ref = [6.2736602345523362, 17.900184232304138, 26.158282220458144]
-        for i in range(3):
-            assert self.bz.get_power_factor()["p"][200][2][i] == approx(ref[i])
+        assert self.bz.get_power_factor()["p"][200][2] == approx(ref)
         assert self.bz.get_power_factor(output="average")["n"][600][4] == approx(411.230962976)
         assert self.bz.get_power_factor(output="average", doping_levels=False, relaxation_time=1e-15)[500][
             459
         ] == approx(6.59277148467)
-        assert self.bz.get_power_factor(output="average", doping_levels=False)[800][61] == approx(
-            2022.67064134
-        )  # TODO: this was originally "eigs"
+        assert self.bz.get_power_factor(output="average", doping_levels=False)[800][61] == approx(2022.67064134)
 
     def test_get_thermal_conductivity(self):
         ref = [2.7719565628862623e-05, 0.00010048046886793946, 0.00015874549392499391]
-        for i in range(3):
-            assert self.bz.get_thermal_conductivity()["p"][300][2][i] == approx(ref[i])
+        assert self.bz.get_thermal_conductivity()["p"][300][2] == approx(ref)
         assert self.bz.get_thermal_conductivity(output="average", relaxation_time=1e-15)["n"][500][0] == approx(
             1.74466575612e-07
         )
@@ -170,27 +157,23 @@ def test_get_thermal_conductivity(self):
 
     def test_get_zt(self):
         ref = [0.097408810215, 0.29335112354, 0.614673998089]
-        for i in range(3):
-            assert self.bz.get_zt()["n"][800][4][i] == approx(ref[i])
+        assert self.bz.get_zt()["n"][800][4] == approx(ref)
         assert self.bz.get_zt(output="average", k_l=0.5)["p"][700][2] == approx(0.0170001879916)
         assert self.bz.get_zt(output="average", doping_levels=False, relaxation_time=1e-15)[300][240] == approx(
             0.0041923533238348342
         )
 
         eigs = self.bz.get_zt(output="eigs", doping_levels=False)[700][65]
         ref_eigs = [0.082420053399668847, 0.29408035502671648, 0.40822061215079392]
-        for idx, val in enumerate(ref_eigs):
-            assert eigs[idx] == approx(val, abs=1e-5)
+        assert eigs == approx(ref_eigs, abs=1e-5)
 
     def test_get_average_eff_mass(self):
         ref = [0.76045816788363574, 0.96181142990667101, 2.9428428773308628]
-        for i in range(3):
-            assert self.bz.get_average_eff_mass()["p"][300][2][i] == approx(ref[i])
+        assert self.bz.get_average_eff_mass()["p"][300][2] == approx(ref)
         ref = [1.1295783824744523, 1.3898454041924351, 5.2459984671977935]
         ref2 = [6.6648842712692078, 31.492540105738343, 37.986369302138954]
-        for i in range(3):
-            assert self.bz.get_average_eff_mass()["n"][600][1][i] == approx(ref[i])
-            assert self.bz.get_average_eff_mass(doping_levels=False)[300][200][i] == approx(ref2[i])
+        assert self.bz.get_average_eff_mass()["n"][600][1] == approx(ref)
+        assert self.bz.get_average_eff_mass(doping_levels=False)[300][200] == approx(ref2)
         ref = [
             [9.61811430e-01, -8.25159596e-19, -4.70319444e-19],
             [-8.25159596e-19, 2.94284288e00, 3.00368916e-18],
@@ -202,12 +185,8 @@ def test_get_average_eff_mass(self):
             [-1.36897140e-17, 8.74169648e-17, 2.21151980e01],
         ]
 
-        for i in range(3):
-            for j in range(3):
-                assert self.bz.get_average_eff_mass(output="tensor")["p"][300][2][i][j] == approx(ref[i][j], abs=1e-4)
-                assert self.bz.get_average_eff_mass(output="tensor", doping_levels=False)[300][500][i][j] == approx(
-                    ref2[i][j], 4
-                )
+        assert self.bz.get_average_eff_mass(output="tensor")["p"][300][2] == approx(ref, abs=1e-4)
+        assert self.bz.get_average_eff_mass(output="tensor", doping_levels=False)[300][500] == approx(ref2, 4)
         assert self.bz.get_average_eff_mass(output="average")["n"][300][2] == approx(1.53769093989, abs=1e-4)
 
     def test_get_carrier_concentration(self):