Forcefield molecular dynamics and forcefield refactor

[esoteric-ephemera]

Completed:

• Add a forcefield maker to perform molecular dynamics with forcefields in ASE, atomate2.forcefields.md.ForceFieldMDMaker
  • The defaults for this class (starting and ending temperature, pressure, Langevin damping, ensemble, thermostat) are matched as closely as possible to the VASP MD maker
  • Units for params are the same between both implementations, see the docstr
  • There's some duplication of the code in atomate2.forcefields.utils.Relaxer, but I don't see a better way to merge those approaches for trajectory observance without a messier implementation
• Specific subclasses of ForceFieldMDMaker for CHGNet, M3GNet, and MACE
• Update matgl dependence to 1.0.0 (needed to use ASE calculator wrapper for M3GNet)
• MD trajectory stored in GridFS. Options to write it to disk either as pymatgen.core.trajectory.Trajectory or ase.io.Trajectory objects
• Refactor Forcefield relax and static jobs to use a common atomate2.forcefields.utils.ase_calculator object. One can either:
  • Load a predefined ASE calculator using one of the atomate2.forcefield.MLFF members,
  • or by specifying a dict with "@module" and "@callable" keys to load a module via MontyDecoder
• Add a force convergence check in atomate2.forcefields.utils.Relaxer which is later passed to ForceFieldTaskDocument to resolve Convergence check for relaxation with forcefields #753.

To Do's:

• Structure of the kwargs entering MDMaker vs. its make function. For example, it may be easier to set the temperature and pressure via MDMaker().make than via MDMaker(temperature=...)

[codecov]

Codecov Report

Attention: Patch coverage is 89.42598% with 35 lines in your changes are missing coverage. Please review.

Project coverage is 76.81%. Comparing base (9f64381) to head (4a88d7b).
Report is 12 commits behind head on main.

❗ Current head 4a88d7b differs from pull request most recent head 78ad999. Consider uploading reports for the commit 78ad999 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #722      +/-   ##
==========================================
+ Coverage   76.64%   76.81%   +0.17%     
==========================================
  Files         114      115       +1     
  Lines        8506     8723     +217     
  Branches     1275     1344      +69     
==========================================
+ Hits         6519     6701     +182     
- Misses       1600     1620      +20     
- Partials      387      402      +15     

Files	Coverage Δ
src/atomate2/common/jobs/utils.py	45.45% <ø> (-33.34%)	⬇️
src/atomate2/common/schemas/phonons.py	94.07% <ø> (ø)
src/atomate2/forcefields/__init__.py	100.00% <100.00%> (ø)
src/atomate2/forcefields/schemas.py	96.42% <97.22%> (+4.53%)	⬆️
src/atomate2/forcefields/jobs.py	92.80% <86.20%> (-1.62%)	⬇️
src/atomate2/forcefields/utils.py	90.50% <94.28%> (+6.16%)	⬆️
src/atomate2/forcefields/md.py	85.00% <85.00%> (ø)

[JaGeo]

@esoteric-ephemera great that you started this. One thing that will likely be really brittle are the trajectories. We need to find a solution for this in the taskdocs before merging this PR. This is already a problem for normal optimizations (okay, our structures are large but still).
We had some discussions on this already in #515

[utf]

Firstly, thanks very much for this @esoteric-ephemera. It looks fantastic.

As I just commented in #515, we have the ability to store trajectories in the datastore for VASP workflows (see materialsproject/emmet#886 for implementation details). I think we should add a similar option for the force field workflow and probably use the datastore by default since ML MD runs are likely to be much longer than AIMD.

[esoteric-ephemera]

Thanks @JaGeo and @utf! I tried to mimic the store_trajectory functionality: by default, only energies, forces, stresses, and magmoms are stored in the TaskDoc ionic steps for FF MD. I can try to add that as a tag to the ForcefieldTaskDocument for consistency.

On this point: for FF MD, one only needs the structure to reconstruct those other four quantities, so maybe storing just the structure for FF MD makes more sense (probably less data-efficient)

[chiang-yuan]

One thing worth mentioning is when I trained MACE MP on MPtrj, the original file compiled by Bowen in JSON format is around 9 times bigger than extxyz. After compression the size difference is even three-fold. pymatgen trajectory has too many repetitive keys and brackets…. ASE trajectory is written in binary so it should be even more efficient than extxyz. We may want to have a separate MDTaskDoc supporting ASE Trajectory binary

[chiang-yuan]

@esoteric-ephemera temperature and pressure schedule tests are working now. Not sure if we need to test all the calculators. It seems redundant to me if we only want to make sure schedule feature is working

