forcefields: allow specifying custom calculator in forcefield makers (#1362)

* Allow ext-load calculator in forcefield makers

* Fix header depth

* Add docs for force_field_name

* Infer calculator_meta from ase_calculator_name

Author: lan496

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,9 +142,10 @@ def make(
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
-            str(self.force_field_name),  # make mypy happy
+            self.ase_calculator_name,
             ase_result,
             self.steps,
+            calculator_meta=self.calculator_meta,
             relax_kwargs=self.relax_kwargs,
             optimizer_kwargs=self.optimizer_kwargs,
             relax_cell=self.relax_cell,
@@ -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,10 +135,11 @@ def make(
             )
 
         return ForceFieldTaskDocument.from_ase_compatible_result(
-            str(self.force_field_name),  # make mypy happy
+            self.ase_calculator_name,
             md_result,
             relax_cell=(self.ensemble == MDEnsemble.npt),
             steps=self.n_steps,
+            calculator_meta=self.calculator_meta,
             relax_kwargs=None,
             optimizer_kwargs=None,
             fix_symmetry=False,

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,