OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation

Wen Yan; Zhi Wang

REVIEW 3 major objections 4 minor 17 references

OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation

T0 review · 3 major / 4 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read The paper claims that any Python callable, PyTorch or NumPy, can serve as the force provider in MD simulation, bypassing TorchScript by handing the C++ engine a pointer to the Python object.

desk verdict Useful plugin with honest benchmarks, but the callback safety story (reference ownership and GIL) is thinner than the 'any Python module' claim. read the letter →

arxiv 2412.18271 v1 pith:6PI6LLXU submitted 2024-12-24 physics.comp-ph

Zhi Wang , Wen Yan This is my paper · ORCID

classification physics.comp-ph

keywords moleculardynamicsPythoncallbackOpenMMpluginPyTorchintegrationtorch.compileCUDAgraphsabinitiomachinelearningforcefields

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper claims that any Python callable — a PyTorch model, a torch.compile module, or a NumPy-backed quantum chemistry routine — can serve directly as the energy-and-force provider for a molecular dynamics simulation, with no TorchScript translation and no need to rewrite the model in C++. It proves this with a plugin that hands the C++ simulation engine a pointer to the Python object and calls it every timestep, using automatic differentiation to get forces when the model returns only energy. In tests on ethanol with a machine-learned force field, every deployment variant conserves the Hamiltonian to fp32 precision, and the torch.compile variant runs about 30 percent faster than the existing C++ TorchScript baseline. The authors further show the same callback works with ab initio quantum chemistry and argue the design ports to other MD engines.

What carries the argument

The callback mechanism: Python's built-in id() gives the integer value of a callable's PyObject pointer; that integer is passed through pybind11 into C++, where it is cast back to a py::object and called with PyObject_Call. A Callable class stores the id, the return type, and the parameters, and the TorchForce and NumPyForce classes expose it as an OpenMM force. The force calculation uses torch.autograd backpropagation when only energy is returned. This single mechanism replaces the static-analysis TorchScript pipeline, which the paper cites as failing for roughly half of real-world models.

What would settle it

Run an NVE simulation in which the Python model object wrapped by the force is deleted from Python and garbage-collected a few steps after force creation; if the callback pointer is not protected, the simulation should segfault, show corrupted energies, or deadlock.

Watch

Extended reading notes

Core claim

The central discovery is that the CPython object model is sufficient, by itself, to couple MD and ML: a Python call's id() yields the PyObject pointer, pybind11 can cast it back into a callable, and the C++ simulation thread can therefore invoke arbitrary Python code in place of a compiled force kernel. Because the force is computed inside Python, PyTorch's autograd can differentiate the energy without explicit force formulas; for quantum codes that supply gradients directly, NumPy arrays carry the same data. The authors validate against OpenMM Torch on a single ethanol molecule with the BAMBOO machine-learned force field, finding energy, force, and trajectory differences at the fp32 limit, and report that torch.compile with CUDA Graph is 8.2 times faster than the baseline in the small-system benchmark.

Load-bearing premise

The whole scheme assumes that the Python object's id(), a raw pointer into the interpreter's memory, stays valid and safely callable from C++ for the entire simulation, but the paper only says the object is kept resident and does not show reference counting, GIL locking, or thread-safety handling.

Editorial extensions

If this is right

Any PyTorch, torch.compile, or NumPy model becomes usable as an MD force field without passing the TorchScript compiler, and the paper reports that about half of real-world models fail TorchScript compilation.
Forces can be obtained by backpropagation through energy-only models, removing the need to hand-derive force expressions.
The same callback pattern works for AIMD by exchanging numpy arrays with quantum chemistry packages, so gradients supplied directly are also supported.
Because the hook is just a Python callable, deployment strategies such as CUDA Graphs, jit.script, and compile can be switched by editing a few lines of Python.
The design ports to other engines that can initialize a Python interpreter, such as Tinker and LAMMPS, according to the authors.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

