Add PhononDos.mae() and PhononBandStructure.has_imaginary_gamma_freq() methods (#3520)

* fix typos and improve code readability

* clarify has_nac doc str

* add PhononBandStructure.has_imaginary_gamma_freq() to check for imaginary modes at the gamma point

* improve tests for phBS.has_imaginary_freq and add new ones for has_imaginary_gamma_freq

* add PhononDos.mae() method to compare two DOSs

* add TestPhononDos.test_mae()

* increase default has_imaginary_freq tol from 1e-5 to 1e-3

Author: janosh

pymatgen/electronic_structure/bandstructure.py

@@ -779,7 +779,7 @@ def get_equivalent_kpoints(self, index):
         TODO: now it uses the label we might want to use coordinates instead
         (in case there was a mislabel)
         """
-        # if the kpoint has no label it can"t have a repetition along the band
+        # if the kpoint has no label it can't have a repetition along the band
         # structure line object
 
         if self.kpoints[index].label is None:

pymatgen/phonon/bandstructure.py

@@ -24,13 +24,13 @@ def get_reasonable_repetitions(n_atoms: int) -> tuple[int, int, int]:
     according to the number of atoms in the system.
     """
     if n_atoms < 4:
-        return (3, 3, 3)
+        return 3, 3, 3
     if 4 <= n_atoms < 15:
-        return (2, 2, 2)
+        return 2, 2, 2
     if 15 <= n_atoms < 50:
-        return (2, 2, 1)
+        return 2, 2, 1
 
-    return (1, 1, 1)
+    return 1, 1, 1
 
 
 def eigenvectors_from_displacements(disp, masses) -> np.ndarray:
@@ -137,8 +137,8 @@ def __init__(
         self.nb_qpoints = len(self.qpoints)
 
         # normalize directions for nac_frequencies and nac_eigendisplacements
-        self.nac_frequencies = []
-        self.nac_eigendisplacements = []
+        self.nac_frequencies: list[tuple[list[float], np.ndarray]] = []
+        self.nac_eigendisplacements: list[tuple[list[float], np.ndarray]] = []
         if nac_frequencies is not None:
             for freq in nac_frequencies:
                 self.nac_frequencies.append(([idx / np.linalg.norm(freq[0]) for idx in freq[0]], freq[1]))
@@ -152,13 +152,29 @@ def min_freq(self) -> tuple[Kpoint, float]:
 
         return self.qpoints[idx[1]], self.bands[idx]
 
-    def has_imaginary_freq(self, tol: float = 1e-5) -> bool:
-        """True if imaginary frequencies are present in the BS."""
+    def has_imaginary_freq(self, tol: float = 1e-3) -> bool:
+        """True if imaginary frequencies are present anywhere in the band structure. Always True if
+        has_imaginary_gamma_freq is True.
+
+        Args:
+            tol: Tolerance for determining if a frequency is imaginary. Defaults to 1e-3.
+        """
         return self.min_freq()[1] + tol < 0
 
+    def has_imaginary_gamma_freq(self, tol: float = 1e-3) -> bool:
+        """Checks if there are imaginary modes at the gamma point.
+
+        Args:
+            tol: Tolerance for determining if a frequency is imaginary. Defaults to 1e-3.
+        """
+        gamma_freqs = self.bands[:, 0]  # frequencies at the Gamma point
+        return any(freq < -tol for freq in gamma_freqs)
+
     @property
     def has_nac(self) -> bool:
-        """True if nac_frequencies are present."""
+        """True if nac_frequencies are present (i.e. the band structure has been
+        calculated taking into account Born-charge-derived non-analytical corrections at Gamma).
+        """
         return len(self.nac_frequencies) > 0
 
     @property
@@ -177,10 +193,10 @@ def get_nac_frequencies_along_dir(self, direction: Sequence) -> np.ndarray | Non
             the frequencies as a numpy array o(3*len(structure), len(qpoints)).
             None if not found.
         """
-        versor = [i / np.linalg.norm(direction) for i in direction]
-        for d, f in self.nac_frequencies:
-            if np.allclose(versor, d):
-                return f
+        versor = [idx / np.linalg.norm(direction) for idx in direction]
+        for dist, freq in self.nac_frequencies:
+            if np.allclose(versor, dist):
+                return freq
 
         return None
 
@@ -195,10 +211,10 @@ def get_nac_eigendisplacements_along_dir(self, direction) -> np.ndarray | None:
             the eigendisplacements as a numpy array of complex numbers with shape
             (3*len(structure), len(structure), 3). None if not found.
         """
-        versor = [i / np.linalg.norm(direction) for i in direction]
-        for d, e in self.nac_eigendisplacements:
-            if np.allclose(versor, d):
-                return e
+        versor = [idx / np.linalg.norm(direction) for idx in direction]
+        for dist, eigen_disp in self.nac_eigendisplacements:
+            if np.allclose(versor, dist):
+                return eigen_disp
 
         return None
 
@@ -426,45 +442,40 @@ def get_equivalent_qpoints(self, index: int) -> list[int]:
         TODO: now it uses the label we might want to use coordinates instead
         (in case there was a mislabel)
         """
-        # if the qpoint has no label it can"t have a repetition along the band
+        # if the qpoint has no label it can't have a repetition along the band
         # structure line object
 
         if self.qpoints[index].label is None:
             return [index]
 
         list_index_qpoints = []
-        for i in range(self.nb_qpoints):
-            if self.qpoints[i].label == self.qpoints[index].label:
-                list_index_qpoints.append(i)
+        for idx in range(self.nb_qpoints):
+            if self.qpoints[idx].label == self.qpoints[index].label:
+                list_index_qpoints.append(idx)
 
         return list_index_qpoints
 
-    def get_branch(self, index: int) -> list[dict]:
-        r"""Returns in what branch(es) is the qpoint. There can be several
-        branches.
+    def get_branch(self, index: int) -> list[dict[str, str | int]]:
+        r"""Returns in what branch(es) is the qpoint. There can be several branches.
 
         Args:
-            index: the qpoint index
+            index (int): the qpoint index
 
         Returns:
-            A list of dictionaries [{"name","start_index","end_index","index"}]
-            indicating all branches in which the qpoint is. It takes into
-            account the fact that one qpoint (e.g., \\Gamma) can be in several
-            branches
+            list[dict[str, str | int]]: [{"name","start_index","end_index","index"}]
+                indicating all branches in which the qpoint is. It takes into
+                account the fact that one qpoint (e.g., \\Gamma) can be in several
+                branches
         """
-        to_return = []
-        for i in self.get_equivalent_qpoints(index):
-            for b in self.branches:
-                if b["start_index"] <= i <= b["end_index"]:
-                    to_return.append(
-                        {
-                            "name": b["name"],
-                            "start_index": b["start_index"],
-                            "end_index": b["end_index"],
-                            "index": i,
-                        }
+        lst = []
+        for pt_idx in self.get_equivalent_qpoints(index):
+            for branch in self.branches:
+                start_idx, end_idx = branch["start_index"], branch["end_index"]
+                if start_idx <= pt_idx <= end_idx:
+                    lst.append(
+                        {"name": branch["name"], "start_index": start_idx, "end_index": end_idx, "index": pt_idx}
                     )
-        return to_return
+        return lst
 
     def write_phononwebsite(self, filename: str | PathLike) -> None:
         """Write a json file for the phononwebsite:
@@ -606,13 +617,13 @@ def from_dict(cls, dct: dict) -> PhononBandStructureSymmLine:
         eigendisplacements = (
             np.array(dct["eigendisplacements"]["real"]) + np.array(dct["eigendisplacements"]["imag"]) * 1j
         )
-        structure = Structure.from_dict(dct["structure"]) if "structure" in dct else None
+        struct = Structure.from_dict(dct["structure"]) if "structure" in dct else None
         return cls(
             dct["qpoints"],
             np.array(dct["bands"]),
             lattice_rec,
             dct["has_nac"],
             eigendisplacements,
             dct["labels_dict"],
-            structure=structure,
+            structure=struct,
         )

pymatgen/phonon/dos.py

@@ -329,6 +329,28 @@ def zero_point_energy(self, structure: Structure | None = None) -> float:
 
         return zpe
 
+    def mae(self, other: PhononDos, two_sided: bool = True) -> float:
+        """Mean absolute error between two DOSs.
+
+        Args:
+            other: Another DOS object.
+            two_sided: Whether to calculate the two-sided MAE meaning interpolate each DOS to the
+                other's frequencies and averaging the two MAEs. Defaults to True.
+
+        Returns:
+            float: Mean absolute error.
+        """
+        # Interpolate other.densities to align with self.frequencies
+        self_interpolated = np.interp(self.frequencies, other.frequencies, other.densities)
+        self_mae = np.abs(self.densities - self_interpolated).mean()
+
+        if two_sided:
+            other_interpolated = np.interp(other.frequencies, self.frequencies, self.densities)
+            other_mae = np.abs(other.densities - other_interpolated).mean()
+            return (self_mae + other_mae) / 2
+
+        return self_mae
+
 
 class CompletePhononDos(PhononDos):
     """This wrapper class defines a total dos, and also provides a list of PDos.

pymatgen/phonon/plotter.py

@@ -285,6 +285,7 @@ def _make_ticks(self, ax: Axes) -> Axes:
             ax.set_xticks(ticks_labels[0])
             ax.set_xticklabels(ticks_labels[1])
 
+        # plot vertical lines at each of the ticks
         for idx, label in enumerate(ticks["label"]):
             if label is not None:
                 ax.axvline(ticks["distance"][idx], color="black")

tests/phonon/test_bandstructure.py

@@ -38,15 +38,27 @@ def test_basic(self):
 
         assert list(self.bs.min_freq()[0].frac_coords) == [0, 0, 0]
         assert self.bs.min_freq()[1] == approx(-0.03700895020)
-        assert self.bs.has_imaginary_freq()
-        assert not self.bs.has_imaginary_freq(tol=0.5)
         assert_allclose(self.bs.asr_breaking(), [-0.0370089502, -0.0370089502, -0.0221388897])
 
         assert self.bs.nb_bands == 6
         assert self.bs.nb_qpoints == 204
 
         assert_allclose(self.bs.qpoints[1].frac_coords, [0.01, 0, 0])
 
+    def test_has_imaginary_freq(self):
+        for tol in (0, 0.01, 0.02, 0.03, 0.04, 0.05):
+            assert self.bs.has_imaginary_freq(tol=tol) == (tol < 0.04)
+
+        for tol in (0, 0.01, 0.02, 0.03, 0.04, 0.05):
+            assert self.bs2.has_imaginary_freq(tol=tol) == (tol < 0.01)
+
+        # test Gamma point imaginary frequency detection
+        for tol in (0, 0.01, 0.02, 0.03, 0.04, 0.05):
+            assert self.bs.has_imaginary_gamma_freq(tol=tol) == (tol < 0.04)
+
+        for tol in (0, 0.01, 0.02, 0.03, 0.04, 0.05):
+            assert self.bs2.has_imaginary_gamma_freq(tol=tol) == (tol < 0.01)
+
     def test_nac(self):
         assert self.bs.has_nac
         assert not self.bs2.has_nac

tests/phonon/test_dos.py

@@ -3,6 +3,7 @@
 import json
 import re
 
+import pytest
 from pytest import approx
 
 from pymatgen.core import Element
@@ -85,6 +86,19 @@ def test_eq(self):
         assert self.dos != 2 * self.dos
         assert 2 * self.dos == self.dos + self.dos
 
+    def test_mae(self):
+        assert self.dos.mae(self.dos) == 0
+        assert self.dos.mae(self.dos + 1) == 1
+        assert self.dos.mae(self.dos - 1) == 1
+        assert self.dos.mae(2 * self.dos) == pytest.approx(0.786546967)
+        assert (2 * self.dos).mae(self.dos) == pytest.approx(0.786546967)
+
+        # test two_sided=False after shifting DOS freqs so MAE requires interpolation
+        dos2 = PhononDos(self.dos.frequencies + 0.01, self.dos.densities)
+        assert self.dos.mae(dos2 + 1, two_sided=False) == pytest.approx(0.999999999)
+        assert self.dos.mae(dos2 - 1, two_sided=False) == pytest.approx(1.00000000000031)
+        assert self.dos.mae(2 * dos2, two_sided=False) == pytest.approx(0.786546967)
+
 
 class TestCompletePhononDos(PymatgenTest):
     def setUp(self):