JDFTXOutfileSlice Durability Improvement (#4418)

* Protecting the final JOutStructure with try/except so in-progress out files can be parsed

* Enabling partial eigstats filling

* Variety of small changes to improve durability of parsing partially finished out files - can return a usable `JDFTXOutfileSlice` as long as electronic minimize has started

* Removing commented out code

* Removing commented out code

* Removing walk-through-helping lines from test

Author: benrich37

src/pymatgen/io/jdftx/_output_utils.py

@@ -336,6 +336,17 @@ def find_all_key(key_input: str, tempfile: list[str], startline: int = 0) -> lis
     return [i for i in range(startline, len(tempfile)) if key_input in tempfile[i]]
 
 
+def _init_dict_from_colon_dump_lines(lines: list[str]):
+    varsdict = {}
+    for line in lines:
+        if ":" in line:
+            lsplit = line.split(":")
+            key = lsplit[0].strip()
+            val = lsplit[1].split()[0].strip()
+            varsdict[key] = val
+    return varsdict
+
+
 def _parse_bandfile_complex(bandfile_filepath: str | Path) -> NDArray[np.complex64]:
     Dtype: TypeAlias = np.complex64
     token_parser = _complex_token_parser

src/pymatgen/io/jdftx/jdftxoutfileslice.py

@@ -25,6 +25,7 @@
 from pymatgen.core.trajectory import Trajectory
 from pymatgen.core.units import Ha_to_eV, ang_to_bohr, bohr_to_ang
 from pymatgen.io.jdftx._output_utils import (
+    _init_dict_from_colon_dump_lines,
     find_all_key,
     find_first_range_key,
     find_key,
@@ -220,6 +221,7 @@ class JDFTXOutfileSlice:
     _electronic_output: ClassVar[list[str]] = [
         "efermi",
         "egap",
+        "optical_egap",
         "emin",
         "emax",
         "homo",
@@ -230,6 +232,7 @@ class JDFTXOutfileSlice:
     ]
     efermi: float | None = None
     egap: float | None = None
+    optical_egap: float | None = None
     emin: float | None = None
     emax: float | None = None
     homo: float | None = None
@@ -666,22 +669,19 @@ def _get_eigstats_varsdict(self, text: list[str], prefix: str | None) -> dict[st
         lines2 = find_all_key("eigStats' ...", text)
         lines3 = [lines1[i] for i in range(len(lines1)) if lines1[i] in lines2]
         if not lines3:
-            varsdict["emin"] = None
-            varsdict["homo"] = None
-            varsdict["efermi"] = None
-            varsdict["lumo"] = None
-            varsdict["emax"] = None
-            varsdict["egap"] = None
+            for key in list(eigstats_keymap.keys()):
+                varsdict[eigstats_keymap[key]] = None
             self.has_eigstats = False
         else:
-            line = lines3[-1]
-            varsdict["emin"] = float(text[line + 1].split()[1]) * Ha_to_eV
-            varsdict["homo"] = float(text[line + 2].split()[1]) * Ha_to_eV
-            varsdict["efermi"] = float(text[line + 3].split()[2]) * Ha_to_eV
-            varsdict["lumo"] = float(text[line + 4].split()[1]) * Ha_to_eV
-            varsdict["emax"] = float(text[line + 5].split()[1]) * Ha_to_eV
-            varsdict["egap"] = float(text[line + 6].split()[2]) * Ha_to_eV
-            self.has_eigstats = True
+            line_start = lines3[-1]
+            line_start_rel_idx = lines1.index(line_start)
+            line_end = lines1[line_start_rel_idx + 1] if len(lines1) >= line_start_rel_idx + 2 else len(lines1) - 1
+            _varsdict = _init_dict_from_colon_dump_lines([text[idx] for idx in range(line_start, line_end)])
+            for key in _varsdict:
+                varsdict[eigstats_keymap[key]] = float(_varsdict[key]) * Ha_to_eV
+            self.has_eigstats = all(eigstats_keymap[key] in varsdict for key in eigstats_keymap) and all(
+                eigstats_keymap[key] is not None for key in eigstats_keymap
+            )
         return varsdict
 
     def _set_eigvars(self, text: list[str]) -> None:
@@ -691,12 +691,8 @@ def _set_eigvars(self, text: list[str]) -> None:
             text (list[str]): Output of read_file for out file.
         """
         eigstats = self._get_eigstats_varsdict(text, self.prefix)
-        self.emin = eigstats["emin"]
-        self.homo = eigstats["homo"]
-        self.efermi = eigstats["efermi"]
-        self.lumo = eigstats["lumo"]
-        self.emax = eigstats["emax"]
-        self.egap = eigstats["egap"]
+        for key, val in eigstats.items():
+            setattr(self, key, val)
         if self.efermi is None:
             if self.mu is None:
                 self.mu = self._get_mu()
@@ -1063,12 +1059,9 @@ def _set_atom_vars(self, text: list[str]) -> None:
         self.atom_elements = atom_elements
         self.atom_elements_int = [Element(x).Z for x in self.atom_elements]
         self.atom_types = atom_types
-        line = find_key("# Ionic positions in", text)
-        if line is not None:
-            line += 1
-            coords = np.array([text[i].split()[2:5] for i in range(line, line + self.nat)], dtype=float)
-            self.atom_coords_final = coords
-            self.atom_coords = coords.copy()
+        if isinstance(self.structure, Structure):
+            self.atom_coords = self.structure.cart_coords
+            self.atom_coords_final = self.structure.cart_coords
 
     def _set_lattice_vars(self, text: list[str]) -> None:
         """Set the lattice variables.
@@ -1246,6 +1239,17 @@ def __str__(self) -> str:
         return pprint.pformat(self)
 
 
+eigstats_keymap = {
+    "eMin": "emin",
+    "HOMO": "homo",
+    "mu": "efermi",
+    "LUMO": "lumo",
+    "eMax": "emax",
+    "HOMO-LUMO gap": "egap",
+    "Optical gap": "optical_egap",
+}
+
+
 def get_pseudo_read_section_bounds(text: list[str]) -> list[list[int]]:
     """Get the boundary line numbers for the pseudopotential read section.
 

src/pymatgen/io/jdftx/joutstructure.py

@@ -584,7 +584,7 @@ def _parse_lattice_lines(self, lattice_lines: list[str]) -> None:
             order (vec i = r[i,:]) and converts from Bohr to Angstroms.
         """
         r = None
-        if len(lattice_lines):
+        if len(lattice_lines) >= 5:
             r = _brkt_list_of_3x3_to_nparray(lattice_lines, i_start=2)
             r = r.T * bohr_to_ang
             self.lattice = Lattice(np.array(r))
@@ -597,7 +597,7 @@ def _parse_strain_lines(self, strain_lines: list[str]) -> None:
             strain tensor. Converts from column-major to row-major order.
         """
         st = None
-        if len(strain_lines):
+        if len(strain_lines) == 4:
             st = _brkt_list_of_3x3_to_nparray(strain_lines, i_start=1)
             st = st.T
         self.strain = st
@@ -616,7 +616,7 @@ def _parse_stress_lines(self, stress_lines: list[str]) -> None:
         # "[Eh/a0^3]" (Hartree per bohr cubed). Check if this changes for direct
         # coordinates.
         st = None
-        if len(stress_lines):
+        if len(stress_lines) == 4:
             st = _brkt_list_of_3x3_to_nparray(stress_lines, i_start=1)
             st = st.T
             st *= Ha_to_eV / (bohr_to_ang**3)
@@ -636,7 +636,7 @@ def _parse_kinetic_stress_lines(self, stress_lines: list[str]) -> None:
         # "[Eh/a0^3]" (Hartree per bohr cubed). Check if this changes for direct
         # coordinates.
         st = None
-        if len(stress_lines):
+        if len(stress_lines) == 4:
             st = _brkt_list_of_3x3_to_nparray(stress_lines, i_start=1)
             st = st.T
             st *= Ha_to_eV / (bohr_to_ang**3)
@@ -659,6 +659,18 @@ def _parse_thermostat_line(self, posns_lines: list[str]) -> None:
         else:
             self.thermostat_velocity = None
 
+    def _check_for_structure_consistency(self, names: list[str]) -> bool:
+        # If JOutStructure was constructed with a reference init_structure
+        if len(self.species):
+            if len(names) != len(self.species):
+                return False
+            _names = list(set(names))
+            _self_names = [s.symbol for s in self.species]
+            for _name in _names:
+                if names.count(_name) != _self_names.count(_name):
+                    return False
+        return True
+
     def _parse_posns_lines(self, posns_lines: list[str]) -> None:
         """Parse positions lines.
 
@@ -673,8 +685,8 @@ def _parse_posns_lines(self, posns_lines: list[str]) -> None:
             the name of the element, and sd is a flag indicating whether the ion is
             excluded from optimization (1) or not (0).
         """
+        self.copy()
         if len(posns_lines):
-            self.remove_sites(list(range(len(self.species))))
             coords_type = posns_lines[0].split("positions in")[1]
             coords_type = coords_type.strip().split()[0].strip()
             _posns: list[NDArray[np.float64]] = []
@@ -697,6 +709,11 @@ def _parse_posns_lines(self, posns_lines: list[str]) -> None:
                 constraint_types.append(constraint_type)
                 constraint_vectors.append(constraint_vector)
                 group_names_list.append(group_names)
+            is_good = self._check_for_structure_consistency(names)
+            if not is_good and len(self.species):
+                # Abort structure updating if we have a pre-existing structure
+                return
+            self.remove_sites(list(range(len(self.species))))
             posns = np.array(_posns)
             if coords_type.lower() != "cartesian":
                 posns = np.dot(posns, self.lattice.matrix)

src/pymatgen/io/jdftx/joutstructures.py

@@ -384,11 +384,18 @@ def _get_joutstructure_list(
     for i, bounds in enumerate(out_bounds):
         if i > 0:
             init_structure = joutstructure_list[-1]
-        joutstructure_list.append(
-            JOutStructure._from_text_slice(
+        joutstructure = None
+        # The final out_slice slice is protected by the try/except block, as this slice has a high
+        # chance of being empty or malformed.
+        try:
+            joutstructure = JOutStructure._from_text_slice(
                 out_slice[bounds[0] : bounds[1]],
                 init_structure=init_structure,
                 opt_type=opt_type,
             )
-        )
+        except (ValueError, IndexError, TypeError, KeyError, AttributeError):
+            if not i == len(out_bounds) - 1:
+                raise
+        if joutstructure is not None:
+            joutstructure_list.append(joutstructure)
     return joutstructure_list

tests/io/jdftx/test_jdftxoutfileslice.py

@@ -208,6 +208,10 @@ def test_as_dict():
     assert isinstance(out_dict, dict)
 
 
+def should_be_parsable_out_slice(out_slice: list[str]):
+    return any("ElecMinimize: Iter:" in line for line in out_slice[::-1])
+
+
 # Make sure all possible exceptions are caught when none_on_error is True
 @pytest.mark.parametrize(("ex_slice"), [(ex_slice1)])
 def test_none_on_partial(ex_slice: list[str]):
@@ -216,4 +220,7 @@ def test_none_on_partial(ex_slice: list[str]):
     for i in range(int(len(ex_slice) / freq)):
         test_slice = ex_slice[: -(i * freq)]
         joutslice = JDFTXOutfileSlice._from_out_slice(test_slice, none_on_error=True)
-        assert isinstance(joutslice, JDFTXOutfileSlice | None)
+        if should_be_parsable_out_slice(test_slice):
+            assert isinstance(joutslice, JDFTXOutfileSlice | None)
+        else:
+            assert isinstance(joutslice, JDFTXOutfileSlice | None)