Forcefield molecular dynamics and forcefield refactor (#722)

* FF MD fully implemented

* change MD defaults to be NVT + Langevin dynamics

* add md defaults that are as consistent with vasp implementation as possible

* correct typing of langevin friction coeff

* Add option to ForceFieldTaskDocument to store trajectory information as pmg object

* Add tests, remove dependence on pickle for trajectory i/o

* fix tests, update matgl dependence

* rename MD forcefield class vars to be consistent with VASP MDSetGenerator

* Add option to explicitly specify thermostat as ASE MolecularDynamics object

* refactoring and add T/P schedule

* implement T/P schedule and attach callback func

* np.nan

* np.nan

* small fixes

* pass ASE dynamics object instead of class, update docstring

* add temp ramp test and make sure upper tria matrix cell

* add stresses to trajeoctroy observer, add pressure schedule test

* avoid `stresses` key not in taskdoc bug

* Refactory TrajectoryObserver, add option to save trajectory as either pmg or ase trajectory objects, add test for trajectory parsing

* Change decomposition to upper triangular cell only when Nose-Hoover requested

* Fix failing tests

* Change MD default TaskDoc ionic step to only store mandatory info

* update pymatgen requirement because of CifParser deprecation

* Significant refactor to all forcefield jobs. Use commmon ASE calc structure, allow loading via MontyDecoder. Add check in relaxation for force convergence, attr in Forcefield taskdoc

* Fix forcefield utils test and lint

* Add revert_dtype env for running forcefield relax and md jobs / undo removal

* refactor nequip jobs, add md default for nequip

* Add GAP and Nequip MD tests, fix arg / kwarg passing in MD

* remove calculator_args / ase_calculator args, revert phonon job change

* remove todo about adding magmoms to forcefield traj observer as that's now implemented

* Remove comments, add option to seed rng for MB velocities, turn on ideal gas stress contribution only when MD trajectory outputs (temp / velocity) are stored

* Ensure CHGNet and M3GNet relax / static makers convert stress to eV/A0**3; add tests for MD NVE ensemble and specifying MolecularDynamics object as input

* fix M3GNet test related to outdated cached model

* Fix dependencies, amend forcefield static option parsing

* linting

* try to fix ci tests caused by torch / dgl incompatibility

* try to fix ci tests caused by torch / dgl incompatibility

* Change: default time_step, units for time_step, initialization of temperature and velocity in TrajectoryObserver, snake case change

* linting / pre-commit

* revert default time_step units to fs

* add missing forcefields to version check in forcefield task document

* allow complex schur decomposition

* more snake case

* Decrease number of steps in tests for speed. Fix (?) temperature / pressure scheduling

* linting

* fix T/P schedule

* linting

* replace dunder __getattribute__ --> getattr in forcefield test_md

---------

Co-authored-by: Yuan Chiang <qaqowotut@gmail.com>
Co-authored-by: Yuan (Cyrus) Chiang <41962462+chiang-yuan@users.noreply.github.com>
Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

.github/workflows/testing.yml

@@ -39,8 +39,7 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-          cache: pip
-          cache-dependency-path: pyproject.toml
+
 
       - name: Install dependencies
         # ERROR: Cannot install atomate2 and atomate2[strict,tests]==0.0.1 because these package versions have conflicting dependencies.

pyproject.toml

@@ -50,9 +50,10 @@ defects = [
 ]
 forcefields = [
     "ase>=3.22.1",
+    "torch<=2.2.1", # incompatibility with dgl if newer versions are used
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
-    "matgl>=0.9.0",
+    "matgl>=1.0.0",
     "quippy-ase>=0.9.14",
 ]
 docs = [
@@ -74,6 +75,7 @@ strict = [
     # must use >= for ase to not uninstall main branch install in CI
     # last known working commit: https://gitlab.com/ase/ase@2bab58f4e
     "ase>=3.22.1",
+    "torch==2.2.1",
     "cclib==1.8.1",
     "chgnet==0.3.5",
     "click==8.1.7",

src/atomate2/common/jobs/utils.py

@@ -40,7 +40,7 @@ def structure_to_conventional(
     structure: Structure, symprec: float = SETTINGS.SYMPREC
 ) -> Structure:
     """
-    Job hat creates a standard conventional structure.
+    Job that creates a standard conventional structure.
 
     Parameters
     ----------

src/atomate2/common/schemas/phonons.py

@@ -128,7 +128,7 @@ class PhononJobDirs(BaseModel):
         None, description="Directory where optimization run was performed."
     )
     taskdoc_run_job_dir: Optional[str] = Field(
-        None, description="Directory where taskdoc was generated."
+        None, description="Directory where task doc was generated."
     )
 
 

src/atomate2/forcefields/__init__.py

@@ -10,4 +10,5 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     GAP = "GAP"
     M3GNet = "M3GNet"
     CHGNet = "CHGNet"
+    Forcefield = "Forcefield"  # default placeholder option
     Nequip = "Nequip"