Merge branch 'master' into update_vasp_outputs

Author: Keian Noori

src/pymatgen/core/periodic_table.py

@@ -236,7 +236,7 @@ def __eq__(self, other: object) -> bool:
 
     def __hash__(self) -> int:
         # multiply Z by 1000 to avoid hash collisions of element N with isotopes of elements N+/-1,2,3...
-        return self.Z * 1000 + self.A if self._is_named_isotope else self.Z
+        return self.Z * 1000 + self.A if self._is_named_isotope else self.Z * 137 * 100
 
     def __repr__(self) -> str:
         return f"Element {self.symbol}"
@@ -1605,7 +1605,7 @@ def get_el_sp(obj: SpeciesLike) -> Element | Species | DummySpecies:
     pass
 
 
-@functools.lru_cache
+@functools.lru_cache(maxsize=1024)
 def get_el_sp(obj: int | SpeciesLike) -> Element | Species | DummySpecies:
     """Utility function to get an Element, Species or DummySpecies from any input.
 

src/pymatgen/io/jdftx/_output_utils.py

@@ -15,6 +15,8 @@
 
 import numpy as np
 
+from pymatgen.electronic_structure.core import Orbital
+
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -476,6 +478,34 @@ def get_proj_tju_from_file(bandfile_filepath: Path | str) -> NDArray[np.float32
     return _parse_bandfile_complex(bandfile_filepath) if is_complex else _parse_bandfile_normalized(bandfile_filepath)
 
 
+def _parse_kptsfrom_bandprojections_file(bandfile_filepath: str | Path) -> tuple[list[float], list[NDArray]]:
+    """Parse kpts from bandprojections file.
+
+    Parse kpts from bandprojections file.
+
+    Args:
+        bandfile_filepath (Path | str): Path to bandprojections file.
+
+    Returns:
+        tuple[list[float], list[np.ndarray[float]]]: Tuple of k-point weights and k-points
+    """
+    wk_list: list[float] = []
+    k_points_list: list[NDArray] = []
+    kpt_lines = []
+    with open(bandfile_filepath) as f:
+        for line in f:
+            if line.startswith("#") and ";" in line:
+                _line = line.split(";")[0].lstrip("#")
+                kpt_lines.append(_line)
+    for line in kpt_lines:
+        k_points = line.split("[")[1].split("]")[0].strip().split()
+        _k_points_floats: list[float] = [float(v) for v in k_points]
+        k_points_list.append(np.array(_k_points_floats))
+        wk = float(line.split("]")[1].strip().split()[0])
+        wk_list.append(wk)
+    return wk_list, k_points_list
+
+
 def _is_complex_bandfile_filepath(bandfile_filepath: str | Path) -> bool:
     """Determine if bandprojections file is complex.
 
@@ -507,6 +537,64 @@ def _is_complex_bandfile_filepath(bandfile_filepath: str | Path) -> bool:
     ["dxy", "dyz", "dz2", "dxz", "dx2-y2"],
     ["fy(3x2-y2)", "fxyz", "fyz2", "fz3", "fxz2", "fz(x2-y2)", "fx(x2-3y2)"],
 ]
+orb_ref_to_o_dict = {
+    "s": int(Orbital.s),
+    "py": int(Orbital.py),
+    "pz": int(Orbital.pz),
+    "px": int(Orbital.px),
+    "dxy": int(Orbital.dxy),
+    "dyz": int(Orbital.dyz),
+    "dz2": int(Orbital.dz2),
+    "dxz": int(Orbital.dxz),
+    "dx2-y2": int(Orbital.dx2),
+    # Keep the f-orbitals arbitrary-ish until they get designated names in pymatgen.
+    orb_ref_list[-1][0]: int(Orbital.f_3),
+    orb_ref_list[-1][1]: int(Orbital.f_2),
+    orb_ref_list[-1][2]: int(Orbital.f_1),
+    orb_ref_list[-1][3]: int(Orbital.f0),
+    orb_ref_list[-1][4]: int(Orbital.f1),
+    orb_ref_list[-1][5]: int(Orbital.f2),
+}
+
+
+def _get_atom_orb_labels_map_dict(bandfile_filepath: Path) -> dict[str, list[str]]:
+    """
+    Return a dictionary mapping each atom symbol to pymatgen-compatible orbital projection string representations.
+
+    Identical to _get_atom_orb_labels_ref_dict, but doesn't include the numbers in the labels.
+
+
+
+    Args:
+        bandfile_filepath (str | Path): The path to the bandfile.
+
+    Returns:
+        dict[str, list[str]]: A dictionary mapping each atom symbol to all atomic orbital projection string
+        representations.
+    """
+    bandfile = read_file(bandfile_filepath)
+    labels_dict: dict[str, list[str]] = {}
+
+    for i, line in enumerate(bandfile):
+        if i > 1:
+            if "#" in line:
+                break
+            lsplit = line.strip().split()
+            sym = lsplit[0]
+            labels_dict[sym] = []
+            lmax = int(lsplit[3])
+            # Would prefer to use "l" rather than "L" here (as uppercase "L" means something else entirely) but
+            # pr*-c*mm*t thinks "l" is an ambiguous variable name.
+            for L in range(lmax + 1):
+                mls = orb_ref_list[L]
+                nshells = int(lsplit[4 + L])
+                for _n in range(nshells):
+                    if nshells > 1:
+                        for ml in mls:
+                            labels_dict[sym].append(f"{ml}")
+                    else:
+                        labels_dict[sym] += mls
+    return labels_dict
 
 
 def _get_atom_orb_labels_ref_dict(bandfile_filepath: Path) -> dict[str, list[str]]:
