Atomate2 OpenMM integration & broader classical MD framework

[orionarcher]

Developed by @xperrylinn and @orionarcher

Summary

This PR builds out support for OpenMM with a framework that could be extended to support other MD codes. Namely LAMMPS, Amber, and Gromacs. Rough visual example here

Core ideas:

• Use the openff.interchange.Interchange object as a core engine-agnostic representation of an MD simulation.
• Define a generic ClassicalMDTaskDocument and generic Interchange creation functions, then hand off responsibility for evolving the simulation to workflows for each MD engine. Currently only OpenMM support is implemented.
• Support low-throughput workflows by making workflows that can easily output to local directories.
• Use prev_task to pass meta-data between jobs. Since OpenMM is programmatic, the core information is not stored in the local directory.
• No InputSets. Since OpenMM is programmatic, InputSets and InputGenerators don't make conceptual sense.

What's implemented:

• Integration of Atomate2 and OpenFF for classical MD simulations
• Support for energy minimization, NPT, NVT, and temperature change jobs
• Implementation of annealing and production workflows
• Serialization and deserialization of OpenFF objects (Molecule, Topology, Interchange, Quantity) with monty
• Utility functions for generating OpenFF Interchange objects and creating molecule specifications

To do / open questions:

• Need to add create conditional imports that fail nicely, where are some examples of this?

Future PR

• Store trajectories in additional_stores, if available
• The core openff.interchange.Interchange object can be larger than 16MB MongoDB doc limit. Need to implement a workaround by zipping the object and/or putting it on S3.
• migrate schema dependencies to emmet
• migrate utility dependencies to pymatgen

Related PRs:

• PR to pymatgen would move most of atomate2.classical_md.utils upstream. PR #3729
• PR to emmet would move atomate2.classical_md.schemas upstream. PR #975

Example usage

mol_specs_dicts = [
    {"smile": "CCO", "count": 10, "name": "ethanol"},
    {"smile": "O", "count": 200, "name": "water"},
]
inter_job = generate_interchange(mol_specs_dicts, mass_density=1)

production_maker = ProductionMaker(
    energy_maker=EnergyMinimizationMaker(),
    npt_maker=NPTMaker(steps=300000, pressure=1),
    anneal_maker=AnnealMaker.from_temps_and_steps(
        steps=1000000, anneal_temp=400, final_temp=300
    ),
    nvt_maker=NVTMaker(steps=5000000),
)

production_flow = production_maker.make(
    inter_job.output.interchange, 
    output_dir="my/directory/"
    prev_task=inter_job.output
)

wf = Flow([inter_job, production_flow])

run_locally(wf)

# will put outputs for the whole flow in "my/directory/"

Additional dependencies introduced

• openmm
• openff-toolkit
• openff-interchange
• packmol
• openbabel

These are all necessary for the classical MD setup and execution workflow.

Checklist

Before a pull request can be merged, the following items must be checked:

• Code is in the standard Python style.
  The easiest way to handle this is to run the following in the correct sequence on
  your local machine. Start with running ruff and ruff format on your new code. This will
  automatically reformat your code to PEP8 conventions and fix many linting issues.
• Doc strings have been added in the Numpy docstring format.
  Run ruff on your code.
• Type annotations are highly encouraged. Run mypy to
  type check your code.
• Tests have been added for any new functionality or bug fixes.
• All linting and tests pass.

Note that the CI system will run all the above checks. But it will be much more
efficient if you already fix most errors prior to submitting the PR. It is highly
recommended that you use the pre-commit hook provided in the repository. Simply run
pre-commit install and a check will be run prior to allowing commits.

[orionarcher]

@utf I believe I've addressed all comments and this is ready to merge.

EDIT: hold off actually, I forgot that I disabled the openmm tests because a new emmet release is needed. Pinging the MP team.

[orionarcher]

I've pinned the emmet pre-release to get tests passing. @utf are you good to merge this once the new emmet version is released and tested against? If so, maybe @janosh can merge once that happens so you don't need to circle back.

[janosh]

huge amount of work here! 👍

[orionarcher]

This is now pinned to the most recent version of emmet and ready to merge @utf.

Your previous suggestion to split out OpenMM and OpenFF was very good. They are now independent and the OpenMM workflows support MLFFs and can return structures to better interoperate with the rest of Atomate2.

[janosh]

@utf i think it could help to have the micromamba-dependent MD CI be its own job or even its own test-md.yml workflow. that would enable only running the MD tests when any MD source files change and would also uncouple the remaining CI from any install/env issues micromamba might encounter

[utf]

I'd be happy with that solution!

[orionarcher]

Using micromamba saves us from manually building enumlib so I favor keeping it in the main CI.

[utf]

This really is fantastic. Thank you very much @orionarcher. @janosh, once you're happy I think we can merge.

[rkingsbury]

Amazing work, thanks for sticking with this @orionarcher !