diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index 985c4d3fb7..0184178643 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -31,51 +31,58 @@ jobs:
- 27017:27017
runs-on: ubuntu-latest
+ defaults:
+ run:
+ shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
strategy:
matrix:
python-version: ["3.9", "3.10", "3.11"]
steps:
- - uses: actions/checkout@v4
+ - name: Check out repo
+ uses: actions/checkout@v4
- - uses: actions/setup-python@v5
- with:
- python-version: ${{ matrix.python-version }}
+ - name: Set up micromamba
+ uses: mamba-org/setup-micromamba@main
+
+ - name: Create mamba environment
+ run: |
+ micromamba create -n a2 python=${{ matrix.python-version }} --yes
+
+ - name: Install uv
+ run: micromamba run -n a2 pip install uv
- - name: Install enumlib
+ - name: Install conda dependencies
run: |
- cd ..
- git clone --recursive https://github.com/msg-byu/enumlib.git
- cd enumlib/symlib/src
- export F90=gfortran
- make
- cd ../../src
- make enum.x
- sudo mv enum.x /usr/local/bin/
- cd ..
- sudo cp aux_src/makeStr.py /usr/local/bin/
- continue-on-error: true # This is not critical to succeed.
+ micromamba install -n a2 -c conda-forge enumlib packmol bader openbabel openff-toolkit==0.16.2 openff-interchange==0.3.22 --yes
- name: Install dependencies
run: |
+ micromamba activate a2
python -m pip install --upgrade pip
mkdir -p ~/.abinit/pseudos
cp -r tests/test_data/abinit/pseudos/ONCVPSP-PBE-SR-PDv0.4 ~/.abinit/pseudos
- pip install .[strict,tests,abinit]
- pip install torch-runstats
- pip install --no-deps nequip==0.5.6
+ uv pip install .[strict,tests,abinit]
+ uv pip install torch-runstats
+ uv pip install --no-deps nequip==0.5.6
- name: Install pymatgen from master if triggered by pymatgen repo dispatch
if: github.event_name == 'repository_dispatch' && github.event.action == 'pymatgen-ci-trigger'
- run: pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
+ run: |
+ micromamba activate a2
+ uv pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
- name: Test Notebooks
- run: pytest --nbmake ./tutorials
+ run: |
+ micromamba activate a2
+ pytest --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb
- name: Test
env:
MP_API_KEY: ${{ secrets.MP_API_KEY }}
- run: pytest --cov=atomate2 --cov-report=xml
+ run: |
+ micromamba activate a2
+ pytest --cov=atomate2 --cov-report=xml
- uses: codecov/codecov-action@v1
if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
diff --git a/docs/user/codes/openmm.md b/docs/user/codes/openmm.md
new file mode 100644
index 0000000000..59271e4db7
--- /dev/null
+++ b/docs/user/codes/openmm.md
@@ -0,0 +1,639 @@
+# Installing Atomate2 from source with OpenMM
+
+```bash
+# setting up our conda environment
+>>> conda create -n atomate2 python=3.11
+>>> conda activate atomate2
+
+# installing atomate2
+>>> pip install git+https://github.com/orionarcher/atomate2.git
+
+# installing classical_md dependencies
+>>> conda install -c conda-forge --file .github/classical_md_requirements.txt
+```
+
+Alternatively, if you anticipate regularly updating
+atomate2 from source (which at this point, you should),
+you can clone the repository and install from source.
+
+``` bash
+# installing atomate2
+>>> git clone https://github.com/orionarcher/atomate2.git
+>>> cd atomate2
+>>> git branch openff
+>>> git checkout openff
+>>> git pull origin openff
+>>> pip install -e .
+```
+
+To test the openmm installation, you can run the following command. If
+you intend to run on GPU, make sure that the tests are passing for CUDA.
+
+```bash
+>>> python -m openmm.testInstallation
+```
+
+# Understanding Atomate2 OpenMM
+
+Atomate2 is really just a collection of jobflow workflows relevant to
+materials science. In all the workflows, we pass our system of interest
+between different jobs to perform the desired simulation. Representing the
+intermediate state of a classical molecular dynamics simulation, however,
+is challenging. While the intermediate representation between stages of
+a periodic DFT simulation can include just the elements, xyz coordinates,
+and box vectors, classical molecular dynamics systems must also include
+velocities and forces. The latter is particularly challenging because
+all MD engines represent forces differently. Rather than implement our
+own representation, we use the `openff.interchange.Interchange` object,
+which catalogs the necessary system properties and interfaces with a
+variety of MD engines. This is the object that we pass between stages of
+a classical MD simulation and it is the starting point of our workflow.
+
+### Setting up the system
+
+The first job we need to create generates the `Interchange` object.
+To specify the system of interest, we use give it the SMILES strings,
+counts, and names (optional) of the molecules we want to include.
+
+
+```python
+from atomate2.openff.core import generate_interchange
+
+mol_specs_dicts = [
+ {"smiles": "O", "count": 200, "name": "water"},
+ {"smiles": "CCO", "count": 10, "name": "ethanol"},
+ {"smiles": "C1=C(C=C(C(=C1O)O)O)C(=O)O", "count": 1, "name": "gallic_acid"},
+]
+
+gallic_interchange_job = generate_interchange(mol_specs_dicts, 1.3)
+```
+
+If you are wondering what arguments are allowed in the dictionaries, check
+out the `create_mol_spec` function in the `atomate2.openff.utils`
+module. Under the hood, this is being called on each mol_spec dict.
+Meaning the code below is functionally identical to the code above.
+
+
+```python
+from atomate2.openff.utils import create_mol_spec
+
+mols_specs = [create_mol_spec(**mol_spec_dict) for mol_spec_dict in mol_specs_dicts]
+
+generate_interchange(mols_specs, 1.3)
+```
+
+In a more complex simulation we might want to scale the ion charges
+and include custom partial charges. An example with a EC:EMC:LiPF6
+electrolyte is shown below. This yields the `elyte_interchange_job`
+object, which we can pass to the next stage of the simulation.
+
+NOTE: It's actually mandatory to include partial charges
+for PF6- here, the built in partial charge method fails.
+
+
+```python
+import numpy as np
+from pymatgen.core.structure import Molecule
+
+
+pf6 = Molecule(
+ ["P", "F", "F", "F", "F", "F", "F"],
+ [
+ [0.0, 0.0, 0.0],
+ [1.6, 0.0, 0.0],
+ [-1.6, 0.0, 0.0],
+ [0.0, 1.6, 0.0],
+ [0.0, -1.6, 0.0],
+ [0.0, 0.0, 1.6],
+ [0.0, 0.0, -1.6],
+ ],
+)
+pf6_charges = np.array([1.34, -0.39, -0.39, -0.39, -0.39, -0.39, -0.39])
+
+mol_specs_dicts = [
+ {"smiles": "C1COC(=O)O1", "count": 100, "name": "EC"},
+ {"smiles": "CCOC(=O)OC", "count": 100, "name": "EMC"},
+ {
+ "smiles": "F[P-](F)(F)(F)(F)F",
+ "count": 50,
+ "name": "PF6",
+ "partial_charges": pf6_charges,
+ "geometry": pf6,
+ "charge_scaling": 0.8,
+ "charge_method": "RESP",
+ },
+ {"smiles": "[Li+]", "count": 50, "name": "Li", "charge_scaling": 0.8},
+]
+
+elyte_interchange_job = generate_interchange(mol_specs_dicts, 1.3)
+```
+
+### Running a basic simulation
+
+To run a production simulation, we will create a production flow,
+link it to our `elyte_interchange_job`, and then run both locally.
+
+In jobflow, jobs and flows are created by
+[Makers](https://materialsproject.github.io/jobflow/tutorials/6-makers.html),
+which can then be linked into more complex flows. Here, `OpenMMFlowMaker` links
+together makers for energy minimization, pressure equilibration, annealing,
+and a nvt simulation. The annealing step is a subflow that saves us from manually
+instantiating three separate jobs.
+
+Finally, we create our production flow and link to the `generate_interchange` job,
+yielding a production ready molecular dynamics workflow.
+
+```python
+from atomate2.openmm.flows.core import OpenMMFlowMaker
+from atomate2.openmm.jobs.core import (
+ EnergyMinimizationMaker,
+ NPTMaker,
+ NVTMaker,
+)
+from jobflow import Flow, run_locally
+
+
+production_maker = OpenMMFlowMaker(
+ name="production_flow",
+ makers=[
+ EnergyMinimizationMaker(traj_interval=10, state_interval=10),
+ NPTMaker(n_steps=100),
+ OpenMMFlowMaker.anneal_flow(n_steps=150),
+ NVTMaker(n_steps=100),
+ ],
+)
+
+production_flow = production_maker.make(
+ elyte_interchange_job.output.interchange,
+ prev_dir=elyte_interchange_job.output.dir_name,
+ output_dir="./tutorial_system",
+)
+
+run_locally(Flow([elyte_interchange_job, production_flow]))
+```
+
+Above, we are running a very short simulation (350 steps total) and reporting out
+the trajectory and state information very frequently. For a more realistic
+simulation, see the "Configuring the Simulation" section below.
+
+When the above code is executed, you should expect to see this within the
+`tutorial_system` directory:
+
+```
+/tutorial_system
+├── state.csv
+├── state2.csv
+├── state3.csv
+├── state4.csv
+├── state5.csv
+├── state6.csv
+├── taskdoc.json
+├── trajectory.dcd
+├── trajectory2.dcd
+├── trajectory3.dcd
+├── trajectory4.dcd
+├── trajectory5.dcd
+├── trajectory6.dcd
+```
+
+Each job saved a separate state and trajectory file. There are 6 because
+the anneal flow creates 3 sub-jobs and the `EnergyMinimizationMaker`
+does not report anything. The `taskdoc.json` file contains the metadata
+for the entire workflow.
+
+Awesome! At this point, we've run a workflow and could start analyzing
+our data. Before we get there though, let's go through some of the
+other simulation options available.
+
+# Digging Deeper
+
+Atomate2 OpenMM supports running a variety of workflows with different
+configurations. Below we dig in to some of the more advanced options.
+
+
+### Configuring the Simulation
+
+Learn more about the configuration of OpenMM simulations
+
+All OpenMM jobs, i.e. anything in `atomate2.openmm.jobs`, inherits
+from the `BaseOpenMMMaker` class. `BaseOpenMMMaker` is highly configurable, you
+can change the timestep, temperature, reporting frequencies, output types, and
+a range of other properties. See the docstring for the full list of options.
+
+Note that when instantiating the `OpenMMFlowMaker` above, we only set the
+`traj_interval` and `state_interval` once, inside `EnergyMinimizationMaker`.
+This is a key feature: all makers will inherit attributes from the previous
+maker if they are not explicitly reset. This allows you to set the timestep
+once and have it apply to all stages of the simulation. The value inheritance
+is as follows: 1) any explicitly set value, 2) the value from the previous
+maker, 3) the default value (as shown below).
+
+
+```python
+from atomate2.openmm.jobs.base import OPENMM_MAKER_DEFAULTS
+
+print(OPENMM_MAKER_DEFAULTS)
+```
+
+```
+{
+ "step_size": 0.001,
+ "temperature": 298,
+ "friction_coefficient": 1,
+ "platform_name": "CPU",
+ "platform_properties": {},
+ "state_interval": 1000,
+ "state_file_name": "state",
+ "traj_interval": 10000,
+ "wrap_traj": False,
+ "report_velocities": False,
+ "traj_file_name": "trajectory",
+ "traj_file_type": "dcd",
+ "embed_traj": False,
+}
+```
+
+Perhaps we want to embed the trajectory in the taskdoc, so that it
+can be saved to the database, but only for our final run so we don't
+waste space. AND we also want to add some tags, so we can identify
+the simulation in our database more easily. Finally, we want to run
+for much longer, more appropriate for a real production workflow.
+
+```python
+production_maker = OpenMMFlowMaker(
+ name="production_flow",
+ tags=["tutorial_production_flow"],
+ makers=[
+ EnergyMinimizationMaker(traj_interval=0),
+ NPTMaker(n_steps=1000000),
+ OpenMMFlowMaker.anneal_flow(n_steps=1500000),
+ NVTMaker(n_steps=5000000, traj_interval=10000, embed_traj=True),
+ ],
+)
+
+production_flow = production_maker.make(
+ elyte_interchange_job.output.interchange,
+ prev_dir=elyte_interchange_job.output.dir_name,
+ output_dir="./tutorial_system",
+)
+
+run_locally(Flow([elyte_interchange_job, production_flow]))
+```
+
+
+
+### Running with Databases
+
+
+Learn to upload your MD data to databases
+
+Before trying this, you should have a basic understanding of JobFlow
+and [Stores](https://materialsproject.github.io/jobflow/stores.html).
+
+To log OpenMM results to a database, you'll need to set up both a MongoStore,
+for taskdocs, and blob storage, for trajectories. Here, I'll show you the
+correct jobflow.yaml file to use the MongoDB storage and MinIO S3 storage
+provided by NERSC. To get this up, you'll need to contact NERSC to get accounts
+on their MongoDB and MinIO services. Then you can follow the instructions in
+the [Stores](https://materialsproject.github.io/jobflow/stores.html) tutorial
+to link jobflow to your databases. Your `jobflow.yaml` should look like this:
+
+```yaml
+JOB_STORE:
+ docs_store:
+ type: MongoStore
+ database: DATABASE
+ collection_name: atomate2_docs # suggested
+ host: mongodb05.nersc.gov
+ port: 27017
+ username: USERNAME
+ password: PASSWORD
+
+ additional_stores:
+ data:
+ type: S3Store
+ index:
+ type: MongoStore
+ database: DATABASE
+ collection_name: atomate2_blobs_index # suggested
+ host: mongodb05.nersc.gov
+ port: 27017
+ username: USERNAME
+ password: PASSWORD
+ key: blob_uuid
+ bucket: oac
+ s3_profile: oac
+ s3_resource_kwargs:
+ verify: false
+ endpoint_url: https://next-gen-minio.materialsproject.org/
+ key: blob_uuid
+```
+
+NOTE: This can work with any MongoDB and S3 storage, not just NERSC's.
+
+As shown in the production example above, you'll need to set the `embed_traj`
+property to `True` in any makers where you want to save the trajectory to
+the database. Otherwise, the trajectory will only be saved locally.
+
+Rather than use `jobflow.yaml`, you could also create the stores in
+Python and pass the stores to the `run_locally` function. This is a bit
+more code, so usually the prior method is preferred.
+
+
+```python
+from jobflow import run_locally, JobStore
+from maggma.stores import MongoStore, S3Store
+
+mongo_info = {
+ "username": "USERNAME",
+ "password": "PASSWORD",
+ "database": "DATABASE",
+ "host": "mongodb05.nersc.gov",
+}
+
+md_doc_store = MongoStore(**mongo_info, collection_name="atomate2_docs")
+
+md_blob_index = MongoStore(
+ **mongo_info,
+ collection_name="atomate2_blobs_index",
+ key="blob_uuid",
+)
+
+md_blob_store = S3Store(
+ index=md_blob_index,
+ bucket="BUCKET",
+ s3_profile="PROFILE",
+ endpoint_url="https://next-gen-minio.materialsproject.org",
+ key="blob_uuid",
+)
+
+# run our previous flow with the new stores
+run_locally(
+ Flow([elyte_interchange_job, production_flow]),
+ store=JobStore(md_doc_store, additional_stores={"data": md_blob_store}),
+ ensure_success=True,
+)
+```
+
+
+### Running on GPUs
+
+
+Learn to accelerate MD simulations with GPUs
+
+
+Running on a GPU is nearly as simple as running on a CPU. The only difference
+is that you need to specify the `platform_properties` argument in the
+`EnergyMinimizationMaker` with the `DeviceIndex` of the GPU you want to use.
+
+
+```python
+production_maker = OpenMMFlowMaker(
+ name="test_production",
+ makers=[
+ EnergyMinimizationMaker(
+ platform_name="CUDA",
+ platform_properties={"DeviceIndex": "0"},
+ ),
+ NPTMaker(),
+ OpenMMFlowMaker.anneal_flow(),
+ NVTMaker(),
+ ],
+)
+```
+
+Some systems (notably perlmutter) have multiple GPUs available on a
+single node. To fully leverage the compute, you'll need to distribute
+4 simulations across the 4 GPUs. A simple way to do this is with MPI.
+
+First you'll need to install mpi4py.
+
+```bash
+>>> conda install mpi4py
+```
+
+Then you can modify and run the following script to distribute the work across the GPUs.
+
+
+```python
+# other imports
+
+from mpi4py import MPI
+
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+
+list_of_mol_spec_lists = []
+# logic to add four mol_spec_lists to list_of_mol_spec_lists
+
+
+flows = []
+for i in range(4):
+ device_index = i
+ mol_specs = list_of_mol_spec_lists[i]
+
+ setup = generate_interchange(mol_specs, 1.0)
+
+ production_maker = OpenMMFlowMaker(
+ name="test_production",
+ makers=[
+ EnergyMinimizationMaker(
+ platform_name="CUDA",
+ platform_properties={"DeviceIndex": str(device_index)},
+ ),
+ NPTMaker(),
+ OpenMMFlowMaker.anneal_flow(),
+ NVTMaker(),
+ ],
+ )
+
+ production_flow = production_maker.make(
+ setup.output.interchange,
+ prev_dir=setup.output.dir_name,
+ output_dir=f"/pscratch/sd/o/oac/openmm_runs/{i}",
+ )
+ flows.append(Flow([setup, production_flow]))
+
+# this script will run four times, each with a different rank, thus distributing the work across the four GPUs.
+run_locally(flows[rank], ensure_success=True)
+```
+
+
+# Analysis with Emmet
+
+For now, you'll need to make sure you have a particular emmet branch installed.
+Later the builders will be integrated into `main`.
+
+```bash
+pip install git+https://github.com/orionarcher/emmet.git@md_builders
+```
+
+### Analyzing Local Data
+
+
+Learn to analyze your data without a database
+
+Emmet will give us a solid head start on analyzing our data even without touching
+a database. Below, we use emmet to create a [MDAnalysis Universe](https://docs.mdanalysis.org/stable/documentation_pages/core/universe.html#module-MDAnalysis.core.universe)
+and a [SolvationAnalysis Solute](https://solvation-analysis.readthedocs.io/en/latest/api/solute.html).
+From here, we can do all sorts of very cool analysis, but that's beyond the
+scope of this tutorial. Consult the tutorials in SolvationAnalysis and MDAnalysis
+for more information.
+
+```python
+from atomate2.openff.core import ClassicalMDTaskDocument
+from emmet.builders.classical_md.utils import create_universe, create_solute
+from openff.interchange import Interchange
+
+ec_emc_taskdoc = ClassicalMDTaskDocument.parse_file("tutorial_system/taskdoc.json")
+interchange = Interchange.parse_raw(ec_emc_taskdoc.interchange)
+mol_specs = ec_emc_taskdoc.mol_specs
+
+u = create_universe(
+ interchange,
+ mol_specs,
+ str("tutorial_system/trajectory5.dcd"),
+ traj_format="DCD",
+)
+
+solute = create_solute(u, solute_name="Li", networking_solvents=["PF6"])
+```
+
+
+### Setting up builders
+
+
+Connect with your databases
+
+If you followed the instructions above to set up a database, you can
+use the `ElectrolyteBuilder` to perform the same analysis as above.
+
+First, we'll need to create the stores where are data is located,
+these should match the stores you used when running your flow.
+
+```python
+from maggma.stores import MongoStore, S3Store
+
+mongo_info = {
+ "username": "USERNAME",
+ "password": "PASSWORD",
+ "database": "DATABASE",
+ "host": "mongodb05.nersc.gov",
+}
+s3_info = {
+ "bucket": "BUCKET",
+ "s3_profile": "PROFILE",
+ "endpoint_url": "https://next-gen-minio.materialsproject.org",
+}
+
+md_docs = MongoStore(**mongo_info, collection_name="atomate2_docs")
+md_blob_index = MongoStore(
+ **mongo_info,
+ collection_name="atomate2_blobs_index",
+ key="blob_uuid",
+)
+md_blob_store = S3Store(
+ **s3_info,
+ index=md_blob_index,
+ key="blob_uuid",
+)
+```
+
+Now we create our Emmet builder and connect to it. We
+will include a query that will only select jobs with
+the tag "tutorial_production_flow" that we used earlier.
+
+```python
+from emmet.builders.classical_md.openmm.core import ElectrolyteBuilder
+
+builder = ElectrolyteBuilder(
+ md_docs, md_blob_store, query={"output.tags": "tutorial_production_flow"}
+)
+builder.connect()
+```
+
+
+Here are some more convenient queries.
+
+Here are some more convenient queries we could use!
+```python
+# query jobs from a specific day
+april_16 = {"completed_at": {"$regex": "^2024-04-16"}}
+may = {"completed_at": {"$regex": "^2024-05"}}
+
+
+# query a particular set of jobs
+job_uuids = [
+ "3d7b4db4-85e5-48a5-9585-07b37910720f",
+ "4202b18f-f156-4705-8ca6-ac2a08093174",
+ "187d9466-c359-4013-9e25-8b4ece6e3ecf",
+]
+my_specific_jobs = {"uuid": {"$in": job_uuids}}
+```
+
+
+
+
+### Analyzing systems individually
+
+
+
+Download and explore systems one-by-one
+
+To analyze a specific system, you'll need the uuid of the taskdoc you want to
+analyze. We can find the uuids of all the taskdocs in our builder by
+retrieving the items and extracting the uuids.
+
+```python
+items = builder.get_items()
+uuids = [item["uuid"] for item in items]
+```
+
+This, however, can quickly get confusing once you have many jobs.
+At this point, I would highly recommend starting to use an application that
+makes it easier to view and navigate MongoDB databases. I recommend
+[Studio3T](https://robomongo.org/) or [DataGrip](https://www.jetbrains.com/datagrip/).
+
+Now we again use our builder to create a `Universe` and `Solute`. This time
+`instatiate_universe` downloads the trajectory, saves it locally, and uses
+it to create a `Universe`.
+
+```python
+# a query that will grab
+tutorial_query = {"tags": "tutorial_production_flow"}
+
+u = builder.instantiate_universe(uuid, "directory/to/store/trajectory")
+
+solute = create_solute(
+ u,
+ solute_name="Li",
+ networking_solvents=["PF6"],
+ fallback_radius=3,
+)
+```
+
+
+### Automated analysis with builders
+
+
+Do it all for all the systems!
+
+Finally, we'll put the H in high-throughput molecular dynamics. Below,
+we create Stores to hold our `SolvationDocs` and `CalculationDocs` and
+execute the builder on all of our jobs!
+
+Later, there will also be `TransportDocs`, `EquilibrationDocs` and more.
+Aggregating most of what you might want to know about an MD simulation.
+
+```python
+solvation_docs = MongoStore(**mongo_info, collection_name="solvation_docs")
+calculation_docs = MongoStore(**mongo_info, collection_name="calculation_docs")
+builder = ElectrolyteBuilder(md_docs, md_blob_store, solvation_docs, calculation_docs)
+
+builder.connect()
+items = builder.get_items()
+processed_docs = builder.process_items(items)
+builder.update_targets(processed_docs)
+```
+
+
diff --git a/pyproject.toml b/pyproject.toml
index d9a5279bcc..294d05fef8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,6 +58,11 @@ forcefields = [
"quippy-ase>=0.9.14",
"sevenn>=0.9.3",
]
+openmm = [
+ "mdanalysis>=2.7.0",
+ "openmm>=8.1.0",
+ "openmm-mdanalysis-reporter>=0.1.0",
+]
docs = [
"FireWorks==2.0.3",
"autodoc_pydantic==2.1.0",
@@ -107,6 +112,9 @@ strict = [
"sevenn==0.9.3.post1",
"torch==2.2.1",
"typing-extensions==4.12.2",
+ "mdanalysis==2.7.0",
+ "openmm==8.1.1",
+ "openmm-mdanalysis-reporter==0.1.0",
]
[project.scripts]
diff --git a/src/atomate2/openff/__init__.py b/src/atomate2/openff/__init__.py
new file mode 100644
index 0000000000..762a41ab61
--- /dev/null
+++ b/src/atomate2/openff/__init__.py
@@ -0,0 +1,70 @@
+"""Module for classical md workflows."""
+
+from openff.interchange import Interchange
+from openff.toolkit.topology import Topology
+from openff.toolkit.topology.molecule import Molecule
+from openff.units import Quantity
+
+
+def openff_mol_as_monty_dict(self: Molecule) -> dict:
+ """Convert a Molecule to a monty dictionary."""
+ mol_dict = self.to_dict()
+ mol_dict["@module"] = "openff.toolkit.topology"
+ mol_dict["@class"] = "Molecule"
+ return mol_dict
+
+
+Molecule.as_dict = openff_mol_as_monty_dict
+
+
+def openff_topology_as_monty_dict(self: Topology) -> dict:
+ """Convert a Topology to a monty dictionary."""
+ top_dict = self.to_dict()
+ top_dict["@module"] = "openff.toolkit.topology"
+ top_dict["@class"] = "Topology"
+ return top_dict
+
+
+Topology.as_dict = openff_topology_as_monty_dict
+
+
+def openff_interchange_as_monty_dict(self: Interchange) -> dict:
+ """Convert an Interchange to a monty dictionary."""
+ int_dict = self.dict()
+ int_dict["@module"] = "openff.interchange"
+ int_dict["@class"] = "Interchange"
+ return int_dict
+
+
+def openff_interchange_from_monty_dict(cls: type[Interchange], d: dict) -> Interchange:
+ """Construct an Interchange from a monty dictionary."""
+ d = d.copy()
+ d.pop("@module", None)
+ d.pop("@class", None)
+ return cls(**d)
+
+
+Interchange.as_dict = openff_interchange_as_monty_dict
+Interchange.from_dict = classmethod(openff_interchange_from_monty_dict)
+
+
+def openff_quantity_as_monty_dict(self: Quantity) -> dict:
+ """Convert a Quantity to a monty dictionary."""
+ q_tuple = self.to_tuple()
+ q_dict = {"magnitude": q_tuple[0], "unit": q_tuple[1]}
+ q_dict["@module"] = "openff.units"
+ q_dict["@class"] = "Quantity"
+ return q_dict
+
+
+def openff_quantity_from_monty_dict(cls: type[Quantity], d: dict) -> Quantity:
+ """Construct a Quantity from a monty dictionary."""
+ d = d.copy()
+ d.pop("@module", None)
+ d.pop("@class", None)
+ q_tuple = (d["magnitude"], d["unit"])
+ return cls.from_tuple(q_tuple)
+
+
+Quantity.as_dict = openff_quantity_as_monty_dict
+Quantity.from_dict = classmethod(openff_quantity_from_monty_dict)
diff --git a/src/atomate2/openff/core.py b/src/atomate2/openff/core.py
new file mode 100644
index 0000000000..9ef7790f3f
--- /dev/null
+++ b/src/atomate2/openff/core.py
@@ -0,0 +1,164 @@
+"""Core jobs for classical MD module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable
+
+import openff.toolkit as tk
+from emmet.core.openff import ClassicalMDTaskDocument, MoleculeSpec
+from emmet.core.vasp.task_valid import TaskState
+from jobflow import Response, job
+from openff.interchange import Interchange
+from openff.interchange.components._packmol import pack_box
+from openff.toolkit import ForceField
+from openff.units import unit
+
+from atomate2.openff.utils import create_mol_spec_list, merge_specs_by_name_and_smiles
+
+
+def openff_job(method: Callable) -> job:
+ """Decorate the ``make`` method of ClassicalMD job makers.
+
+ This is a thin wrapper around :obj:`~jobflow.core.job.Job` that configures common
+ settings for all ClassicalMD jobs. Namely, configures the output schema to be a
+ :obj:`.ClassicalMDTaskDocument`.
+
+ Any makers that return classical md jobs (not flows) should decorate the ``make``
+ method with @openff_job. For example:
+
+ .. code-block:: python
+
+ class MyClassicalMDMaker(BaseOpenMMMaker):
+ @openff_job
+ def make(structure):
+ # code to run OpenMM job.
+ pass
+
+ Parameters
+ ----------
+ method : callable
+ A BaseVaspMaker.make method. This should not be specified directly and is
+ implied by the decorator.
+
+ Returns
+ -------
+ callable
+ A decorated version of the make function that will generate jobs.
+ """
+ return job(
+ method,
+ output_schema=ClassicalMDTaskDocument,
+ data=["interchange", "traj_blob"],
+ )
+
+
+@openff_job
+def generate_interchange(
+ input_mol_specs: list[MoleculeSpec | dict],
+ mass_density: float,
+ force_field: str = "openff_unconstrained-2.1.1.offxml",
+ pack_box_kwargs: dict = None,
+ tags: list[str] = None,
+) -> Response:
+ """Generate an OpenFF Interchange object from a list of molecule specifications.
+
+ This function takes a list of molecule specifications (either as
+ MoleculeSpec objects or dictionaries), a target mass density, and
+ optional force field and box packing parameters. It processes the molecule
+ specifications, packs them into a box using the specified mass density, and
+ creates an OpenFF Interchange object using the specified force field.
+
+ If you'd like to have multiple distinct input geometries, you
+ can pass multiple mol_specs with the same name and SMILES string.
+ After packing the box, they will be merged into a single mol_spec
+ and treated as a single component in the resulting system.
+
+ Parameters
+ ----------
+ input_mol_specs : List[Union[MoleculeSpec, dict]]
+ A list of molecule specifications, either as MoleculeSpec objects or
+ dictionaries that can be passed to `create_mol_spec` to create
+ MoleculeSpec objects. See the `create_mol_spec` function
+ for details on the expected format of the dictionaries.
+ mass_density : float
+ The target mass density for packing the molecules into
+ a box, kg/L.
+ force_field : str, optional
+ The name of the force field to use for creating the
+ Interchange object. This is passed directly to openff.toolkit.ForceField.
+ Default is "openff_unconstrained-2.1.1.offxml".
+ pack_box_kwargs : Dict, optional
+ Additional keyword arguments to pass to the
+ toolkit.interchange.components._packmol.pack_box. Default is an empty dict.
+ tags : List[str], optional
+ A list of tags to attach to the task document.
+
+ Returns
+ -------
+ ClassicalMDTaskDocument
+ A task document containing the generated OpenFF Interchange
+ object, molecule specifications, and force field information.
+
+ Notes
+ -----
+ - The function assumes that all dictionaries in the mol_specs list can be used to
+ create valid MoleculeSpec objects.
+ - The function sorts the molecule specifications based on their SMILES string
+ and name before packing the box.
+ - The function uses the merge_specs_by_name_and_smiles function to merge molecule
+ specifications with the same name and SMILES string.
+ - The function currently does not support passing a list of force fields due to
+ limitations in the OpenFF Toolkit.
+ """
+ mol_specs = create_mol_spec_list(input_mol_specs)
+
+ pack_box_kwargs = pack_box_kwargs or {}
+ topology = pack_box(
+ molecules=[tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs],
+ number_of_copies=[spec.count for spec in mol_specs],
+ mass_density=mass_density * unit.grams / unit.milliliter,
+ **pack_box_kwargs,
+ )
+
+ mol_specs = merge_specs_by_name_and_smiles(mol_specs)
+
+ # TODO: ForceField doesn't currently support iterables, fix this
+ # force_field: str | Path | List[str | Path] = "openff_unconstrained-2.1.1.offxml",
+
+ # valid FFs: https://github.com/openforcefield/openff-forcefields
+ ff_object = ForceField(force_field)
+
+ interchange = Interchange.from_smirnoff(
+ force_field=ff_object,
+ topology=topology,
+ charge_from_molecules=[
+ tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs
+ ],
+ allow_nonintegral_charges=True,
+ )
+
+ # currently not needed because ForceField isn't correctly supporting iterables
+ # coerce force_field to a str or list of str
+ # if not isinstance(force_field, list):
+ # force_field = [force_field]
+ # ff_list = [ff.name if isinstance(ff, Path) else ff for ff in force_field]
+ # force_field_names = ff_list if len(force_field) > 1 else ff_list[0]
+
+ interchange_json = interchange.json()
+
+ dir_name = Path.cwd()
+
+ task_doc = ClassicalMDTaskDocument(
+ dir_name=str(dir_name),
+ state=TaskState.SUCCESS,
+ interchange=interchange_json,
+ mol_specs=mol_specs,
+ force_field=force_field,
+ tags=tags,
+ )
+
+ with open(dir_name / "taskdoc.json", "w") as file:
+ file.write(task_doc.json())
+
+ return Response(output=task_doc)
diff --git a/src/atomate2/openff/utils.py b/src/atomate2/openff/utils.py
new file mode 100644
index 0000000000..92838a971f
--- /dev/null
+++ b/src/atomate2/openff/utils.py
@@ -0,0 +1,306 @@
+"""Utility functions for classical md subpackage."""
+
+from __future__ import annotations
+
+import copy
+import re
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+import openff.toolkit as tk
+from emmet.core.openff import MoleculeSpec
+from pymatgen.core import Element, Molecule
+from pymatgen.io.openff import create_openff_mol
+
+if TYPE_CHECKING:
+ import pathlib
+
+
+def create_mol_spec(
+ smiles: str,
+ count: int,
+ name: str = None,
+ charge_scaling: float = 1,
+ charge_method: str = None,
+ geometry: Molecule | str | pathlib.Path = None,
+ partial_charges: list[float] = None,
+) -> MoleculeSpec:
+ """Create a MoleculeSpec from a SMILES string and other parameters.
+
+ Constructs an OpenFF Molecule using create_openff_mol and creates a MoleculeSpec
+ with the specified parameters.
+
+ Parameters
+ ----------
+ smiles : str
+ The SMILES string of the molecule.
+ count : int
+ The number of molecules to create.
+ name : str, optional
+ The name of the molecule. If not provided, defaults to the SMILES string.
+ charge_scaling : float, optional
+ The scaling factor for partial charges. Default is 1.
+ charge_method : str, optional
+ The charge method to use if partial charges are not provided. If not specified,
+ defaults to "custom" if partial charges are provided, else "am1bcc".
+ geometry : Union[pymatgen.core.Molecule, str, Path], optional
+ The geometry to use for adding conformers. Can be a Pymatgen Molecule, file path
+ or None.
+ partial_charges : List[float], optional
+ A list of partial charges to assign, or None to use the charge method.
+
+ Returns
+ -------
+ MoleculeSpec
+ The created MoleculeSpec
+ """
+ if charge_method is None:
+ charge_method = "custom" if partial_charges is not None else "am1bcc"
+
+ openff_mol = create_openff_mol(
+ smiles,
+ geometry,
+ charge_scaling,
+ partial_charges,
+ charge_method,
+ )
+
+ # create mol_spec
+ return MoleculeSpec(
+ name=(name or smiles),
+ count=count,
+ charge_scaling=charge_scaling,
+ charge_method=charge_method,
+ openff_mol=openff_mol.to_json(),
+ )
+
+
+def create_mol_spec_list(
+ input_mol_specs: list[MoleculeSpec | dict],
+) -> list[MoleculeSpec]:
+ """
+ Coerce and sort a MoleculeSpecs and dicts to MoleculeSpecs.
+
+ Will sort alphabetically based on concatenated smiles and name.
+
+ Parameters
+ ----------
+ input_mol_specs : list[dict | MoleculeSpec]
+ List of dicts or MoleculeSpecs to coerce and sort.
+
+ Returns
+ -------
+ List[MoleculeSpec]
+ List of MoleculeSpecs sorted by smiles and name.
+ """
+ mol_specs = []
+
+ for spec in input_mol_specs:
+ if isinstance(spec, dict):
+ mol_specs.append(create_mol_spec(**spec))
+ elif isinstance(spec, MoleculeSpec):
+ mol_specs.append(copy.deepcopy(spec))
+ else:
+ raise TypeError(
+ f"item in mol_specs is a {type(spec)}, but mol_specs "
+ f"must be a list of dicts or MoleculeSpec"
+ )
+
+ mol_specs.sort(
+ key=lambda x: tk.Molecule.from_json(x.openff_mol).to_smiles() + x.name
+ )
+
+ return mol_specs
+
+
+def merge_specs_by_name_and_smiles(mol_specs: list[MoleculeSpec]) -> list[MoleculeSpec]:
+ """Merge MoleculeSpecs with the same name and SMILES string.
+
+ Groups MoleculeSpecs by their name and SMILES string, and merges the counts of specs
+ with matching name and SMILES. Returns a list of unique MoleculeSpecs.
+
+ Parameters
+ ----------
+ mol_specs : List[MoleculeSpec]
+ A list of MoleculeSpecs to merge.
+
+ Returns
+ -------
+ List[MoleculeSpec]
+ A list of merged MoleculeSpecs with unique name and SMILES combinations.
+ """
+ mol_specs = copy.deepcopy(mol_specs)
+ merged_spec_dict: dict[tuple[str, str], MoleculeSpec] = {}
+ for spec in mol_specs:
+ key = (tk.Molecule.from_json(spec.openff_mol).to_smiles(), spec.name)
+ if key in merged_spec_dict:
+ merged_spec_dict[key].count += spec.count
+ else:
+ merged_spec_dict[key] = spec
+ return list(merged_spec_dict.values())
+
+
+def calculate_elyte_composition(
+ solvents: dict[str, float],
+ salts: dict[str, float],
+ solvent_densities: dict = None,
+ solvent_ratio_dimension: Literal["mass", "volume"] = "mass",
+) -> dict[str, float]:
+ """Calculate the normalized mass ratios of an electrolyte solution.
+
+ Parameters
+ ----------
+ solvents : dict
+ Dictionary of solvent SMILES strings and their relative unit fraction.
+ salts : dict
+ Dictionary of salt SMILES strings and their molarities.
+ solvent_densities : dict
+ Dictionary of solvent SMILES strings and their densities (g/ml).
+ solvent_ratio_dimension: optional, str
+ Whether the solvents are included with a ratio of "mass" or "volume"
+
+ Returns
+ -------
+ dict
+ A dictionary containing the normalized mass ratios of molecules in
+ the electrolyte solution.
+ """
+ # Check if all solvents have corresponding densities
+ solvent_densities = solvent_densities or {}
+ if set(solvents) > set(solvent_densities):
+ raise ValueError("solvent_densities must contain densities for all solvents.")
+
+ # convert masses to volumes so we can normalize volume
+ if solvent_ratio_dimension == "mass":
+ solvents = {
+ smile: mass / solvent_densities[smile] for smile, mass in solvents.items()
+ }
+
+ # normalize volume ratios
+ total_vol = sum(solvents.values())
+ solvent_volumes = {smile: vol / total_vol for smile, vol in solvents.items()}
+
+ # Convert volume ratios to mass ratios using solvent densities
+ mass_ratio = {
+ smile: vol * solvent_densities[smile] for smile, vol in solvent_volumes.items()
+ }
+
+ # Calculate the molecular weights of the solvent
+ masses = {el.Z: el.atomic_mass for el in Element}
+ salt_mws = {}
+ for smile in salts:
+ mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
+ salt_mws[smile] = sum([masses[atom.atomic_number] for atom in mol.atoms])
+
+ # Convert salt mole ratios to mass ratios
+ salt_mass_ratio = {
+ salt: molarity * salt_mws[salt] / 1000 for salt, molarity in salts.items()
+ }
+
+ # Combine solvent and salt mass ratios
+ combined_mass_ratio = {**mass_ratio, **salt_mass_ratio}
+
+ # Calculate the total mass
+ total_mass = sum(combined_mass_ratio.values())
+
+ # Normalize the mass ratios
+ return {species: mass / total_mass for species, mass in combined_mass_ratio.items()}
+
+
+def counts_from_masses(species: dict[str, float], n_mol: int) -> dict[str, float]:
+ """Calculate the number of mols needed to yield a given mass ratio.
+
+ Parameters
+ ----------
+ species : list of str
+ Dictionary of species SMILES strings and their relative mass fractions.
+ n_mol : float
+ Total number of mols. Returned array will sum to near n_mol.
+
+
+ Returns
+ -------
+ numpy.ndarray
+ n_mols: Number of each SMILES needed for the given mass ratio.
+ """
+ masses = {el.Z: el.atomic_mass for el in Element}
+
+ mol_weights = []
+ for smile in species:
+ mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
+ mol_weights.append(sum([masses[atom.atomic_number] for atom in mol.atoms]))
+
+ mol_ratio = np.array(list(species.values())) / np.array(mol_weights)
+ mol_ratio /= sum(mol_ratio)
+ return {
+ smile: int(np.round(ratio * n_mol)) for smile, ratio in zip(species, mol_ratio)
+ }
+
+
+def counts_from_box_size(
+ species: dict[str, float], side_length: float, density: float = 0.8
+) -> dict[str, float]:
+ """Calculate the number of molecules needed to fill a box.
+
+ Parameters
+ ----------
+ species : dict of str, float
+ Dictionary of species SMILES strings and their relative mass fractions.
+ side_length : int
+ Side length of the cubic simulation box in nm.
+ density : int, optional
+ Density of the system in g/cm^3. Default is 1 g/cm^3.
+
+ Returns
+ -------
+ dict of str, float
+ Number of each species needed to fill the box with the given density.
+ """
+ masses = {el.Z: el.atomic_mass for el in Element}
+
+ na = 6.02214076e23
+ volume = (side_length * 1e-7) ** 3 # Convert from nm3 to cm^3
+ total_mass = volume * density # grams
+
+ # Calculate molecular weights
+ mol_weights = []
+ for smile in species:
+ mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
+ mol_weights.append(sum([masses[atom.atomic_number] for atom in mol.atoms]))
+ mean_mw = np.mean(mol_weights)
+ n_mol = (total_mass / mean_mw) * na
+
+ # Calculate the number of moles needed for each species
+ mol_ratio = np.array(list(species.values())) / np.array(mol_weights)
+ mol_ratio /= sum(mol_ratio)
+
+ # Convert moles to number of molecules
+ return {
+ smile: int(np.round(ratio * n_mol))
+ for smile, ratio in zip(species.keys(), mol_ratio)
+ }
+
+
+def create_mol_dicts(
+ counts: dict[str, float],
+ ion_charge_scaling: float,
+ name_lookup: dict[str, str] = None,
+ xyz_charge_lookup: dict[str, tuple] = None,
+) -> list[dict]:
+ """Create lists of mol specs from just counts. Still rudimentary."""
+ spec_dicts = []
+ for smile, count in counts.items():
+ spec_dict = {
+ "smile": smile,
+ "count": count,
+ "name": name_lookup.get(smile, smile),
+ }
+ if re.search(r"[+-]", smile):
+ spec_dict["charge_scaling"] = ion_charge_scaling
+ xyz_charge = xyz_charge_lookup.get(smile)
+ if xyz_charge is not None:
+ spec_dict["geometry"] = xyz_charge[0]
+ spec_dict["partial_charges"] = xyz_charge[1]
+ spec_dict["charge_method"] = "RESP"
+ spec_dicts.append(spec_dict)
+ return spec_dicts
diff --git a/src/atomate2/openmm/__init__.py b/src/atomate2/openmm/__init__.py
new file mode 100644
index 0000000000..52c075b2ad
--- /dev/null
+++ b/src/atomate2/openmm/__init__.py
@@ -0,0 +1 @@
+"""Core classes for OpenMM module."""
diff --git a/src/atomate2/openmm/flows/__init__.py b/src/atomate2/openmm/flows/__init__.py
new file mode 100644
index 0000000000..734b8e4d2a
--- /dev/null
+++ b/src/atomate2/openmm/flows/__init__.py
@@ -0,0 +1,3 @@
+"""Core Flows for OpenMM module."""
+
+from atomate2.openmm.flows.core import OpenMMFlowMaker
diff --git a/src/atomate2/openmm/flows/core.py b/src/atomate2/openmm/flows/core.py
new file mode 100644
index 0000000000..15ca62fb62
--- /dev/null
+++ b/src/atomate2/openmm/flows/core.py
@@ -0,0 +1,232 @@
+"""Core flows for OpenMM module."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from emmet.core.openmm import Calculation, OpenMMInterchange, OpenMMTaskDocument
+from jobflow import Flow, Job, Response
+from monty.json import MontyDecoder, MontyEncoder
+
+from atomate2.openmm.jobs.base import openmm_job
+from atomate2.openmm.jobs.core import NVTMaker, TempChangeMaker
+from atomate2.openmm.utils import create_list_summing_to
+
+if TYPE_CHECKING:
+ from openff.interchange import Interchange
+
+ from atomate2.openmm.jobs.base import BaseOpenMMMaker
+
+
+def _get_calcs_reversed(job: Job | Flow) -> list[Calculation | list]:
+ """Unwrap a nested list of calcs from jobs or flows."""
+ if isinstance(job, Flow):
+ return [_get_calcs_reversed(sub_job) for sub_job in job.jobs]
+ return job.output.calcs_reversed
+
+
+def _flatten_calcs(nested_calcs: list) -> list[Calculation]:
+ """Flattening nested calcs."""
+ flattened = []
+ for item in nested_calcs:
+ if isinstance(item, list):
+ flattened.extend(_flatten_calcs(item))
+ else:
+ flattened.append(item)
+ return flattened
+
+
+@openmm_job
+def collect_outputs(
+ prev_dir: str,
+ tags: list[str] | None,
+ job_uuids: list[str],
+ calcs_reversed: list[Calculation | list],
+ task_type: str,
+) -> Response:
+ """Reformat the output of the OpenMMFlowMaker into a OpenMMTaskDocument."""
+ with open(Path(prev_dir) / "taskdoc.json") as file:
+ task_dict = json.load(file, cls=MontyDecoder)
+ task_doc = OpenMMTaskDocument.model_validate(task_dict)
+
+ # this must be done here because we cannot unwrap the calcs
+ # when they are an output reference
+ calcs = _flatten_calcs(calcs_reversed)
+ calcs.reverse()
+ task_doc.calcs_reversed = calcs
+ task_doc.tags = tags
+ task_doc.job_uuids = job_uuids
+ task_doc.task_type = task_type
+
+ with open(Path(task_doc.dir_name) / "taskdoc.json", "w") as file:
+ json.dump(task_doc.model_dump(), file, cls=MontyEncoder)
+
+ return Response(output=task_doc)
+
+
+@dataclass
+class OpenMMFlowMaker:
+ """Run a production simulation.
+
+ This flexible flow links together any flows of OpenMM jobs in
+ a linear sequence.
+
+ Attributes
+ ----------
+ name : str
+ The name of the production job. Default is "production".
+ tags : list[str]
+ Tags to apply to the final job. Will only be applied if collect_jobs is True.
+ makers: list[BaseOpenMMMaker]
+ A list of makers to string together.
+ collect_outputs : bool
+ If True, a final job is added that collects all jobs into a single
+ task document.
+ """
+
+ name: str = "flexible"
+ tags: list[str] = field(default_factory=list)
+ makers: list[BaseOpenMMMaker | OpenMMFlowMaker] = field(default_factory=list)
+ collect_outputs: bool = True
+ final_task_type: str = "collect"
+
+ def make(
+ self,
+ interchange: Interchange | OpenMMInterchange | str,
+ prev_dir: str | None = None,
+ ) -> Flow:
+ """Run the production simulation using the provided Interchange object.
+
+ Parameters
+ ----------
+ interchange : Interchange
+ The Interchange object containing the system
+ to simulate.
+ prev_task : Optional[ClassicalMDTaskDocument]
+ The directory of the previous task.
+ output_dir : Optional[Union[str, Path]]
+ The directory to write reporter files to.
+
+ Returns
+ -------
+ Flow
+ A Flow object containing the OpenMM jobs for the simulation.
+ """
+ if len(self.makers) == 0:
+ raise ValueError("At least one maker must be included")
+
+ jobs: list = []
+ job_uuids: list = []
+ calcs_reversed = []
+ for maker in self.makers:
+ job = maker.make(
+ interchange=interchange,
+ prev_dir=prev_dir,
+ )
+ interchange = job.output.interchange
+ prev_dir = job.output.dir_name
+ jobs.append(job)
+
+ # collect the uuids and calcs for the final collect job
+ if isinstance(job, Flow):
+ job_uuids.extend(job.job_uuids)
+ else:
+ job_uuids.append(job.uuid)
+ calcs_reversed.append(_get_calcs_reversed(job))
+
+ if self.collect_outputs:
+ collect_job = collect_outputs(
+ prev_dir,
+ tags=self.tags or None,
+ job_uuids=job_uuids,
+ calcs_reversed=calcs_reversed,
+ task_type=self.final_task_type,
+ )
+ jobs.append(collect_job)
+
+ return Flow(
+ jobs,
+ output=collect_job.output,
+ )
+ return Flow(
+ jobs,
+ output=job.output,
+ )
+
+ @classmethod
+ def anneal_flow(
+ cls,
+ name: str = "anneal",
+ tags: list[str] | None = None,
+ anneal_temp: int = 400,
+ final_temp: int = 298,
+ n_steps: int | tuple[int, int, int] = 1500000,
+ temp_steps: int | tuple[int, int, int] | None = None,
+ job_names: tuple[str, str, str] = ("raise temp", "hold temp", "lower temp"),
+ **kwargs,
+ ) -> OpenMMFlowMaker:
+ """Create an AnnealMaker from the specified temperatures, steps, and job names.
+
+ Parameters
+ ----------
+ name : str, optional
+ The name of the annealing job. Default is "anneal".
+ tags : list[str], optional
+ Tags to apply to the final job.
+ anneal_temp : int, optional
+ The annealing temperature. Default is 400.
+ final_temp : int, optional
+ The final temperature after annealing. Default is 298.
+ n_steps : int or Tuple[int, int, int], optional
+ The number of steps for each stage of annealing.
+ If an integer is provided, it will be divided into three equal parts.
+ If a tuple of three integers is provided, each value represents the
+ steps for the corresponding stage. Default is 1500000.
+ temp_steps : int or Tuple[int, int, int], optional
+ The number of temperature steps for raising and
+ lowering the temperature. If an integer is provided, it will be used
+ for both stages. If a tuple of three integers is provided, each value
+ represents the temperature steps for the corresponding stage.
+ Default is None and all jobs will automatically determine temp_steps.
+ job_names : Tuple[str, str, str], optional
+ The names for the jobs in each stage of annealing.
+ Default is ("raise temp", "hold temp", "lower temp").
+ **kwargs
+ Additional keyword arguments to be passed to the job makers.
+
+ Returns
+ -------
+ AnnealMaker
+ An AnnealMaker instance with the specified parameters.
+ """
+ if isinstance(n_steps, int):
+ n_steps = tuple(create_list_summing_to(n_steps, 3))
+ if isinstance(temp_steps, int) or temp_steps is None:
+ temp_steps = (temp_steps, temp_steps, temp_steps)
+
+ raise_temp_maker = TempChangeMaker(
+ n_steps=n_steps[0],
+ name=job_names[0],
+ temperature=anneal_temp,
+ temp_steps=temp_steps[0],
+ **kwargs,
+ )
+ nvt_maker = NVTMaker(
+ n_steps=n_steps[1], name=job_names[1], temperature=anneal_temp, **kwargs
+ )
+ lower_temp_maker = TempChangeMaker(
+ n_steps=n_steps[2],
+ name=job_names[2],
+ temperature=final_temp,
+ temp_steps=temp_steps[2],
+ **kwargs,
+ )
+ return cls(
+ name=name,
+ tags=tags,
+ makers=[raise_temp_maker, nvt_maker, lower_temp_maker],
+ collect_outputs=False,
+ )
diff --git a/src/atomate2/openmm/jobs/__init__.py b/src/atomate2/openmm/jobs/__init__.py
new file mode 100644
index 0000000000..a3a085ad75
--- /dev/null
+++ b/src/atomate2/openmm/jobs/__init__.py
@@ -0,0 +1,8 @@
+"""Definitions of OpenMM jobs."""
+
+from atomate2.openmm.jobs.core import (
+ EnergyMinimizationMaker,
+ NPTMaker,
+ NVTMaker,
+ TempChangeMaker,
+)
diff --git a/src/atomate2/openmm/jobs/base.py b/src/atomate2/openmm/jobs/base.py
new file mode 100644
index 0000000000..01ae0bb136
--- /dev/null
+++ b/src/atomate2/openmm/jobs/base.py
@@ -0,0 +1,623 @@
+"""The base class for OpenMM simulation makers."""
+
+from __future__ import annotations
+
+import copy
+import json
+import time
+import warnings
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, NoReturn
+
+from emmet.core.openmm import (
+ Calculation,
+ CalculationInput,
+ CalculationOutput,
+ OpenMMInterchange,
+ OpenMMTaskDocument,
+)
+from jobflow import Maker, Response, job
+from mdareporter.mdareporter import MDAReporter
+from monty.json import MontyDecoder, MontyEncoder
+from openmm import Integrator, LangevinMiddleIntegrator, Platform, XmlSerializer
+from openmm.app import StateDataReporter
+from openmm.unit import angstrom, kelvin, picoseconds
+from pymatgen.core import Structure
+
+from atomate2.openmm.utils import increment_name, task_reports
+
+if TYPE_CHECKING:
+ from openmm.app.simulation import Simulation
+
+
+try:
+ # so we can load OpenMM Interchange created by openmmml
+ import openmmml
+except ImportError:
+ openmmml = None
+
+try:
+ from openff.interchange import Interchange
+except ImportError:
+
+ class Interchange: # type: ignore[no-redef]
+ """Dummy class for failed imports of Interchange."""
+
+ def model_validate(self, _: str) -> None:
+ """Parse raw is the first method called on the Interchange object."""
+ raise ImportError(
+ "openff-interchange must be installed for OpenMM makers to"
+ "to support OpenFF Interchange objects."
+ )
+
+
+OPENMM_MAKER_DEFAULTS = {
+ "step_size": 0.001,
+ "temperature": 298,
+ "friction_coefficient": 1,
+ "platform_name": "CPU",
+ "platform_properties": {},
+ "state_interval": 1000,
+ "state_file_name": "state",
+ "traj_interval": 10000,
+ "wrap_traj": False,
+ "report_velocities": False,
+ "traj_file_name": "trajectory",
+ "traj_file_type": "dcd",
+ "embed_traj": False,
+ "save_structure": False,
+}
+
+
+def openmm_job(method: Callable) -> job:
+ """Decorate the ``make`` method of ClassicalMD job makers.
+
+ This is a thin wrapper around :obj:`~jobflow.core.job.Job` that configures common
+ settings for all ClassicalMD jobs. Namely, configures the output schema to be a
+ :obj:`.ClassicalMDTaskDocument`.
+
+ Any makers that return classical md jobs (not flows) should decorate the ``make``
+ method with @openff_job. For example:
+
+ .. code-block:: python
+
+ class MyClassicalMDMaker(BaseOpenMMMaker):
+ @openff_job
+ def make(structure):
+ # code to run OpenMM job.
+ pass
+
+ Parameters
+ ----------
+ method : callable
+ A BaseVaspMaker.make method. This should not be specified directly and is
+ implied by the decorator.
+
+ Returns
+ -------
+ callable
+ A decorated version of the make function that will generate jobs.
+ """
+ return job(
+ method,
+ output_schema=OpenMMTaskDocument,
+ data=["interchange", "traj_blob"],
+ )
+
+
+@dataclass
+class BaseOpenMMMaker(Maker):
+ """Base class for OpenMM simulation makers.
+
+ This class provides a foundation for creating OpenMM simulation
+ makers. It includes common attributes and methods for setting up,
+ running, and closing OpenMM simulations. Subclasses can override
+ the run_openmm method to define specific simulation logic.
+
+ In general, any missing values will be taken from the
+ previous task, if possible, and the default values defined in
+ atomate2.openmm.OPENMM_MAKER_DEFAULTS, if not.
+
+ Attributes
+ ----------
+ name : str
+ The name of the OpenMM job.
+ tags : Optional[List[str]]
+ Tags for the OpenMM job.
+ n_steps : Optional[int]
+ The number of simulation steps to run.
+ step_size : Optional[float]
+ The size of each simulation step (picoseconds).
+ temperature : Optional[float]
+ The simulation temperature (kelvin).
+ friction_coefficient : Optional[float]
+ The friction coefficient for the
+ integrator (inverse picoseconds).
+ platform_name : Optional[str]
+ The name of the OpenMM platform to use, passed to
+ Interchange.to_openmm_simulation.
+ platform_properties : Optional[dict]
+ Properties for the OpenMM platform,
+ passed to Interchange.to_openmm_simulation.
+ state_interval : Optional[int]
+ The interval for saving simulation state.
+ To record no state, set to 0.
+ state_file_name : Optional[str]
+ The name of the state file to save.
+ traj_interval : Optional[int]
+ The interval for saving trajectory frames. To record
+ no trajectory, set to 0.
+ wrap_traj : Optional[bool]
+ Whether to wrap trajectory coordinates.
+ report_velocities : Optional[bool]
+ Whether to report velocities in the trajectory file.
+ traj_file_name : Optional[str]
+ The name of the trajectory file to save.
+ traj_file_type : Optional[str]
+ The type of trajectory file to save. Supports any output format
+ supported by MDAnalysis.
+ embed_traj : Optional[bool]
+ Whether to embed the trajectory in the task document.
+ save_structure : Optional[bool]
+ Whether to save the final structure in the task document.
+ """
+
+ name: str = "base openmm job"
+ tags: list[str] | None = field(default=None)
+ n_steps: int | None = field(default=None)
+ step_size: float | None = field(default=None)
+ temperature: float | None = field(default=None)
+ friction_coefficient: float | None = field(default=None)
+ platform_name: str | None = field(default=None)
+ platform_properties: dict | None = field(default=None)
+ state_interval: int | None = field(default=None)
+ state_file_name: str | None = field(default=None)
+ traj_interval: int | None = field(default=None)
+ wrap_traj: bool | None = field(default=None)
+ report_velocities: bool | None = field(default=None)
+ traj_file_name: str | None = field(default=None)
+ traj_file_type: str | None = field(default=None)
+ embed_traj: bool | None = field(default=None)
+ save_structure: bool | None = field(default=None)
+
+ @openmm_job
+ def make(
+ self,
+ interchange: Interchange | OpenMMInterchange | str,
+ prev_dir: str | None = None,
+ ) -> Response:
+ """Run an OpenMM calculation.
+
+ This method sets up an OpenMM simulation, runs the simulation based
+ on the specific job logic defined in run_openmm, and closes the
+ simulation. It returns a response containing the output task document.
+
+ Parameters
+ ----------
+ interchange : Union[Interchange, OpenMMInterchange, str]
+ An Interchange object, OpenMMInterchange object or byte encoded equivalent.
+ prev_task : Optional[str]
+ The directory of the previous task.
+
+ Returns
+ -------
+ Response
+ A response object containing the output task document.
+ """
+ if prev_dir:
+ with open(Path(prev_dir) / "taskdoc.json") as file:
+ task_dict = json.load(file, cls=MontyDecoder)
+ prev_task = OpenMMTaskDocument.model_validate(task_dict)
+ else:
+ prev_task = None
+
+ interchange = self._load_interchange(interchange)
+
+ dir_name = Path.cwd()
+
+ sim = self._create_simulation(interchange, prev_task)
+
+ self._add_reporters(sim, dir_name, prev_task)
+
+ # Run the simulation
+ start = time.time()
+ self.run_openmm(sim)
+ elapsed_time = time.time() - start
+
+ self._update_interchange(interchange, sim, prev_task)
+
+ structure = self._create_structure(sim, prev_task)
+
+ task_doc = self._create_task_doc(
+ interchange, structure, elapsed_time, dir_name, prev_task
+ )
+
+ # leaving the MDAReporter makes the builders fail
+ for _ in range(len(sim.reporters)):
+ reporter = sim.reporters.pop()
+ del reporter
+ del sim
+
+ # write out task_doc json to output dir
+ with open(dir_name / "taskdoc.json", "w") as file:
+ json.dump(task_doc.model_dump(), file, cls=MontyEncoder)
+
+ return Response(output=task_doc)
+
+ def _load_interchange(
+ self, interchange: Interchange | OpenMMInterchange | str
+ ) -> Interchange | OpenMMInterchange:
+ """Load an Interchange object from a JSON string or bytes.
+
+ This method loads an Interchange object from a JSON string or bytes
+ representation. It is used to convert interchange data to an Interchange
+ object for use in the OpenMM simulation.
+
+ Parameters
+ ----------
+ interchange : Union[Interchange, str]
+ An Interchange object or a JSON string of the Interchange object.
+
+ Returns
+ -------
+ Interchange
+ The loaded Interchange object.
+ """
+ if isinstance(interchange, str):
+ try:
+ interchange = Interchange.parse_raw(interchange)
+ except: # noqa: E722
+ # parse with openmm instead
+ interchange = OpenMMInterchange.model_validate_json(interchange)
+ else:
+ interchange = copy.deepcopy(interchange)
+ return interchange
+
+ def _add_reporters(
+ self,
+ sim: Simulation,
+ dir_name: Path,
+ prev_task: OpenMMTaskDocument | None = None,
+ ) -> None:
+ """Add reporters to the OpenMM simulation.
+
+ This method adds DCD and state reporters to the OpenMM
+ simulation based on the specified intervals and settings.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ dir_name : Path
+ The directory to save the reporter output files.
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document.
+ """
+ has_steps = self._resolve_attr("n_steps", prev_task) > 0
+ # add trajectory reporter
+ traj_interval = self._resolve_attr("traj_interval", prev_task)
+ traj_file_name = self._resolve_attr("traj_file_name", prev_task)
+ traj_file_type = self._resolve_attr("traj_file_type", prev_task)
+ report_velocities = self._resolve_attr("report_velocities", prev_task)
+
+ if has_steps & (traj_interval > 0):
+ writer_kwargs = {}
+ # these are the only file types that support velocities
+ if traj_file_type in ["h5md", "nc", "ncdf"]:
+ writer_kwargs["velocities"] = report_velocities
+ writer_kwargs["forces"] = False
+ elif report_velocities and traj_file_type not in ["trr"]:
+ raise ValueError(
+ f"File type {traj_file_type} does not support velocities as"
+ f"of MDAnalysis 2.7.0. Select another file type"
+ f"or do not attempt to report velocities."
+ )
+
+ traj_file = dir_name / f"{traj_file_name}.{traj_file_type}"
+
+ if traj_file.exists() and task_reports(prev_task, "traj"):
+ self.traj_file_name = increment_name(traj_file_name)
+
+ # TODO: MDA 2.7.0 has a bug that prevents velocity reporting
+ # this is a stop gap measure before MDA 2.8.0 is released
+ kwargs = dict(
+ file=str(dir_name / f"{self.traj_file_name}.{traj_file_type}"),
+ reportInterval=traj_interval,
+ enforcePeriodicBox=self._resolve_attr("wrap_traj", prev_task),
+ )
+ if report_velocities:
+ # assert package version
+
+ kwargs["writer_kwargs"] = writer_kwargs
+ warnings.warn(
+ "Reporting velocities is only supported with the"
+ "development version of MDAnalysis, >= 2.8.0, "
+ "proceed with caution.",
+ stacklevel=1,
+ )
+ traj_reporter = MDAReporter(**kwargs)
+
+ sim.reporters.append(traj_reporter)
+
+ # add state reporter
+ state_interval = self._resolve_attr("state_interval", prev_task)
+ state_file_name = self._resolve_attr("state_file_name", prev_task)
+ if has_steps & (state_interval > 0):
+ state_file = dir_name / f"{state_file_name}.csv"
+ if state_file.exists() and task_reports(prev_task, "state"):
+ self.state_file_name = increment_name(state_file_name)
+
+ state_reporter = StateDataReporter(
+ file=f"{dir_name / self.state_file_name}.csv",
+ reportInterval=state_interval,
+ step=True,
+ potentialEnergy=True,
+ kineticEnergy=True,
+ totalEnergy=True,
+ temperature=True,
+ volume=True,
+ density=True,
+ )
+ sim.reporters.append(state_reporter)
+
+ def run_openmm(self, simulation: Simulation) -> NoReturn:
+ """Abstract method for running the OpenMM simulation.
+
+ This method should be implemented by subclasses to
+ define the specific simulation logic. It takes an
+ OpenMM simulation object and evolves the simulation.
+
+ Parameters
+ ----------
+ simulation : Simulation
+ The OpenMM simulation object.
+
+ Raises
+ ------
+ NotImplementedError
+ If the method is not implemented by a subclass.
+ """
+ raise NotImplementedError(
+ "`run_openmm` should be implemented by each child class."
+ )
+
+ def _resolve_attr(
+ self,
+ attr: str,
+ prev_task: OpenMMTaskDocument | None = None,
+ add_defaults: dict | None = None,
+ ) -> Any:
+ """Resolve an attribute and set its value.
+
+ This method retrieves the value of an attribute from the current maker,
+ previous task input, or a default value (in that order of priority). It
+ sets the attribute on the current maker and returns the resolved value.
+
+ Default values are defined in `OPENMM_MAKER_DEFAULTS`.
+
+ Parameters
+ ----------
+ attr : str
+ The name of the attribute to resolve.
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document.
+ add_defaults : Optional[dict]
+ Additional default values to use,
+ overrides `OPENMM_MAKER_DEFAULTS`.
+
+ Returns
+ -------
+ Any
+ The resolved attribute value.
+ """
+ prev_task = prev_task or OpenMMTaskDocument()
+
+ # retrieve previous CalculationInput through multiple Optional fields
+ if prev_task.calcs_reversed:
+ prev_input = prev_task.calcs_reversed[0].input
+ else:
+ prev_input = None
+
+ defaults = {**OPENMM_MAKER_DEFAULTS, **(add_defaults or {})}
+
+ if getattr(self, attr, None) is not None:
+ attr_value = getattr(self, attr)
+ elif getattr(prev_input, attr, None) is not None:
+ attr_value = getattr(prev_input, attr)
+ else:
+ attr_value = defaults.get(attr)
+
+ setattr(self, attr, attr_value)
+ return getattr(self, attr)
+
+ def _create_integrator(
+ self,
+ prev_task: OpenMMTaskDocument | None = None,
+ ) -> Integrator:
+ """Create an OpenMM integrator.
+
+ This method creates a Langevin middle integrator based on the
+ resolved temperature, friction coefficient, and step size.
+
+ Parameters
+ ----------
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document.
+
+ Returns
+ -------
+ LangevinMiddleIntegrator
+ The created OpenMM integrator.
+ """
+ return LangevinMiddleIntegrator(
+ self._resolve_attr("temperature", prev_task) * kelvin,
+ self._resolve_attr("friction_coefficient", prev_task) / picoseconds,
+ self._resolve_attr("step_size", prev_task) * picoseconds,
+ )
+
+ def _create_simulation(
+ self,
+ interchange: Interchange | OpenMMInterchange,
+ prev_task: OpenMMTaskDocument | None = None,
+ ) -> Simulation:
+ """Create an OpenMM simulation.
+
+ This method creates an OpenMM simulation using the provided Interchange object,
+ the get_integrator method, and the platform and platform_properties attributes.
+
+ Parameters
+ ----------
+ interchange : Interchange
+ The Interchange object containing the MD data.
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document.
+
+ Returns
+ -------
+ Simulation
+ The created OpenMM simulation object.
+ """
+ integrator = self._create_integrator(prev_task)
+ platform = Platform.getPlatformByName(
+ self._resolve_attr("platform_name", prev_task)
+ )
+ platform_properties = self._resolve_attr("platform_properties", prev_task)
+
+ return interchange.to_openmm_simulation(
+ integrator,
+ platform=platform,
+ platformProperties=platform_properties,
+ )
+
+ def _update_interchange(
+ self,
+ interchange: Interchange | OpenMMInterchange,
+ sim: Simulation,
+ prev_task: OpenMMTaskDocument | None = None,
+ ) -> None:
+ """Update the Interchange object with the current simulation state.
+
+ This method updates the positions, velocities, and box vectors of the
+ Interchange object based on the current state of the OpenMM simulation.
+
+ Parameters
+ ----------
+ interchange : Interchange
+ The Interchange object to update.
+ sim : Simulation
+ The OpenMM simulation object.
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document.
+ """
+ state = sim.context.getState(
+ getPositions=True,
+ getVelocities=True,
+ enforcePeriodicBox=self._resolve_attr("wrap_traj", prev_task),
+ )
+ if isinstance(interchange, Interchange):
+ interchange.positions = state.getPositions(asNumpy=True)
+ interchange.velocities = state.getVelocities(asNumpy=True)
+ interchange.box = state.getPeriodicBoxVectors(asNumpy=True)
+ elif isinstance(interchange, OpenMMInterchange):
+ interchange.state = XmlSerializer.serialize(state)
+
+ def _create_structure(
+ self, sim: Simulation, prev_task: OpenMMTaskDocument | None = None
+ ) -> Structure | None:
+ """Create a pymatgen Structure from the OpenMM simulation.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ """
+ if not self._resolve_attr("save_structure", prev_task):
+ return None
+
+ state = sim.context.getState(
+ getPositions=True,
+ getVelocities=True,
+ enforcePeriodicBox=self._resolve_attr("wrap_traj", prev_task),
+ )
+
+ return Structure(
+ lattice=state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(angstrom),
+ species=[atom.element.symbol for atom in sim.topology.atoms()],
+ coords=state.getPositions(asNumpy=True).value_in_unit(angstrom),
+ coords_are_cartesian=True,
+ )
+
+ def _create_task_doc(
+ self,
+ interchange: Interchange | OpenMMInterchange,
+ structure: Structure | None,
+ elapsed_time: float | None = None,
+ dir_name: Path | None = None,
+ prev_task: OpenMMTaskDocument | None = None,
+ ) -> OpenMMTaskDocument:
+ """Create a task document for the OpenMM job.
+
+ This method creates an OpenMMTaskDocument based on the current
+ maker attributes, previous task document, and simulation results.
+
+ Parameters
+ ----------
+ interchange : Interchange
+ The updated Interchange object.
+ structure : Structure
+ The final structure of the simulation.
+ elapsed_time : Optional[float]
+ The elapsed time of the simulation. Default is None.
+ dir_name : Optional[Path]
+ The directory where the output files are saved.
+ Default is None.
+ prev_task : Optional[OpenMMTaskDocument]
+ The previous task document. Default is None.
+
+ Returns
+ -------
+ OpenMMTaskDocument
+ The created task document.
+ """
+ maker_attrs = copy.deepcopy(dict(vars(self)))
+ job_name = maker_attrs.pop("name")
+ tags = maker_attrs.pop("tags")
+ state_file_name = self._resolve_attr("state_file_name", prev_task)
+ traj_file_name = self._resolve_attr("traj_file_name", prev_task)
+ traj_file_type = self._resolve_attr("traj_file_type", prev_task)
+ calc = Calculation(
+ dir_name=str(dir_name),
+ has_openmm_completed=True,
+ input=CalculationInput(**maker_attrs),
+ output=CalculationOutput.from_directory(
+ dir_name,
+ f"{state_file_name}.csv",
+ f"{traj_file_name}.{traj_file_type}",
+ elapsed_time=elapsed_time,
+ embed_traj=self._resolve_attr("embed_traj", prev_task),
+ ),
+ completed_at=str(datetime.now(tz=timezone.utc)),
+ task_name=job_name,
+ calc_type=self.__class__.__name__,
+ )
+
+ prev_task = prev_task or OpenMMTaskDocument()
+
+ interchange_json = interchange.json()
+ # interchange_bytes = interchange_json.encode("utf-8")
+
+ return OpenMMTaskDocument(
+ tags=tags,
+ dir_name=str(dir_name),
+ state="successful",
+ calcs_reversed=[calc],
+ interchange=interchange_json,
+ mol_specs=prev_task.mol_specs,
+ structure=structure,
+ force_field=prev_task.force_field,
+ task_name=calc.task_name,
+ task_type="test",
+ last_updated=datetime.now(tz=timezone.utc),
+ )
diff --git a/src/atomate2/openmm/jobs/core.py b/src/atomate2/openmm/jobs/core.py
new file mode 100644
index 0000000000..821d1953eb
--- /dev/null
+++ b/src/atomate2/openmm/jobs/core.py
@@ -0,0 +1,220 @@
+"""Core OpenMM jobs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+from openmm import Integrator, LangevinMiddleIntegrator, MonteCarloBarostat
+from openmm.unit import atmosphere, kelvin, kilojoules_per_mole, nanometer, picoseconds
+
+from atomate2.openmm.jobs.base import BaseOpenMMMaker
+from atomate2.openmm.utils import create_list_summing_to
+
+if TYPE_CHECKING:
+ from emmet.core.openmm import OpenMMTaskDocument
+ from openmm.app import Simulation
+
+
+@dataclass
+class EnergyMinimizationMaker(BaseOpenMMMaker):
+ """A maker class for performing energy minimization using OpenMM.
+
+ This class inherits from BaseOpenMMMaker, only new attributes are documented.
+ n_steps must be 0.
+
+ Attributes
+ ----------
+ name : str
+ The name of the energy minimization job.
+ Default is "energy minimization".
+ tolerance : float
+ The energy tolerance for minimization. Default is 10 kj/nm.
+ max_iterations : int
+ The maximum number of minimization iterations.
+ Default is 0, which means no maximum.
+ """
+
+ name: str = "energy minimization"
+ n_steps: int = 0
+ tolerance: float = 10
+ max_iterations: int = 0
+
+ def run_openmm(self, sim: Simulation) -> None:
+ """Run the energy minimization with OpenMM.
+
+ This method performs energy minimization on the molecular system using
+ the OpenMM simulation package. It minimizes the energy of the system
+ based on the specified tolerance and maximum number of iterations.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ """
+ if self.n_steps != 0:
+ raise ValueError("Energy minimization should have 0 steps.")
+
+ # Minimize the energy
+ sim.minimizeEnergy(
+ tolerance=self.tolerance * kilojoules_per_mole / nanometer,
+ maxIterations=self.max_iterations,
+ )
+
+
+@dataclass
+class NPTMaker(BaseOpenMMMaker):
+ """A maker class for performing NPT (isothermal-isobaric) simulations using OpenMM.
+
+ This class inherits from BaseOpenMMMaker, only new attributes are documented.
+
+ Attributes
+ ----------
+ name : str
+ The name of the NPT simulation job. Default is "npt simulation".
+ n_steps : int
+ The number of simulation steps. Default is 1,000,000.
+ pressure : float
+ The pressure of the simulation in atmospheres.
+ Default is 1 atm.
+ pressure_update_frequency : int
+ The number of steps between pressure update attempts.
+ """
+
+ name: str = "npt simulation"
+ n_steps: int = 1_000_000
+ pressure: float = 1
+ pressure_update_frequency: int = 10
+
+ def run_openmm(self, sim: Simulation) -> None:
+ """Evolve the simulation for self.n_steps in the NPT ensemble.
+
+ This adds a Monte Carlo barostat to the system to put it into NPT, runs the
+ simulation for the specified number of steps, and then removes the barostat.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ """
+ # Add barostat to system
+ context = sim.context
+ system = context.getSystem()
+
+ barostat_force_index = system.addForce(
+ MonteCarloBarostat(
+ self.pressure * atmosphere,
+ sim.context.getIntegrator().getTemperature(),
+ self.pressure_update_frequency,
+ )
+ )
+
+ # Re-init the context after adding thermostat to System
+ context.reinitialize(preserveState=True)
+
+ # Run the simulation
+ sim.step(self.n_steps)
+
+ # Remove thermostat and update context
+ system.removeForce(barostat_force_index)
+ context.reinitialize(preserveState=True)
+
+
+@dataclass
+class NVTMaker(BaseOpenMMMaker):
+ """A maker class for performing NVT (canonical ensemble) simulations using OpenMM.
+
+ This class inherits from BaseOpenMMMaker, only new attributes are documented.
+
+ Attributes
+ ----------
+ name : str
+ The name of the NVT simulation job. Default is "nvt simulation".
+ n_steps : int
+ The number of simulation steps. Default is 1,000,000.
+ """
+
+ name: str = "nvt simulation"
+ n_steps: int = 1_000_000
+
+ def run_openmm(self, sim: Simulation) -> None:
+ """Evolve the simulation with OpenMM for self.n_steps.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ """
+ # Run the simulation
+ sim.step(self.n_steps)
+
+
+@dataclass
+class TempChangeMaker(BaseOpenMMMaker):
+ """A maker class for performing simulations with temperature changes using OpenMM.
+
+ This class inherits from BaseOpenMMMaker and provides
+ functionality for running simulations with temperature
+ changes using the OpenMM simulation package.
+
+ Attributes
+ ----------
+ name : str
+ The name of the temperature change job. Default is "temperature change".
+ n_steps : int
+ The total number of simulation steps. Default is 1000000.
+ temp_steps : Optional[int]
+ The number of steps over which the temperature is raised, by
+ default will be set to steps / 10000.
+ starting_temperature : Optional[float]
+ The starting temperature of the simulation.
+ If not provided it will inherit from the previous task.
+ """
+
+ name: str = "temperature change"
+ n_steps: int = 1_000_000
+ temp_steps: int | None = None
+ starting_temperature: float | None = None
+
+ def run_openmm(self, sim: Simulation) -> None:
+ """Evolve the simulation while gradually changing the temperature.
+
+ self.temperature is the final temperature. self.temp_steps
+ determines how many gradiations there are between the starting and
+ final temperature. At each gradiation, the system is evolved for a
+ number of steps such that the total number of steps is self.n_steps.
+
+ Parameters
+ ----------
+ sim : Simulation
+ The OpenMM simulation object.
+ """
+ integrator = sim.context.getIntegrator()
+
+ temps = np.linspace(
+ self.starting_temperature, self.temperature, self.temp_steps
+ )
+ steps = create_list_summing_to(self.n_steps, self.temp_steps)
+ for temp, n_steps in zip(temps, steps):
+ integrator.setTemperature(temp * kelvin)
+ sim.step(n_steps)
+
+ def _create_integrator(
+ self, prev_task: OpenMMTaskDocument | None = None
+ ) -> Integrator:
+ # we resolve this here because prev_task is available
+ temp_steps_default = (self.n_steps // 10000) or 1
+ self.temp_steps = self._resolve_attr(
+ "temp_steps", prev_task, add_defaults={"temp_steps": temp_steps_default}
+ )
+
+ # we do this dance so _resolve_attr takes its value from the previous task
+ temp_holder, self.temperature = self.temperature, None
+ self.starting_temperature = self._resolve_attr("temperature", prev_task)
+ self.temperature = temp_holder or self._resolve_attr("temperature", prev_task)
+ return LangevinMiddleIntegrator(
+ self.starting_temperature * kelvin,
+ self._resolve_attr("friction_coefficient", prev_task) / picoseconds,
+ self._resolve_attr("step_size", prev_task) * picoseconds,
+ )
diff --git a/src/atomate2/openmm/jobs/generate.py b/src/atomate2/openmm/jobs/generate.py
new file mode 100644
index 0000000000..b32ecc542a
--- /dev/null
+++ b/src/atomate2/openmm/jobs/generate.py
@@ -0,0 +1,325 @@
+"""Utilities for working with the OPLS forcefield in OpenMM."""
+
+from __future__ import annotations
+
+import copy
+import io
+import re
+import warnings
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from xml.etree.ElementTree import tostring
+
+import numpy as np
+from emmet.core.openff import MoleculeSpec
+from emmet.core.openmm import OpenMMInterchange, OpenMMTaskDocument
+from emmet.core.vasp.task_valid import TaskState
+from jobflow import Response
+from openmm import Context, LangevinMiddleIntegrator, System, XmlSerializer
+from openmm.app import PME, ForceField
+from openmm.app.pdbfile import PDBFile
+from openmm.unit import kelvin, picoseconds
+from pymatgen.core import Element
+from pymatgen.io.openff import get_atom_map
+
+from atomate2.openff.utils import create_mol_spec, merge_specs_by_name_and_smiles
+from atomate2.openmm.jobs.base import openmm_job
+
+try:
+ import openff.toolkit as tk
+ from openff.interchange.components._packmol import pack_box
+ from openff.units import unit
+except ImportError as e:
+ raise ImportError(
+ "Using the atomate2.openmm.generate "
+ "module requires the openff-toolkit package."
+ ) from e
+
+
+class XMLMoleculeFF:
+ """A class for manipulating XML files representing OpenMM-compatible forcefields."""
+
+ def __init__(self, xml_string: str) -> None:
+ """Create an XMLMoleculeFF object from a string version of the XML file."""
+ self.tree = ET.parse(io.StringIO(xml_string)) # noqa: S314
+
+ root = self.tree.getroot()
+ canonical_order = {}
+ for i, atom in enumerate(root.findall(".//Residues/Residue/Atom")):
+ canonical_order[atom.attrib["type"]] = i
+
+ non_to_res_map = {}
+ for i, atom in enumerate(root.findall(".//NonbondedForce/Atom")):
+ non_to_res_map[i] = canonical_order[atom.attrib["type"]]
+ # self._res_to_non.append(canonical_order[atom.attrib["type"]])
+ # invert map, change to list
+ self._res_to_non = [
+ k for k, v in sorted(non_to_res_map.items(), key=lambda item: item[1])
+ ]
+ self._non_to_res = list(non_to_res_map.values())
+
+ def __str__(self) -> str:
+ """Return the a string version of the XML file."""
+ return tostring(self.tree.getroot(), encoding="unicode")
+
+ def increment_types(self, increment: str) -> None:
+ """Increment the type names in the XMLMoleculeFF object.
+
+ This method is needed because LigParGen will reuse type names
+ in XML files, then causing an error in OpenMM. We differentiate
+ the types with this method.
+ """
+ root_type = [
+ (".//AtomTypes/Type", "name"),
+ (".//AtomTypes/Type", "class"),
+ (".//Residues/Residue/Atom", "type"),
+ (".//HarmonicBondForce/Bond", "class"),
+ (".//HarmonicAngleForce/Angle", "class"),
+ (".//PeriodicTorsionForce/Proper", "class"),
+ (".//PeriodicTorsionForce/Improper", "class"),
+ (".//NonbondedForce/Atom", "type"),
+ ]
+ for xpath, type_stub in root_type:
+ for element in self.tree.getroot().findall(xpath):
+ for key in element.attrib:
+ if type_stub in key:
+ element.attrib[key] += increment
+
+ def to_openff_molecule(self) -> tk.Molecule:
+ """Convert the XMLMoleculeFF to an openff_toolkit Molecule."""
+ if sum(self.partial_charges) > 1e-3:
+ # TODO: update message
+ warnings.warn("Formal charges not considered.", stacklevel=1)
+
+ p_table = {e.symbol: e.number for e in Element}
+ openff_mol = tk.Molecule()
+ for atom in self.tree.getroot().findall(".//Residues/Residue/Atom"):
+ symbol = re.match(r"^[A-Za-z]+", atom.attrib["name"]).group()
+ atomic_number = p_table[symbol]
+ openff_mol.add_atom(atomic_number, formal_charge=0, is_aromatic=False)
+
+ for bond in self.tree.getroot().findall(".//Residues/Residue/Bond"):
+ openff_mol.add_bond(
+ int(bond.attrib["from"]),
+ int(bond.attrib["to"]),
+ bond_order=1,
+ is_aromatic=False,
+ )
+
+ openff_mol.partial_charges = self.partial_charges * unit.elementary_charge
+
+ return openff_mol
+
+ @property
+ def partial_charges(self) -> np.ndarray:
+ """Get the partial charges from the XMLMoleculeFF object."""
+ atoms = self.tree.getroot().findall(".//NonbondedForce/Atom")
+ charges = np.array([float(atom.attrib["charge"]) for atom in atoms])
+ return charges[self._res_to_non]
+
+ @partial_charges.setter
+ def partial_charges(self, partial_charges: np.ndarray) -> None:
+ for i, atom in enumerate(self.tree.getroot().findall(".//NonbondedForce/Atom")):
+ charge = partial_charges[self._non_to_res[i]]
+ atom.attrib["charge"] = str(charge)
+
+ def assign_partial_charges(self, mol_or_method: tk.Molecule | str) -> None:
+ """Assign partial charges to the XMLMoleculeFF object.
+
+ Parameters
+ ----------
+ mol_or_method : Union[tk.Molecule, str]
+ If a molecule is provided, it must have partial charges assigned.
+ If a string is provided, openff_toolkit.Molecule.assign_partial_charges
+ will be used to generate the partial charges.
+ """
+ if isinstance(mol_or_method, str):
+ openff_mol = self.to_openff_molecule()
+ openff_mol.assign_partial_charges(mol_or_method)
+ mol_or_method = openff_mol
+ self_mol = self.to_openff_molecule()
+ isomorphic, atom_map = get_atom_map(mol_or_method, self_mol)
+ mol_charges = mol_or_method.partial_charges[list(atom_map.values())].magnitude
+ self.partial_charges = mol_charges
+
+ def to_file(self, file: str | Path) -> None:
+ """Write the XMLMoleculeFF object to an XML file."""
+ self.tree.write(file, encoding="utf-8")
+
+ @classmethod
+ def from_file(cls, file: str | Path) -> XMLMoleculeFF:
+ """Create an XMLMoleculeFF object from an XML file."""
+ with open(file) as f:
+ xml_str = f.read()
+ return cls(xml_str)
+
+
+def create_system_from_xml(
+ topology: tk.Topology,
+ xml_mols: list[XMLMoleculeFF],
+) -> System:
+ """Create an OpenMM system from a list of molecule specifications and XML files."""
+ io_files = []
+ for i, xml in enumerate(xml_mols):
+ xml_copy = copy.deepcopy(xml)
+ xml_copy.increment_types(f"_{i}")
+ io_files.append(io.StringIO(str(xml_copy)))
+
+ ff = ForceField(io_files[0])
+ for i, xml in enumerate(io_files[1:]): # type: ignore[assignment]
+ ff.loadFile(xml, resname_prefix=f"_{i + 1}")
+
+ return ff.createSystem(topology.to_openmm(), nonbondedMethod=PME)
+
+
+@openmm_job
+def generate_openmm_interchange(
+ input_mol_specs: list[MoleculeSpec | dict],
+ mass_density: float,
+ ff_xmls: list[str],
+ xml_method_and_scaling: tuple[str, float] = None,
+ pack_box_kwargs: dict = None,
+ tags: list[str] = None,
+) -> Response:
+ """Generate an OpenMM Interchange object from a list of molecule specifications.
+
+ This function takes a list of molecule specifications (either as
+ MoleculeSpec objects or dictionaries), a target mass density, and
+ optional force field and box packing parameters. It processes the molecule
+ specifications, packs them into a box using the specified mass density, and
+ creates an OpenFF Interchange object using the specified force field.
+
+ If you'd like to have multiple distinct input geometries, you
+ can pass multiple mol_specs with the same name and SMILES string.
+ After packing the box, they will be merged into a single mol_spec
+ and treated as a single component in the resulting system.
+
+ Parameters
+ ----------
+ input_mol_specs : List[Union[MoleculeSpec, dict]]
+ A list of molecule specifications, either as MoleculeSpec objects or
+ dictionaries that can be passed to `create_mol_spec` to create
+ MoleculeSpec objects. See the `create_mol_spec` function
+ for details on the expected format of the dictionaries.
+ mass_density : float
+ The target mass density for packing the molecules into
+ a box, kg/L.
+ ff_xmls : List[str]
+ A list of force field XML strings, these should be the raw text
+ of the XML files. The order of the XML strings
+ must match the order of the input_mol_specs.
+ xml_method_and_scaling : Tuple[str, float], optional
+ A tuple containing the charge method and scaling factor to use for
+ the partial charges in the xml. If this is not set, partial charges
+ will be generated by openff toolkit.
+ pack_box_kwargs : Dict, optional
+ Additional keyword arguments to pass to the
+ toolkit.interchange.components._packmol.pack_box. Default is an empty dict.
+ tags : List[str], optional
+ A list of tags to attach to the task document.
+
+ Returns
+ -------
+ ClassicalMDTaskDocument
+ A task document containing the generated OpenFF Interchange
+ object, molecule specifications, and force field information.
+
+ Notes
+ -----
+ - The function assumes that all dictionaries in the mol_specs list can be used to
+ create valid MoleculeSpec objects.
+ - The function sorts the molecule specifications based on their SMILES string
+ and name before packing the box.
+ - The function uses the merge_specs_by_name_and_smiles function to merge molecule
+ specifications with the same name and SMILES string.
+ """
+ mol_specs = []
+ for spec in input_mol_specs:
+ if isinstance(spec, dict):
+ mol_specs.append(create_mol_spec(**spec))
+ elif isinstance(spec, MoleculeSpec):
+ mol_specs.append(copy.deepcopy(spec))
+ else:
+ raise TypeError(
+ f"item in mol_specs is a {type(spec)}, but mol_specs "
+ f"must be a list of dicts or MoleculeSpec"
+ )
+
+ xml_mols = [XMLMoleculeFF(xml) for xml in ff_xmls]
+ if len(mol_specs) != len(xml_mols):
+ raise ValueError(
+ "The number of molecule specifications and XML files must match."
+ )
+
+ for mol_spec, xml_mol in zip(mol_specs, xml_mols):
+ openff_mol = tk.Molecule.from_json(mol_spec.openff_mol)
+ xml_openff_mol = xml_mol.to_openff_molecule()
+ is_isomorphic, atom_map = get_atom_map(openff_mol, xml_openff_mol)
+ if not is_isomorphic:
+ raise ValueError(
+ "The mol_specs and ff_xmls must index identical molecules."
+ )
+ if xml_method_and_scaling:
+ charge_method, charge_scaling = xml_method_and_scaling
+ mol_spec.charge_method = charge_method
+ mol_spec.charge_scaling = charge_scaling
+ openff_mol.partial_charges = xml_openff_mol.partial_charges
+ mol_spec.openff_mol = openff_mol.to_json()
+ else:
+ xml_mol.assign_partial_charges(openff_mol)
+
+ mol_specs.sort(
+ key=lambda x: tk.Molecule.from_json(x.openff_mol).to_smiles() + x.name
+ )
+ mol_specs = merge_specs_by_name_and_smiles(mol_specs)
+
+ pack_box_kwargs = pack_box_kwargs or {}
+ topology = pack_box(
+ molecules=[tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs],
+ number_of_copies=[spec.count for spec in mol_specs],
+ mass_density=mass_density * unit.grams / unit.milliliter,
+ **pack_box_kwargs,
+ )
+
+ system = create_system_from_xml(topology, xml_mols)
+
+ # these values don't actually matter because integrator is only
+ # used to generate the state
+ integrator = LangevinMiddleIntegrator(
+ 298 * kelvin, 1 / picoseconds, 1 * picoseconds
+ )
+ context = Context(system, integrator)
+ context.setPositions(topology.get_positions().magnitude / 10)
+ state = context.getState(getPositions=True)
+
+ with io.StringIO() as s:
+ PDBFile.writeFile(
+ topology.to_openmm(), np.zeros(shape=(topology.n_atoms, 3)), file=s
+ )
+ s.seek(0)
+ pdb = s.read()
+
+ interchange = OpenMMInterchange(
+ system=XmlSerializer.serialize(system),
+ state=XmlSerializer.serialize(state),
+ topology=pdb,
+ )
+
+ # TODO: fix all jsons
+ interchange_json = interchange.json()
+
+ dir_name = Path.cwd()
+
+ task_doc = OpenMMTaskDocument(
+ dir_name=str(dir_name),
+ state=TaskState.SUCCESS,
+ interchange=interchange_json,
+ mol_specs=mol_specs,
+ force_field="opls", # TODO: change to flexible value
+ tags=tags,
+ )
+
+ with open(dir_name / "taskdoc.json", "w") as file:
+ file.write(task_doc.json())
+
+ return Response(output=task_doc)
diff --git a/src/atomate2/openmm/utils.py b/src/atomate2/openmm/utils.py
new file mode 100644
index 0000000000..5ab5e2b151
--- /dev/null
+++ b/src/atomate2/openmm/utils.py
@@ -0,0 +1,134 @@
+"""Utilities for working with the OPLS forcefield in OpenMM."""
+
+from __future__ import annotations
+
+import re
+import tempfile
+import time
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from emmet.core.openmm import OpenMMTaskDocument
+
+
+def download_opls_xml(
+ names_smiles: dict[str, str], output_dir: str | Path, overwrite_files: bool = False
+) -> None:
+ """Download an OPLS-AA/M XML file from the LigParGen website using Selenium."""
+ try:
+ from selenium import webdriver
+ from selenium.webdriver.chrome.service import Service as ChromeService
+ from selenium.webdriver.common.by import By
+ from webdriver_manager.chrome import ChromeDriverManager
+
+ except ImportError:
+ warnings.warn(
+ "The `selenium` package is not installed. "
+ "It's required to run the opls web scraper.",
+ stacklevel=1,
+ )
+
+ # Initialize the Chrome driver
+ driver = webdriver.Chrome(service=ChromeService(ChromeDriverManager().install()))
+
+ for name, smiles in names_smiles.items():
+ final_file = Path(output_dir) / f"{name}.xml"
+ if final_file.exists() and not overwrite_files:
+ continue
+ try:
+ # Specify the directory where you want to download files
+ with tempfile.TemporaryDirectory() as tmpdir:
+ download_dir = tmpdir
+
+ # Set up Chrome options
+ chrome_options = webdriver.ChromeOptions()
+ prefs = {"download.default_directory": download_dir}
+ chrome_options.add_experimental_option("prefs", prefs)
+
+ # Initialize Chrome with the options
+ driver = webdriver.Chrome(options=chrome_options)
+
+ # Open the first webpage
+ driver.get("https://zarbi.chem.yale.edu/ligpargen/")
+
+ # Find the SMILES input box and enter the SMILES code
+ smiles_input = driver.find_element(By.ID, "smiles")
+ smiles_input.send_keys(smiles)
+
+ # Find and click the "Submit Molecule" button
+ submit_button = driver.find_element(
+ By.XPATH,
+ '//button[@type="submit" and contains(text(), "Submit Molecule")]',
+ )
+ submit_button.click()
+
+ # Wait for the second page to load
+ # time.sleep(2) # Adjust this delay as needed based on the loading time
+
+ # Find and click the "XML" button under Downloads and OpenMM
+ xml_button = driver.find_element(
+ By.XPATH, '//input[@type="submit" and @value="XML"]'
+ )
+ xml_button.click()
+
+ # Wait for the file to download
+ time.sleep(0.3) # Adjust as needed based on the download time
+
+ file = next(Path(tmpdir).iterdir())
+ # copy downloaded file to output_file using os
+ Path(file).rename(final_file)
+
+ except Exception as e: # noqa: BLE001
+ warnings.warn(
+ f"{name} ({smiles}) failed to download because an error occurred: {e}",
+ stacklevel=1,
+ )
+
+ driver.quit()
+
+
+def create_list_summing_to(total_sum: int, n_pieces: int) -> list:
+ """Create a NumPy array with n_pieces elements that sum up to total_sum.
+
+ Divides total_sum by n_pieces to determine the base value for each element.
+ Distributes the remainder evenly among the elements.
+
+ Parameters
+ ----------
+ total_sum : int
+ The desired sum of the array elements.
+ n_pieces : int
+ The number of elements in the array.
+
+ Returns
+ -------
+ numpy.ndarray
+ A 1D NumPy array with n_pieces elements summing up to total_sum.
+ """
+ div, mod = total_sum // n_pieces, total_sum % n_pieces
+ return [div + 1] * mod + [div] * (n_pieces - mod)
+
+
+def increment_name(file_name: str) -> str:
+ """Increment the count in a file name."""
+ # logic to increment count on file name
+ re_match = re.search(r"(\d*)$", file_name)
+ position = re_match.start(1)
+ new_count = int(re_match.group(1) or 1) + 1
+ return f"{file_name[:position]}{new_count}"
+
+
+def task_reports(task: OpenMMTaskDocument, traj_or_state: str = "traj") -> bool:
+ """Check if a task reports trajectories or states."""
+ if not task.calcs_reversed:
+ return False
+ calc_input = task.calcs_reversed[0].input
+ if traj_or_state == "traj":
+ report_freq = calc_input.traj_interval
+ elif traj_or_state == "state":
+ report_freq = calc_input.state_interval
+ else:
+ raise ValueError("traj_or_state must be 'traj' or 'state'")
+ return calc_input.n_steps >= report_freq
diff --git a/tests/common/schemas/test_cclib.py b/tests/common/schemas/test_cclib.py
index bb34fd84bf..7be2bbdd4c 100644
--- a/tests/common/schemas/test_cclib.py
+++ b/tests/common/schemas/test_cclib.py
@@ -14,7 +14,8 @@
cclib = None
-@pytest.mark.skipif(cclib is None, reason="requires cclib to be installed")
+# @pytest.mark.skipif(cclib is None, reason="requires cclib to be installed")
+@pytest.mark.skip(reason="cclib is not working in CI")
def test_cclib_taskdoc(test_dir):
p = test_dir / "schemas"
diff --git a/tests/openff_md/__init__.py b/tests/openff_md/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/openff_md/conftest.py b/tests/openff_md/conftest.py
new file mode 100644
index 0000000000..7375e82f33
--- /dev/null
+++ b/tests/openff_md/conftest.py
@@ -0,0 +1,93 @@
+import numpy as np
+import openff.toolkit as tk
+import pytest
+from jobflow import run_locally
+from openff.interchange import Interchange
+from openff.interchange.components._packmol import pack_box
+from openff.toolkit import ForceField
+from openff.units import unit
+
+from atomate2.openff.utils import create_mol_spec, merge_specs_by_name_and_smiles
+
+
+@pytest.fixture
+def run_job(tmp_path):
+ def run_job(job):
+ response_dict = run_locally(job, ensure_success=True, root_dir=tmp_path)
+ return list(response_dict.values())[-1][1].output
+
+ return run_job
+
+
+@pytest.fixture
+def mol_specs_small():
+ return [
+ create_mol_spec("CCO", 10, name="ethanol", charge_method="mmff94"),
+ create_mol_spec("O", 20, name="water", charge_method="mmff94"),
+ ]
+
+
+@pytest.fixture
+def openff_data(test_dir):
+ return test_dir / "openff"
+
+
+@pytest.fixture
+def mol_files(openff_data):
+ geo_dir = openff_data / "molecule_charge_files"
+ return {
+ "CCO_xyz": str(geo_dir / "CCO.xyz"),
+ "CCO_charges": str(geo_dir / "CCO.npy"),
+ "FEC_r_xyz": str(geo_dir / "FEC-r.xyz"),
+ "FEC_s_xyz": str(geo_dir / "FEC-s.xyz"),
+ "FEC_charges": str(geo_dir / "FEC.npy"),
+ "PF6_xyz": str(geo_dir / "PF6.xyz"),
+ "PF6_charges": str(geo_dir / "PF6.npy"),
+ "Li_charges": str(geo_dir / "Li.npy"),
+ "Li_xyz": str(geo_dir / "Li.xyz"),
+ }
+
+
+@pytest.fixture
+def mol_specs_salt(mol_files):
+ charges = np.load(mol_files["PF6_charges"])
+ return [
+ create_mol_spec("CCO", 10, name="ethanol", charge_method="mmff94"),
+ create_mol_spec("O", 20, name="water", charge_method="mmff94"),
+ create_mol_spec("[Li+]", 5, name="li", charge_method="mmff94"),
+ create_mol_spec(
+ "F[P-](F)(F)(F)(F)F",
+ 5,
+ name="pf6",
+ partial_charges=charges,
+ geometry=mol_files["PF6_xyz"],
+ ),
+ ]
+
+
+@pytest.fixture(scope="package")
+def interchange():
+ o = create_mol_spec("O", 300, charge_method="mmff94")
+ cco = create_mol_spec("CCO", 10, charge_method="mmff94")
+ cco2 = create_mol_spec("CCO", 20, name="cco2", charge_method="mmff94")
+ mol_specs = [o, cco, cco2]
+ mol_specs.sort(
+ key=lambda x: tk.Molecule.from_json(x.openff_mol).to_smiles() + x.name
+ )
+
+ topology = pack_box(
+ molecules=[tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs],
+ number_of_copies=[spec.count for spec in mol_specs],
+ mass_density=0.8 * unit.grams / unit.milliliter,
+ )
+
+ mol_specs = merge_specs_by_name_and_smiles(mol_specs)
+
+ return Interchange.from_smirnoff(
+ force_field=ForceField("openff_unconstrained-2.1.1.offxml"),
+ topology=topology,
+ charge_from_molecules=[
+ tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs
+ ],
+ allow_nonintegral_charges=True,
+ )
diff --git a/tests/openff_md/test_core.py b/tests/openff_md/test_core.py
new file mode 100644
index 0000000000..77679d9a1b
--- /dev/null
+++ b/tests/openff_md/test_core.py
@@ -0,0 +1,64 @@
+from emmet.core.openff import ClassicalMDTaskDocument, MoleculeSpec
+from openff.interchange import Interchange
+
+from atomate2.openff.core import generate_interchange
+
+
+def test_generate_interchange(mol_specs_small, run_job):
+ mass_density = 1
+ force_field = "openff_unconstrained-2.1.1.offxml"
+ mol_specs = mol_specs_small
+
+ job = generate_interchange(mol_specs, mass_density, force_field)
+ task_doc = run_job(job)
+
+ assert isinstance(task_doc, ClassicalMDTaskDocument)
+ assert task_doc.force_field == force_field
+
+ interchange = Interchange.parse_raw(task_doc.interchange)
+ assert isinstance(interchange, Interchange)
+
+ topology = interchange.topology
+ assert topology.n_molecules == 30
+ assert topology.n_atoms == 150
+ assert topology.n_bonds == 120
+
+ molecule_specs = task_doc.mol_specs
+ assert len(molecule_specs) == 2
+ assert all(isinstance(spec, MoleculeSpec) for spec in molecule_specs)
+ assert molecule_specs[0].name == "ethanol"
+ assert molecule_specs[0].count == 10
+ assert molecule_specs[1].name == "water"
+ assert molecule_specs[1].count == 20
+
+ # TODO: debug issue with ForceField accepting iterables of FFs
+
+ # Test with mol_specs as a list of dicts
+ mol_specs_dicts = [
+ {"smiles": "CCO", "count": 10, "name": "ethanol", "charge_method": "mmff94"},
+ {"smiles": "O", "count": 20, "name": "water", "charge_method": "mmff94"},
+ ]
+ job = generate_interchange(mol_specs_dicts, mass_density, force_field)
+ task_doc = run_job(job)
+ molecule_specs = task_doc.mol_specs
+ assert len(molecule_specs) == 2
+ assert molecule_specs[0].name == "ethanol"
+ assert molecule_specs[0].count == 10
+ assert molecule_specs[1].name == "water"
+ assert molecule_specs[1].count == 20
+
+
+def test_generate_interchange_salt(mol_specs_salt, run_job):
+ mass_density = 1
+ force_field = "openff_unconstrained-2.1.1.offxml"
+ mol_specs = mol_specs_salt
+
+ job = generate_interchange(mol_specs, mass_density, force_field)
+ task_doc = run_job(job)
+
+ molecule_specs = task_doc.mol_specs
+ assert len(molecule_specs) == 4
+ assert molecule_specs[1].name == "ethanol"
+ assert molecule_specs[1].count == 10
+ assert molecule_specs[2].name == "water"
+ assert molecule_specs[2].count == 20
diff --git a/tests/openff_md/test_utils.py b/tests/openff_md/test_utils.py
new file mode 100644
index 0000000000..e3591cab43
--- /dev/null
+++ b/tests/openff_md/test_utils.py
@@ -0,0 +1,255 @@
+import numpy as np
+import openff.toolkit as tk
+import pymatgen
+import pytest
+from emmet.core.openff import MoleculeSpec
+from openff.interchange import Interchange
+from openff.toolkit.topology import Topology
+from openff.toolkit.topology.molecule import Molecule
+from openff.units import Quantity
+from pymatgen.analysis.graphs import MoleculeGraph
+from pymatgen.io.openff import (
+ add_conformer,
+ assign_partial_charges,
+ create_openff_mol,
+ get_atom_map,
+ infer_openff_mol,
+ mol_graph_to_openff_mol,
+)
+
+from atomate2.openff.utils import (
+ counts_from_box_size,
+ counts_from_masses,
+ create_mol_spec,
+ merge_specs_by_name_and_smiles,
+)
+
+
+def test_molgraph_to_openff_pf6(mol_files):
+ """transform a water MoleculeGraph to a OpenFF water molecule"""
+ pf6_mol = pymatgen.core.Molecule.from_file(mol_files["PF6_xyz"])
+ pf6_mol.set_charge_and_spin(charge=-1)
+ pf6_molgraph = MoleculeGraph.with_edges(
+ pf6_mol,
+ {
+ (0, 1): {"weight": 1},
+ (0, 2): {"weight": 1},
+ (0, 3): {"weight": 1},
+ (0, 4): {"weight": 1},
+ (0, 5): {"weight": 1},
+ (0, 6): {"weight": 1},
+ },
+ )
+
+ pf6_openff_1 = tk.Molecule.from_smiles("F[P-](F)(F)(F)(F)F")
+
+ pf6_openff_2 = mol_graph_to_openff_mol(pf6_molgraph)
+ assert pf6_openff_1 == pf6_openff_2
+
+
+def test_molgraph_to_openff_cco(mol_files):
+ from pymatgen.analysis.local_env import OpenBabelNN
+
+ cco_pmg = pymatgen.core.Molecule.from_file(mol_files["CCO_xyz"])
+ cco_molgraph = MoleculeGraph.with_local_env_strategy(cco_pmg, OpenBabelNN())
+
+ cco_openff_1 = mol_graph_to_openff_mol(cco_molgraph)
+
+ cco_openff_2 = tk.Molecule.from_smiles("CCO")
+ cco_openff_2.assign_partial_charges("mmff94")
+
+ assert cco_openff_1 == cco_openff_2
+
+
+@pytest.mark.parametrize(
+ "xyz_path, smiles, map_values",
+ [
+ ("CCO_xyz", "CCO", [0, 1, 2, 3, 4, 5, 6, 7, 8]),
+ ("FEC_r_xyz", "O=C1OC[C@@H](F)O1", [0, 1, 2, 3, 4, 6, 7, 9, 8, 5]),
+ ("FEC_s_xyz", "O=C1OC[C@H](F)O1", [0, 1, 2, 3, 4, 6, 7, 9, 8, 5]),
+ ("PF6_xyz", "F[P-](F)(F)(F)(F)F", [1, 0, 2, 3, 4, 5, 6]),
+ ],
+)
+def test_get_atom_map(xyz_path, smiles, map_values, mol_files):
+ mol = pymatgen.core.Molecule.from_file(mol_files[xyz_path])
+ inferred_mol = infer_openff_mol(mol)
+ openff_mol = tk.Molecule.from_smiles(smiles)
+ isomorphic, atom_map = get_atom_map(inferred_mol, openff_mol)
+ assert isomorphic
+ assert map_values == list(atom_map.values())
+
+
+@pytest.mark.parametrize(
+ "xyz_path, n_atoms, n_bonds",
+ [
+ ("CCO_xyz", 9, 8),
+ ("FEC_r_xyz", 10, 10),
+ ("FEC_s_xyz", 10, 10),
+ ("PF6_xyz", 7, 6),
+ ],
+)
+def test_infer_openff_mol(xyz_path, n_atoms, n_bonds, mol_files):
+ mol = pymatgen.core.Molecule.from_file(mol_files[xyz_path])
+ openff_mol = infer_openff_mol(mol)
+ assert isinstance(openff_mol, tk.Molecule)
+ assert openff_mol.n_atoms == n_atoms
+ assert openff_mol.n_bonds == n_bonds
+
+
+def test_add_conformer(mol_files):
+ openff_mol = tk.Molecule.from_smiles("CCO")
+ geometry = pymatgen.core.Molecule.from_file(mol_files["CCO_xyz"])
+ openff_mol, atom_map = add_conformer(openff_mol, geometry)
+ assert openff_mol.n_conformers == 1
+ assert list(atom_map.values()) == list(range(openff_mol.n_atoms))
+
+
+def test_assign_partial_charges(mol_files):
+ openff_mol = tk.Molecule.from_smiles("CCO")
+ geometry = pymatgen.core.Molecule.from_file(mol_files["CCO_xyz"])
+ openff_mol, atom_map = add_conformer(openff_mol, geometry)
+ partial_charges = np.load(mol_files["CCO_charges"])
+ openff_mol = assign_partial_charges(openff_mol, atom_map, "am1bcc", partial_charges)
+ assert np.allclose(openff_mol.partial_charges.magnitude, partial_charges)
+
+
+def test_create_openff_mol(mol_files):
+ smiles = "CCO"
+ geometry = mol_files["CCO_xyz"]
+ partial_charges = np.load(mol_files["CCO_charges"])
+ openff_mol = create_openff_mol(smiles, geometry, 1.0, partial_charges, "am1bcc")
+ assert isinstance(openff_mol, tk.Molecule)
+ assert openff_mol.n_atoms == 9
+ assert openff_mol.n_bonds == 8
+ assert np.allclose(openff_mol.partial_charges.magnitude, partial_charges)
+
+
+def test_create_mol_spec(mol_files):
+ smiles, count, name, geometry = "CCO", 10, "ethanol", mol_files["CCO_xyz"]
+ partial_charges = np.load(mol_files["CCO_charges"])
+ mol_spec = create_mol_spec(
+ smiles, count, name, 1.0, "am1bcc", geometry, partial_charges
+ )
+ assert isinstance(mol_spec, MoleculeSpec)
+ assert mol_spec.name == name
+ assert mol_spec.count == count
+ assert mol_spec.charge_scaling == 1.0
+ assert mol_spec.charge_method == "am1bcc"
+ assert isinstance(tk.Molecule.from_json(mol_spec.openff_mol), tk.Molecule)
+
+
+def test_merge_specs_by_name_and_smiles(mol_files):
+ smiles1, count1, name1, geometry1 = "CCO", 5, "ethanol", mol_files["CCO_xyz"]
+ partial_charges1 = np.load(mol_files["CCO_charges"])
+ mol_spec1 = create_mol_spec(
+ smiles1, count1, name1, 1.0, "am1bcc", geometry1, partial_charges1
+ )
+
+ smiles2, count2, name2, geometry2 = "CCO", 8, "ethanol", mol_files["CCO_xyz"]
+ partial_charges2 = np.load(mol_files["CCO_charges"])
+ mol_spec2 = create_mol_spec(
+ smiles2, count2, name2, 1.0, "am1bcc", geometry2, partial_charges2
+ )
+
+ mol_specs = [mol_spec1, mol_spec2]
+ merged_specs = merge_specs_by_name_and_smiles(mol_specs)
+ assert len(merged_specs) == 1
+ assert merged_specs[0].count == count1 + count2
+
+ mol_specs[1].name = "ethanol2"
+ merged_specs = merge_specs_by_name_and_smiles(mol_specs)
+ assert len(merged_specs) == 2
+ assert merged_specs[0].name == name1
+ assert merged_specs[0].count == count1
+ assert merged_specs[1].name == "ethanol2"
+ assert merged_specs[1].count == count2
+
+
+def test_openff_mol_as_from_monty_dict():
+ mol = Molecule.from_smiles("CCO")
+ mol_dict = mol.as_dict()
+ reconstructed_mol = Molecule.from_dict(mol_dict)
+
+ assert mol.to_smiles() == reconstructed_mol.to_smiles()
+ assert mol.n_atoms == reconstructed_mol.n_atoms
+ assert mol.n_bonds == reconstructed_mol.n_bonds
+ assert mol.n_angles == reconstructed_mol.n_angles
+ assert mol.n_propers == reconstructed_mol.n_propers
+ assert mol.n_impropers == reconstructed_mol.n_impropers
+
+
+def test_openff_topology_as_from_monty_dict():
+ topology = Topology.from_molecules([Molecule.from_smiles("CCO")])
+ topology_dict = topology.as_dict()
+ reconstructed_topology = Topology.from_dict(topology_dict)
+
+ assert topology.n_molecules == reconstructed_topology.n_molecules
+ assert topology.n_atoms == reconstructed_topology.n_atoms
+ assert topology.n_bonds == reconstructed_topology.n_bonds
+ assert topology.box_vectors == reconstructed_topology.box_vectors
+
+
+def test_openff_interchange_as_from_monty_dict(interchange):
+ # interchange = Interchange.from_smirnoff("openff-2.0.0.offxml", "CCO")
+ interchange_dict = interchange.as_dict()
+ reconstructed_interchange = Interchange.from_dict(interchange_dict)
+
+ assert np.all(interchange.positions == reconstructed_interchange.positions)
+ assert np.all(interchange.velocities == reconstructed_interchange.velocities)
+ assert np.all(interchange.box == reconstructed_interchange.box)
+
+ assert interchange.mdconfig == reconstructed_interchange.mdconfig
+
+ topology = interchange.topology
+ reconstructed_topology = reconstructed_interchange.topology
+
+ assert topology.n_molecules == reconstructed_topology.n_molecules
+ assert topology.n_atoms == reconstructed_topology.n_atoms
+ assert topology.n_bonds == reconstructed_topology.n_bonds
+ assert np.all(topology.box_vectors == reconstructed_topology.box_vectors)
+
+
+def test_openff_quantity_as_from_monty_dict():
+ quantity = Quantity(1.0, "kilocalorie / mole")
+ quantity_dict = quantity.as_dict()
+ reconstructed_quantity = Quantity.from_dict(quantity_dict)
+
+ assert quantity.magnitude == reconstructed_quantity.magnitude
+ assert quantity.units == reconstructed_quantity.units
+ assert quantity == reconstructed_quantity
+
+
+def test_calculate_elyte_composition():
+ from atomate2.openff.utils import calculate_elyte_composition, counts_from_masses
+
+ vol_ratio = {"O": 0.5, "CCO": 0.5}
+ salts = {"[Li+]": 1.0, "F[P-](F)(F)(F)(F)F": 1.0}
+ solvent_densities = {"O": 1.0, "CCO": 0.8}
+
+ comp_dict = calculate_elyte_composition(
+ vol_ratio, salts, solvent_densities, "volume"
+ )
+ counts = counts_from_masses(comp_dict, 100)
+ assert sum(counts.values()) == 100
+
+ mol_ratio = {
+ "[Li+]": 0.00616,
+ "F[P-](F)(F)(F)(F)F": 0.128,
+ "C1COC(=O)O1": 0.245, # EC
+ "CCOC(=O)OC": 0.462, # EMC
+ "CC#N": 0.130,
+ "FC1COC(=O)O1": 0.028,
+ }
+ counts2 = counts_from_masses(mol_ratio, 1000)
+ assert np.allclose(sum(counts2.values()), 1000, atol=5)
+
+
+def test_counts_calculators():
+ mass_fractions = {"O": 0.5, "CCO": 0.5}
+
+ counts_size = counts_from_box_size(mass_fractions, 3)
+ counts_number = counts_from_masses(mass_fractions, 324)
+
+ assert 200 < sum(counts_size.values()) < 500
+
+ assert counts_size == counts_number
diff --git a/tests/openmm_md/__init__.py b/tests/openmm_md/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/openmm_md/conftest.py b/tests/openmm_md/conftest.py
new file mode 100644
index 0000000000..4bb1d843a3
--- /dev/null
+++ b/tests/openmm_md/conftest.py
@@ -0,0 +1,56 @@
+import openff.toolkit as tk
+import pytest
+from jobflow import run_locally
+from openff.interchange import Interchange
+from openff.interchange.components._packmol import pack_box
+from openff.toolkit import ForceField
+from openff.units import unit
+
+from atomate2.openff.utils import create_mol_spec, merge_specs_by_name_and_smiles
+
+
+@pytest.fixture
+def run_job(tmp_path):
+ def run_job(job):
+ response_dict = run_locally(job, ensure_success=True, root_dir=tmp_path)
+ return list(response_dict.values())[-1][1].output
+
+ return run_job
+
+
+@pytest.fixture
+def openmm_data(test_dir):
+ return test_dir / "openmm"
+
+
+@pytest.fixture(scope="package")
+def interchange():
+ o = create_mol_spec("O", 300, charge_method="mmff94")
+ cco = create_mol_spec("CCO", 10, charge_method="mmff94")
+ cco2 = create_mol_spec("CCO", 20, name="cco2", charge_method="mmff94")
+ mol_specs = [o, cco, cco2]
+ mol_specs.sort(
+ key=lambda x: tk.Molecule.from_json(x.openff_mol).to_smiles() + x.name
+ )
+
+ topology = pack_box(
+ molecules=[tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs],
+ number_of_copies=[spec.count for spec in mol_specs],
+ mass_density=0.8 * unit.grams / unit.milliliter,
+ )
+
+ mol_specs = merge_specs_by_name_and_smiles(mol_specs)
+
+ return Interchange.from_smirnoff(
+ force_field=ForceField("openff_unconstrained-2.1.1.offxml"),
+ topology=topology,
+ charge_from_molecules=[
+ tk.Molecule.from_json(spec.openff_mol) for spec in mol_specs
+ ],
+ allow_nonintegral_charges=True,
+ )
+
+
+@pytest.fixture
+def output_dir(test_dir):
+ return test_dir / "classical_md" / "output_dir"
diff --git a/tests/openmm_md/flows/__init__.py b/tests/openmm_md/flows/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/openmm_md/flows/test_core.py b/tests/openmm_md/flows/test_core.py
new file mode 100644
index 0000000000..b43d95103a
--- /dev/null
+++ b/tests/openmm_md/flows/test_core.py
@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+from emmet.core.openmm import OpenMMTaskDocument
+from jobflow import Flow
+from MDAnalysis import Universe
+from monty.json import MontyDecoder
+from openff.interchange import Interchange
+
+from atomate2.openmm.flows.core import OpenMMFlowMaker
+from atomate2.openmm.jobs import EnergyMinimizationMaker, NPTMaker, NVTMaker
+
+
+def test_anneal_maker(interchange, run_job):
+ # Create an instance of AnnealMaker with custom parameters
+ anneal_maker = OpenMMFlowMaker.anneal_flow(
+ name="test_anneal",
+ anneal_temp=500,
+ final_temp=300,
+ n_steps=30,
+ temp_steps=1,
+ job_names=("heat", "hold", "cool"),
+ platform_name="CPU",
+ )
+
+ # Run the AnnealMaker flow
+ anneal_flow = anneal_maker.make(interchange)
+
+ task_doc = run_job(anneal_flow)
+
+ # Check the output task document
+ assert isinstance(task_doc, OpenMMTaskDocument)
+ assert task_doc.state == "successful"
+ assert len(task_doc.calcs_reversed) == 1
+
+ # Check the individual jobs in the flow
+ raise_temp_job = anneal_flow.jobs[0]
+ assert raise_temp_job.maker.name == "heat"
+ assert raise_temp_job.maker.n_steps == 10
+ assert raise_temp_job.maker.temperature == 500
+ assert raise_temp_job.maker.temp_steps == 1
+
+ nvt_job = anneal_flow.jobs[1]
+ assert nvt_job.maker.name == "hold"
+ assert nvt_job.maker.n_steps == 10
+ assert nvt_job.maker.temperature == 500
+
+ lower_temp_job = anneal_flow.jobs[2]
+ assert lower_temp_job.maker.name == "cool"
+ assert lower_temp_job.maker.n_steps == 10
+ assert lower_temp_job.maker.temperature == 300
+ assert lower_temp_job.maker.temp_steps == 1
+
+
+# @pytest.mark.skip("Reporting to HDF5 is broken in MDA upstream.")
+def test_hdf5_writing(interchange, run_job):
+ # Create an instance of AnnealMaker with custom parameters
+ import MDAnalysis
+ from packaging.version import Version
+
+ if Version(MDAnalysis.__version__) < Version("2.8.0"):
+ return
+
+ anneal_maker = OpenMMFlowMaker.anneal_flow(
+ name="test_anneal",
+ n_steps=3,
+ temp_steps=1,
+ platform_name="CPU",
+ traj_file_type="h5md",
+ report_velocities=True,
+ traj_interval=1,
+ )
+
+ # Run the AnnealMaker flow
+ anneal_maker.collect_outputs = True
+ anneal_flow = anneal_maker.make(interchange)
+
+ task_doc = run_job(anneal_flow)
+
+ calc_output_names = [calc.output.traj_file for calc in task_doc.calcs_reversed]
+ assert len(list(Path(task_doc.dir_name).iterdir())) == 5
+ assert set(calc_output_names) == {
+ "trajectory3.h5md",
+ "trajectory2.h5md",
+ "trajectory.h5md",
+ }
+
+
+def test_collect_outputs(interchange, run_job):
+ # Create an instance of ProductionMaker with custom parameters
+ production_maker = OpenMMFlowMaker(
+ name="test_production",
+ tags=["test"],
+ makers=[
+ EnergyMinimizationMaker(max_iterations=1, save_structure=True),
+ NVTMaker(n_steps=5),
+ ],
+ collect_outputs=True,
+ )
+
+ # Run the ProductionMaker flow
+ production_flow = production_maker.make(interchange)
+ run_job(production_flow)
+
+
+def test_flow_maker(interchange, run_job):
+ # Create an instance of ProductionMaker with custom parameters
+ production_maker = OpenMMFlowMaker(
+ name="test_production",
+ tags=["test"],
+ makers=[
+ EnergyMinimizationMaker(max_iterations=1),
+ NPTMaker(n_steps=5, pressure=1.0, state_interval=1, traj_interval=1),
+ OpenMMFlowMaker.anneal_flow(anneal_temp=400, final_temp=300, n_steps=5),
+ NVTMaker(n_steps=5),
+ ],
+ )
+
+ # Run the ProductionMaker flow
+ production_flow = production_maker.make(interchange)
+ task_doc = run_job(production_flow)
+
+ # Check the output task document
+ assert isinstance(task_doc, OpenMMTaskDocument)
+ assert task_doc.state == "successful"
+ assert len(task_doc.calcs_reversed) == 6
+ assert task_doc.calcs_reversed[-1].task_name == "energy minimization"
+ assert task_doc.calcs_reversed[0].task_name == "nvt simulation"
+ assert task_doc.tags == ["test"]
+ assert len(task_doc.job_uuids) == 6
+ assert task_doc.job_uuids[0] is not None
+
+ # Check the individual jobs in the flow
+ energy_job = production_flow.jobs[0]
+ assert isinstance(energy_job.maker, EnergyMinimizationMaker)
+
+ npt_job = production_flow.jobs[1]
+ assert isinstance(npt_job.maker, NPTMaker)
+ assert npt_job.maker.n_steps == 5
+ assert npt_job.maker.pressure == 1.0
+
+ anneal_flow = production_flow.jobs[2]
+ assert isinstance(anneal_flow, Flow)
+ assert anneal_flow.jobs[0].maker.temperature == 400
+ assert anneal_flow.jobs[2].maker.temperature == 300
+
+ nvt_job = production_flow.jobs[3]
+ assert isinstance(nvt_job.maker, NVTMaker)
+ assert nvt_job.maker.n_steps == 5
+
+ # Test length of state attributes in calculation output
+ calc_output = task_doc.calcs_reversed[0].output
+ assert len(calc_output.steps_reported) == 5
+
+ all_steps = [calc.output.steps_reported for calc in task_doc.calcs_reversed]
+ assert all_steps == [
+ [1, 2, 3, 4, 5],
+ [1],
+ [1, 2],
+ [1, 2],
+ [1, 2, 3, 4, 5],
+ None,
+ ]
+ # Test that the state interval is respected
+ assert calc_output.steps_reported == list(range(1, 6))
+ assert calc_output.traj_file == "trajectory5.dcd"
+ assert calc_output.state_file == "state5.csv"
+
+ interchange = Interchange.parse_raw(task_doc.interchange)
+ topology = interchange.to_openmm_topology()
+ u = Universe(topology, str(Path(task_doc.dir_name) / "trajectory5.dcd"))
+
+ assert len(u.trajectory) == 5
+
+
+def test_traj_blob_embed(interchange, run_job, tmp_path):
+ nvt = NVTMaker(n_steps=2, traj_interval=1, embed_traj=True)
+
+ # Run the ProductionMaker flow
+ nvt_job = nvt.make(interchange)
+ task_doc = run_job(nvt_job)
+
+ interchange = Interchange.parse_raw(task_doc.interchange)
+ topology = interchange.to_openmm_topology()
+ u = Universe(topology, str(Path(task_doc.dir_name) / "trajectory.dcd"))
+
+ assert len(u.trajectory) == 2
+
+ calc_output = task_doc.calcs_reversed[0].output
+ assert calc_output.traj_blob is not None
+
+ # Write the bytes back to a file
+ with open(tmp_path / "doc_trajectory.dcd", "wb") as f:
+ f.write(bytes.fromhex(calc_output.traj_blob))
+
+ u2 = Universe(topology, str(tmp_path / "doc_trajectory.dcd"))
+
+ assert np.all(u.atoms.positions == u2.atoms.positions)
+
+ with open(Path(task_doc.dir_name) / "taskdoc.json") as file:
+ task_dict = json.load(file, cls=MontyDecoder)
+ task_doc_parsed = OpenMMTaskDocument.model_validate(task_dict)
+
+ parsed_output = task_doc_parsed.calcs_reversed[0].output
+
+ assert parsed_output.traj_blob == calc_output.traj_blob
+
+
+@pytest.mark.skip("for local testing and debugging")
+def test_fireworks(interchange):
+ # Create an instance of ProductionMaker with custom parameters
+
+ production_maker = OpenMMFlowMaker(
+ name="test_production",
+ tags=["test"],
+ makers=[
+ EnergyMinimizationMaker(max_iterations=1),
+ NPTMaker(n_steps=5, pressure=1.0, state_interval=1, traj_interval=1),
+ OpenMMFlowMaker.anneal_flow(anneal_temp=400, final_temp=300, n_steps=5),
+ NVTMaker(n_steps=5),
+ ],
+ )
+
+ interchange_json = interchange.json()
+ # interchange_bytes = interchange_json.encode("utf-8")
+
+ # Run the ProductionMaker flow
+ production_flow = production_maker.make(interchange_json)
+
+ from fireworks import LaunchPad
+ from jobflow.managers.fireworks import flow_to_workflow
+
+ wf = flow_to_workflow(production_flow)
+
+ lpad = LaunchPad.auto_load()
+ lpad.add_wf(wf)
+
+ # from fireworks.core.rocket_launcher import launch_rocket
+ #
+ # launch_rocket(lpad)
diff --git a/tests/openmm_md/jobs/__init__.py b/tests/openmm_md/jobs/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/openmm_md/jobs/test_base.py b/tests/openmm_md/jobs/test_base.py
new file mode 100644
index 0000000000..6b120a9f83
--- /dev/null
+++ b/tests/openmm_md/jobs/test_base.py
@@ -0,0 +1,221 @@
+import copy
+
+import numpy as np
+import pytest
+from emmet.core.openff import ClassicalMDTaskDocument
+from emmet.core.openmm import Calculation, CalculationInput, OpenMMTaskDocument
+from jobflow import Flow, Job
+from mdareporter import MDAReporter
+from openmm.app import Simulation, StateDataReporter
+from openmm.openmm import LangevinMiddleIntegrator
+from openmm.unit import kelvin, picoseconds
+from pymatgen.core import Structure
+
+from atomate2.openff.core import generate_interchange
+from atomate2.openmm.jobs.base import BaseOpenMMMaker
+
+
+def test_add_reporters(interchange, tmp_path):
+ maker = BaseOpenMMMaker(
+ traj_interval=100, state_interval=50, wrap_traj=True, n_steps=1
+ )
+ sim = maker._create_simulation(interchange) # noqa: SLF001
+ dir_name = tmp_path / "test_output"
+ dir_name.mkdir()
+
+ maker._add_reporters(sim, dir_name) # noqa: SLF001
+
+ assert len(sim.reporters) == 2
+ assert isinstance(sim.reporters[0], MDAReporter)
+ next_dcd = sim.reporters[0].describeNextReport(sim)
+ assert next_dcd[0] == 100 # steps until next report
+ assert next_dcd[5] is True # enforce periodic boundaries
+ assert isinstance(sim.reporters[1], StateDataReporter)
+ next_state = sim.reporters[1].describeNextReport(sim)
+ assert next_state[0] == 50 # steps until next report
+
+
+def test_resolve_attr():
+ maker = BaseOpenMMMaker(temperature=301, friction_coefficient=2)
+ prev_task = ClassicalMDTaskDocument(
+ calcs_reversed=[Calculation(input=CalculationInput(step_size=0.002))]
+ )
+
+ assert maker._resolve_attr("temperature") == 301 # noqa: SLF001
+ assert maker._resolve_attr("friction_coefficient") == 2 # noqa: SLF001
+ assert maker._resolve_attr("step_size", prev_task) == 0.002 # noqa: SLF001
+ assert maker._resolve_attr("platform_name") == "CPU" # noqa: SLF001
+
+
+def test_create_integrator():
+ maker = BaseOpenMMMaker(temperature=300, friction_coefficient=2, step_size=0.002)
+ integrator = maker._create_integrator() # noqa: SLF001
+
+ assert isinstance(integrator, LangevinMiddleIntegrator)
+ assert integrator.getTemperature() == 300 * kelvin
+ assert integrator.getFriction() == 2 / picoseconds
+ assert integrator.getStepSize() == 0.002 * picoseconds
+
+
+def test_create_simulation(interchange):
+ maker = BaseOpenMMMaker()
+
+ sim = maker._create_simulation(interchange) # noqa: SLF001
+
+ assert isinstance(sim, Simulation)
+ assert isinstance(sim.integrator, LangevinMiddleIntegrator)
+ assert sim.context.getPlatform().getName() == "CPU"
+
+
+def test_update_interchange(interchange):
+ interchange = copy.deepcopy(interchange)
+ maker = BaseOpenMMMaker(wrap_traj=True)
+ sim = maker._create_simulation(interchange) # noqa: SLF001
+ start_positions = interchange.positions
+ start_velocities = interchange.velocities
+ start_box = interchange.box
+
+ # Run the simulation for one step
+ sim.step(1)
+
+ maker._update_interchange(interchange, sim, None) # noqa: SLF001
+
+ assert interchange.positions.shape == start_positions.shape
+ assert interchange.velocities.shape == (1170, 3)
+
+ assert np.any(interchange.positions != start_positions)
+ assert np.any(interchange.velocities != start_velocities)
+ assert np.all(interchange.box == start_box)
+
+
+def test_create_task_doc(interchange, tmp_path):
+ maker = BaseOpenMMMaker(n_steps=1000, temperature=300)
+ dir_name = tmp_path / "test_output"
+ dir_name.mkdir()
+
+ task_doc = maker._create_task_doc( # noqa: SLF001
+ interchange,
+ None,
+ elapsed_time=10.5,
+ dir_name=dir_name,
+ )
+
+ assert isinstance(task_doc, OpenMMTaskDocument)
+ assert task_doc.dir_name == str(dir_name)
+ assert task_doc.state == "successful"
+ assert len(task_doc.calcs_reversed) == 1
+ assert task_doc.calcs_reversed[0].input.n_steps == 1000
+ assert task_doc.calcs_reversed[0].input.temperature == 300
+ assert task_doc.calcs_reversed[0].output.elapsed_time == 10.5
+
+
+def test_make(interchange, tmp_path, run_job):
+ # Create an instance of BaseOpenMMMaker
+ maker = BaseOpenMMMaker(
+ n_steps=1000,
+ step_size=0.002,
+ platform_name="CPU",
+ state_interval=100,
+ traj_interval=50,
+ temperature=300,
+ friction_coefficient=1,
+ save_structure=True,
+ )
+
+ # monkey patch to allow running the test without openmm
+
+ def do_nothing(self, sim):
+ pass
+
+ BaseOpenMMMaker.run_openmm = do_nothing
+
+ # Call the make method
+ base_job = maker.make(interchange)
+ assert isinstance(base_job, Job)
+
+ task_doc = run_job(base_job)
+
+ # Assert the specific values in the task document
+ assert isinstance(task_doc, OpenMMTaskDocument)
+ assert task_doc.state == "successful"
+ # assert task_doc.dir_name == str(tmp_path)
+ assert len(task_doc.calcs_reversed) == 1
+ assert isinstance(task_doc.structure, Structure)
+
+ # Assert the calculation details
+ calc = task_doc.calcs_reversed[0]
+ # assert calc.dir_name == str(tmp_path)
+ assert calc.has_openmm_completed is True
+ assert calc.input.n_steps == 1000
+ assert calc.input.step_size == 0.002
+ assert calc.input.platform_name == "CPU"
+ assert calc.input.state_interval == 100
+ assert calc.input.traj_interval == 50
+ assert calc.input.temperature == 300
+ assert calc.input.friction_coefficient == 1
+ assert calc.output is not None
+ assert calc.completed_at is not None
+ assert calc.task_name == "base openmm job"
+ assert calc.calc_type == "BaseOpenMMMaker"
+
+
+def test_make_w_velocities(interchange, run_job):
+ # monkey patch to allow running the test without openmm
+ def do_nothing(self, sim):
+ pass
+
+ BaseOpenMMMaker.run_openmm = do_nothing
+
+ maker1 = BaseOpenMMMaker(
+ n_steps=1000,
+ report_velocities=True,
+ )
+
+ with pytest.raises(RuntimeError):
+ run_job(maker1.make(interchange))
+ # run_job(base_job)
+
+ import MDAnalysis
+ from packaging.version import Version
+
+ if Version(MDAnalysis.__version__) < Version("2.8.0"):
+ return
+
+ maker2 = BaseOpenMMMaker(
+ n_steps=1000,
+ report_velocities=True,
+ traj_file_type="h5md",
+ )
+
+ base_job = maker2.make(interchange)
+ task_doc = run_job(base_job)
+
+ # Assert the calculation details
+ calc = task_doc.calcs_reversed[0]
+ assert calc.input.report_velocities is True
+
+
+def test_make_from_prev(run_job):
+ mol_specs_dicts = [
+ {"smiles": "CCO", "count": 50, "name": "ethanol", "charge_method": "mmff94"},
+ {"smiles": "O", "count": 300, "name": "water", "charge_method": "mmff94"},
+ ]
+ inter_job = generate_interchange(mol_specs_dicts, 1)
+
+ # Create an instance of BaseOpenMMMaker
+ maker = BaseOpenMMMaker(n_steps=10)
+
+ # monkey patch to allow running the test without openmm
+ def do_nothing(self, sim):
+ pass
+
+ BaseOpenMMMaker.run_openmm = do_nothing
+
+ # Call the make method
+ base_job = maker.make(
+ inter_job.output.interchange, prev_dir=inter_job.output.dir_name
+ )
+
+ task_doc = run_job(Flow([inter_job, base_job]))
+
+ assert task_doc.mol_specs is not None
diff --git a/tests/openmm_md/jobs/test_core.py b/tests/openmm_md/jobs/test_core.py
new file mode 100644
index 0000000000..cfa3b5173e
--- /dev/null
+++ b/tests/openmm_md/jobs/test_core.py
@@ -0,0 +1,58 @@
+import numpy as np
+from openff.interchange import Interchange
+
+from atomate2.openmm.jobs import (
+ EnergyMinimizationMaker,
+ NPTMaker,
+ NVTMaker,
+ TempChangeMaker,
+)
+
+
+def test_energy_minimization_maker(interchange, run_job):
+ maker = EnergyMinimizationMaker(max_iterations=1)
+ base_job = maker.make(interchange)
+ task_doc = run_job(base_job)
+ new_interchange = Interchange.parse_raw(task_doc.interchange)
+
+ assert np.any(new_interchange.positions != interchange.positions)
+
+
+def test_npt_maker(interchange, run_job):
+ maker = NPTMaker(n_steps=10, pressure=0.1, pressure_update_frequency=1)
+ base_job = maker.make(interchange)
+ task_doc = run_job(base_job)
+ new_interchange = Interchange.parse_raw(task_doc.interchange)
+
+ # test that coordinates and box size has changed
+ assert np.any(new_interchange.positions != interchange.positions)
+ assert np.any(new_interchange.box != interchange.box)
+
+
+def test_nvt_maker(interchange, run_job):
+ maker = NVTMaker(n_steps=10, state_interval=1)
+ base_job = maker.make(interchange)
+ task_doc = run_job(base_job)
+ new_interchange = Interchange.parse_raw(task_doc.interchange)
+
+ # test that coordinates have changed
+ assert np.any(new_interchange.positions != interchange.positions)
+
+ # Test length of state attributes in calculation output
+ calc_output = task_doc.calcs_reversed[0].output
+ assert len(calc_output.steps_reported) == 10
+
+ # Test that the state interval is respected
+ assert calc_output.steps_reported == list(range(1, 11))
+
+
+def test_temp_change_maker(interchange, run_job):
+ maker = TempChangeMaker(n_steps=10, temperature=310, temp_steps=10)
+ base_job = maker.make(interchange)
+ task_doc = run_job(base_job)
+ new_interchange = Interchange.parse_raw(task_doc.interchange)
+
+ # test that coordinates have changed and starting temperature is present and correct
+ assert np.any(new_interchange.positions != interchange.positions)
+ assert task_doc.calcs_reversed[0].input.temperature == 310
+ assert task_doc.calcs_reversed[0].input.starting_temperature == 298
diff --git a/tests/openmm_md/jobs/test_generate.py b/tests/openmm_md/jobs/test_generate.py
new file mode 100644
index 0000000000..e7784a6319
--- /dev/null
+++ b/tests/openmm_md/jobs/test_generate.py
@@ -0,0 +1,190 @@
+import numpy as np
+import openff.toolkit as tk
+from emmet.core.openmm import OpenMMInterchange
+from jobflow import Flow
+from openff.interchange.components._packmol import pack_box
+from openff.units import unit
+from openmm import XmlSerializer
+
+from atomate2.openff.utils import create_mol_spec
+from atomate2.openmm.jobs import EnergyMinimizationMaker
+from atomate2.openmm.jobs.base import BaseOpenMMMaker
+from atomate2.openmm.jobs.generate import (
+ XMLMoleculeFF,
+ create_system_from_xml,
+ generate_openmm_interchange,
+)
+
+
+def test_create_system_from_xml(openmm_data):
+ # load strings of xml files into dict
+ ff_xmls = [
+ XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CCO.xml"),
+ XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml"),
+ ]
+
+ # uncomment to regenerate data
+ # download_opls_xml("CCO", opls_xmls / "CCO.xml")
+ # download_opls_xml("CO", opls_xmls / "CO.xml")
+
+ mol_specs = [
+ {"smiles": "CCO", "count": 10},
+ {"smiles": "CO", "count": 20},
+ ]
+
+ topology = pack_box(
+ molecules=[tk.Molecule.from_smiles(spec["smiles"]) for spec in mol_specs],
+ number_of_copies=[spec["count"] for spec in mol_specs],
+ mass_density=0.8 * unit.grams / unit.milliliter,
+ )
+
+ create_system_from_xml(topology, ff_xmls)
+
+
+def test_xml_molecule_from_file(openmm_data):
+ xml_mol = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
+
+ assert isinstance(str(xml_mol), str)
+ assert len(str(xml_mol)) > 100
+
+
+def test_to_openff_molecule(openmm_data):
+ xml_mol = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
+
+ mol = xml_mol.to_openff_molecule()
+ assert len(mol.atoms) == 6
+ assert len(mol.bonds) == 5
+
+
+def test_assign_partial_charges_w_mol(openmm_data):
+ xml_mol = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
+
+ openff_mol = tk.Molecule()
+
+ atom_c00 = openff_mol.add_atom(6, 0, is_aromatic=False)
+ atom_h02 = openff_mol.add_atom(1, 0, is_aromatic=False)
+ atom_h03 = openff_mol.add_atom(1, 0, is_aromatic=False)
+ atom_h04 = openff_mol.add_atom(1, 0, is_aromatic=False)
+ atom_h05 = openff_mol.add_atom(1, 0, is_aromatic=False)
+ atom_o01 = openff_mol.add_atom(8, 0, is_aromatic=False)
+
+ openff_mol.add_bond(atom_c00, atom_o01, bond_order=1, is_aromatic=False)
+ openff_mol.add_bond(atom_c00, atom_h02, bond_order=1, is_aromatic=False)
+ openff_mol.add_bond(atom_c00, atom_h03, bond_order=1, is_aromatic=False)
+ openff_mol.add_bond(atom_c00, atom_h04, bond_order=1, is_aromatic=False)
+ openff_mol.add_bond(atom_o01, atom_h05, bond_order=1, is_aromatic=False)
+
+ openff_mol.assign_partial_charges("mmff94")
+
+ xml_mol.assign_partial_charges(openff_mol)
+ assert xml_mol.partial_charges[0] > 0.2 # C
+ assert xml_mol.partial_charges[1] < -0.3 # O
+ assert xml_mol.partial_charges[5] > 0.1 # alcohol H
+
+
+def test_assign_partial_charges_w_method(openmm_data):
+ xml_mol = XMLMoleculeFF.from_file(openmm_data / "opls_xml_files" / "CO.xml")
+ xml_mol.assign_partial_charges("mmff94")
+ assert xml_mol.partial_charges[0] > 0.2 # C
+ assert xml_mol.partial_charges[1] < -0.3 # O
+ assert xml_mol.partial_charges[5] > 0.1 # alcohol H
+
+
+def test_generate_openmm_interchange(openmm_data, run_job):
+ mol_specs = [
+ create_mol_spec("CCO", 10, name="ethanol", charge_method="mmff94"),
+ create_mol_spec("CO", 300, name="water", charge_method="mmff94"),
+ ]
+
+ ff_xmls = [
+ (openmm_data / "opls_xml_files" / "CCO.xml").read_text(),
+ (openmm_data / "opls_xml_files" / "CO.xml").read_text(),
+ ]
+
+ job = generate_openmm_interchange(
+ mol_specs, 1.0, ff_xmls, xml_method_and_scaling=("cm1a-lbcc", 1.14)
+ )
+ task_doc = run_job(job)
+ molecule_specs = task_doc.mol_specs
+ assert len(molecule_specs) == 2
+ assert molecule_specs[0].name == "ethanol"
+ assert molecule_specs[0].count == 10
+ assert molecule_specs[1].name == "water"
+ assert molecule_specs[1].count == 300
+ co = tk.Molecule.from_json(molecule_specs[1].openff_mol)
+
+ assert np.allclose(
+ co.partial_charges.magnitude,
+ np.array([-0.0492, -0.5873, 0.0768, 0.0768, 0.0768, 0.4061]), # from file
+ )
+
+
+def test_make_from_prev(openmm_data, run_job):
+ mol_specs = [
+ create_mol_spec("CCO", 10, name="ethanol", charge_method="mmff94"),
+ create_mol_spec("CO", 300, name="water", charge_method="mmff94"),
+ ]
+
+ ff_xmls = [
+ (openmm_data / "opls_xml_files" / "CCO.xml").read_text(),
+ (openmm_data / "opls_xml_files" / "CO.xml").read_text(),
+ ]
+ inter_job = generate_openmm_interchange(mol_specs, 1, ff_xmls)
+
+ # Create an instance of BaseOpenMMMaker
+ maker = BaseOpenMMMaker(n_steps=10)
+
+ # monkey patch to allow running the test without openmm
+ def do_nothing(self, sim):
+ pass
+
+ BaseOpenMMMaker.run_openmm = do_nothing
+
+ # Call the make method
+ base_job = maker.make(
+ inter_job.output.interchange, prev_dir=inter_job.output.dir_name
+ )
+
+ task_doc = run_job(Flow([inter_job, base_job]))
+
+ assert task_doc.mol_specs is not None
+
+
+def test_evolve_simulation(openmm_data, run_job):
+ mol_specs = [
+ create_mol_spec("CCO", 10, name="ethanol", charge_method="mmff94"),
+ create_mol_spec("CO", 300, name="water", charge_method="mmff94"),
+ ]
+
+ ff_xmls = [
+ (openmm_data / "opls_xml_files" / "CCO.xml").read_text(),
+ (openmm_data / "opls_xml_files" / "CO.xml").read_text(),
+ ]
+ inter_job = generate_openmm_interchange(mol_specs, 1, ff_xmls)
+
+ task_doc = run_job(inter_job)
+
+ # test that opls charges are not being used
+ co = tk.Molecule.from_json(task_doc.mol_specs[1].openff_mol)
+ assert not np.allclose(
+ co.partial_charges.magnitude,
+ np.array([-0.5873, -0.0492, 0.0768, 0.0768, 0.4061, 0.0768]), # from file
+ )
+
+ interchange_str = task_doc.interchange
+ interchange = OpenMMInterchange.parse_raw(interchange_str)
+
+ initial_state = XmlSerializer.deserialize(interchange.state)
+ initial_position = initial_state.getPositions(asNumpy=True)
+
+ maker = EnergyMinimizationMaker(max_iterations=1)
+ min_job = maker.make(interchange)
+
+ task_doc2 = run_job(min_job)
+ interchange_str2 = task_doc2.interchange
+ interchange2 = OpenMMInterchange.parse_raw(interchange_str2)
+
+ final_state = XmlSerializer.deserialize(interchange2.state)
+ final_position = final_state.getPositions(asNumpy=True)
+
+ assert not (final_position == initial_position).all()
diff --git a/tests/openmm_md/test_utils.py b/tests/openmm_md/test_utils.py
new file mode 100644
index 0000000000..9fd07726c6
--- /dev/null
+++ b/tests/openmm_md/test_utils.py
@@ -0,0 +1,29 @@
+import pytest
+
+from atomate2.openmm.utils import download_opls_xml, increment_name
+
+
+@pytest.mark.skip("annoying test")
+def test_download_xml(tmp_path):
+ pytest.importorskip("selenium")
+
+ download_opls_xml("CCO", tmp_path / "CCO.xml")
+
+ assert (tmp_path / "CCO.xml").exists()
+
+
+def test_increment_file_name():
+ test_cases = [
+ ("report", "report2"),
+ ("report123", "report124"),
+ ("report.123", "report.124"),
+ ("report-123", "report-124"),
+ ("report-dcd", "report-dcd2"),
+ ("report.123.dcd", "report.123.dcd2"),
+ ]
+
+ for file_name, expected_output in test_cases:
+ result = increment_name(file_name)
+ assert (
+ result == expected_output
+ ), f"Failed for case: {file_name}. Expected: {expected_output}, Got: {result}"
diff --git a/tests/test_data/openff/geometries/trimer.pdb b/tests/test_data/openff/geometries/trimer.pdb
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/test_data/openff/geometries/trimer.txt b/tests/test_data/openff/geometries/trimer.txt
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/test_data/openff/molecule_charge_files/CCO.npy b/tests/test_data/openff/molecule_charge_files/CCO.npy
new file mode 100644
index 0000000000..4724234e78
Binary files /dev/null and b/tests/test_data/openff/molecule_charge_files/CCO.npy differ
diff --git a/tests/test_data/openff/molecule_charge_files/CCO.xyz b/tests/test_data/openff/molecule_charge_files/CCO.xyz
new file mode 100644
index 0000000000..403720cd49
--- /dev/null
+++ b/tests/test_data/openff/molecule_charge_files/CCO.xyz
@@ -0,0 +1,11 @@
+9
+
+C 1.000000 1.000000 0.000000
+C -0.515000 1.000000 0.000000
+O -0.999000 1.000000 1.335000
+H 1.390000 1.001000 -1.022000
+H 1.386000 0.119000 0.523000
+H 1.385000 1.880000 0.526000
+H -0.907000 0.118000 -0.516000
+H -0.897000 1.894000 -0.501000
+H -0.661000 0.198000 1.768000
diff --git a/tests/test_data/openff/molecule_charge_files/FEC-r.xyz b/tests/test_data/openff/molecule_charge_files/FEC-r.xyz
new file mode 100644
index 0000000000..f94a8923ee
--- /dev/null
+++ b/tests/test_data/openff/molecule_charge_files/FEC-r.xyz
@@ -0,0 +1,12 @@
+10
+
+O 1.000000 1.000000 0.000000
+C -0.219000 1.000000 0.000000
+O -0.984000 1.000000 1.133000
+C -2.322000 0.780000 0.720000
+C -2.300000 1.205000 -0.711000
+H -3.034000 0.686000 -1.332000
+F -2.507000 2.542000 -0.809000
+O -0.983000 0.948000 -1.128000
+H -3.008000 1.375000 1.328000
+H -2.544000 -0.285000 0.838000
diff --git a/tests/test_data/openff/molecule_charge_files/FEC-s.xyz b/tests/test_data/openff/molecule_charge_files/FEC-s.xyz
new file mode 100644
index 0000000000..af492afc65
--- /dev/null
+++ b/tests/test_data/openff/molecule_charge_files/FEC-s.xyz
@@ -0,0 +1,12 @@
+10
+
+O 1.000000 1.000000 0.000000
+C -0.219000 1.000000 0.000000
+O -0.981000 1.000000 1.133000
+C -2.323000 0.828000 0.723000
+C -2.305000 1.254000 -0.707000
+H -2.567000 2.305000 -0.862000
+F -3.125000 0.469000 -1.445000
+O -0.983000 1.001000 -1.127000
+H -2.991000 1.447000 1.328000
+H -2.610000 -0.222000 0.848000
diff --git a/tests/test_data/openff/molecule_charge_files/FEC.npy b/tests/test_data/openff/molecule_charge_files/FEC.npy
new file mode 100644
index 0000000000..016912f7d6
Binary files /dev/null and b/tests/test_data/openff/molecule_charge_files/FEC.npy differ
diff --git a/tests/test_data/openff/molecule_charge_files/FEC_bad.npy b/tests/test_data/openff/molecule_charge_files/FEC_bad.npy
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/test_data/openff/molecule_charge_files/Li.npy b/tests/test_data/openff/molecule_charge_files/Li.npy
new file mode 100644
index 0000000000..b87c6ecebb
Binary files /dev/null and b/tests/test_data/openff/molecule_charge_files/Li.npy differ
diff --git a/tests/test_data/openff/molecule_charge_files/Li.xyz b/tests/test_data/openff/molecule_charge_files/Li.xyz
new file mode 100644
index 0000000000..7f08d77c84
--- /dev/null
+++ b/tests/test_data/openff/molecule_charge_files/Li.xyz
@@ -0,0 +1,3 @@
+1
+
+Li 0.0 0.0 0.0
diff --git a/tests/test_data/openff/molecule_charge_files/PF6.npy b/tests/test_data/openff/molecule_charge_files/PF6.npy
new file mode 100644
index 0000000000..1a4c723b70
Binary files /dev/null and b/tests/test_data/openff/molecule_charge_files/PF6.npy differ
diff --git a/tests/test_data/openff/molecule_charge_files/PF6.xyz b/tests/test_data/openff/molecule_charge_files/PF6.xyz
new file mode 100644
index 0000000000..3f3df87fa8
--- /dev/null
+++ b/tests/test_data/openff/molecule_charge_files/PF6.xyz
@@ -0,0 +1,9 @@
+7
+
+P 0.0 0.0 0.0
+F 1.6 0.0 0.0
+F -1.6 0.0 0.0
+F 0.0 1.6 0.0
+F 0.0 -1.6 0.0
+F 0.0 0.0 1.6
+F 0.0 0.0 -1.6
diff --git a/tests/test_data/openmm/opls_xml_files/CCO.xml b/tests/test_data/openmm/opls_xml_files/CCO.xml
new file mode 100644
index 0000000000..6e8c4ea08b
--- /dev/null
+++ b/tests/test_data/openmm/opls_xml_files/CCO.xml
@@ -0,0 +1,88 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/tests/test_data/openmm/opls_xml_files/CO.xml b/tests/test_data/openmm/opls_xml_files/CO.xml
new file mode 100644
index 0000000000..4381bd9348
--- /dev/null
+++ b/tests/test_data/openmm/opls_xml_files/CO.xml
@@ -0,0 +1,56 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/tutorials/openmm_tutorial.ipynb b/tutorials/openmm_tutorial.ipynb
new file mode 100644
index 0000000000..8cfb3baf9c
--- /dev/null
+++ b/tutorials/openmm_tutorial.ipynb
@@ -0,0 +1,420 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "0",
+ "metadata": {
+ "jupyter": {
+ "outputs_hidden": false
+ }
+ },
+ "source": [
+ "### Installing Atomate2 From Source with OpenMM\n",
+ "\n",
+ "```bash\n",
+ "# setting up our conda environment\n",
+ ">>> conda create -n atomate2 python=3.11\n",
+ ">>> conda activate atomate2\n",
+ "\n",
+ "# installing atomate2\n",
+ ">>> pip install 'git+https://github.com/orionarcher/atomate2.git#egg=atomate2[classical_md]'\n",
+ "\n",
+ "# installing classical_md dependencies\n",
+ ">>> conda install -c conda-forge --file .github/classical_md_requirements.txt\n",
+ "```\n",
+ "\n",
+ "Alternatively, if you anticipate regularly updating atomate2 from source (which at this point, you should), you can clone the repository and install from source.\n",
+ "\n",
+ "``` bash\n",
+ "# installing atomate2\n",
+ ">>> git clone https://github.com/orionarcher/atomate2.git\n",
+ ">>> cd atomate2\n",
+ ">>> git branch openff\n",
+ ">>> git checkout openff\n",
+ ">>> git pull origin openff\n",
+ ">>> pip install -e '.[classical_md]'\n",
+ "```\n",
+ "\n",
+ "To test the openmm installation, you can run the following command. If you intend to run on GPU, make sure that the tests are passing for CUDA.\n",
+ "\n",
+ "```bash\n",
+ ">>> python -m openmm.testInstallation\n",
+ "```\n",
+ "\n",
+ "```bash\n",
+ "pip uninstall pymongo\n",
+ "pip uninstall bson\n",
+ "pip install pymongo\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1",
+ "metadata": {},
+ "source": [
+ "### Understanding Atomate2 OpenMM \n",
+ "\n",
+ "Atomate2 is really just a collection of jobflow workflows relevant to materials science. In all the workflows, we pass our system of interest between different jobs to perform the desired simulation. Representing the intermediate state of a classical molecular dynamics simulation, however, is challenging. While the intermediate representation between stages of a periodic DFT simulation can include just the elements, xyz coordinates, and box vectors, classical molecular dynamics systems must also include velocities and forces. The latter is particularly challenging because all MD engines represent forces differently. Rather than implement our own representation, we use the `openff.interchange.Interchange` object, which catalogs the necessary system properties and interfaces with a variety of MD engines. This is the object that we pass between stages of a classical MD simulation and it is the starting point of our workflow."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2",
+ "metadata": {},
+ "source": [
+ "### Pouring a Glass of Wine\n",
+ "\n",
+ "The first job we need to create generates the `Interchange` object. To specify the system of interest, we use give it the SMILES strings, counts, and names (optional) of the molecules we want to include."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3",
+ "metadata": {
+ "jupyter": {
+ "outputs_hidden": false
+ }
+ },
+ "outputs": [],
+ "source": [
+ "from atomate2.openff.core import generate_interchange\n",
+ "\n",
+ "mol_specs_dicts = [\n",
+ " {\"smiles\": \"O\", \"count\": 200, \"name\": \"water\"},\n",
+ " {\"smiles\": \"CCO\", \"count\": 10, \"name\": \"ethanol\"},\n",
+ " {\"smiles\": \"C1=C(C=C(C(=C1O)O)O)C(=O)O\", \"count\": 1, \"name\": \"gallic_acid\"},\n",
+ "]\n",
+ "\n",
+ "gallic_interchange_job = generate_interchange(mol_specs_dicts, 1.3)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4",
+ "metadata": {},
+ "source": [
+ "If you are wondering what arguments are allowed in the dictionaries, check out the `create_mol_spec` function in the `atomate2.openff.utils` module. Under the hood, this is being called on each mol_spec dict. Meaning the code below is functionally identical to the code above."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from atomate2.openff.utils import create_mol_spec\n",
+ "\n",
+ "mols_specs = [create_mol_spec(**mol_spec_dict) for mol_spec_dict in mol_specs_dicts]\n",
+ "\n",
+ "generate_interchange(mols_specs, 1.3)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "6",
+ "metadata": {},
+ "source": [
+ "In a more complex simulation we might want to scale the ion charges and include custom partial charges. An example with the Gen2 electrolyte is shown below. This yields the `elyte_interchange_job` object, which we can pass to the next stage of the simulation.\n",
+ "\n",
+ "NOTE: It's actually mandatory to include partial charges for PF6- here, the built in partial charge method fails."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "7",
+ "metadata": {
+ "jupyter": {
+ "outputs_hidden": false
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "from pymatgen.core.structure import Molecule\n",
+ "\n",
+ "pf6 = Molecule(\n",
+ " [\"P\", \"F\", \"F\", \"F\", \"F\", \"F\", \"F\"],\n",
+ " [\n",
+ " [0.0, 0.0, 0.0],\n",
+ " [1.6, 0.0, 0.0],\n",
+ " [-1.6, 0.0, 0.0],\n",
+ " [0.0, 1.6, 0.0],\n",
+ " [0.0, -1.6, 0.0],\n",
+ " [0.0, 0.0, 1.6],\n",
+ " [0.0, 0.0, -1.6],\n",
+ " ],\n",
+ ")\n",
+ "pf6_charges = np.array([1.34, -0.39, -0.39, -0.39, -0.39, -0.39, -0.39])\n",
+ "\n",
+ "mol_specs_dicts = [\n",
+ " {\"smiles\": \"C1COC(=O)O1\", \"count\": 100, \"name\": \"EC\"},\n",
+ " {\"smiles\": \"CCOC(=O)OC\", \"count\": 100, \"name\": \"EMC\"},\n",
+ " {\n",
+ " \"smiles\": \"F[P-](F)(F)(F)(F)F\",\n",
+ " \"count\": 50,\n",
+ " \"name\": \"PF6\",\n",
+ " \"partial_charges\": pf6_charges,\n",
+ " \"geometry\": pf6,\n",
+ " \"charge_scaling\": 0.8,\n",
+ " \"charge_method\": \"RESP\",\n",
+ " },\n",
+ " {\"smiles\": \"[Li+]\", \"count\": 50, \"name\": \"Li\", \"charge_scaling\": 0.8},\n",
+ "]\n",
+ "\n",
+ "elyte_interchange_job = generate_interchange(mol_specs_dicts, 1.3)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "8",
+ "metadata": {},
+ "source": [
+ "### The basic simulation\n",
+ "\n",
+ "To run a production simulation, we will create a production flow, link it to our `elyte_interchange_job`, and then run both locally.\n",
+ "\n",
+ "In jobflow, jobs and flows are created by [Makers](https://materialsproject.github.io/jobflow/tutorials/6-makers.html), which can then be linked into more complex flows. The production maker links together makers for energy minimization, pressure equilibration, annealing, and a nvt simulation. The annealing is itself a flow flow that links together nvt and tempchange makers (it uses the `anneal_flow` method to save us from creating three more jobs manually). When linked up the `generate_interchange` job this yields a production ready molecular dynamics workflow.\n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from jobflow import Flow, run_locally\n",
+ "\n",
+ "from atomate2.openmm.flows.core import OpenMMFlowMaker\n",
+ "from atomate2.openmm.jobs.core import EnergyMinimizationMaker, NPTMaker, NVTMaker\n",
+ "\n",
+ "production_maker = OpenMMFlowMaker(\n",
+ " name=\"production_flow\",\n",
+ " makers=[\n",
+ " EnergyMinimizationMaker(traj_interval=10, state_interval=10),\n",
+ " NPTMaker(n_steps=100),\n",
+ " OpenMMFlowMaker.anneal_flow(n_steps=150),\n",
+ " NVTMaker(n_steps=100),\n",
+ " ],\n",
+ ")\n",
+ "\n",
+ "production_flow = production_maker.make(\n",
+ " elyte_interchange_job.output.interchange,\n",
+ " prev_dir=elyte_interchange_job.output.dir_name,\n",
+ ")\n",
+ "\n",
+ "run_locally(\n",
+ " Flow([elyte_interchange_job, production_flow]), root_dir=\"./tutorial_system\"\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "10",
+ "metadata": {},
+ "source": [
+ "When the above code is executed, you should expect to see something like this:\n",
+ "\n",
+ "```\n",
+ "/tutorial_system\n",
+ "├── state.csv\n",
+ "├── state2.csv\n",
+ "├── state3.csv\n",
+ "├── state4.csv\n",
+ "├── state5.csv\n",
+ "├── state6.csv\n",
+ "├── taskdoc.json\n",
+ "├── trajectory.dcd\n",
+ "├── trajectory2.dcd\n",
+ "├── trajectory3.dcd\n",
+ "├── trajectory4.dcd\n",
+ "├── trajectory5.dcd\n",
+ "├── trajectory6.dcd\n",
+ "```\n",
+ "\n",
+ "We see that each job saved a separate state and trajectory file. There are 6 because the `AnnealMaker` creates 3 sub-jobs and the `EnergyMinimizationMaker` does not report anything. We also see a `taskdoc.json` file, which contains the metadata for the entire workflow. This is needed when we later want to do downstream analysis in `emmet`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "11",
+ "metadata": {},
+ "source": [
+ "### Configuring the Simulation\n",
+ "\n",
+ "All OpenMM jobs, i.e. anything in `atomate2.openmm.jobs`, inherits from the `BaseOpenMMMaker` class. `BaseOpenMMMaker` is highly configurable, you can change the timestep, temperature, reporting frequencies, output types, and a range of other properties. See the docstring for the full list of options.\n",
+ "\n",
+ "Note that when instantiating the `ProductionMaker` above, we only set the `traj_interval` and `state_interval` once, inside `EnergyMinimizationMaker`. This is a key feature: all makers will inherit attributes from the previous maker if they are not explicitly reset. This allows you to set the timestep once and have it apply to all stages of the simulation. More explicitly, the value inheritance is as follows: 1) any explictly set value, 2) the value from the previous maker, 3) the default value, shown below."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "12",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from atomate2.openmm.jobs.base import OPENMM_MAKER_DEFAULTS\n",
+ "\n",
+ "OPENMM_MAKER_DEFAULTS"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "13",
+ "metadata": {},
+ "source": [
+ "Perhaps we want to record a trajectory with velocities but only for the final NVT run. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "14",
+ "metadata": {},
+ "source": [
+ "### Running with Databases\n",
+ "\n",
+ "Before trying this, you should have a basic understanding of JobFlow and [Stores](https://materialsproject.github.io/jobflow/stores.html).\n",
+ "\n",
+ "To log OpenMM results to a database, you'll need to set up both a MongoStore, for taskdocs, and blob storage, for trajectories. Here, I'll show you the correct jobflow.yaml file to use the MongoDB storage and MinIO S3 storage provided by NERSC. To get this up, you'll need to contact NERSC to get accounts on their MongoDB and MinIO services. Then you can follow the instructions in the [Stores](https://materialsproject.github.io/jobflow/stores.html) tutorial to link jobflow to your databases. Your `jobflow.yaml` should look like this:\n",
+ "\n",
+ "```yaml\n",
+ "JOB_STORE:\n",
+ " docs_store:\n",
+ " type: MongoStore\n",
+ " database: DATABASE\n",
+ " collection_name: atomate2_docs # suggested\n",
+ " host: mongodb05.nersc.gov\n",
+ " port: 27017\n",
+ " username: USERNAME\n",
+ " password: PASSWORD\n",
+ "\n",
+ " additional_stores:\n",
+ " data:\n",
+ " type: S3Store\n",
+ " index:\n",
+ " type: MongoStore\n",
+ " database: DATABASE\n",
+ " collection_name: atomate2_blobs_index # suggested\n",
+ " host: mongodb05.nersc.gov\n",
+ " port: 27017\n",
+ " username: USERNAME\n",
+ " password: PASSWORD\n",
+ " key: blob_uuid\n",
+ " bucket: oac\n",
+ " s3_profile: oac\n",
+ " s3_resource_kwargs:\n",
+ " verify: false\n",
+ " endpoint_url: https://next-gen-minio.materialsproject.org/\n",
+ " key: blob_uuid\n",
+ "```\n",
+ "\n",
+ "NOTE: This can work with any MongoDB and S3 storage, not just NERSC's.\n",
+ "\n",
+ "Rather than use `jobflow.yaml`, you could also create the stores in Python and pass the stores to the `run_locally` function. This is shown below for completeness but the prior method is usually recommended.\n",
+ "\n",
+ "\n",
+ "```python\n",
+ "from jobflow import run_locally, JobStore\n",
+ "from maggma.stores import MongoStore, S3Store, MemoryStore\n",
+ "\n",
+ "md_doc_store = MongoStore(\n",
+ " username=\"USERNAME\",\n",
+ " password=\"PASSWORD\",\n",
+ " database=\"DATABASE\",\n",
+ " collection_name=\"atomate2_docs\", # suggested\n",
+ " host=\"mongodb05.nersc.gov\",\n",
+ " port=27017,\n",
+ ")\n",
+ "\n",
+ "md_blob_index = MongoStore(\n",
+ " username=\"USERNAME\",\n",
+ " password=\"PASSWORD\",\n",
+ " database=\"DATABASE\",\n",
+ " collection_name=\"atomate2_blobs_index\", # suggested\n",
+ " host=\"mongodb05.nersc.gov\",\n",
+ " port=27017,\n",
+ " key=\"blob_uuid\",\n",
+ ")\n",
+ "\n",
+ "md_blob_store = S3Store(\n",
+ " index=md_blob_index,\n",
+ " bucket=\"BUCKET\",\n",
+ " s3_profile=\"PROFILE\",\n",
+ " endpoint_url=\"https://next-gen-minio.materialsproject.org\",\n",
+ " key=\"blob_uuid\",\n",
+ ")\n",
+ "\n",
+ "wf = [] # set up whatever workflow you'd like to run\n",
+ "\n",
+ "# run the flow with our custom store\n",
+ "run_locally(\n",
+ " wf,\n",
+ " store=JobStore(md_doc_store, additional_stores={\"data\": md_blob_store}),\n",
+ " ensure_success=True,\n",
+ ")\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "15",
+ "metadata": {},
+ "source": [
+ "### Running on GPU(s)\n",
+ "\n",
+ "Running on a GPU is nearly as simple as running on a CPU. The only difference is that you need to specify the `platform_properties` argument in the `EnergyMinimizationMaker` with the `DeviceIndex` of the GPU you want to use."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "16",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "production_maker = OpenMMFlowMaker(\n",
+ " name=\"production_flow\",\n",
+ " makers=[\n",
+ " EnergyMinimizationMaker(\n",
+ " platform_name=\"CUDA\",\n",
+ " platform_properties={\"DeviceIndex\": \"0\"},\n",
+ " ),\n",
+ " NPTMaker(n_steps=100),\n",
+ " OpenMMFlowMaker.anneal_flow(n_steps=150),\n",
+ " NVTMaker(n_steps=100),\n",
+ " ],\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "17",
+ "metadata": {},
+ "source": [
+ "To run on a system with multiple GPUs, the 'DeviceIndex' can be changed to a different number for each job."
+ ]
+ }
+ ],
+ "metadata": {
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.8"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}