tests/io/aims use numpy.testing.assert_allclose and pytest.MonkeyPatch (#3575)

* tests/io/aims use numpy.testing.assert_allclose for better error msgs and pytest.MonkeyPatch for auto-reverted side-effects

* remove superfluous np.all() in asserts

* fix pytest.raises match regex escape

* rename d->dct

Author: janosh

examples/aims_io/FHI-aims-example.ipynb

@@ -7,15 +7,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from pymatgen.io.aims.inputs import AimsGeometryIn, AimsCube, AimsControlIn\n",
-    "from pymatgen.io.aims.outputs import AimsOutput\n",
-    "\n",
-    "from pymatgen.core import Structure, Lattice\n",
+    "from pathlib import Path\n",
+    "from subprocess import check_call\n",
     "\n",
     "import numpy as np\n",
+    "from numpy.testing import assert_allclose\n",
     "\n",
-    "from pathlib import Path\n",
-    "from subprocess import check_call\n"
+    "from pymatgen.core import Lattice, Structure\n",
+    "from pymatgen.io.aims.inputs import AimsControlIn, AimsCube, AimsGeometryIn\n",
+    "from pymatgen.io.aims.outputs import AimsOutput"
    ]
   },
   {
@@ -29,7 +29,7 @@
     "AIMS_CMD = \"aims.x\"\n",
     "AIMS_OUTPUT = \"aims.out\"\n",
     "AIMS_SD = \"species_dir\"\n",
-    "AIMS_TEST_DIR = \"../../tests/io/aims/species_directory/light/\"\n"
+    "AIMS_TEST_DIR = \"../../tests/io/aims/species_directory/light/\""
    ]
   },
   {
@@ -41,12 +41,10 @@
    "source": [
     "# Create test structure\n",
     "structure = Structure(\n",
-    "    lattice=Lattice(\n",
-    "        np.array([[0, 2.715, 2.715],[2.715, 0, 2.715], [2.715, 2.715, 0]])\n",
-    "    ),\n",
+    "    lattice=Lattice(np.array([[0, 2.715, 2.715], [2.715, 0, 2.715], [2.715, 2.715, 0]])),\n",
     "    species=[\"Si\", \"Si\"],\n",
-    "    coords=np.array([np.zeros(3), np.ones(3) *  0.25])\n",
-    ")\n"
+    "    coords=np.array([np.zeros(3), np.ones(3) * 0.25]),\n",
+    ")"
    ]
   },
   {
@@ -78,9 +76,9 @@
     "\n",
     "# Cube file output controlled by the AimsCube class\n",
     "cont_in[\"cubes\"] = [\n",
-    "    AimsCube(\"total_density\", origin=[0,0,0], points=[11, 11, 11]),\n",
+    "    AimsCube(\"total_density\", origin=[0, 0, 0], points=[11, 11, 11]),\n",
     "    AimsCube(\"eigenstate_density 1\", origin=[0, 0, 0], points=[11, 11, 11]),\n",
-    "]\n"
+    "]"
    ]
   },
   {
@@ -91,11 +89,11 @@
    "outputs": [],
    "source": [
     "# Write the input files\n",
-    "workdir = Path.cwd() / \"workdir/\"\n",
-    "workdir.mkdir(exist_ok=True)\n",
+    "work_dir = Path.cwd() / \"workdir/\"\n",
+    "work_dir.mkdir(exist_ok=True)\n",
     "\n",
-    "geo_in.write_file(workdir, overwrite=True)\n",
-    "cont_in.write_file(structure, workdir, overwrite=True)\n"
+    "geo_in.write_file(work_dir, overwrite=True)\n",
+    "cont_in.write_file(structure, work_dir, overwrite=True)"
    ]
   },
   {
@@ -106,8 +104,8 @@
    "outputs": [],
    "source": [
     "# Run the calculation\n",
-    "with open(f\"{workdir}/{AIMS_OUTPUT}\", \"w\") as outfile:\n",
-    "    aims_run = check_call([AIMS_CMD], cwd=workdir, stdout=outfile)\n"
+    "with open(f\"{work_dir}/{AIMS_OUTPUT}\", \"w\") as outfile:\n",
+    "    aims_run = check_call([AIMS_CMD], cwd=work_dir, stdout=outfile)"
    ]
   },
   {
@@ -118,22 +116,14 @@
    "outputs": [],
    "source": [
     "# Read the aims output file and the final relaxed geometry\n",
-    "outputs = AimsOutput.from_outfile(f\"{workdir}/{AIMS_OUTPUT}\")\n",
-    "relaxed_structure = AimsGeometryIn.from_file(f\"{workdir}/geometry.in.next_step\")\n",
+    "outputs = AimsOutput.from_outfile(f\"{work_dir}/{AIMS_OUTPUT}\")\n",
+    "relaxed_structure = AimsGeometryIn.from_file(f\"{work_dir}/geometry.in.next_step\")\n",
     "\n",
     "# Check the results\n",
     "assert outputs.get_results_for_image(-1).lattice == relaxed_structure.structure.lattice\n",
-    "assert np.all(outputs.get_results_for_image(-1).frac_coords == relaxed_structure.structure.frac_coords)\n",
-    "\n",
-    "assert np.allclose(\n",
-    "    outputs.get_results_for_image(-1).properties[\"stress\"],\n",
-    "    outputs.stress\n",
-    ")\n",
-    "\n",
-    "assert np.allclose(\n",
-    "    outputs.get_results_for_image(-1).site_properties[\"force\"],\n",
-    "    outputs.forces\n",
-    ")\n"
+    "assert_allclose(outputs.get_results_for_image(-1).frac_coords, relaxed_structure.structure.frac_coords)\n",
+    "assert_allclose(outputs.get_results_for_image(-1).properties[\"stress\"], outputs.stress)\n",
+    "assert_allclose(outputs.get_results_for_image(-1).site_properties[\"force\"], outputs.forces)"
    ]
   }
  ],