[janosh]

this looks like a float comparison (except for structure)? if so, should this use pytest.approx (and getattr)?

[esoteric-ephemera]

Changed this to use approx when appropriate. The intent was that two runs should be identical when run either explicitly specifying the dynamics object, or implicitly via str, but that is maybe too optimistic

I tried using getattr, but the forcefield task doc has no defined getattr method, hence the dunder

I can write these out explictly if that's preferable to the dunder. Or try to add a getattr

I see what you mean, replaced the __getattribute__ calls with getattr

[janosh]

code-wise this looks great! thanks again @esoteric-ephemera 👍
only remaining concern is test time. tests on the mlff_md branch take 17 min vs 12 on the main branch. can we reduce the number of steps or system size in tests?

[esoteric-ephemera]

@janosh:

only remaining concern is test time. tests on the mlff_md branch take 17 min vs 12 on the main branch. can we reduce the number of steps or system size in tests?

Tried to reduce the time for this by reducing the number of steps in the temperature / pressure MD tests, as well as only testing for CHGNet (MACE is 3x slow for these tests)

@chiang-yuan: the temperature and pressure scheduling didn't seem right to me. When I ran the tests, the temperature schedule gave me [300, 3000, 3000,...,3000], when it should've been a "smooth" (linear) ramp from 300 to 3000 K, no?

I've modified the logic for the scheduling in a new function ForceFieldMDMaker._interpolate_quantity, can you take a look and let me know if this is what you had in mind?

[chiang-yuan]

Thanks @esoteric-ephemera !! Sorry I confuse the behavior of numpy interp and scipy interp1d. Thanks for the nice catch! I slightly changed a bit the syntax for schedule interpolation and make it seems (at least to me) explainable

[esoteric-ephemera]

Thanks @esoteric-ephemera !! Sorry I confuse the behavior of numpy interp and scipy interp1d. Thanks for the nice catch! I slightly changed a bit the syntax for schedule interpolation and make it seems (at least to me) explainable

Perfect thanks! This looks much cleaner than my edit

[utf]

Hi @esoteric-ephemera. Thanks again for all the edits. Is this good to go now?

I'll go through and see if there are any incompatibilities. QuantumChemist seems to have found a missing dir_name which they patched in #791, so there may be other missing attrs

Did you have a chance to do this?

[esoteric-ephemera]

Hi @esoteric-ephemera. Thanks again for all the edits. Is this good to go now?

I'll go through and see if there are any incompatibilities. QuantumChemist seems to have found a missing dir_name which they patched in #791, so there may be other missing attrs

Did you have a chance to do this?

Confirmed that the following works without issue:

from emmet.core.tasks import TaskDoc

forcefield_taskdoc: ForceFieldTaskDocument = <output of some forcefield job>
vasp_style_taskdoc = TaskDoc(**forcefield_taskdoc.model_dump())

There are some optional fields that aren't populated (calcs_reversed, dir_name, etc.) but those aren't so important given that MD trajectories are being stored in GridFS.

For relax / static jobs, is there interest in having a calcs_reversed field on the forcefield task doc?

[utf]

Thanks for checking - calcs_reversed is specific to the calculator used and won't be part of the universal task document (when we get to addressing #741). As long as the other fields match we should be good to go.

[esoteric-ephemera]

Then yes, I think we're good to go!

[janosh]

i'll go ahead and merge since everyone seems happy with this.

thanks so much @esoteric-ephemera, very excited to have MLFF MD makers in atomate2! 🎉

[utf]

Fantastic, thanks everyone who contributed to this PR and especially @esoteric-ephemera and @chiang-yuan!

[JaGeo]

Great!

Short question to the ML potential crowd: wouldn't a LAMMPS interface be even more beneficial? Especially for MD? I am aware of the atomate2-LAMMPS addon but it currently does not look ready to use (and is even archived: https://github.com/Matgenix/atomate2-lammps).

[janosh]

never used LAMMPS myself, I think OpenMM might be the better fit for GPU-accelerated ML potentials based on what I heard from @orionarcher

[orionarcher]

I've found OpenMM to be MUCH easier to use than LAMMPS and seen it run much faster on GPUs. Though I don't know how many popular ML potentials have been implemented, OpenMM has a plugin for integration with Pytorch.

Caveat: I have only used OpenMM for classical MD and not ML potentials

[JaGeo]

@orionarcher thanks! From the ML community side I always hear the LAMMPS wish. Just wanted to see if there are plans or opinions.

(Cc @ml-evs?)