More concise test_from_boundary_distance (#3242)

* make test_from_boundary_distance more concise

* snake_case vars

---------

Co-authored-by: Shyue Ping Ong <shyuep@users.noreply.github.com>

pymatgen/analysis/chemenv/utils/scripts_utils.py

@@ -245,23 +245,23 @@ def compute_environments(chemenv_configuration):
                 input_source = input('Enter materials project id (e.g. "mp-1902") : ')
             from pymatgen.ext.matproj import MPRester
 
-            a = MPRester()
-            structure = a.get_structure_by_material_id(input_source)
+            mpr = MPRester()
+            structure = mpr.get_structure_by_material_id(input_source)
         lgf.setup_structure(structure)
         print(f"Computing environments for {structure.composition.reduced_formula} ... ")
         se = lgf.compute_structure_environments(maximum_distance_factor=max_dist_factor)
         print("Computing environments finished")
         while True:
             test = input(
-                "See list of environments determined for each (unequivalent) site ? "
+                "See list of environments determined for each (inequivalent) site ? "
                 '("y" or "n", "d" with details, "g" to see the grid) : '
             )
             strategy = default_strategy
             if test in ["y", "d", "g"]:
                 strategy.set_structure_environments(se)
-                for eqslist in se.equivalent_sites:
-                    site = eqslist[0]
-                    isite = se.structure.index(site)
+                for equiv_list in se.equivalent_sites:
+                    site = equiv_list[0]
+                    site_idx = se.structure.index(site)
                     try:
                         if strategy.uniquely_determines_coordination_environments:
                             ces = strategy.get_site_coordination_environments(site)
@@ -275,34 +275,33 @@ def compute_environments(chemenv_configuration):
                         continue
                     comp = site.species
                     # ce = strategy.get_site_coordination_environment(site)
+                    reduced_formula = comp.get_reduced_formula_and_factor()[0]
                     if strategy.uniquely_determines_coordination_environments:
                         ce = ces[0]
                         if ce is None:
                             continue
-                        thecg = allcg.get_geometry_from_mp_symbol(ce[0])
-                        mystring = (
-                            f"Environment for site #{isite} {comp.get_reduced_formula_and_factor()[0]}"
-                            f" ({comp}) : {thecg.name} ({ce[0]})\n"
+                        the_cg = allcg.get_geometry_from_mp_symbol(ce[0])
+                        msg = (
+                            f"Environment for site #{site_idx} {reduced_formula}"
+                            f" ({comp}) : {the_cg.name} ({ce[0]})\n"
                         )
                     else:
-                        mystring = (
-                            f"Environments for site #{isite} {comp.get_reduced_formula_and_factor()[0]} ({comp}) : \n"
-                        )
+                        msg = f"Environments for site #{site_idx} {reduced_formula} ({comp}) : \n"
                         for ce in ces:
                             cg = allcg.get_geometry_from_mp_symbol(ce[0])
                             csm = ce[1]["other_symmetry_measures"]["csm_wcs_ctwcc"]
-                            mystring += f" - {cg.name} ({cg.mp_symbol}): {ce[2]:.2%} (csm : {csm:2f})\n"
+                            msg += f" - {cg.name} ({cg.mp_symbol}): {ce[2]:.2%} (csm : {csm:2f})\n"
                     if (
                         test in ["d", "g"]
                         and strategy.uniquely_determines_coordination_environments
-                        and thecg.mp_symbol != UNCLEAR_ENVIRONMENT_SYMBOL
+                        and the_cg.mp_symbol != UNCLEAR_ENVIRONMENT_SYMBOL
                     ):
-                        mystring += "  <Continuous symmetry measures>  "
-                        mingeoms = se.ce_list[isite][thecg.coordination_number][0].minimum_geometries()
-                        for mingeom in mingeoms:
-                            csm = mingeom[1]["other_symmetry_measures"]["csm_wcs_ctwcc"]
-                            mystring += f"{mingeom[0]} : {csm:.2f}       "
-                    print(mystring)
+                        msg += "  <Continuous symmetry measures>  "
+                        min_geoms = se.ce_list[site_idx][the_cg.coordination_number][0].minimum_geometries()
+                        for min_geom in min_geoms:
+                            csm = min_geom[1]["other_symmetry_measures"]["csm_wcs_ctwcc"]
+                            msg += f"{min_geom[0]} : {csm:.2f}       "
+                    print(msg)
             if test == "g":
                 while True:
                     test = input(
@@ -312,10 +311,10 @@ def compute_environments(chemenv_configuration):
                     try:
                         indices = [int(x) for x in test.split()]
                         print(str(indices))
-                        for isite in indices:
-                            if isite < 0:
+                        for site_idx in indices:
+                            if site_idx < 0:
                                 raise IndexError
-                            se.plot_environments(isite)
+                            se.plot_environments(site_idx)
                         break
                     except ValueError:
                         print("This is not a valid site")

tests/io/vasp/test_sets.py

@@ -64,6 +64,13 @@
 NO_PSP_DIR = SETTINGS.get("PMG_VASP_PSP_DIR") is None
 skip_if_no_psp_dir = mark.skipif(NO_PSP_DIR, reason="PMG_VASP_PSP_DIR is not set.")
 
+_dummy_structure = Structure(
+    [1, 0, 0, 0, 1, 0, 0, 0, 1],
+    ["I"],
+    [[0, 0, 0]],
+    site_properties={"magmom": [[0, 0, 1]]},
+)
+
 
 @pytest.mark.parametrize(
     "input_set",
@@ -1821,47 +1828,39 @@ def tearDown(self):
 
     def test_ipa(self):
         prev_run = f"{TEST_FILES_DIR}/absorption/static"
-        absorptionipa = MPAbsorptionSet.from_prev_calc(prev_calc_dir=prev_run, copy_wavecar=True, mode="IPA")
-        absorptionipa.write_input(self.tmp)
+        absorption_ipa = MPAbsorptionSet.from_prev_calc(prev_calc_dir=prev_run, copy_wavecar=True, mode="IPA")
+        absorption_ipa.write_input(self.tmp)
         assert os.path.isfile(os.path.join(self.tmp, "WAVECAR"))
-        assert absorptionipa.incar["NBANDS"] == 32
-        assert absorptionipa.incar["ALGO"] == "Exact"
-        assert absorptionipa.incar["LOPTICS"]
+        assert absorption_ipa.incar["NBANDS"] == 32
+        assert absorption_ipa.incar["ALGO"] == "Exact"
+        assert absorption_ipa.incar["LOPTICS"]
 
         # test override_from_prev_calc
-        absorptionipa = MPAbsorptionSet(_dummy_structure, copy_wavecar=True, mode="IPA")
-        absorptionipa.override_from_prev_calc(prev_calc_dir=prev_run)
-        absorptionipa.write_input(self.tmp)
+        absorption_ipa = MPAbsorptionSet(_dummy_structure, copy_wavecar=True, mode="IPA")
+        absorption_ipa.override_from_prev_calc(prev_calc_dir=prev_run)
+        absorption_ipa.write_input(self.tmp)
         assert os.path.isfile(os.path.join(self.tmp, "WAVECAR"))
-        assert absorptionipa.incar["NBANDS"] == 32
-        assert absorptionipa.incar["ALGO"] == "Exact"
-        assert absorptionipa.incar["LOPTICS"]
+        assert absorption_ipa.incar["NBANDS"] == 32
+        assert absorption_ipa.incar["ALGO"] == "Exact"
+        assert absorption_ipa.incar["LOPTICS"]
 
     def test_rpa(self):
         prev_run = f"{TEST_FILES_DIR}/absorption/ipa"
-        absorptionrpa = MPAbsorptionSet.from_prev_calc(prev_run, copy_wavecar=True, mode="RPA")
-        absorptionrpa.write_input(self.tmp)
+        absorption_rpa = MPAbsorptionSet.from_prev_calc(prev_run, copy_wavecar=True, mode="RPA")
+        absorption_rpa.write_input(self.tmp)
         assert os.path.isfile(os.path.join(self.tmp, "WAVECAR"))
         assert os.path.isfile(os.path.join(self.tmp, "WAVEDER"))
-        assert absorptionrpa.incar["NOMEGA"] == 1000
-        assert absorptionrpa.incar["NBANDS"] == 48
-        assert absorptionrpa.incar["ALGO"] == "CHI"
+        assert absorption_rpa.incar["NOMEGA"] == 1000
+        assert absorption_rpa.incar["NBANDS"] == 48
+        assert absorption_rpa.incar["ALGO"] == "CHI"
 
         # test override_from_prev_calc
         prev_run = f"{TEST_FILES_DIR}/absorption/ipa"
-        absorptionrpa = MPAbsorptionSet(_dummy_structure, copy_wavecar=True, mode="RPA")
-        absorptionrpa.override_from_prev_calc(prev_calc_dir=prev_run)
-        absorptionrpa.write_input(self.tmp)
+        absorption_rpa = MPAbsorptionSet(_dummy_structure, copy_wavecar=True, mode="RPA")
+        absorption_rpa.override_from_prev_calc(prev_calc_dir=prev_run)
+        absorption_rpa.write_input(self.tmp)
         assert os.path.isfile(os.path.join(self.tmp, "WAVECAR"))
         assert os.path.isfile(os.path.join(self.tmp, "WAVEDER"))
-        assert absorptionrpa.incar["NOMEGA"] == 1000
-        assert absorptionrpa.incar["NBANDS"] == 48
-        assert absorptionrpa.incar["ALGO"] == "CHI"
-
-
-_dummy_structure = Structure(
-    [1, 0, 0, 0, 1, 0, 0, 0, 1],
-    ["I"],
-    [[0, 0, 0]],
-    site_properties={"magmom": [[0, 0, 1]]},
-)
+        assert absorption_rpa.incar["NOMEGA"] == 1000
+        assert absorption_rpa.incar["NBANDS"] == 48
+        assert absorption_rpa.incar["ALGO"] == "CHI"

tests/transformations/test_standard_transformations.py

@@ -160,57 +160,41 @@ def test_from_scaling_factors(self):
         assert len(struct) == 4 * functools.reduce(lambda a, b: a * b, scale_factors)
 
     def test_from_boundary_distance(self):
-        struct_cubic = Structure.from_spacegroup("Pm-3m", [[4, 0, 0], [0, 4, 0], [0, 0, 4]], ["H"], [[0, 0, 0]])
-        trafo_allow_rotate_true = SupercellTransformation.from_boundary_distance(
-            structure=struct_cubic, min_boundary_dist=9, allow_rotation=True
-        )
-        trafo_allow_rotate_false = SupercellTransformation.from_boundary_distance(
-            structure=struct_cubic, min_boundary_dist=9, allow_rotation=False
-        )
-        struct_super_allow_rotate_true = trafo_allow_rotate_true.apply_transformation(struct_cubic.copy())
-        struct_super_allow_rotate_false = trafo_allow_rotate_false.apply_transformation(struct_cubic.copy())
-        min_expand_allow_rotate_true = np.int8(
-            9 / np.array([struct_super_allow_rotate_true.lattice.d_hkl(plane) for plane in np.eye(3)])
-        )
-        min_expand_allow_rotate_false = np.int8(
-            9 / np.array([struct_super_allow_rotate_false.lattice.d_hkl(plane) for plane in np.eye(3)])
-        )
-        assert len(struct_super_allow_rotate_true) == 14
-        assert len(struct_super_allow_rotate_false) == 27
-        assert np.count_nonzero(min_expand_allow_rotate_true) == 0
-        assert np.count_nonzero(min_expand_allow_rotate_false) == 0
+        struct_cubic = Structure.from_spacegroup("Pm-3m", 4 * np.eye(3), ["H"], [[0, 0, 0]])
 
         for struct in [struct_cubic, self.struct]:
-            for min_boundary_dist in range(5, 16, 2):
-                trafo_allow_rotate_true = SupercellTransformation.from_boundary_distance(
-                    structure=struct, min_boundary_dist=min_boundary_dist, allow_rotation=True
+            for min_dist in range(6, 19, 4):
+                trafo = SupercellTransformation.from_boundary_distance(
+                    structure=struct, min_boundary_dist=min_dist, allow_rotation=False
                 )
-                trafo_allow_rotate_false = SupercellTransformation.from_boundary_distance(
-                    structure=struct, min_boundary_dist=min_boundary_dist, allow_rotation=False
+                trafo_allow_rotate = SupercellTransformation.from_boundary_distance(
+                    structure=struct, min_boundary_dist=min_dist, allow_rotation=True
                 )
-                struct_super_allow_rotate_true = trafo_allow_rotate_true.apply_transformation(struct.copy())
-                struct_super_allow_rotate_false = trafo_allow_rotate_false.apply_transformation(struct.copy())
+                struct_super = trafo.apply_transformation(struct.copy())
+                struct_super_allow_rotate = trafo_allow_rotate.apply_transformation(struct.copy())
+
+                if min_dist == 9 and struct is struct_cubic:
+                    assert len(struct_super_allow_rotate) == 14
+                    assert len(struct_super) == 27
+
                 min_expand_allow_rotate_true = np.int8(
-                    min_boundary_dist
-                    / np.array([struct_super_allow_rotate_true.lattice.d_hkl(plane) for plane in np.eye(3)])
+                    min_dist / np.array([struct_super_allow_rotate.lattice.d_hkl(plane) for plane in np.eye(3)])
                 )
                 min_expand_allow_rotate_false = np.int8(
-                    min_boundary_dist
-                    / np.array([struct_super_allow_rotate_false.lattice.d_hkl(plane) for plane in np.eye(3)])
+                    min_dist / np.array([struct_super.lattice.d_hkl(plane) for plane in np.eye(3)])
                 )
-                assert len(struct_super_allow_rotate_true) <= len(struct_super_allow_rotate_false)
-                assert np.count_nonzero(min_expand_allow_rotate_true) == 0
-                assert np.count_nonzero(min_expand_allow_rotate_false) == 0
 
-        max_atoms, min_boundary_dist = 10, 9
+                assert sum(min_expand_allow_rotate_true != 0) == 0
+                assert sum(min_expand_allow_rotate_false != 0) == 0
+                assert len(struct_super_allow_rotate) <= len(struct_super)
+
+        max_atoms = 10
         with pytest.raises(
             RuntimeError,
             match=f"{max_atoms=} exceeded while trying to solve for supercell. "
-            f"You can try lowering {min_boundary_dist=}",
+            f"You can try lowering min_boundary_dist=6",
         ):
-            SupercellTransformation.from_boundary_distance(
-                structure=struct_cubic, min_boundary_dist=min_boundary_dist, allow_rotation=False, max_atoms=max_atoms
-            )
+            SupercellTransformation.from_boundary_distance(structure=self.struct, max_atoms=max_atoms)
 
 
 class TestOxidationStateDecorationTransformation(unittest.TestCase):