@@ -153,7 +143,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,

pymatgen/analysis/ferroelectricity/polarization.py

@@ -104,7 +104,7 @@ def get_total_ionic_dipole(structure, zval_dict):
 
 
 class PolarizationLattice(Structure):
-    """Why is a Lattice inheriting a structure? This is ridiculous."""
+    """TODO Why is a Lattice inheriting a structure? This is ridiculous."""
 
     def get_nearest_site(self, coords, site, r=None):
         """
@@ -144,14 +144,7 @@ class Polarization:
     electron Angstroms along the three lattice directions (a,b,c).
     """
 
-    def __init__(
-        self,
-        p_elecs,
-        p_ions,
-        structures,
-        p_elecs_in_cartesian=True,
-        p_ions_in_cartesian=False,
-    ):
+    def __init__(self, p_elecs, p_ions, structures, p_elecs_in_cartesian=True, p_ions_in_cartesian=False):
         """
         p_elecs: np.array of electronic contribution to the polarization with shape [N, 3]
         p_ions: np.array of ionic contribution to the polarization with shape [N, 3]
@@ -270,7 +263,7 @@ def get_same_branch_polarization_data(self, convert_to_muC_per_cm2=True, all_in_
         lattices = [s.lattice for s in self.structures]
         volumes = np.array([s.lattice.volume for s in self.structures])
 
-        L = len(p_elec)
+        n_elecs = len(p_elec)
 
         e_to_muC = -1.6021766e-13
         cm2_to_A2 = 1e16
@@ -282,38 +275,41 @@ def get_same_branch_polarization_data(self, convert_to_muC_per_cm2=True, all_in_
             # Convert the total polarization
             p_tot = np.multiply(units.T[:, np.newaxis], p_tot)
             # adjust lattices
-            for i in range(L):
-                lattice = lattices[i]
-                lattices[i] = Lattice.from_parameters(*(np.array(lattice.lengths) * units.ravel()[i]), *lattice.angles)
+            for idx in range(n_elecs):
+                lattice = lattices[idx]
+                lattices[idx] = Lattice.from_parameters(
+                    *(np.array(lattice.lengths) * units.ravel()[idx]), *lattice.angles
+                )
         #  convert polarizations to polar lattice
         elif convert_to_muC_per_cm2 and all_in_polar:
             abc = [lattice.abc for lattice in lattices]
             abc = np.array(abc)  # [N, 3]
             p_tot /= abc  # e * Angstroms to e
             p_tot *= abc[-1] / volumes[-1] * e_to_muC * cm2_to_A2  # to muC / cm^2
-            for i in range(L):
+            for idx in range(n_elecs):
                 lattice = lattices[-1]  # Use polar lattice
                 # Use polar units (volume)
-                lattices[i] = Lattice.from_parameters(*(np.array(lattice.lengths) * units.ravel()[-1]), *lattice.angles)
+                lattices[idx] = Lattice.from_parameters(
+                    *(np.array(lattice.lengths) * units.ravel()[-1]), *lattice.angles
+                )
 
         d_structs = []
         sites = []
-        for i in range(L):
-            lattice = lattices[i]
-            frac_coord = np.divide(np.array([p_tot[i]]), np.array(lattice.lengths))
+        for idx in range(n_elecs):
+            lattice = lattices[idx]
+            frac_coord = np.divide(np.array([p_tot[idx]]), np.array(lattice.lengths))
             d = PolarizationLattice(lattice, ["C"], [np.array(frac_coord).ravel()])
             d_structs.append(d)
             site = d[0]
             # Adjust nonpolar polarization to be closest to zero.
             # This is compatible with both a polarization of zero or a half quantum.
-            prev_site = [0, 0, 0] if i == 0 else sites[-1].coords
+            prev_site = [0, 0, 0] if idx == 0 else sites[-1].coords
             new_site = d.get_nearest_site(prev_site, site)
             sites.append(new_site[0])
 
         adjust_pol = []
-        for s, d in zip(sites, d_structs):
-            lattice = d.lattice
-            adjust_pol.append(np.multiply(s.frac_coords, np.array(lattice.lengths)).ravel())
+        for site, d in zip(sites, d_structs):
+            adjust_pol.append(np.multiply(site.frac_coords, np.array(d.lattice.lengths)).ravel())
         return np.array(adjust_pol)
 
     def get_lattice_quanta(self, convert_to_muC_per_cm2=True, all_in_polar=True):
@@ -324,7 +320,7 @@ def get_lattice_quanta(self, convert_to_muC_per_cm2=True, all_in_polar=True):
         lattices = [s.lattice for s in self.structures]
         volumes = np.array([struct.volume for struct in self.structures])
 
-        L = len(self.structures)
+        n_structs = len(self.structures)
 
         e_to_muC = -1.6021766e-13
         cm2_to_A2 = 1e16
@@ -334,13 +330,17 @@ def get_lattice_quanta(self, convert_to_muC_per_cm2=True, all_in_polar=True):
         # convert polarizations and lattice lengths prior to adjustment
         if convert_to_muC_per_cm2 and not all_in_polar:
             # adjust lattices
-            for i in range(L):
-                lattice = lattices[i]
-                lattices[i] = Lattice.from_parameters(*(np.array(lattice.lengths) * units.ravel()[i]), *lattice.angles)
+            for idx in range(n_structs):
+                lattice = lattices[idx]
+                lattices[idx] = Lattice.from_parameters(
+                    *(np.array(lattice.lengths) * units.ravel()[idx]), *lattice.angles
+                )
         elif convert_to_muC_per_cm2 and all_in_polar:
-            for i in range(L):
+            for idx in range(n_structs):
                 lattice = lattices[-1]
-                lattices[i] = Lattice.from_parameters(*(np.array(lattice.lengths) * units.ravel()[-1]), *lattice.angles)
+                lattices[idx] = Lattice.from_parameters(
+                    *(np.array(lattice.lengths) * units.ravel()[-1]), *lattice.angles
+                )
 
         return np.array([np.array(latt.lengths) for latt in lattices])
 

pymatgen/analysis/local_env.py

@@ -2286,12 +2286,12 @@ def __init__(self, types, parameters=None, cutoff=-10.0) -> None:
 
         self._comp_azi = False
         self._params = []
-        for i, typ in enumerate(self._types):
-            d = deepcopy(default_op_params[typ]) if default_op_params[typ] is not None else None
-            if parameters is None or parameters[i] is None:
-                self._params.append(d)
+        for idx, typ in enumerate(self._types):
+            dct = deepcopy(default_op_params[typ]) if default_op_params[typ] is not None else None
+            if parameters is None or parameters[idx] is None:
+                self._params.append(dct)
             else:
-                self._params.append(deepcopy(parameters[i]))
+                self._params.append(deepcopy(parameters[idx]))
 
         self._computerijs = self._computerjks = self._geomops = False
         self._geomops2 = self._boops = False
@@ -3923,20 +3923,20 @@ def get_nn_data(self, structure: Structure, n: int, length=None):
             for entry in nn:
                 r2 = _get_radius(entry["site"])
                 if r1 > 0 and r2 > 0:
-                    d = r1 + r2
+                    dist = r1 + r2
                 else:
                     warnings.warn(
                         "CrystalNN: cannot locate an appropriate radius, "
                         "covalent or atomic radii will be used, this can lead "
                         "to non-optimal results."
                     )
-                    d = _get_default_radius(structure[n]) + _get_default_radius(entry["site"])
+                    dist = _get_default_radius(structure[n]) + _get_default_radius(entry["site"])
 
                 dist = np.linalg.norm(structure[n].coords - entry["site"].coords)
                 dist_weight: float = 0
 
-                cutoff_low = d + self.distance_cutoffs[0]
-                cutoff_high = d + self.distance_cutoffs[1]
+                cutoff_low = dist + self.distance_cutoffs[0]
+                cutoff_high = dist + self.distance_cutoffs[1]
 
                 if dist <= cutoff_low:
                     dist_weight = 1

pymatgen/analysis/molecule_matcher.py

@@ -402,9 +402,9 @@ def _align_heavy_atoms(mol1, mol2, vmol1, vmol2, ilabel1, ilabel2, eq_atoms):
                 a2 = aligned_mol2.GetAtom(c2)
                 for c1 in candidates1:
                     a1 = canon_mol1.GetAtom(c1)
-                    d = a1.GetDistance(a2)
-                    if d < distance:
-                        distance = d
+                    dist = a1.GetDistance(a2)
+                    if dist < distance:
+                        distance = dist
                         canon_idx = c1
                 canon_label2[c2 - 1] = canon_idx
                 candidates1.remove(canon_idx)
@@ -462,9 +462,9 @@ def _align_hydrogen_atoms(mol1, mol2, heavy_indices1, heavy_indices2):
             a2 = cmol2.GetAtom(h2)
             for h1 in hydrogen_label1:
                 a1 = cmol1.GetAtom(h1)
-                d = a1.GetDistance(a2)
-                if d < distance:
-                    distance = d
+                dist = a1.GetDistance(a2)
+                if dist < distance:
+                    distance = dist
                     idx = h1
             hydrogen_label2.append(idx)
             hydrogen_label1.remove(idx)
@@ -828,11 +828,11 @@ def kabsch(P: np.ndarray, Q: np.ndarray):
         V, _S, WT = np.linalg.svd(C)
 
         # Getting the sign of the det(V*Wt) to decide whether
-        d = np.linalg.det(np.dot(V, WT))
+        det = np.linalg.det(np.dot(V, WT))
 
         # And finally calculating the optimal rotation matrix R
         # we need to correct our rotation matrix to ensure a right-handed coordinate system.
-        return np.dot(np.dot(V, np.diag([1, 1, d])), WT)
+        return np.dot(np.dot(V, np.diag([1, 1, det])), WT)
 
 
 class BruteForceOrderMatcher(KabschMatcher):

pymatgen/analysis/piezo_sensitivity.py

@@ -377,9 +377,9 @@ def get_unstable_FCM(self, max_force=1):
                 D[3 * op[1] : 3 * op[1] + 3, 3 * op[0] : 3 * op[0] + 3] = np.zeros([3, 3])
 
                 for symop in op[4]:
-                    tempfcm = D[3 * op[2] : 3 * op[2] + 3, 3 * op[3] : 3 * op[3] + 3]
-                    tempfcm = symop.transform_tensor(tempfcm)
-                    D[3 * op[0] : 3 * op[0] + 3, 3 * op[1] : 3 * op[1] + 3] += tempfcm
+                    temp_fcm = D[3 * op[2] : 3 * op[2] + 3, 3 * op[3] : 3 * op[3] + 3]
+                    temp_fcm = symop.transform_tensor(temp_fcm)
+                    D[3 * op[0] : 3 * op[0] + 3, 3 * op[1] : 3 * op[1] + 3] += temp_fcm
 
                 if len(op[4]) != 0:
                     D[3 * op[0] : 3 * op[0] + 3, 3 * op[1] : 3 * op[1] + 3] = D[

pymatgen/io/aims/inputs.py

@@ -306,14 +306,13 @@ def __post_init__(self) -> None:
             ValueError: If any of the inputs is invalid
         """
         split_type = self.type.split()