Inference: A practical production wrapper would need explicit reference counting or a kept-alive registry plus a documented GIL policy; without those, the raw pointer that makes the method flexible is a memory-safety hazard in threaded or long-running simulations.
Inference: The performance comparison is dominated by kernel-launch overhead on a single molecule, so for larger systems the Python-callback overhead may shrink relative to compute and the reported 8.2 times speedup should not be expected to generalize without retesting.
Inference: The same id()-casting trick could generalize beyond MD to other C++ simulation frameworks that accept external potential terms, but each port would need the same interpreter-lifecycle care.
Inference: A testable extension is replacing the synchronous per-step callback with a batched async queue, which could keep Python and GPU work overlapped without changing the force interface.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Short version: this is a useful engineering paper with a real gap between what it demonstrates and what it claims. The plugin works, the benchmarks are honest, and the idea of using id() to pass a Python callable into OpenMM's C++ force object is clever. But the mechanism as described is not as general as advertised: it relies on the caller keeping the Python object alive and says nothing about GIL handling, which matters if forces are evaluated on worker threads. Those are fixable with documentation and tests, but they should not be hand-waved.

What's new: OpenMM-Torch and others require TorchScript or specialized operators. This plugin bypasses that by storing the PyObject pointer (via id()) and invoking the callable directly, which means torch.compile and plain NumPy functions work without extra compilation. The paper validates this on a single ethanol molecule with eight deployment strategies; energy drift, trajectory differences, and force errors all come out at fp32 precision level, and the speed table is informative. The AIMD demo with GPU4PySCF shows the same callback works for an external quantum chemistry code, which is a nice proof of flexibility.

The soft spots: (1) Reference ownership. The paper says 'ensuring that the object remains resident in memory' but never says who guarantees that. id() returns a raw address, not a strong reference. If the user's model variable goes out of scope, the pointer dangles and the simulation segfaults. The example works only because model42 stays alive for the whole script. A robust design would keep a reference in the Callable object. (2) GIL. PyObject_Call from C++ needs PyGILState_Ensure/Release unless the thread already holds the GIL. OpenMM's multi-threaded CPU platform evaluates forces from worker threads; the paper never mentions this. It may work in practice on the CUDA platform, but 'any Python module' is not established. (3) Portability to Tinker/LAMMPS is speculation based on 'implementation experience', not a port. (4) The 50% TorchScript failure rate appears without a citation. These are not fatal to the core result, but the authors should address 1 and 2 in the text and ideally add a test that deletes the model and runs on the multi-threaded CPU platform.

Who this is for: researchers who want to drop a PyTorch model into OpenMM without fighting TorchScript. It deserves a serious referee; the code is open, the evidence is reproducible, and the safety concerns are fixable. I would send it out with a request that the authors clarify reference lifetime and thread safety.

Referee Report

3 major / 4 minor

Summary. The paper presents OpenMM-Python-Force, an OpenMM plugin that lets a molecular dynamics simulation obtain energies and forces from an arbitrary Python callable via a callback mechanism. The Python object's id() is passed to C++, reinterpreted as a PyObject pointer, and invoked with pybind11. The authors validate the approach with NVE simulations of a single ethanol molecule using the BAMBOO machine-learned force field under eight deployment strategies (native PyTorch, torch.jit.script, torch.compile, with and without CUDA Graphs), reporting energy and force agreement at fp32 precision, and demonstrating that torch.compile plus CUDA Graphs is fastest. They also show an AIMD example with PySCF/GPU4PySCF via a NumPyForce, and argue the mechanism can be ported to other MD engines such as Tinker and LAMMPS.

Significance. If the callback mechanism is robust, the plugin is a valuable and timely contribution: it lets OpenMM users integrate arbitrary PyTorch or NumPy models without TorchScript compatibility, which is a real practical bottleneck. The numerical validation is internally consistent, the baseline comparison against OpenMM Torch is appropriate, and the source code is openly available. The performance tables are clear and the AIMD example demonstrates flexibility beyond machine-learned potentials. However, the paper's central 'general solution' claim depends on two unverified assumptions: that the Python object remains valid for the whole simulation, and that callbacks are thread-safe with respect to the GIL. These issues, plus the speculative portability claim, need to be addressed before the contribution can be accepted as a general mechanism.

