forcefields: allow specifying custom calculator in forcefield makers

docs/user/codes/forcefields.md

@@ -3,6 +3,19 @@
 # Machine Learning forcefields / interatomic potentials
 
 `atomate2` includes an interface to a few common machine learning interatomic potentials (MLIPs), also known variously as machine learning forcefields (MLFFs), or foundation potentials (FPs) for universal variants.
+
+Most of `Maker` classes using the forcefields inherit from `atomate2.forcefields.utils.ForceFieldMixin` to specify which forcefield to use.
+The `ForceFieldMixin` mixin provides the following configurable parameters:
+
+- `force_field_name`: Name of the forcefield to use.
+- `calculator_kwargs`: Keyword arguments to pass to the corresponding ASE calculator.
+
+These parameters are passed to `atomate2.forcefields.utils.ase_calculator()` to instantiate the appropriate ASE calculator.
+
+The `force_field_name` should be either one of predefined `atomate2.forcefields.utils.MLFF` (or its string equivalent) or a dictionary decodable as a class or function for ASE calculator as follows.
+
+## Using predefined forcefields supported via `atomate2.forcefields.utils.MLFF`
+
 Support is provided for the following models, which can be selected using `atomate2.forcefields.utils.MLFF`, as shown in the table below.
 **You need only install packages for the forcefields you wish to use.**
 