@@ -623,6 +711,28 @@ def _get_orb_label(ion: str, idx: int, orb: str) -> str:
     return f"{ion}#{idx + 1}({orb})"
 
 
+def _get_u_to_oa_map(bandfile_filepath: Path) -> list[tuple[int, int]]:
+    """
+    Return a list, where the u'th element is a tuple of the atomic orbital index and the ion index.
+
+    Args:
+        bandfile_filepath (str | Path): The path to the bandfile.
+
+    Returns:
+        list[tuple[int, int]]: A list, where the u'th element is a tuple of the atomic orbital index and the ion index.
+    """
+    map_labels_dict = _get_atom_orb_labels_map_dict(bandfile_filepath)
+    atom_count_list = _get_atom_count_list(bandfile_filepath)
+    u_to_oa_map = []
+    a = 0
+    for ion, ion_count in atom_count_list:
+        for _i in range(ion_count):
+            for orb in map_labels_dict[ion]:
+                u_to_oa_map.append((orb_ref_to_o_dict[orb], a))
+            a += 1
+    return u_to_oa_map
+
+
 def _get_orb_label_list(bandfile_filepath: Path) -> tuple[str, ...]:
     """
     Return a tuple of all atomic orbital projection string representations.

src/pymatgen/io/jdftx/generic_tags.py

@@ -70,6 +70,56 @@ def validate_value_type(self, tag: str, value: Any, try_auto_type_fix: bool = Fa
             tuple[str, bool, Any]: The tag, whether the value is of the correct type, and the possibly fixed value.
         """
 
+    def is_equal_to(self, val1: Any | list[Any], obj2: AbstractTag, val2: Any | list[Any]) -> bool:
+        """Check if the two values are equal.
+
+        Args:
+            val1 (Any): The value of this tag object.
+            obj2 (AbstractTag): The other tag object.
+            val2 (Any): The value of the other tag object.
+
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+        if self.can_repeat:
+            if not obj2.can_repeat:
+                return False
+            val1 = val1 if isinstance(val1, list) else [val1]
+            val2 = val2 if isinstance(val2, list) else [val2]
+            if len(val1) != len(val2):
+                return False
+            return all(True in [self._is_equal_to(v1, obj2, v2) for v2 in val2] for v1 in val1)
+        return self._is_equal_to(val1, obj2, val2)
+
+    @abstractmethod
+    def _is_equal_to(self, val1: Any, obj2: AbstractTag, val2: Any) -> bool:
+        """Check if the two values are equal.
+
+        Used to check if the two values are equal. Assumes val1 and val2 are single elements.
+
+        Args:
+            val1 (Any): The value of this tag object.
+            obj2 (AbstractTag): The other tag object.
+            val2 (Any): The value of the other tag object.
+
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+
+    def _is_same_tagtype(
+        self,
+        obj2: AbstractTag,
+    ) -> bool:
+        """Check if the two values are equal.
+
+        Args:
+            obj2 (AbstractTag): The other tag object.
+
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+        return isinstance(self, type(obj2))
+
     def _validate_value_type(
         self, type_check: type, tag: str, value: Any, try_auto_type_fix: bool = False
     ) -> tuple[str, bool, Any]:
@@ -258,6 +308,19 @@ def validate_value_type(self, tag: str, value: Any, try_auto_type_fix: bool = Fa
         """
         return self._validate_value_type(bool, tag, value, try_auto_type_fix=try_auto_type_fix)
 
+    def _is_equal_to(self, val1: Any, obj2: AbstractTag, val2: Any) -> bool:
+        """Check if the two values are equal.
+
+        Args:
+            val1 (Any): The value of this tag object.
+            obj2 (AbstractTag): The other tag object.
+            val2 (Any): The value of the other tag object.
+
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+        return self._is_same_tagtype(obj2) and val1 == val2
+
     def raise_value_error(self, tag: str, value: str) -> None:
         """Raise a ValueError for the value string.
 
