Add DynamicOpenMMFlowMaker for dynamic OpenMM Simulations
#1115
Conversation
…of TYPE_CHECKING to use as callable in default_factory
…id shared mutable defaults
…d placeholder should_continue function
orionarcher
left a comment
orionarcher
left a comment
There was a problem hiding this comment.
Looks good so far, just a couple comments
src/atomate2/openmm/flows/core.py
Outdated
| @dataclass | ||
| class DynamicOpenMMFlowMaker(Maker): |
There was a problem hiding this comment.
How about putting this in a folder called dynamic instead of core?
There was a problem hiding this comment.
If you mean a dynamic.py file in the same flows folder, I agree.
There was a problem hiding this comment.
Yep, thats what I meant!
src/atomate2/openmm/flows/core.py
Outdated
| @openmm_job | ||
| def default_should_continue( |
There was a problem hiding this comment.
test this by itself too
src/atomate2/openmm/flows/core.py
Outdated
| ( | ||
| list[OpenMMTaskDocument], | ||
| int, | ||
| int, | ||
| str, | ||
| float | None, | ||
| float, | ||
| float, | ||
| ) | ||
| should_continue: Callable[ |
There was a problem hiding this comment.
What's happening here? very weird lint maybe?
There was a problem hiding this comment.
This was unused/redundant code that should've been deleted, thanks! Also, I've fixed the mypy errors using a Protocol, so now all mypy tests pass.
src/atomate2/openmm/flows/core.py
Outdated
| task_docs: list[OpenMMTaskDocument], | ||
| stage_index: int, | ||
| max_stages: int, | ||
| physical_property: str = "potential_energy", |
There was a problem hiding this comment.
You could also make this take a callable that acts on a list of task docs, that'd be a bit more flexible that locking in a couple different physical parameters.
There was a problem hiding this comment.
I agree, ideally we want should_continue to solely depend on task_docs. I'd like to think about this some more, since there may be a physical_property not yet included in the Calculation (RMSD, RDFs, etc.), but can still be derived from trajectory information. So I'd propose keeping this as an MVP and generalizing in a future PR.
|
thanks @orionarcher, I've addressed all comments + any pre-commit/pytests are working on my end |
orionarcher
left a comment
orionarcher
left a comment
There was a problem hiding this comment.
Thank you for addressing my comments, I think this looks good.
|
Thanks @orionarcher and @janosh! At least from my side of things, I don't have any more changes to make, so please let me know if there's anything else you need from my end to merge this and #1111. |
…sproject#1115) * created DynamicOpenMMFlowMaker; moved BaseOpenMMMaker import outside of TYPE_CHECKING to use as callable in default_factory * use of lambda function is more appropriate for default_factory to avoid shared mutable defaults * added dynamic flow logic and apply_flow_control instace classes; added placeholder should_continue function * restructured dynamic_flow and removed apply_flow_control * removed dynamic_collect_outputs * added and passed DynamicOpenMMFlowMaker tests; fixed dynamic flow logic * undo incorrect change * created dynamic.py * created test_should_continue * mypy_extensions passes mypy, but adds additional dependency * mypy_extensions passes mypy, but adds additional dependency * replaced mypy_extensions with protocol, mypy tests pass. --------- Co-authored-by: Shehan M Parmar <[email protected]>
3 participants
Summary
DynamicOpenMMFlowMakerinsrc/atomate2/openmm/flows/dynamic.pydefault_should_continuefunctions insrc/atomate2/openmm/flows/dynamic.pytest_dynamic_flow_makerintests/openmm_md/flows/test_dynamic.pyThis PR resolves the final enhancement proposed in #1102. Currently, subclasses to
BaseOpenMMMakerincludeEnergyMinimizationMaker,NPTMaker,NVTMaker, andTempChangeMaker. This PR adds theDynamicOpenMMFlowMakerclass, suitable for high-throughput classical MD workflows. Equilibration timescales are highly system- and model-dependent for condensed-phase systems even with notable structural or compositional similarities (see ionic liquids, polymers, etc.). Now there is support for on-the-fly OpenMM trajectory propagation. Below is a simple water equilibration:Based on
default_should_continue, equilibration is deemed complete once the potential energy vs time's rate of change plateaus. A newNPTMakeris instantiated dynamically, creating equilibration stages. The water test case below, the simulation converges once the decay rate falls below thethreshold=1e-3:DynamicOpenMMFlowMaker Job Logs Output
TODO (if any)
I've runpre-commit run --allwith only the following issue unresolvedEDIT: Resolved with Protocol solution (See ShouldContinueProtocol)
Future Work
The scope of this PR is a straightforward, extensible DynamicOpenMMFlowMaker. For more complex workflows, future enhancements would include:
traj_intervalatcollect_jobPhaseclasses to easily stitch together intentionally separate, dynamic flows (i.e.,EquilibrationPhase,ProductionPhase, etc.)default_should_continuewith broader physical parameter support (e.g., integrate with pymbar, RED, etc.)Checklist
The easiest way to handle this is to run the following in the correct sequence on
your local machine. Start with running
ruffandruff formaton your new code. This willautomatically reformat your code to PEP8 conventions and fix many linting issues.
Run ruff on your code.
type check your code.
Note that the CI system will run all the above checks. But it will be much more
efficient if you already fix most errors prior to submitting the PR. It is highly
recommended that you use the pre-commit hook provided in the repository. Simply run
pre-commit installand a check will be run prior to allowing commits.