@@ -20,3 +33,18 @@ Support is provided for the following models, which can be selected using `atoma
 | Neuroevolution Potential (NEP) | `NEP` | [10.1103/PhysRevB.104.104309](https://doi.org/10.1103/PhysRevB.104.104309) | Relies on `calorine` package |
 | Neural Equivariant Interatomic Potentials (Nequip) | `Nequip` | [10.1038/s41467-022-29939-5](https://doi.org/10.1038/s41467-022-29939-5) | Relies on the `nequip` package |
 | SevenNet | `SevenNet` | [10.1021/acs.jctc.4c00190](https://doi.org/10.1021/acs.jctc.4c00190) | Relies on the `sevenn` package |
+
+## Using custom forcefields by dictionary
+
+`force_field_name` also accepts a MSONable dictionary for specifying a custom ASE calculator class or function [^calculator-meta-type-annotation].
+For example, a `Job` created with the following code snippet instantiates `chgnet.model.dynamics.CHGNetCalculator` as the ASE calculator:
+```python
+job = ForceFieldStaticMaker(
+    force_field_name={
+        "@module": "chgnet.model.dynamics",
+        "@callable": "CHGNetCalculator",
+    }
+).make(structure)
+```
+
+[^calculator-meta-type-annotation]: In this context, the type annotation of the decoded dict should be either `Type[Calculator]` or `Callable[..., Calculator]`, where `Calculator` is from `ase.calculators.calculator`.

docs/user/key_concepts_overview.md

@@ -103,7 +103,7 @@ An `InputSet` is a convenient way to provide a collection of input data for one
 The [pymatgen](https://github.com/materialsproject/pymatgen) class `InputSet` is a core class to manage and write the input files for the several computational codes to a file location the user specifies.
 There are predefined "recipes" for generating `InputSets` tailored to specific tasks like structural relaxation or the band structure calculation and more, that are provided as `InputGenerator` classes.
 
-## Technical Aspects
+### Technical Aspects
 
 The `InputSet` objects posses the `write_input()` method that is used to write all the necessary files.
 

src/atomate2/forcefields/flows/approx_neb.py

@@ -66,20 +66,19 @@ def get_charge_density(
     @classmethod
     def from_force_field_name(
         cls,
-        force_field_name: str | MLFF,
+        force_field_name: str | MLFF | dict,
         **kwargs,
     ) -> Self:
         """
         Create an ApproxNEB flow from a forcefield name.
 
         Parameters
         ----------
-        force_field_name : str or .MLFF
+        force_field_name : str or .MLFF or dict
             The name of the force field.
         **kwargs
             Additional kwargs to pass to ApproxNEB
 
-
         Returns
         -------
         MLFFApproxNebFromEndpointsMaker

src/atomate2/forcefields/flows/elastic.py

@@ -103,7 +103,7 @@ def prev_calc_dir_argname(self) -> str | None:
     @classmethod
     def from_force_field_name(
         cls,
-        force_field_name: str | MLFF,
+        force_field_name: str | MLFF | dict,
         mlff_kwargs: dict | None = None,
         **kwargs,
     ) -> Self:
@@ -112,7 +112,7 @@ def from_force_field_name(
 
         Parameters
         ----------
-        force_field_name : str or .MLFF
+        force_field_name : str or .MLFF or dict
             The name of the force field.
         mlff_kwargs : dict or None (default)
             kwargs to pass to `ForceFieldRelaxMaker`.

src/atomate2/forcefields/flows/eos.py

@@ -59,7 +59,7 @@ class ForceFieldEosMaker(CommonEosMaker):
     @classmethod
     def from_force_field_name(
         cls,
-        force_field_name: str | MLFF,
+        force_field_name: str | MLFF | dict,
         relax_initial_structure: bool = True,
         **kwargs,
     ) -> Self:
@@ -68,7 +68,7 @@ def from_force_field_name(
 
         Parameters
         ----------
-        force_field_name : str or .MLFF
+        force_field_name : str or .MLFF or dict
             The name of the force field.
         relax_initial_structure: bool = True
             Whether to relax the initial structure before performing an EOS fit.

src/atomate2/forcefields/flows/phonons.py

@@ -158,10 +158,15 @@ def mlff(self) -> MLFF:
         """The MLFF enum corresponding to the force field name."""
         return self.phonon_displacement_maker.mlff
 
+    @property
+    def ase_calculator_name(self) -> str:
+        """The name of the ASE calculator used in this flow."""
+        return self.phonon_displacement_maker.ase_calculator_name
+
     @classmethod
     def from_force_field_name(
         cls,
-        force_field_name: str | MLFF,
+        force_field_name: str | MLFF | dict,
         relax_initial_structure: bool = True,
         **kwargs,
     ) -> Self:
@@ -170,14 +175,13 @@ def from_force_field_name(
 
         Parameters
         ----------
-        force_field_name : str or .MLFF
+        force_field_name : str or .MLFF or dict
             The name of the force field.
         relax_initial_structure: bool = True
             Whether to relax the initial structure before performing an EOS fit.
         **kwargs
             Additional kwargs to pass to PhononMaker
 
-
         Returns
         -------
         PhononMaker

src/atomate2/forcefields/flows/qha.py

@@ -92,7 +92,7 @@ def prev_calc_dir_argname(self) -> None:
     @classmethod
     def from_force_field_name(
         cls,
-        force_field_name: str | MLFF,
+        force_field_name: str | MLFF | dict,
         relax_initial_structure: bool = True,
         run_eos_flow: bool = True,
         **kwargs,
@@ -102,7 +102,7 @@ def from_force_field_name(
 
         Parameters
         ----------
-        force_field_name : str or .MLFF
+        force_field_name : str or .MLFF or dict
             The name of the force field.
         relax_initial_structure: bool = True
             Whether to relax the initial structure before performing an EOS fit.
@@ -111,7 +111,6 @@ def from_force_field_name(
         **kwargs
             Additional kwargs to pass to ForceFieldEosMaker
 
-
         Returns
         -------
         ForceFieldQhaMaker

src/atomate2/forcefields/jobs.py

@@ -73,7 +73,7 @@ class ForceFieldRelaxMaker(ForceFieldMixin, AseRelaxMaker):
     ----------
     name : str
         The job name.
-    force_field_name : str or .MLFF
+    force_field_name : str or .MLFF or dict
         The name of the force field.
     relax_cell : bool = True
         Whether to allow the cell shape/volume to change during relaxation.
@@ -106,7 +106,7 @@ class ForceFieldRelaxMaker(ForceFieldMixin, AseRelaxMaker):
     """
 
     name: str = "Force field relax"
-    force_field_name: str | MLFF = MLFF.Forcefield
+    force_field_name: str | MLFF | dict = MLFF.Forcefield
     relax_cell: bool = True
     fix_symmetry: bool = False
     symprec: float | None = 1e-2
@@ -142,7 +142,8 @@ def make(
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
-            str(self.force_field_name),  # make mypy happy
+            self.ase_calculator_name,
+            self.calculator_meta,
             ase_result,
             self.steps,
             relax_kwargs=self.relax_kwargs,
@@ -170,7 +171,7 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker):
     ----------
     name : str
         The job name.
-    force_field_name : str or .MLFF
+    force_field_name : str or .MLFF or dict
         The name of the force field.
     calculator_kwargs : dict
         Keyword arguments that will get passed to the ASE calculator.
@@ -180,7 +181,7 @@ class ForceFieldStaticMaker(ForceFieldRelaxMaker):
     """
 
     name: str = "Force field static"
-    force_field_name: str | MLFF = MLFF.Forcefield
+    force_field_name: str | MLFF | dict = MLFF.Forcefield
     relax_cell: bool = False
     steps: int = 1
     relax_kwargs: dict = field(default_factory=dict)

src/atomate2/forcefields/md.py

@@ -43,7 +43,7 @@ class ForceFieldMDMaker(ForceFieldMixin, AseMDMaker):
     ----------
     name : str
         The name of the MD Maker
-    force_field_name : str or .MLFF
+    force_field_name : str or .MLFF or dict
         The name of the forcefield (for provenance)
     time_step : float | None = None.
         The timestep of the MD run in fs.
@@ -135,7 +135,8 @@ def make(
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
-            str(self.force_field_name),  # make mypy happy
+            self.ase_calculator_name,
+            self.calculator_meta,
             md_result,
             relax_cell=(self.ensemble == MDEnsemble.npt),
             steps=self.n_steps,

src/atomate2/forcefields/neb.py

@@ -24,7 +24,7 @@ class ForceFieldNebFromImagesMaker(ForceFieldMixin, AseNebFromImagesMaker):
     """Run NEB with an ML forcefield using ASE."""
 
     name: str = "Forcefield NEB from images"
-    force_field_name: str | MLFF = MLFF.Forcefield
+    force_field_name: str | MLFF | dict = MLFF.Forcefield
 
     @job(data=_FORCEFIELD_DATA_OBJECTS, schema=NebResult)
     def make(
@@ -49,7 +49,7 @@ class ForceFieldNebFromEndpointsMaker(ForceFieldMixin, AseNebFromEndpointsMaker)
     """Run NEB with an ML forcefield using ASE."""
 
     name: str = "Forcefield NEB from endpoints"
-    force_field_name: str | MLFF = MLFF.Forcefield
+    force_field_name: str | MLFF | dict = MLFF.Forcefield
 
     @job(data=_FORCEFIELD_DATA_OBJECTS, schema=NebResult)
     def make(
@@ -69,19 +69,21 @@ def make(
         return self._run_ase_safe(images=images, prev_dir=prev_dir)
 
     @classmethod
-    def from_force_field_name(cls, force_field_name: str | MLFF, **kwargs) -> Self:
+    def from_force_field_name(
+        cls, force_field_name: str | MLFF | dict, **kwargs
+    ) -> Self:
         """Create a force field NEB job from its name.
 
         Parameters
         ----------
-        force_field_name : str or MLFF
-            The name of the forcefield. Should be a valid MLFF member.
+        force_field_name : str or MLFF or dict
+            The name of the forcefield.
         **kwargs
             kwargs to pass to ForceFieldNebFromEndpointsMaker.
         """
         endpoint_relax_maker = ForceFieldRelaxMaker(force_field_name=force_field_name)
         return cls(
-            name=f"{force_field_name} NEB from endpoints maker",
+            name=f"{endpoint_relax_maker.mlff.name} NEB from endpoints maker",
             endpoint_relax_maker=endpoint_relax_maker,
             force_field_name=force_field_name,
             **kwargs,

src/atomate2/forcefields/schemas.py

@@ -7,6 +7,7 @@
 from emmet.core.types.enums import StoreTrajectoryOption
 from pydantic import BaseModel, Field
 from pymatgen.core import Molecule
+from typing_extensions import assert_never
 
 from atomate2.ase.schemas import (
     AseMoleculeTaskDoc,
@@ -17,6 +18,7 @@
     _task_doc_translation_keys,
 )
 from atomate2.forcefields import MLFF
+from atomate2.forcefields.utils import _load_calc_cls
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -91,6 +93,7 @@ class ForceFieldTaskDocument(AseStructureTaskDoc, ForceFieldMeta):
     def from_ase_compatible_result(
         cls,
         ase_calculator_name: str,
+        calculator_meta: MLFF | dict,
         result: AseResult,
         steps: int,
         relax_kwargs: dict = None,
@@ -114,6 +117,8 @@ def from_ase_compatible_result(
         ----------
         ase_calculator_name : str
             Name of the ASE calculator used.
+        calculator_meta : MLFF or dict
+            Metadata about the calculator used.
         result : AseResult
             The output results from the task.
         fix_symmetry : bool
@@ -151,9 +156,24 @@ def from_ase_compatible_result(
         ff_kwargs = {
             "forcefield_name": task_document_kwargs.get(
                 "forcefield_name", ase_calculator_name
-            )
+            ),
         }
 
+        if pkg_name := _get_pkg_name(calculator_meta):
+            import importlib.metadata
+
+            ff_kwargs["forcefield_version"] = importlib.metadata.version(pkg_name)
+
+        return (
+            ForceFieldMoleculeTaskDocument
+            if isinstance(result.final_mol_or_struct, Molecule)
+            else cls
+        ).from_ase_task_doc(ase_task_doc, **ff_kwargs)
+
+
+def _get_pkg_name(calculator_meta: MLFF | dict) -> str | None:
+    """Get the package name for a given force field."""
+    if isinstance(calculator_meta, MLFF):
         # map force field name to its package name
         model_to_pkg_map = {
             MLFF.M3GNet: "matgl",
@@ -168,16 +188,8 @@ def from_ase_compatible_result(
             MLFF.MATPES_PBE: "matgl",
             MLFF.MATPES_R2SCAN: "matgl",
         }
-
-        if pkg_name := {str(k): v for k, v in model_to_pkg_map.items()}.get(
-            ff_kwargs["forcefield_name"]
-        ):
-            import importlib.metadata
-
-            ff_kwargs["forcefield_version"] = importlib.metadata.version(pkg_name)
-
-        return (
-            ForceFieldMoleculeTaskDocument
-            if isinstance(result.final_mol_or_struct, Molecule)
-            else cls
-        ).from_ase_task_doc(ase_task_doc, **ff_kwargs)
+        return model_to_pkg_map.get(calculator_meta)
+    if isinstance(calculator_meta, dict):
+        calc_cls = _load_calc_cls(calculator_meta)
+        return calc_cls.__module__.split(".")[0]
+    assert_never(calculator_meta)