@@ -335,6 +398,23 @@ def validate_value_type(self, tag: str, value: Any, try_auto_type_fix: bool = Fa
         """
         return self._validate_value_type(str, tag, value, try_auto_type_fix=try_auto_type_fix)
 
+    def _is_equal_to(self, val1: Any, obj2: AbstractTag, val2: Any) -> bool:
+        """Check if the two values are equal.
+
+        Args:
+            val1 (Any): The value of this tag object.
+            obj2 (AbstractTag): The other tag object.
+            val2 (Any): The value of the other tag object.
+
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+        if self._is_same_tagtype(obj2):
+            if not all(isinstance(x, str) for x in (val1, val2)):
+                raise ValueError("Both values must be strings for StrTag comparison")
+            return val1.strip() == val2.strip()
+        return False
+
     def read(self, tag: str, value: str) -> str:
         """Read the value string for this tag.
 
@@ -379,6 +459,8 @@ class AbstractNumericTag(AbstractTag):
     ub: float | None = None  # upper bound
     lb_incl: bool = True  # lower bound inclusive
     ub_incl: bool = True  # upper bound inclusive
+    eq_atol: float = 1.0e-8  # absolute tolerance for equality check
+    eq_rtol: float = 1.0e-5  # relative tolerance for equality check
 
     def val_is_within_bounds(self, value: float) -> bool:
         """Check if the value is within the bounds.
@@ -425,6 +507,22 @@ def validate_value_bounds(
             return False, self.get_invalid_value_error_str(tag, value)
         return True, ""
 
+    def _is_equal_to(self, val1, obj2, val2):
+        """Check if the two values are equal.
+
+        Used to check if the two values are equal. Doesn't need to be redefined for IntTag and FloatTag.
+
+        Args:
+            val1 (Any): The value of this tag object.
+            obj2 (AbstractTag): The other tag object.
+            val2 (Any): The value of the other tag object.
+            rtol (float, optional): Relative tolerance. Defaults to 1.e-5.
+            atol (float, optional): Absolute tolerance. Defaults to 1.e-8.
+        Returns:
+            bool: True if the two tag object/value pairs are equal, False otherwise.
+        """
+        return self._is_same_tagtype(obj2) and np.isclose(val1, val2, rtol=self.eq_rtol, atol=self.eq_atol)
+
 
 @dataclass
 class IntTag(AbstractNumericTag):
@@ -620,6 +718,10 @@ def get_token_len(self) -> int:
         """
         return self._get_token_len()
 
+    def _is_equal_to(self, val1, obj2, val2):
+        return True  # TODO: We still need to actually implement initmagmom as a multi-format tag
+        # raise NotImplementedError("equality not yet implemented for InitMagMomTag")
+
 
 @dataclass
 class TagContainer(AbstractTag):
@@ -1013,6 +1115,28 @@ def get_dict_representation(self, tag: str, value: list) -> dict | list[dict]:
         list_value = self._make_str_for_dict(tag, value)
         return self.read(tag, list_value)
 
+    def _is_equal_to(self, val1, obj2, val2):
+        """Check if the two values are equal.
+
+        Return False if (checked in following order)
+        - obj2 is not a TagContainer
+        - all of val1's subtags are not in val2
+        - val1 and val2 are not the same length (different number of subtags)
+        - at least one subtag in val1 is not equal to the corresponding subtag in val2
+        """
+        if self._is_same_tagtype(obj2):
+            if isinstance(val1, dict) and isinstance(val2, dict):
+                if all(subtag in val2 for subtag in val1) and (len(list(val1.keys())) == len(list(val2.keys()))):
+                    for subtag, subtag_type in self.subtags.items():
+                        if (subtag in val1) and (
+                            not subtag_type.is_equal_to(val1[subtag], obj2.subtags[subtag], val2[subtag])
+                        ):
+                            return False
+                    return True
+                return False
+            raise ValueError("Values must be in dictionary format for TagContainer comparison")
+        return False
+
 
 # TODO: Write StructureDefferedTagContainer back in (commented out code block removed
 # on 11/4/24) and make usable for tags like initial-magnetic-moments
@@ -1162,6 +1286,9 @@ def get_token_len(self) -> int:
         """
         raise NotImplementedError("This method is not supposed to be called directly on MultiformatTag objects!")
 
+    def _is_equal_to(self, val1, obj2, val2):
+        raise NotImplementedError("This method is not supposed to be called directly on MultiformatTag objects!")
+
 
 @dataclass
 class BoolTagContainer(TagContainer):