Merge pull request #141 from materialsproject/static-gen-run_types

Generate static run-types for IDEs

emmet-core/emmet/core/vasp/calc_types/__init__.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+try:
+    import emmet.core.vasp.calc_types.enums
+except ImportError:
+    import emmet.core.vasp.calc_types.generate
+
+from emmet.core.vasp.calc_types.enums import RunType, TaskType, CalcType
+from emmet.core.vasp.calc_types.utils import run_type, task_type, calc_type

emmet-core/emmet/core/vasp/calc_types/generate.py

@@ -0,0 +1,98 @@
+""" Module to define various calculation types as Enums for VASP """
+import datetime
+from itertools import groupby, product
+from pathlib import Path
+from typing import Dict, Iterator, List
+
+import bson
+import numpy as np
+from monty.json import MSONable
+from monty.serialization import loadfn
+from pydantic import BaseModel
+from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher
+from pymatgen.core.structure import Structure
+from typing_extensions import Literal
+
+from emmet.core import SETTINGS
+from emmet.core.utils import ValueEnum
+
+_RUN_TYPE_DATA = loadfn(str(Path(__file__).parent.joinpath("run_types.yaml").resolve()))
+_TASK_TYPES = [
+    "NSCF Line",
+    "NSCF Uniform",
+    "Dielectric",
+    "DFPT",
+    "DFPT Dielectric",
+    "NMR Nuclear Shielding",
+    "NMR Electric Field Gradient",
+    "Static",
+    "Structure Optimization",
+    "Deformation",
+]
+
+_RUN_TYPES = (
+    [
+        rt
+        for functional_class in _RUN_TYPE_DATA
+        for rt in _RUN_TYPE_DATA[functional_class]
+    ]
+    + [
+        f"{rt}+U"
+        for functional_class in _RUN_TYPE_DATA
+        for rt in _RUN_TYPE_DATA[functional_class]
+    ]
+    + ["LDA", "LDA+U"]
+)
+
+
+def get_enum_source(enum_name, doc, items):
+    header = f"""
+class {enum_name}(ValueEnum):
+    \"\"\" {doc} \"\"\"\n
+"""
+    items = [f'    {const} = "{val}"' for const, val in items.items()]
+
+    return header + "\n".join(items)
+
+
+run_type_enum = get_enum_source(
+    "RunType",
+    "VASP calculation run types",
+    dict(
+        {
+            "_".join(rt.split()).replace("+", "_").replace("-", "_"): rt
+            for rt in _RUN_TYPES
+        }
+    ),
+)
+task_type_enum = get_enum_source(
+    "TaskType",
+    "VASP calculation task types",
+    {"_".join(tt.split()): tt for tt in _TASK_TYPES},
+)
+calc_type_enum = get_enum_source(
+    "CalcType",
+    "VASP calculation types",
+    {
+        f"{'_'.join(rt.split()).replace('+','_').replace('-','_')}_{'_'.join(tt.split())}": f"{rt} {tt}"
+        for rt, tt in product(_RUN_TYPES, _TASK_TYPES)
+    },
+)
+
+
+with open(Path(__file__).parent / "enums.py", "w") as f:
+    f.write(
+        """\"\"\"
+Autogenerated Enums for VASP RunType, TaskType, and CalcType
+Do not edit this by hand. Edit generate.py or run_types.yaml instead
+\"\"\"
+from emmet.core.utils import ValueEnum
+
+"""
+    )
+    f.write(run_type_enum)
+    f.write("\n\n")
+    f.write(task_type_enum)
+    f.write("\n\n")
+    f.write(calc_type_enum)
+    f.write("\n")

emmet-core/emmet/core/vasp/calc_types.py → emmet-core/emmet/core/vasp/calc_types/utils.py

@@ -14,52 +14,9 @@
 from typing_extensions import Literal
 
 from emmet.core import SETTINGS
-from emmet.core.utils import ValueEnum
+from emmet.core.vasp.calc_types.enums import RunType, TaskType, CalcType
 
 _RUN_TYPE_DATA = loadfn(str(Path(__file__).parent.joinpath("run_types.yaml").resolve()))
-_TASK_TYPES = [
-    "NSCF Line",
-    "NSCF Uniform",
-    "Dielectric",
-    "DFPT",
-    "DFPT Dielectric",
-    "NMR Nuclear Shielding",
-    "NMR Electric Field Gradient",
-    "Static",
-    "Structure Optimization",
-    "Deformation",
-]
-
-_RUN_TYPES = (
-    [
-        rt
-        for functional_class in _RUN_TYPE_DATA
-        for rt in _RUN_TYPE_DATA[functional_class]
-    ]
-    + [
-        f"{rt}+U"
-        for functional_class in _RUN_TYPE_DATA
-        for rt in _RUN_TYPE_DATA[functional_class]
-    ]
-    + ["LDA", "LDA+U"]
-)
-
-RunType = ValueEnum(  # type: ignore
-    "RunType", dict({"_".join(rt.split()).replace("+", "_"): rt for rt in _RUN_TYPES})
-)
-RunType.__doc__ = "VASP calculation run types"
-
-TaskType = ValueEnum("TaskType", {"_".join(tt.split()): tt for tt in _TASK_TYPES})  # type: ignore
-TaskType.__doc__ = "VASP calculation task types"
-
-CalcType = ValueEnum(  # type: ignore
-    "CalcType",
-    {
-        f"{'_'.join(rt.split()).replace('+','_')}_{'_'.join(tt.split())}": f"{rt} {tt}"
-        for rt, tt in product(_RUN_TYPES, _TASK_TYPES)
-    },
-)
-CalcType.__doc__ = "VASP calculation types"
 
 
 def run_type(parameters: Dict) -> RunType:

emmet-core/emmet/stubs/utils.py

@@ -154,7 +154,7 @@ def MSONable_to_pydantic(monty_cls: type, pydantic_model=None):
             for field_name, field_type in _type_hints
         }
 
-    model = create_model(monty_cls.__name__, **monty_props, **props)
+    model = create_model(monty_cls.__name__, **monty_props, **props)  # type: ignore
     if hasattr(monty_cls, "__doc__"):
         setattr(model, "__doc__", monty_cls.__doc__)
     return model

tests/emmet-core/test_structuremetadata.py

@@ -1,5 +1,5 @@
 import pytest
-from pymatgen import Lattice
+from pymatgen import Lattice, Element
 
 from emmet.core.structure import StructureMetadata
 from emmet.core.symmetry import CrystalSystem, SymmetryData
@@ -22,13 +22,16 @@ def test_symmetry(structure):
     assert symm_doc.symbol == "Pm-3m"
     assert symm_doc.crystal_system == CrystalSystem.cubic
 
+    assert symm_doc.dict()["crystal_system"] == CrystalSystem.cubic
+    assert str(symm_doc.dict()["crystal_system"]) == "Cubic"
+
 
 def test_structure_metadata(structure):
 
     meta_doc = StructureMetadata.from_structure(structure)
 
     assert meta_doc.nsites == 1
-    assert meta_doc.elements == ["Fe"]
+    assert meta_doc.elements == [Element.Fe]
     assert meta_doc.nelements == 1
     assert meta_doc.formula_pretty == "Fe"
     assert meta_doc.formula_anonymous == "A"