Merge pull request #112 from benrich37/missing-unit-conversion

Slightly cleaner nan stripping

Author: benrich37

src/pymatgen/io/jdftx/inputs.py

@@ -21,14 +21,7 @@ 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,
-    FloatTag,
-    MultiformatTag,
-    TagContainer,
-)
+from pymatgen.io.jdftx.generic_tags import AbstractTag, BoolTagContainer, DumpTagContainer, MultiformatTag, TagContainer
 from pymatgen.io.jdftx.jdftxinfile_default_inputs import default_inputs
 from pymatgen.io.jdftx.jdftxinfile_master_format import (
     __PHONON_TAGS__,
@@ -958,105 +951,84 @@ 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_infile_of_nans(infile: JDFTXInfile) -> JDFTXInfile:
+    infile_cleaned = infile.copy()
+    popmaps = get_popmaps(infile_cleaned)
+    apply_popmaps(infile_cleaned, popmaps)
+    return infile_cleaned
 
 
-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 hasnan(val: Any) -> bool:
+    if type(val) in [float, int]:
+        return np.isnan(val)
+    if isinstance(val, list):
+        for v in val:
+            if hasnan(v):
+                return True
+    elif isinstance(val, dict):
+        for v in val.values():
+            if hasnan(v):
+                return True
+    return False
 
 
-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
+def get_nanmaps(val: Any) -> list[list[str | int]]:
+    nanmaps: list[list[str | int]] = []
+    generator = None
+    if type(val) in [float, int]:
+        if hasnan(val):
+            nanmaps.append([])
+    elif isinstance(val, dict):
+        generator = [(k, v) for k, v in val.items()]
+    elif isinstance(val, list):
+        generator = [(i, v) for i, v in enumerate(val)]
+    elif type(val in [str]):
+        pass
+    else:
+        raise ValueError(f"Unexpected type {type(val)} of value {val}")
+    if generator is not None:
+        for k, v in generator:
+            submaps = get_nanmaps(v)
+            for submap in submaps:
+                add_map = [k]
+                add_map.extend(submap)
+                nanmaps.append(add_map)
+    return nanmaps
+
+
+def get_popmaps(infile: JDFTXInfile) -> list[list[str | int]]:
+    popmaps = []
+    for tag in infile:
+        nanmaps = get_nanmaps(infile[tag])
+        for nanmap in nanmaps:
+            add_map = [tag]
+            add_map.extend(nanmap)
+            popmaps.append(add_map)
+    return popmaps
+
+
+def apply_popmaps(edit_infile: JDFTXInfile, popmaps: list[list[str | int]]) -> None:
+    for popmap in popmaps:
+        popmap_use = popmap
+        popmap_destination = edit_infile
+        to = get_tag_object_on_val(popmap[0], edit_infile[popmap[0]])
+        if isinstance(to, TagContainer):
+            subtag = to
+            popmap_destination = popmap_destination[popmap[0]]
+            popmap_route = popmap[1:]
+            popmap_use = [popmap[0]]
+            step = None
+            for i, step in enumerate(popmap_route):
+                if isinstance(step, str):
+                    subtag = subtag.subtags[step]
+                    if not subtag.optional:
+                        step = popmap_route[i - 1]
+                        break
+                popmap_use.append(step)
+            popmap_destination = edit_infile
+            for step in popmap_use[:-1]:
+                popmap_destination = popmap_destination[step]
+            popmap_destination.pop(popmap_use[-1], None)
 
 
 @dataclass

tests/io/jdftx/test_jdftxinfile.py

@@ -9,7 +9,12 @@
 
 from pymatgen.core.structure import Site, Structure
 from pymatgen.core.units import bohr_to_ang
-from pymatgen.io.jdftx.inputs import JDFTXInfile, JDFTXStructure, selective_dynamics_site_prop_to_jdftx_interpretable
+from pymatgen.io.jdftx.inputs import (
+    JDFTXInfile,
+    JDFTXStructure,
+    clean_infile_of_nans,
+    selective_dynamics_site_prop_to_jdftx_interpretable,
+)
 from pymatgen.io.jdftx.jdftxinfile_default_inputs import antoinePvap, default_inputs
 from pymatgen.io.jdftx.jdftxinfile_master_format import get_tag_object
 
@@ -615,3 +620,21 @@ def test_jdftxinfile_comparison():
 
 def test_antoine_pvap():
     assert_same_value(antoinePvap(298, 7.31549, 1794.88, -34.764), 1.06736e-10)
+
+
+def test_nan_stripping():
+    jif = JDFTXInfile.from_file(ex_infile1_fname)
+    # Single optional subtag being nan should only remove that subtag
+    jif["elec-cutoff"]["EcutRho"] = np.nan
+    jif = clean_infile_of_nans(jif)
+    assert "EcutRho" not in jif["elec-cutoff"]
+    # Single required subtag being nan should remove the whole tag
+    jif["latt-move-scale"]["s0"] = np.nan
+    jif = clean_infile_of_nans(jif)
+    assert "latt-move-scale" not in jif
+    jif.pop("fluid-solvent")
+    jif.read_line("fluid-solvent H2O 55.338 ScalarEOS epsBulk nan")
+    assert "epsBulk" in jif["fluid-solvent"][0]
+    jif = clean_infile_of_nans(jif)
+    assert "epsBulk" not in jif["fluid-solvent"][0]
+    print(jif)