Merge pull request #110 from benrich37/faster-parsing

Misc edits

Author: benrich37

src/pymatgen/io/jdftx/inputs.py

@@ -21,7 +21,14 @@ class is written.
 from pymatgen.core import Lattice, Structure
 from pymatgen.core.periodic_table import Element
 from pymatgen.core.units import bohr_to_ang
-from pymatgen.io.jdftx.generic_tags import AbstractTag, BoolTagContainer, DumpTagContainer, MultiformatTag, TagContainer
+from pymatgen.io.jdftx.generic_tags import (
+    AbstractTag,
+    BoolTagContainer,
+    DumpTagContainer,
+    FloatTag,
+    MultiformatTag,
+    TagContainer,
+)
 from pymatgen.io.jdftx.jdftxinfile_default_inputs import default_inputs
 from pymatgen.io.jdftx.jdftxinfile_master_format import (
     __PHONON_TAGS__,
@@ -45,14 +52,15 @@ class is written.
 
 __author__ = "Jacob Clary, Ben Rich"
 
+
 # TODO: Add check for whether all ions have or lack velocities.
 # TODO: Add default value filling like JDFTx does.
 # TODO: Add more robust checking for if two repeatable tag values represent the
 # same information. This is likely fixed by implementing filling of default values.
 # TODO: Incorporate something to collapse repeated dump tags of the same frequency
 # into a single value.
-
-
+# TODO: Add a method to strip all tags that contain their default values for simpler
+# files written (especially when a `JDFTXInfile` is created from `JDFTXOutfileSlice`)
 class JDFTXInfile(dict, MSONable):
     """Class for reading/writing JDFtx input files.
 
@@ -270,6 +278,32 @@ def from_jdftxstructure(
         jstr = jdftxstructure.get_str()
         return cls.from_str(jstr)
 
+    def read_line(
+        self,
+        line: str,
+        validate_value_boundaries: bool = True,
+        autofix: bool = True,
+        overwrite_nonrepeatable: bool = True,
+    ) -> None:
+        """Read a single line and update the JDFTXInfile object.
+
+        Convenience method for reading a single line and updating the JDFTXInfile object.
+
+        Args:
+            line (str): Line to read.
+        """
+        line = line.strip()
+        tag_object, tag, value = self._preprocess_line(line)
+        if not tag_object.can_repeat and overwrite_nonrepeatable and tag in self:
+            del self[tag]
+        processed_value = tag_object.read(tag, value)
+        _params = self.as_dict(skip_module_keys=True)
+        _params = self._store_value(_params, tag_object, tag, processed_value)
+        self.update(_params)
+        self.validate_tags(try_auto_type_fix=autofix, error_on_failed_fix=True)
+        if validate_value_boundaries:
+            self.validate_boundaries()
+
     @classmethod
     def from_str(
         cls,
@@ -397,7 +431,9 @@ def copy(self) -> JDFTXInfile:
         Returns:
             JDFTXInfile: Copy of the JDFTXInfile object.
         """
-        return type(self)(self)
+        # Wasn't working before
+        # return type(self)(self)
+        return self.from_dict(self.as_dict(skip_module_keys=True), validate_value_boundaries=False)
 
     def get_text_list(self) -> list[str]:
         """Get a list of strings representation of the JDFTXInfile.
@@ -426,14 +462,24 @@ def get_text_list(self) -> list[str]:
                 text.append("")
         return text
 
-    def write_file(self, filename: PathLike) -> None:
+    # TODO: JDFTXInfile can accept nan for values, as this is occasionally what is stored
+    # for unused variables, but JDFTx has no way read nan for an input value. All subtags
+    # with nan values should be removed before writing to file.
+    # TODO: Detect for and warn for tags that can be used together but likely shouldn't be,
+    # ie (ion-width being 0 while fluid is not None)
+    def write_file(self, filename: PathLike, strip_nan: bool = False) -> None:
         """Write JDFTXInfile to an in file.
 
         Args:
             filename (PathLike): Filename to write to.
+            strip_nan (bool, optional): Whether to strip all subtags with nan values before writing.
+                                         Defaults to False. WARNING - VERY JANKY RIGHT NOW
         """
+        write_infile = self
+        if strip_nan:
+            write_infile = clean_infile_of_nans(self)
         with open(filename, mode="w") as file:
-            file.write(str(self))
+            file.write(str(write_infile))
 
     @classmethod
     def to_jdftxstructure(
@@ -912,6 +958,107 @@ def movescale_array_to_selective_dynamics_site_prop(movescale: ArrayLike[int | f
     return selective_dynamics
 
 
+# def _strip_nans(infile: JDFTXInfile) -> JDFTXInfile:
+#     for k, v in infile.items():
+#         tag_object = get_tag_object_on_val(k, v)
+#         if isinstance(v, dict):
+#             infile[k] = _strip_nans(v)
+#         if isinstance(v, float) and np.isnan(v):
+#             infile[k] = None
+#         elif isinstance(v, list):
+#             infile[k] = [x for x in v if not (isinstance(x, float) and np.isnan(x))]
+#         elif isinstance(v, dict):
+#             infile[k] = _strip_nans(v)
+
+# def _has_nans(value: float | list) -> bool:
+#     if isinstance(value, float) and np.isnan(value):
+#         return True
+#     elif isinstance(value, list):
+#         return any(_has_nans(x) for x in value)
+#     return False
+
+
+def _isnan(x):
+    try:
+        return np.isnan(x)
+    except TypeError:
+        return False
+
+
+def _check_tagcontainer_for_nan(tag_container: TagContainer, val_dict: dict):
+    hasnans = []
+    for kk, vv in val_dict.items():
+        tag = tag_container.subtags[kk]
+        if not isinstance(tag, TagContainer):
+            if _isnan(vv):
+                print(f"Tag {kk} has nan value")
+                hasnans.append(kk)
+        elif not tag.can_repeat:
+            _hasnans = _check_tagcontainer_for_nan(tag, vv)
+            if len(_hasnans) > 0:
+                hasnans.append({kk: _hasnans})
+    return hasnans
+
+
+def has_nan_in_required_subtag(tag_container: TagContainer, val_dict: dict):
+    for kk, vv in val_dict.items():
+        tag = tag_container.subtags[kk]
+        if not isinstance(tag, TagContainer):
+            if _isnan(vv) and not tag.optional:
+                return True
+        elif not tag.can_repeat and has_nan_in_required_subtag(tag, vv):
+            return True
+    return False
+
+
+def clean_tagcontainer_of_nans(tag_container: TagContainer, val_dict: dict):
+    subtags_to_delete = []
+    for kk, vv in val_dict.items():
+        tag = tag_container.subtags[kk]
+        if not isinstance(tag, TagContainer):
+            if _isnan(vv):
+                print(f"Removing tag {kk} with nan value")
+                subtags_to_delete.append(kk)
+        elif not tag.can_repeat:
+            if has_nan_in_required_subtag(tag, vv) and tag_container.optional:
+                subtags_to_delete.append(kk)
+            else:
+                clean_tagcontainer_of_nans(tag, vv)
+                if len(tag.subtags) == 0:
+                    print(f"Removing empty tag container {kk}")
+                    subtags_to_delete.append(kk)
+            clean_tagcontainer_of_nans(tag, vv)
+    for kk in subtags_to_delete:
+        val_dict.pop(kk)
+
+
+def clean_infile_of_nans(infile: JDFTXInfile) -> JDFTXInfile:
+    hasnans = []
+    for k, v in infile.items():
+        tag = get_tag_object_on_val(k, v)
+        if isinstance(tag, FloatTag) and _isnan(v):
+            print(f"Tag {k} has nan value")
+        elif isinstance(tag, TagContainer) and not tag.can_repeat:
+            _hasnans = _check_tagcontainer_for_nan(tag, v)
+            if len(_hasnans) > 0:
+                hasnans.append({k: _hasnans})
+    infile_cleaned = JDFTXInfile.from_dict(infile.as_dict())
+
+    # infile_cleaned = JDFTXInfile(infile)
+
+    for h in hasnans:
+        if isinstance(h, str):
+            infile_cleaned.pop(h)
+        elif isinstance(h, dict):
+            for _h in h:
+                tag = get_tag_object_on_val(_h, infile_cleaned[_h])
+                if _isnan(infile_cleaned[_h]):
+                    infile_cleaned.pop(_h)
+                else:
+                    clean_tagcontainer_of_nans(tag, infile_cleaned[_h])
+    return infile_cleaned
+
+
 @dataclass
 class JDFTXStructure(MSONable):
     """Object for representing the data in JDFTXStructure tags.

src/pymatgen/io/jdftx/jdftxinfile_default_inputs.py

@@ -7,6 +7,7 @@
 # One example is fluid-minimize, which has different convergence thresholds and max iterations depending
 # on the algorithm specified. For these tags, a second set of default values which can map partially
 # filled tagcontainers to the set as filled by JDFTx is needed.
+# TODO: Make sure ion-width, which changes based on 'fluid' tag, is handled correctly.
 default_inputs = {
     "basis": "kpoint-dependent",
     "coords-type": "Lattice",

src/pymatgen/io/jdftx/jdftxinfile_master_format.py

@@ -574,7 +574,7 @@
                     write_tagname=False,
                     optional=False,
                 ),
-                "origin": TagContainer(
+                "center": TagContainer(
                     allow_list_representation=True,
                     optional=True,
                     subtags={

src/pymatgen/io/jdftx/jdftxoutfileslice.py

@@ -291,7 +291,7 @@ class JDFTXOutfileSlice:
     converged: bool | None = None
     structure: Structure | None = None
     initial_structure: Structure | None = None
-    trajectory: Trajectory | None = None
+    _trajectory: Trajectory | None = None
     electronic_output: dict | None = None
     eopt_type: str | None = None
     elecmindata: JElSteps | None = None
@@ -336,10 +336,10 @@ def _get_mu(self) -> None | float:
     def _from_out_slice(
         cls,
         text: list[str],
+        skim_levels: list[str],
+        skip_props: list[str],
         is_bgw: bool = False,
         none_on_error: bool = True,
-        skim_levels: list[str] | None = None,
-        skip_props: list[str] | None = None,
     ) -> JDFTXOutfileSlice | None:
         """
         Read slice of out file into a JDFTXOutfileSlice instance.
@@ -356,16 +356,14 @@ def _from_out_slice(
         instance = cls()
         instance.is_bgw = is_bgw
         try:
-            instance._from_out_slice_init_all(text, skim_levels=skim_levels, skip_props=skip_props)
+            instance._from_out_slice_init_all(text, skim_levels, skip_props)
         except (ValueError, IndexError, TypeError, KeyError, AttributeError):
             if none_on_error:
                 return None
             raise
         return instance
 
-    def _from_out_slice_init_all(
-        self, text: list[str], skim_levels: list[str] | None = None, skip_props: list[str] | None = None
-    ) -> None:
+    def _from_out_slice_init_all(self, text: list[str], skim_levels: list[str], skip_props: list[str]) -> None:
         self._set_internal_infile(text)
         # self._set_min_settings(text)
         self._set_geomopt_vars(text)
@@ -470,12 +468,23 @@ def _set_trajectory(self) -> None:
             structures = [slc.structure for slc in self.jstrucs]
             constant_lattice = self.constant_lattice if self.constant_lattice is not None else False
             frame_properties = [slc.properties for slc in self.jstrucs]
-            self.trajectory = Trajectory.from_structures(
+            self._trajectory = Trajectory.from_structures(
                 structures=structures,
                 constant_lattice=constant_lattice,
                 frame_properties=frame_properties,
             )
 
+    @property
+    def trajectory(self) -> Trajectory | None:
+        """Return pymatgen trajectory object.
+
+        Returns:
+            Trajectory: pymatgen Trajectory object containing intermediate Structure's of outfile slice calculation.
+        """
+        if self._trajectory is None:
+            self._set_trajectory()
+        return self._trajectory
+
     def _set_electronic_output(self) -> None:
         """Return a dictionary with all relevant electronic information.
 
@@ -884,9 +893,7 @@ def _get_initial_structure(self, text: list[str]) -> Structure | None:
             raise ValueError("Provided out file slice's inputs preamble does not contain input structure data.")
         return init_struc
 
-    def _set_jstrucs(
-        self, text: list[str], skim_levels: list[str] | None = None, skip_props: list[str] | None = None
-    ) -> None:
+    def _set_jstrucs(self, text: list[str], skim_levels: list[str], skip_props: list[str]) -> None:
         """Set the jstrucs class variable.
 
         Set the JStructures object to jstrucs from the out file text and all class attributes initialized from jstrucs.
@@ -903,17 +910,19 @@ def _set_jstrucs(
             expected_etype = "G"
         self.jstrucs = JOutStructures._from_out_slice(
             text,
+            # skim_levels,
+            # skip_props,
             opt_type=self.geom_opt_label,
             init_struc=self.initial_structure,
             is_md=self.is_md,
             expected_etype=expected_etype,
             skim_levels=skim_levels,
             skip_props=skip_props,
         )
-        if self.etype is None:
-            self.etype = self.jstrucs[-1].etype
         if self.jstrucs is not None:
-            self._set_trajectory()
+            if self.etype is None:
+                self.etype = self.jstrucs[-1].etype
+            # self._set_trajectory()
             self.mu = self._get_mu()
             for var in _jofs_atr_from_jstrucs:
                 setattr(self, var, getattr(self.jstrucs, var))

src/pymatgen/io/jdftx/joutstructure.py

@@ -108,7 +108,7 @@ class JOutStructure(Structure):
     elec_grad_k: float | np.float64 | None = None
     elec_alpha: float | np.float64 | None = None
     elec_linmin: float | np.float64 | None = None
-    structure: Structure | None = None
+    _structure: Structure | None = None
     is_md: bool = False
     # thermostat_velocity: np.ndarray | None = None
     _velocities: list[NDArray[np.float64] | None] | None = None
@@ -386,7 +386,6 @@ def _from_text_slice(
             instance._parse_lattice_lines(line_collections["lattice"]["lines"])
             # Posns must be parsed before forces and lowdin analysis so that they can be stored in site_properties
             cur_species = instance._parse_posns_lines(line_collections["posns"]["lines"], cur_species)
-
         if "forces" not in skip_props:
             instance._parse_forces_lines(line_collections["forces"]["lines"])
         if "lowdin" not in skip_props:
@@ -406,9 +405,16 @@ def _from_text_slice(
         # Set relevant properties in self.properties
         instance._fill_properties()
         # Done last in case of any changes to site-properties
-        instance._init_structure(cur_species)
+        # instance._init_structure(cur_species)
         return instance
 
+    @property
+    def structure(self) -> Structure:
+        """Return structure attribute."""
+        if self._structure is None:
+            self._init_structure()
+        return self._structure
+
     def _init_e_sp_backup(self) -> None:
         """Initialize self.e with coverage for single-point calculations."""
         err_str = None
@@ -967,11 +973,11 @@ def _fill_properties(self) -> None:
             "strain": self.strain,
         }
 
-    def _init_structure(self, cur_species: Sequence[Element | Species]) -> None:
+    def _init_structure(self) -> None:
         """Initialize structure attribute."""
-        self.structure = Structure(
+        self._structure = Structure(
             lattice=self.lattice,
-            species=cur_species,
+            species=self.species,
             coords=self.cart_coords,
             site_properties=self.site_properties,
             coords_are_cartesian=True,