Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics

This repository implements learning Hamiltonian Flow Maps (HFM) for molecular dynamics simulations and other Hamiltonian systems.

Our learned Hamiltonian Flow Map enables stable simulations at significantly larger timesteps (right) than standard Velocity Verlet integration (left), which becomes inaccurate or unstable already at comparatively small step sizes. Consequently, our model explores conformational space more rapidly and identifies all relevant modes more efficiently than the baseline.

Interactive Simulator

See Hamiltonian Flow Maps in action. Run large-timestep molecular dynamics of Paracetamol directly in your browser. Adjust the timestep and watch the learned flow map propagate the molecule's phase-space state in real time.

Project Structure

• configs/: Hydra configuration files for models and simulation environments.
• hfm/: Core library source code, including the main application scripts.
• tests/: Unit tests for simulations and conservation filters.
• examples/: Minimal working examples for toy potentials and simple systems.

Installation

All required dependencies are managed via pyproject.toml. It is recommended to use uv for fast environment setup:

# Clone the repository and navigate to the root
cd hfm

# Setup the environment and verify installation
uv run python -c "import hfm; print('Setup successful')"

Short intro

See our standalone toy example for a quick intro into training and evaluating Hamiltonian Flow Maps, and a toy example with our code infrastructure.

Datasets

MD17 and MD22 datasets can be downloaded directly from the SGDML website. The gravity dataset can be generated using the potential we provide or the original SEGNN code base.

Training Models

To train a new model, edit the parameters in configs/config.yaml and start a single run from the hfm subdirectory:

cd hfm
uv run python app.py

You can use our custom resolvers for ASE units directly in all Hydra configs. For example, to define a duration of 0.1 nanoseconds in ASE units use: ${ns:0.1}

Running Simulations

Simulations are managed via simbench.py. Use the configuration files in configs/simbench.yaml to define your environment (integrators like NVE/NVT and simulation filters).

cd hfm
uv run python simbench.py

• Use sim_env/{your_config_file} to specify a simulation environment consisting of integrators and simulation filters.
• The main config (simbench.yaml) can be adapted to add/remove logged metrics or adjust constants.
• In simbench.yaml, specify the HFM model path and the path to the force field for energy correction.

You can download pretrained model files directly from huggingface.

Once you downloaded the data files and models you can run simulations, for example for paracetamol:

uv run python simbench.py load_models.sim={your_download_dir}/paracetamol/mean_flowparams_last.pkl data_module=md17_paracetamol globals.T_equilibrium=500 globals.gt_integration_timestep_fs=1.0

Testing

Use discover on the test directory to run all tests (run command from root directory):

uv run python -m unittest discover tests/

Alternatively run a single test using:

uv run python -m unittest tests/test_nvt_simulation.py

Contributing

Feel free to open an issue if you encounter any problems or have questions.

Citation

If you find our work useful, please cite:

@inproceedings{ripken2026learning,
  title = {Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics},
  author = {Ripken, Winfried and
            Plainer, Michael and
            Lied, Gregor and
            Frank, Thorben and
            Unke, Oliver T. and
            Chmiela, Stefan and
            No{\'e}, Frank and
            M{\"u}ller, Klaus-Robert},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML 2026)},
  year = {2026},
  publisher = {PMLR}
}