major comments (3)

[Section 2 (Callback mechanism)] The sentence 'ensuring that the object remains resident in memory' asserts a necessary condition but does not specify the mechanism that guarantees it. The callback passes only id(model42), a raw integer, to the Callable class; nothing described in the paper performs Py_INCREF or otherwise takes a strong reference. If the Python variable model42 is deleted or goes out of scope, the numeric pointer obtained from id() becomes dangling, and the simulation can crash or produce undefined behavior. In the given example, model42 is a global in the same script, so it is kept alive by accident of script layout rather than by a documented guarantee. Furthermore, the paper does not describe any GIL management (e.g., PyGILState_Ensure/Release or py::gil_scoped_acquire) for force evaluation, even though OpenMM may call the force from C++ worker threads, for instance on the multi-threaded CPU platform. Please specify how the plugin manages reference ownership and thread state, and add tests that delete the Python callable during a simulation and that run on the CPU platform.
[Section 3.3 (Extensibility to Other MD Engines)] The portability claim is not supported by evidence. The paper states that adding a 'callback Python energy term' to Tinker or LAMMPS 'would require comparable code modifications' and that Python interpreter initialization 'would necessitate only minimal additional changes,' but no implementation, patch, benchmark, or proof-of-concept for any other engine is provided. The analogy to Tinker9's Fortran runtime initialization is conceptual only. Please either provide a minimal implementation for at least one other engine or revise the claim to state that cross-engine portability is a design expectation that has not yet been tested.
[Section 3.1 (Example: Ethanol)] The numerical validation is limited to a single ethanol molecule in vacuum, 100 NVE steps, and the CUDA mixed-precision platform. This is sufficient to demonstrate accuracy for that setup, but it does not exercise the callback mechanism under conditions where the safety concerns from Section 2 would surface, such as the multi-threaded CPU platform, long simulations, or callables constructed inside a function and then garbage collected. As a result, the 'general solution' claim in the abstract and conclusion is stronger than what the tests establish. Please add at least one CPU-platform run and, if possible, a test that deliberately destroys the Python callable to verify the reference-ownership behavior.

minor comments (4)

[Abstract] The word 'throughtorch.Tensor' should be 'through torch.Tensor'.
[Section 1 (Introduction)] The statement that 'approximately 50% of real-world models fail to compile successfully using torch.jit.script' is given without a citation; please provide a source or remove the specific percentage.
[Section 3.1 (Errors in Forces)] The text reports force RMSD values but does not state their units; the caption of Figure 4 says 'in kJ/mol and nm', which is not a standard unit for force. Please state the units explicitly, e.g., kJ/mol/nm, in both the text and the figure caption.
[Section 3.2 (AIMD Simulation)] The claim that 'the overhead from device-to-device data transfer and floating-point conversion proves negligible' is presented without profiling data; please add a brief justification or a timing measurement to support this statement.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: the callback mechanism is validated against an external baseline and does not derive its conclusions from its own assumptions.

full rationale

The paper is an engineering report rather than a derivation. Its central claim is that a Python callable can be invoked from OpenMM through the CPython C API by passing id(model) as a PyObject pointer and calling PyObject_Call. This is an implementation mechanism, not a claim derived from a fitted parameter or from the benchmark results. Numerical correctness is established by comparing energies and forces against the external OpenMM Torch baseline, and the performance conclusions compare torch.compile and CUDA graphs against that same baseline. The BAMBOO machine-learned force field is used only as a test model, and while some authors overlap with the paper, no load-bearing argument relies on BAMBOO's predictions being correct; the same is true for the Tinker9 and GPU4PySCF references, which serve as examples of portability and as an AIMD demonstration. The acknowledged risks around Python object lifetime, reference counting, and GIL/thread safety are important correctness concerns, but they are not circularity: the paper does not assume the safety of the callback in order to prove the callback is safe. No equation in the paper reduces to another by construction, and no fitted input is relabeled as a prediction. The work is self-contained as a performance and compatibility study, so the circularity score is 0.