+        cube_type = split_type[0]
         if split_type[0] in ALLOWED_AIMS_CUBE_TYPES:
             if len(split_type) > 1:
-                msg = f"Cube of type {split_type[0]} can not have a state associated with it"
-                raise ValueError(msg)
+                raise ValueError(f"{cube_type=} can not have a state associated with it")
         elif split_type[0] in ALLOWED_AIMS_CUBE_TYPES_STATE:
             if len(split_type) != 2:
-                msg = f"Cube of type {split_type[0]} must have a state associated with it"
-                raise ValueError(msg)
+                raise ValueError(f"{cube_type=} must have a state associated with it")
         else:
             raise ValueError("Cube type undefined")
 

pymatgen/io/vasp/outputs.py

@@ -1931,7 +1931,7 @@ def __init__(self, filename):
                 self.read_pattern({key: rf"{key}\s+=\s+([\d\-\.]+)"})
             if not self.data[key]:
                 continue
-            final_energy_contribs[key] = sum(float(f) for f in self.data[key][-1])
+            final_energy_contribs[key] = sum(map(float, self.data[key][-1]))
         self.final_energy_contribs = final_energy_contribs
 
     def read_pattern(self, patterns, reverse=False, terminate_on_match=False, postprocess=str):
@@ -2053,7 +2053,7 @@ def read_electrostatic_potential(self):
 
         pattern = {"radii": r"the test charge radii are((?:\s+[\.\-\d]+)+)"}
         self.read_pattern(pattern, reverse=True, terminate_on_match=True, postprocess=str)
-        self.sampling_radii = [float(f) for f in self.data["radii"][0][0].split()]
+        self.sampling_radii = [*map(float, self.data["radii"][0][0].split())]
 
         header_pattern = r"\(the norm of the test charge is\s+[\.\-\d]+\)"
         table_pattern = r"((?:\s+\d+\s*[\.\-\d]+)+)"
@@ -2064,7 +2064,7 @@ def read_electrostatic_potential(self):
 
         pots = re.findall(r"\s+\d+\s*([\.\-\d]+)+", pots)
 
-        self.electrostatic_potential = [float(f) for f in pots]
+        self.electrostatic_potential = [*map(float, pots)]
 
     @staticmethod
     def _parse_sci_notation(line):