Proposal: Extending Jax MD with Monte Carlo Capabilities and Bonded Potentials

[stefanhiemer]

Dear Jax MD maintainers,

I am writing to inquire about the possibility of extending the capabilities of Jax MD, specifically by incorporating additional Monte Carlo (MC) methods and bonded potentials into the codebase. I work with molecular systems, glasses and high entropy alloys and think that MC would greatly help my research, this project and the general community.

My plans would be to first introduce standard methods (e. g. Widom insertion), standard MC moves (translation, rotation, etc. pp.) and then move to the more advanced things like various configurational-bias moves and some non-standard ensembles. I need bonded potentials as I want to simulate molecules or glasses with well defined chemical bonds, so I aim to implement the simple standard potentials there too.

I am planning to start working on these extensions towards the end of the year, depending some grant decisions as I am currently transitioning from PhD student to postdoc.

Thank you for your time and consideration. I look forward to hearing your thoughts on this matter.

Best regards
Stefan Hiemer

[abhijeetgangan]

Hi Stefan,

I would also be interested in extending JAX MD to incorporate Monte Carlo methods. Is it possible to reach out to you via email?

Best regards,
Abhijeet Gangan

[stefanhiemer]

Hi Abhijeet,

my contacts are there
https://www.matsim.techfak.uni-erlangen.de/staff/stefan-hiemer.shtml

Regards
Stefan

[cpgoodri]

I don't want to speak for the main developers, but it sounds like you are actually suggesting two separate extensions: 1) more bounded potentials, and 2) a suite of MC functionality. Is there a reason I don't realize that these are actually more intricacy connected?

There are some bond potential functionality, though its not very robust. At some point I wrote a generic 3-body interface, but I don't think it ever got merged with the main branch.

Regarding the MC functionality, you elaborated a bit in the original post about what these functionalities would be, but I'm wondering if you could expand a bit. Or perhaps better is there a MC package out there with built-out and documented functionality that could be used as a reference?

Do you have a sense of how well these MC functions (or at least standard implementations of them) would play with basic JAX transformations like jit and grad?

[stefanhiemer]

Hi there,

regarding your first question:

MC and bonded potentials are connected as many moves used especially for proteins or polymers make explicit use of the bonded-nonbonded separation to (re-)grow molecules/proteins. Just take a look at the figures in this publication and you will probably understand it by intuition:
Chameleon: A generalized, connectivity altering software for tackling properties of realistic polymer systems

What type of MC functionality I would like to add:

I would like to add them in different stages as to see how fast I can implement stuff and what I have to rewrite in Jax as

1. Standard MC moves

   • random translation/rotation,
   • Widom insertion (muVT)
   • box volume change (NPT)
     Parts of this are already implemented in LAMMPS (https://docs.lammps.org/Packages_details.html#pkg-mc), so one can readily compare the implementations with lammps

2. Moves specific to crystals like lattice switch Monte Carlo which we can compare this with monteswitch (https://github.com/tomlunderwood/monteswitch). There need to be added some tricks to more efficiently sample different lattices.

3. Moves under constant tension/stress along the lines of these publications:

R. F. Najafabadi, “Monte carlo simulation of structural and mechanical properties of crystal and bicrystal systems at finite temperature,” Ph.D. dissertation, Massachusetts Institute of Technology, 198

R. Najafabadi and S. Yip, “Observation of finite-temperature bain transformation (fcc→ r bcc) in monte carlo simulation of iron,” Scripta metallurgica, vol. 17, no. 10, pp. 1199–1204, 1983
Unfortunately I am not aware if this specific method has been implemented somewhere.

If you do not want to read those entirely I recommend to read the chapter 5.5 in Daan Frenkel's standard textbook.

4. Ensembles to study quasistatic fracture with a pseudo grandcanonical ensemble where the chemical potential is used to pick off bonds, so this method also relies very much on bonded potential (at least in its vanilla form):

J. P. Nilmeier, G. E. Crooks, D. D. Minh, and J. D. Chodera, “Nonequilibrium candidate monte carlo is an efficient tool for equilibrium simulation,” Proceedings of the National Academy of Sciences, vol. 108, no. 45, pp. E1009–E1018, 2011.

T. Mulla, S. Moeini, K. Ioannidou, R.-M. Pellenq, and F.-J. Ulm, “Phase diagram of brittle fracture in the semi-grand-canonical ensemble,” Physical Review E, vol. 103, no. 1, p. 013003, 2021

I am not aware that any package has implemented this already, but I may be able to contact the authors via common connections.

6. Configurational bias moves for molecule/atom insertion and regrowth. This can be readily compared with Towhee and should enable Jax Md to deal with molecular systems and generate complicated geometries in the first place. A lot of those methods are already implemented in Towhee, but that code is entirely serial and is more or less dead. Marcus Martin still maintains it, but no further development takes place. I want to expand these methods later to other systems like glasses. If you want to know more about this, check either the chapter in Daan Frenkel's book or this essay by Marcus Martin https://towhee.sourceforge.net/algorithm/cbmc.html

We can compare here the implementations with Towhee.

Regarding JIT and automatic differentiation, I see no reason why these should fail in the MC case. The only conflicts with Jax MDs current architecture that I can image is the parallelization that might be suboptimal with MC. But that we have to check which is why I want to implement simple things at first, benchmark to check the performance and then propose changes if necessary. Personally I am not to worried that much I have to admit.

Have I answered your questions? Or do you need more details?

Regards
Stefan

[charles-liang999]

Dear Stefan,

This is a very interesting proposal! Any updates on this "Proposal: Extending Jax MD with Monte Carlo Capabilities and Bonded Potentials"?

Best,
Charles

[stefanhiemer]

Hi Charles,

no progress, but I've got a grant where I will implement some MC use in Jax MD. The prject will start in march and I am currently in the process of moving to another country and finishing stuff up. That's why I have not done any programming yet.

@charles-liang999 Anything specific from my list that you are especially interested? Maybe I can implement it first then and have you then as a tester.

Regards
Stefan

[charles-liang999]

@stefanhiemer Cool! I am interested in standard MC moves, including NVT MC and the ones listed.

Best,
Charles