Assumptions & free parameters 0 free parameters · 3 assumptions · 0 invented entities

The plugin introduces no fitted parameters or new physical entities. It rests on domain assumptions about CPython object lifetime and interpreter safety that are plausible but not fully specified, and on the correctness of PyTorch autograd for the force computation.

assumptions (3)

domain assumption The PyObject pointer captured via id() remains valid and the object is kept alive for the entire simulation duration.
Section 2 states 'ensuring that the object remains resident in memory' but does not describe how reference counting or garbage collection is prevented. If the object is collected, the pointer becomes dangling.
domain assumption The Python interpreter can be safely called from the C++ MD engine without explicit GIL management or thread-affinity controls.
The paper does not discuss the Global Interpreter Lock. If the callback is invoked from a non-main worker thread without acquiring the GIL, the simulation may deadlock or crash.
domain assumption For TorchForce, automatic differentiation through torch.compile yields the correct forces for any model.
The authors test one simple quadratic model and BAMBOO. They rely on PyTorch autograd to compute forces, which is generally sound but is not proven for arbitrary models in this paper.

how reviews work

0 comments

Cite this review

Pith. "Pith review of OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation." pith.science (2026). https://pith.science/paper/6PI6LLXU

@misc{pith2026241218271,
  author       = {Pith},
  title        = {Pith review of: OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/6PI6LLXU}},
  note         = {Machine review of arXiv:2412.18271}
}

read the original abstract

We present OpenMM-Python-Force, a plugin designed to extend OpenMM's functionality by enabling integration of energy and force calculations from external Python programs via a callback mechanism. During molecular dynamics simulations, data exchange can be implemented through torch.Tensor or numpy.ndarray, depending on the specific use case. This enhancement significantly expands OpenMM's capabilities, facilitating seamless integration of accelerated Python modules within molecular dynamics simulations. This approach represents a general solution that can be adapted to other molecular dynamics engines beyond OpenMM. The source code is openly available at https://github.com/bytedance/OpenMM-Python-Force.

Figures

Figures reproduced from arXiv: 2412.18271 by the authors.

