bandstructure attribute for JDFTOutputs - initializes to BandStructure if successfully stores 'eigenvals' and 'kpts'. Initializes projections in BandStructure if successfully stores 'bandprojections'. (#4413)

Author: benrich37

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/outputs.py

@@ -16,15 +16,24 @@ class is written.
 import numpy as np
 from monty.dev import deprecated
 
+from pymatgen.core import Lattice
+from pymatgen.core.units import Ha_to_eV
+from pymatgen.electronic_structure.bandstructure import BandStructure
+from pymatgen.electronic_structure.core import Spin
 from pymatgen.io.jdftx._output_utils import (
     _get_nbands_from_bandfile_filepath,
     _get_orb_label_list,
+    _get_u_to_oa_map,
+    _parse_kptsfrom_bandprojections_file,
     get_proj_tju_from_file,
+    orb_ref_list,
     read_outfile_slices,
 )
 from pymatgen.io.jdftx.jdftxoutfileslice import JDFTXOutfileSlice
 
 if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
     from pymatgen.core.structure import Structure
     from pymatgen.core.trajectory import Trajectory
     from pymatgen.io.jdftx.inputs import JDFTXInfile
@@ -68,6 +77,8 @@ class is written.
 implemented_store_vars = (
     "bandProjections",
     "eigenvals",
+    "kpts",
+    "bandstructure",
 )
 
 
@@ -104,17 +115,22 @@ class JDFTXOutputs:
             0-based index will appear mimicking a principle quantum number (ie "0px" for first shell and "1px" for
             second shell). The actual principal quantum number is not stored in the JDFTx output files and must be
             inferred by the user.
+        kpts (list[np.ndarray]): A list of the k-points used in the calculation. Each k-point is a 3D numpy array.
+        wk_list (list[np.ndarray]): A list of the weights for the k-points used in the calculation.
     """
 
     calc_dir: str | Path = field(init=True)
     outfile_name: str | Path | None = field(init=True)
     store_vars: list[str] = field(default_factory=list, init=True)
     paths: dict[str, Path] = field(init=False)
     outfile: JDFTXOutfile = field(init=False)
-    bandProjections: np.ndarray | None = field(init=False)
-    eigenvals: np.ndarray | None = field(init=False)
+    bandProjections: NDArray | None = field(init=False)
+    eigenvals: NDArray | None = field(init=False)
+    kpts: list[NDArray] | None = field(init=False)
+    wk_list: list[NDArray] | None = field(init=False)
     # Misc metadata for interacting with the data
     orb_label_list: tuple[str, ...] | None = field(init=False)
+    bandstructure: BandStructure | None = field(init=False)
 
     @classmethod
     def from_calc_dir(
@@ -136,13 +152,21 @@ def from_calc_dir(
         Returns:
             JDFTXOutputs: The JDFTXOutputs object.
         """
-        if store_vars is None:
-            store_vars = []
+        store_vars = cls._check_store_vars(store_vars)
         return cls(calc_dir=Path(calc_dir), store_vars=store_vars, outfile_name=outfile_name)
 
     def __post_init__(self):
         self._init_paths()
         self._store_vars()
+        self._init_bandstructure()
+
+    def _check_store_vars(store_vars: list[str] | None) -> list[str]:
+        if store_vars is None:
+            return []
+        if "bandstructure" in store_vars:
+            store_vars += ["kpts", "eigenvals"]
+            store_vars.pop(store_vars.index("bandstructure"))
+        return list(set(store_vars))
 
     def _init_paths(self):
         self.paths = {}
@@ -208,7 +232,101 @@ def _store_eigenvals(self):
         if "eigenvals" in self.paths:
             self.eigenvals = np.fromfile(self.paths["eigenvals"])
             nstates = int(len(self.eigenvals) / self.outfile.nbands)
-            self.eigenvals = self.eigenvals.reshape(nstates, self.outfile.nbands)
+            self.eigenvals = self.eigenvals.reshape(nstates, self.outfile.nbands) * Ha_to_eV
+
+    def _check_kpts(self):
+        if "bandProjections" in self.paths and self.paths["bandProjections"].exists():
+            # TODO: Write kpt file inconsistency checking
+            return
+        if "kPts" in self.paths and self.paths["kPts"].exists():
+            raise NotImplementedError("kPts file parsing not yet implemented.")
+        raise RuntimeError("No k-point data found in JDFTx output files.")
+
+    def _store_kpts(self):
+        if "bandProjections" in self.paths and self.paths["bandProjections"].exists():
+            wk_list, kpts_list = _parse_kptsfrom_bandprojections_file(self.paths["bandProjections"])
+            self.kpts = kpts_list
+            self.wk_list = wk_list
+
+    def _init_bandstructure(self):
+        if True in [v is None for v in [self.kpts, self.eigenvals]]:
+            return
+        kpoints = np.array(self.kpts)
+        eigenvals = self._get_pmg_eigenvals()
+        projections = None
+        if self.bandProjections is not None:
+            projections = self._get_pmg_projections()
+        self.bandstructure = BandStructure(
+            kpoints,
+            eigenvals,
+            Lattice(self.outfile.structure.lattice.reciprocal_lattice.matrix),
+            self.outfile.mu,
+            projections=projections,
+            structure=self.outfile.structure,
+        )
+
+    def _get_lmax(self) -> tuple[str | None, int | None]:
+        """Get the maximum l quantum number and projection array orbital length.
+
+        Returns:
+            tuple[str | None, int | None]: The maximum l quantum number and projection array orbital length.
+                Both are None if no bandProjections file is available.
+        """
+        if self.orb_label_list is None:
+            return None, None
+        orbs = [label.split("(")[1].split(")")[0] for label in self.orb_label_list]
+        if orb_ref_list[-1][0] in orbs:
+            return "f", 16
+        if orb_ref_list[-2][0] in orbs:
+            return "d", 9
+        if orb_ref_list[-3][0] in orbs:
+            return "p", 4
+        if orb_ref_list[-4][0] in orbs:
+            return "s", 1
+        raise ValueError("Unrecognized orbital labels in orb_label_list.")
+
+    def _get_pmg_eigenvals(self) -> dict | None:
+        if self.eigenvals is None:
+            return None
+        _e_skj = self.eigenvals.copy().reshape(self.outfile.nspin, -1, self.outfile.nbands)
+        _e_sjk = np.swapaxes(_e_skj, 1, 2)
+        spins = [Spin.up, Spin.down]
+        eigenvals = {}
+        for i in range(self.outfile.nspin):
+            eigenvals[spins[i]] = _e_sjk[i]
+        return eigenvals
+
+    def _get_pmg_projections(self) -> dict | None:
+        """Return pymatgen-compatible projections dictionary.
+
+        Converts the bandProjections array to a pymatgen-compatible dictionary
+
+        Returns:
+            dict | None:
+        """
+        lmax, norbmax = self._get_lmax()
+        if norbmax is None:
+            return None
+        if self.orb_label_list is None:
+            return None
+        _proj_tju = self.bandProjections.copy()
+        if _proj_tju.dtype is np.complex64:
+            # Convert <orb|band(state)> to |<orb|band(state)>|^2
+            _proj_tju = np.abs(_proj_tju) ** 2
+        # Convert to standard datatype - using np.real to suppress warnings
+        proj_tju = np.array(np.real(_proj_tju), dtype=float)
+        proj_skju = proj_tju.reshape([self.outfile.nspin, -1, self.outfile.nbands, len(self.orb_label_list)])
+        proj_sjku = np.swapaxes(proj_skju, 1, 2)
+        nspin, nbands, nkpt, nproj = proj_sjku.shape
+        u_to_oa_map = _get_u_to_oa_map(self.paths["bandProjections"])
+        projections = {}
+        spins = [Spin.up, Spin.down]
+        for i in range(nspin):
+            projections[spins[i]] = np.zeros([nbands, nkpt, norbmax, len(self.outfile.structure)])
+            # TODO: Consider jitting this loop
+            for u in range(nproj):
+                projections[spins[i]][:, :, *u_to_oa_map[u]] += proj_sjku[i, :, :, u]
+        return projections
 
 
 _jof_atr_from_last_slice = (

tests/io/jdftx/test_jdftxoutput.py

@@ -19,6 +19,8 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
+_implemented_store_vars = [v for v in implemented_store_vars if v != "bandstructure"]
+
 
 @pytest.mark.parametrize(
     ("calc_dir", "known_paths"),
@@ -34,6 +36,7 @@ def test_known_paths(calc_dir: Path, known_paths: dict):
 @pytest.mark.parametrize(
     ("calc_dir", "store_vars"),
     [
+        (n2_ex_calc_dir, ["bandProjections", "eigenvals", "kpts"]),
         (n2_ex_calc_dir, ["bandProjections", "eigenvals"]),
         (n2_ex_calc_dir, ["bandProjections"]),
         (n2_ex_calc_dir, []),
@@ -42,14 +45,35 @@ def test_known_paths(calc_dir: Path, known_paths: dict):
 def test_store_vars(calc_dir: Path, store_vars: list[str]):
     """Test that the stored variables are correct."""
     jo = JDFTXOutputs.from_calc_dir(calc_dir, store_vars=store_vars)
-    for var in implemented_store_vars:
+    for var in _implemented_store_vars:
         assert hasattr(jo, var)
         if var in store_vars:
             assert getattr(jo, var) is not None
         else:
             assert getattr(jo, var) is None
 
 
+@pytest.mark.parametrize(
+    ("calc_dir", "store_vars"),
+    [
+        (n2_ex_calc_dir, ["bandProjections", "eigenvals", "kpts", "bandstructure"]),
+        (n2_ex_calc_dir, ["bandProjections", "bandstructure"]),
+        (n2_ex_calc_dir, ["bandstructure"]),
+    ],
+)
+def test_bandstructure_construction(calc_dir: Path, store_vars: list[str]):
+    """Test that the stored variables are correct."""
+    jo = JDFTXOutputs.from_calc_dir(calc_dir, store_vars=store_vars)
+    assert hasattr(jo, "bandstructure")
+    required_for_bandstructure = ["eigenvals", "kpts"]
+    for var in required_for_bandstructure:
+        assert hasattr(jo, var)
+    if "bandProjections" in store_vars:
+        assert len(jo.bandstructure.projections)
+    else:
+        assert not len(jo.bandstructure.projections)
+
+
 @pytest.mark.parametrize(
     ("calc_dir", "known_metadata"),
     [