forcefields: Standardize MLFF handling and maker naming (#1360)

* Standardize MLFF handling and maker naming

* Introduce mlff property in tutorial

Author: lan496

src/atomate2/forcefields/flows/approx_neb.py

@@ -9,8 +9,8 @@
 from typing_extensions import Self
 
 from atomate2.common.flows.approx_neb import ApproxNebFromEndpointsMaker
-from atomate2.forcefields import MLFF, _get_formatted_ff_name
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker
+from atomate2.forcefields.utils import MLFF
 
 
 @dataclass
@@ -84,15 +84,11 @@ def from_force_field_name(
         -------
         MLFFApproxNebFromEndpointsMaker
         """
-        force_field_name = _get_formatted_ff_name(force_field_name)
-        kwargs.update(
-            image_relax_maker=ForceFieldRelaxMaker(
-                force_field_name=force_field_name, relax_cell=False
-            ),
+        image_relax_maker = ForceFieldRelaxMaker(
+            force_field_name=force_field_name, relax_cell=False
         )
+        kwargs.update(image_relax_maker=image_relax_maker)
         return cls(
-            name=(
-                f"{force_field_name.split('MLFF.')[-1]} ApproxNEB from endpoints Maker"
-            ),
+            name=(f"{image_relax_maker.mlff.name} ApproxNEB from endpoints Maker"),
             **kwargs,
         )

src/atomate2/forcefields/flows/elastic.py

@@ -7,14 +7,15 @@
 
 from atomate2 import SETTINGS
 from atomate2.common.flows.elastic import BaseElasticMaker
-from atomate2.forcefields import MLFF, _get_formatted_ff_name
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker
 
 if TYPE_CHECKING:
     from typing import Any
 
     from typing_extensions import Self
 
+    from atomate2.forcefields import MLFF
+
 # default options for the forcefield makers in ElasticMaker
 _DEFAULT_RELAX_KWARGS: dict[str, Any] = {
     "force_field_name": "CHGNet",
@@ -125,19 +126,20 @@ def from_force_field_name(
         default_kwargs: dict[str, Any] = {
             **_DEFAULT_RELAX_KWARGS,
             **(mlff_kwargs or {}),
-            "force_field_name": _get_formatted_ff_name(force_field_name),
+            "force_field_name": force_field_name,
         }
+        bulk_relax_maker = ForceFieldRelaxMaker(
+            relax_cell=True,
+            **default_kwargs,
+        )
         kwargs.update(
-            bulk_relax_maker=ForceFieldRelaxMaker(
-                relax_cell=True,
-                **default_kwargs,
-            ),
+            bulk_relax_maker=bulk_relax_maker,
             elastic_relax_maker=ForceFieldRelaxMaker(
                 relax_cell=False,
                 **default_kwargs,
             ),
         )
         return cls(
-            name=f"{str(force_field_name).split('MLFF.')[-1]} elastic",
+            name=f"{bulk_relax_maker.mlff.name} elastic",
             **kwargs,
         )

src/atomate2/forcefields/flows/eos.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING
 
 from atomate2.common.flows.eos import CommonEosMaker
-from atomate2.forcefields import _get_formatted_ff_name
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker
 
 if TYPE_CHECKING:
@@ -81,22 +80,19 @@ def from_force_field_name(
         -------
         ForceFieldEosMaker
         """
-        force_field_name = _get_formatted_ff_name(force_field_name)
+        eos_relax_maker = ForceFieldRelaxMaker(
+            force_field_name=force_field_name, relax_cell=False
+        )
         kwargs.update(
             initial_relax_maker=(
                 ForceFieldRelaxMaker(force_field_name=force_field_name)
                 if relax_initial_structure
                 else None
-            )
-        )
-
-        kwargs.update(
-            eos_relax_maker=ForceFieldRelaxMaker(
-                force_field_name=force_field_name, relax_cell=False
             ),
+            eos_relax_maker=eos_relax_maker,
             static_maker=None,
         )
         return cls(
-            name=f"{force_field_name.split('MLFF.')[-1]} EOS Maker",
+            name=f"{eos_relax_maker.mlff.name} EOS Maker",
             **kwargs,
         )

src/atomate2/forcefields/flows/phonons.py

@@ -7,8 +7,8 @@
 
 from atomate2 import SETTINGS
 from atomate2.common.flows.phonons import BasePhononMaker
-from atomate2.forcefields import _get_formatted_ff_name
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
+from atomate2.forcefields.utils import MLFF
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -153,6 +153,11 @@ def prev_calc_dir_argname(self) -> None:
         """
         return
 
+    @property
+    def mlff(self) -> MLFF:
+        """The MLFF enum corresponding to the force field name."""
+        return self.phonon_displacement_maker.mlff
+
     @classmethod
     def from_force_field_name(
         cls,
@@ -177,24 +182,22 @@ def from_force_field_name(
         -------
         PhononMaker
         """
-        force_field_name = _get_formatted_ff_name(force_field_name)
-
+        static_energy_maker = ForceFieldStaticMaker(force_field_name=force_field_name)
         kwargs.update(
             bulk_relax_maker=(
                 ForceFieldRelaxMaker(
                     force_field_name=force_field_name, relax_kwargs={"fmax": 1e-5}
                 )
                 if relax_initial_structure
                 else None
-            )
-        )
-        kwargs.update(
-            static_energy_maker=ForceFieldStaticMaker(
-                force_field_name=force_field_name
             ),
+            static_energy_maker=static_energy_maker,
             phonon_displacement_maker=ForceFieldStaticMaker(
                 force_field_name=force_field_name
             ),
             born_maker=None,
         )
-        return cls(name=f"{force_field_name.split('MLFF.')[-1]} Phonon Maker", **kwargs)
+        return cls(
+            name=(f"{static_energy_maker.mlff.name} Phonon Maker"),
+            **kwargs,
+        )

src/atomate2/forcefields/flows/qha.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING, Literal
 
 from atomate2.common.flows.qha import CommonQhaMaker
-from atomate2.forcefields import _get_formatted_ff_name
 from atomate2.forcefields.flows.phonons import PhononMaker
 from atomate2.forcefields.jobs import ForceFieldRelaxMaker
 
@@ -117,16 +116,12 @@ def from_force_field_name(
         -------
         ForceFieldQhaMaker
         """
-        force_field_name = _get_formatted_ff_name(force_field_name)
         kwargs.update(
             initial_relax_maker=(
                 ForceFieldRelaxMaker(force_field_name=force_field_name)
                 if relax_initial_structure
                 else None
-            )
-        )
-
-        kwargs.update(
+            ),
             eos_relax_maker=(
                 ForceFieldRelaxMaker(
                     force_field_name=force_field_name,
@@ -135,13 +130,14 @@ def from_force_field_name(
                 )
                 if run_eos_flow
                 else None
-            )
+            ),
+        )
+        phonon_maker = PhononMaker.from_force_field_name(
+            force_field_name=force_field_name, relax_initial_structure=False
         )
         return cls(
-            phonon_maker=PhononMaker.from_force_field_name(
-                force_field_name=force_field_name, relax_initial_structure=False
-            ),
-            name=f"{force_field_name.split('MLFF.')[-1]} QHA Maker",
+            phonon_maker=phonon_maker,
+            name=f"{phonon_maker.mlff.name} QHA Maker",
             **kwargs,
         )
 

src/atomate2/forcefields/utils.py

@@ -12,6 +12,7 @@
 from ase.units import Bohr
 from ase.units import GPa as _GPa_to_eV_per_A3
 from monty.json import MontyDecoder
+from typing_extensions import deprecated
 
 if TYPE_CHECKING:
     from collections.abc import Generator
@@ -75,27 +76,34 @@ def _missing_(cls, value: Any) -> Any:
 }
 
 
-def _get_formatted_ff_name(force_field_name: str | MLFF) -> str:
-    """
-    Get the standardized force field name.
+def _get_standardized_mlff(force_field_name: str | MLFF) -> MLFF:
+    """Get the standardized force field name.
 
     Parameters
     ----------
     force_field_name : str or .MLFF
         The name of the force field
+        For str, accept both with and without the `MLFF.` prefix.
 
     Returns
     -------
-    str : the name of the forcefield from MLFF
+    MLFF: the name of the forcefield
     """
     if isinstance(force_field_name, str):
         # ensure `force_field_name` uses enum format
+        if force_field_name.startswith("MLFF."):
+            force_field_name = force_field_name.split("MLFF.")[-1]
+
         if force_field_name in MLFF.__members__:
             force_field_name = MLFF[force_field_name]
         elif force_field_name in [v.value for v in MLFF]:
             force_field_name = MLFF(force_field_name)
-    force_field_name = str(force_field_name)
-    if force_field_name in {"MLFF.MACE", "MACE"}:
+        else:
+            raise ValueError(
+                f"force_field_name={force_field_name} is not a valid MLFF name."
+            )
+
+    if force_field_name == MLFF.MACE:
         warnings.warn(
             "Because the default MP-trained MACE model is constantly evolving, "
             "we no longer recommend using `MACE` or `MLFF.MACE` to specify "
@@ -108,6 +116,24 @@ def _get_formatted_ff_name(force_field_name: str | MLFF) -> str:
     return force_field_name
 
 
+@deprecated("Use _get_standardized_mlff instead.")
+def _get_formatted_ff_name(force_field_name: str | MLFF) -> str:
+    """
+    Get the standardized force field name.
+
+    Parameters
+    ----------
+    force_field_name : str or .MLFF
+        The name of the force field
+
+    Returns
+    -------
+    str : the name of the forcefield from MLFF
+    """
+    force_field_name = _get_standardized_mlff(force_field_name)
+    return str(force_field_name)
+
+
 @dataclass
 class ForceFieldMixin:
     """Mix-in class for force-fields.
@@ -134,18 +160,17 @@ def __post_init__(self) -> None:
         if hasattr(super(), "__post_init__"):
             super().__post_init__()  # type: ignore[misc]
 
-        self.force_field_name = _get_formatted_ff_name(self.force_field_name)
+        mlff = _get_standardized_mlff(self.force_field_name)
+        self.force_field_name: str = str(mlff)  # Narrow-down type for mypy
 
         # Pad calculator_kwargs with default values, but permit user to override them
         self.calculator_kwargs = {
-            **_DEFAULT_CALCULATOR_KWARGS.get(
-                MLFF(self.force_field_name.split("MLFF.")[-1]), {}
-            ),
+            **_DEFAULT_CALCULATOR_KWARGS.get(mlff, {}),
             **self.calculator_kwargs,
         }
 
         if not self.task_document_kwargs.get("force_field_name"):
-            self.task_document_kwargs["force_field_name"] = str(self.force_field_name)
+            self.task_document_kwargs["force_field_name"] = self.force_field_name
 
     def _run_ase_safe(self, *args, **kwargs) -> AseResult:
         if not hasattr(self, "run_ase"):
@@ -159,10 +184,15 @@ def _run_ase_safe(self, *args, **kwargs) -> AseResult:
     def calculator(self) -> Calculator:
         """ASE calculator, can be overwritten by user."""
         return ase_calculator(
-            str(self.force_field_name),  # make mypy happy
+            self.force_field_name,
             **self.calculator_kwargs,
         )
 
+    @property
+    def mlff(self) -> MLFF:
+        """The MLFF enum corresponding to the force field name."""
+        return MLFF(str(self.force_field_name).split("MLFF.")[-1])
+
 
 def ase_calculator(
     calculator_meta: str | MLFF | dict, **kwargs: Any

tutorials/force_fields/phonon_workflow.ipynb

@@ -134,16 +134,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "PhononMaker.from_force_field_name(force_field_name=\"MACE_MP_0B3\")\n",
-    "run_locally(flow, create_folders=True, raise_immediately=True, root_dir=tmp_dir)"
+    "maker = PhononMaker.from_force_field_name(force_field_name=\"MACE_MP_0B3\")"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "11",
    "metadata": {},
    "source": [
-    "Now, we clean up the temporary directory that we made. In reality, you might want to keep this data."
+    "We can confirm that the specified force field is being used via `PhononMaker.mlff` property:"
    ]
   },
   {
@@ -152,6 +151,26 @@
    "id": "12",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "from atomate2.forcefields.utils import MLFF\n",
+    "\n",
+    "assert maker.mlff == MLFF.MACE_MP_0B3  # noqa: S101"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "Now, we clean up the temporary directory that we made. In reality, you might want to keep this data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import shutil\n",
     "\n",
@@ -163,14 +182,14 @@
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.6"
+   "pygments_lexer": "ipython3",
+   "version": "3.11.14"
   }
  },
  "nbformat": 4,