**Figure 2.** Evolution of system Hamiltonians over 100 time-steps for different [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗

**Figure 3.** Comparison of energies across deployment strategies: unsigned differ [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗

Figures from the paper (1 more)

**Figure 4.** Figure 4: The unsigned differences in potential energies and root mean square [PITH_FULL_IMAGE:figures/full_fig_p009_4.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

17 extracted references · 13 canonical work pages

[1]

P.; Abreu, C

Eastman, P.; Galvelis, R.; Pel \'a ez, R. P.; Abreu, C. R. A.; Farr, S. E.; Gallicchio, E.; Gorenko, A.; Henry, M. M.; Hu, F.; Huang, J.; Kr \"a mer, A.; Michel, J.; Mitchell, J. A.; Pande, V. S.; Rodrigues, J. P.; Rodriguez-Guerra , J.; Simmonett, A. C.; Singh, S.; Swails, J.; Turner, P.; Wang, Y.; Zhang, I.; Chodera, J. D.; De Fabritiis, G.; Markland, T...

work page 2024
[2]

https://github.com/openmm/openmm-torch

OpenMM Torch . https://github.com/openmm/openmm-torch

work page
[3]

K.; Maher, B.; Pan, Y.; Puhrsch, C.; Reso, M.; Saroufim, M.; Siraichi, M

Ansel, J.; Yang, E.; He, H.; Gimelshein, N.; Jain, A.; Voznesensky, M.; Bao, B.; Bell, P.; Berard, D.; Burovski, E.; Chauhan, G.; Chourdia, A.; Constable, W.; Desmaison, A.; DeVito, Z.; Ellison, E.; Feng, W.; Gong, J.; Gschwind, M.; Hirsh, B.; Huang, S.; Kalambarkar, K.; Kirsch, L.; Lazos, M.; Lezcano, M.; Liang, Y.; Liang, J.; Lu, Y.; Luk, C. K.; Maher, ...

work page 2024
[4]

https://github.com/openmm/NNPOps

NNPOps . https://github.com/openmm/NNPOps

work page
[5]

P.; Simeon, G.; Galvelis, R.; Mirarchi, A.; Eastman, P.; Doerr, S.; Th \"o lke, P.; Markland, T

Pelaez, R. P.; Simeon, G.; Galvelis, R.; Mirarchi, A.; Eastman, P.; Doerr, S.; Th \"o lke, P.; Markland, T. E.; De Fabritiis, G. TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations . Journal of Chemical Theory and Computation 2024, 20, 4076--4087

work page 2024
[6]

Pybind11 -- Seamless Operability between C ++11 and Python

Jakob, W.; Rhinelander, J.; Moldovan, D. Pybind11 -- Seamless Operability between C ++11 and Python . 2017; https://github.com/pybind/pybind11

work page 2017
[7]

https://www.swig.org

SWIG . https://www.swig.org

work page
[8]

BAMBOO : A Predictive and Transferable Machine Learning Force Field Framework for Liquid Electrolyte Development

Gong, S.; Zhang, Y.; Mu, Z.; Pu, Z.; Wang, H.; Yu, Z.; Chen, M.; Zheng, T.; Wang, Z.; Chen, L.; Wu, X.; Shi, S.; Gao, W.; Yan, W.; Xiang, L. BAMBOO : A Predictive and Transferable Machine Learning Force Field Framework for Liquid Electrolyte Development. 2024; https://arxiv.org/abs/2404.07181

arXiv 2024

Show all 17 references

[9]

S.; Bogdanov, N

Sun, Q.; Zhang, X.; Banerjee, S.; Bao, P.; Barbry, M.; Blunt, N. S.; Bogdanov, N. A.; Booth, G. H.; Chen, J.; Cui, Z.-H.; Eriksen, J. J.; Gao, Y.; Guo, S.; Hermann, J.; Hermes, M. R.; Koh, K.; Koval, P.; Lehtola, S.; Li, Z.; Liu, J.; Mardirossian, N.; McClain, J. D.; Motta, M....

2020
[10]

Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework

Wu, X.; Sun, Q.; Pu, Z.; Zheng, T.; Ma, W.; Yan, W.; Yu, X.; Wu, Z.; Huo, M.; Li, X.; Ren, W.; Gong, S.; Zhang, Y.; Gao, W. Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework . 2024; http://arxiv.org/abs/2404.09452

2024 arXiv
[11]

Li, R.; Sun, Q.; Zhang, X.; Chan, G. K.-L. Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework . 2024; https://arxiv.org/abs/2407.09700

2024 arXiv
[12]

J.; Marques, M

Lehtola, S.; Steigemann, C.; Oliveira, M. J.; Marques, M. A. Recent Developments in Libxc --- A Comprehensive Library of Functionals for Density Functional Theory. SoftwareX 2018, 7, 1--5

2018
[13]

A.; Wang, Z.; Lu, C.; Laury, M

Rackers, J. A.; Wang, Z.; Lu, C.; Laury, M. L.; Lagard \`e re, L.; Schnieders, M. J.; Piquemal, J.-P.; Ren, P.; Ponder, J. W. Tinker 8: Software Tools for Molecular Design . Journal of Chemical Theory and Computation 2018, 14, 5273--5289

2018
[14]

P.; Aktulga, H

Thompson, A. P.; Aktulga, H. M.; Berger, R.; Bolintineanu, D. S.; Brown, W. M.; Crozier, P. S.; In 'T Veld, P. J.; Kohlmeyer, A.; Moore, S. G.; Nguyen, T. D.; Shan, R.; Stevens, M. J.; Tranchida, J.; Trott, C.; Plimpton, S. J. LAMMPS - a Flexible Simulation Tool for Particle-B...

2022
[15]

Wang, Z.; Ponder, J. W. Tinker9: Next Generation of Tinker with GPU Support . https://github.com/TinkerTools/tinker9
[16]

https://docs.python.org/3/c-api/index.html

Python/ C API Reference Manual . https://docs.python.org/3/c-api/index.html
[17]

universal MLFF

Pybind11 Documentation: Embedding the Interpreter. https://pybind11.readthedocs.io/en/stable/advanced/embedding.html mcitethebibliography main.tex0000664000000000000000000006116314732465642011250 0ustar rootroot [letterpaper,journal=jctcce,manuscript=article] achemso geometry ...

2024 doi arXiv

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.

[1] [1]

P.; Abreu, C

Eastman, P.; Galvelis, R.; Pel \'a ez, R. P.; Abreu, C. R. A.; Farr, S. E.; Gallicchio, E.; Gorenko, A.; Henry, M. M.; Hu, F.; Huang, J.; Kr \"a mer, A.; Michel, J.; Mitchell, J. A.; Pande, V. S.; Rodrigues, J. P.; Rodriguez-Guerra , J.; Simmonett, A. C.; Singh, S.; Swails, J.; Turner, P.; Wang, Y.; Zhang, I.; Chodera, J. D.; De Fabritiis, G.; Markland, T...

work page 2024

[2] [2]

https://github.com/openmm/openmm-torch

OpenMM Torch . https://github.com/openmm/openmm-torch

work page

[3] [3]

K.; Maher, B.; Pan, Y.; Puhrsch, C.; Reso, M.; Saroufim, M.; Siraichi, M

Ansel, J.; Yang, E.; He, H.; Gimelshein, N.; Jain, A.; Voznesensky, M.; Bao, B.; Bell, P.; Berard, D.; Burovski, E.; Chauhan, G.; Chourdia, A.; Constable, W.; Desmaison, A.; DeVito, Z.; Ellison, E.; Feng, W.; Gong, J.; Gschwind, M.; Hirsh, B.; Huang, S.; Kalambarkar, K.; Kirsch, L.; Lazos, M.; Lezcano, M.; Liang, Y.; Liang, J.; Lu, Y.; Luk, C. K.; Maher, ...

work page 2024

[4] [4]

https://github.com/openmm/NNPOps

NNPOps . https://github.com/openmm/NNPOps

work page

[5] [5]

P.; Simeon, G.; Galvelis, R.; Mirarchi, A.; Eastman, P.; Doerr, S.; Th \"o lke, P.; Markland, T

Pelaez, R. P.; Simeon, G.; Galvelis, R.; Mirarchi, A.; Eastman, P.; Doerr, S.; Th \"o lke, P.; Markland, T. E.; De Fabritiis, G. TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations . Journal of Chemical Theory and Computation 2024, 20, 4076--4087

work page 2024

[6] [6]

Pybind11 -- Seamless Operability between C ++11 and Python

Jakob, W.; Rhinelander, J.; Moldovan, D. Pybind11 -- Seamless Operability between C ++11 and Python . 2017; https://github.com/pybind/pybind11

work page 2017

[7] [7]

https://www.swig.org

SWIG . https://www.swig.org

work page

[8] [8]

BAMBOO : A Predictive and Transferable Machine Learning Force Field Framework for Liquid Electrolyte Development

Gong, S.; Zhang, Y.; Mu, Z.; Pu, Z.; Wang, H.; Yu, Z.; Chen, M.; Zheng, T.; Wang, Z.; Chen, L.; Wu, X.; Shi, S.; Gao, W.; Yan, W.; Xiang, L. BAMBOO : A Predictive and Transferable Machine Learning Force Field Framework for Liquid Electrolyte Development. 2024; https://arxiv.org/abs/2404.07181

arXiv 2024