Introduction to Molly.jl

Joe Greener

MRC Laboratory of Molecular Biology

http://jgreener64.github.io

​

​

​

• Introduce molecular dynamics for biomolecules.

​

• Describe the Julia package Molly.jl and why it exists.

​

• Show examples of custom interactions, simulators and loggers.

​

• Briefly introduce differentiable molecular simulation.

​

• Show some of the features of PyMOL.

Outline

• MSc+PhD, Imperial College London (2013-2017)
• Prediction of allosteric sites on proteins

​

• Postdoc, University College London (2017-2021)
• Protein structure prediction and design

​

• Independent researcher, MRC-LMB (2021-)
• Differentiable molecular simulation

​

• Open source contributions
• BioStructures.jl, Molly.jl, Biopython

My scientific history

• Set up a physical system, define the rules that determine the forces and press play.

​

• Numerical integration of F = ma. Often very computationally expensive, there are few shortcuts.

​

• MD has helped us understand many processes in chemistry and biology.

​

• Ever-improving compute resources, innovative machine learning approaches and better protein structure prediction mean these methods will continue to develop.

Molecular dynamics (MD)

• Some success in folding small, fast-folding proteins and understanding folding pathways.
• Commonly used for mechanistic interpretation in experimental studies. Can be used to inform experiments in a virtuous cycle.
• Can access information hard to acquire by experiments, e.g. atomic-level resolution and allosteric pathways.
• In the last few years binding free energy methods have become popular in pharma companies for improving small molecule binders.

Molecular dynamics successes

From Lindorff-Larsen et al. 2011

• At the CASP protein structure prediction assessment MD has consistently failed to predict the structure of standard proteins.
• Disordered proteins too compact in many force fields, protein aggregation depends more on force field choice than sequence.
• Enzymatic reactions require special treatment (QM-MM) as bonds can’t break in molecular mechanics force fields.
• Not routine for experimentalists the way AlphaFold has become, partly due to expertise and compute required.

Molecular dynamics failures

• The key concern is whether the question of interest can be answered on the time scale available for study.
• Changes in flexibility, and conformation for small biomolecules, can be studied on the scale of ns to μs.
• Comparing multiple simulations often works well, e.g. testing mutations to inform wet lab experiments.
• Enhanced sampling methods can be used to study larger systems and work particularly well when appropriate collective variables are available.

Where does molecular dynamics work well?

AlphaFold database model of TrbA from E. coli

MD useful for working out how this region behaves

Force field problem

Sampling problem

Correct functional form?

Correct parameters?

QM / MM / coarse-grained

Ordered/disordered proteins

Ligand parameters

The problems with molecular dynamics for biomolecules

From Robustelli et al. 2018

Barnase-barstar complex

Dissociation time ~1 day

~10¹⁹ simulation steps!

• Moore’s law is holding in the GPU space so the available simulation length is always increasing.
• Software is continuously developing to make better use of modern GPUs and parallelisation.
• Force fields are improving for disordered proteins, nucleic acids and small molecule binders. MLIPs can achieve quantum-level accuracy.
• Free energy perturbation methods are improving and becoming widely used in pharma.
• AlphaFold has dramatically increased the number of proteins available for simulation.

Recent developments in MD

From Mu et al. 2021

• A lot of work recently has shown that neural networks are able to give quantum-level accuracy at much faster speeds.
• These methods are trained to match forces/energies on quantum datasets but often show issues over longer simulations.
• The field is addressing problems like long-range interactions and worse performance than classical force fields (>10x slower).
• MLIPs are promising, but are not immediately applicable to biological systems which are large and require long simulation times.

Machine learning interaction potentials (MLIPs)

From Musaelian et al. 2023

• There are a number of excellent, very mature software packages for MD.

​

• Most packages have one or more fast kernels to run simulations (C, C++, CUDA, Fortran), and a layer on top to interact with (binaries, Python, config files).

​

• Most packages are hard to interact with all the way down and have a mixed ability to customise.

​

• No mature MD package has support for differentiable simulations.

Software for molecular simulation

• A pure Julia implementation of MD. One language all the way down. Simulation scripts are Julia scripts.

​

• Closest in design to OpenMM, which has a Python API (but multiple kernels under the hood).

​

• User-defined potentials, simulators etc. are easy to define and as fast as built-in features. Everything is defined imperatively in Julia.

​

• Differentiable with Zygote.jl. Gradients match finite differencing for reverse and forward mode AD.

​

• Can simulate standard proteins with the trajectories matching OpenMM.

Foldit1 (PDB ID 6MRR) simulated with Molly.jl in the a99SB-ILDN force field with explicit solvent (not shown).

Molly.jl

• Non-bonded interactions: Lennard-Jones, Coulomb (plus reaction field), gravity, soft sphere, Mie
• Bonded interactions: harmonic bonds and angles, Morse/FENE bonds, cosine angles, periodic torsion angles
• Read in OpenMM and Gromacs force field files and coordinate files using Chemfiles.jl
• Implements AtomsBase.jl interface
• Verlet, velocity Verlet, Störmer-Verlet and flexible Langevin integrators
• Steepest descent energy minimisation
• Andersen, Berendsen and rescaling thermostats
• Monte Carlo barostat

Features

• Periodic and infinite boundary conditions in a cubic/triclinic box
• Flexible loggers to track arbitrary properties throughout simulations
• Cutoff algorithms for non-bonded interactions
• Various neighbour list implementations to speed up calculation of non-bonded forces, including use of CellListMap.jl
• Implicit solvent GBSA methods
• Unitful.jl compatible
• Some analysis functions, e.g. RDF
• Basic visualisation with GLMakie.jl
• Runs on CPU (threaded) or GPU

Features continued

• Constrained bonds and angles (work in progress)
• Particle mesh Ewald summation
• Other temperature/pressure coupling methods
• Full virial calculation
• System preparation - solvent box, add hydrogens etc.
• Domain decomposition algorithms
• Alchemical free energy calculations
• API stability
• High test coverage
• High performance

Missing features

• Easy to add new features
• One language throughout so fairly easy to understand
• Differentiable simulations are fast

​

​

​

​

• Lacks maturity, testing and features of well-established software
• GPU and multithreaded performance not great

​

Advantages

Disadvantages

Can define components individually

Foldit1 (PDB ID 6MRR) simulated with Molly.jl in the a99SB-ILDN force field with explicit solvent (not shown).

Can setup like OpenMM, but runs in one language

• Unitful.jl
• StaticArrays.jl
• Zygote.jl

Demonstration of useful Julia packages

• Pairwise: between all (or most) pairs of atoms, e.g. Mie potential.
• Specific: between specific atoms due to topology, e.g. bonds and angles in water.
• General: all coordinates in, all forces out, e.g. MLIP or DFT.

Three types of interactions in Molly

Custom interactions

​

Custom simulators and couplers

​

Custom loggers

​

Examples

Peptide (shown with VMD)

Training to reproduce a familiar logo

Diatomic gas

Solar system

Making and breaking bonds

Plotting potential energies

• Development currently very active
• High-performance GPU kernels (Jaydev Singh Rao)
• Bond and angle constraints (Ethan Meitz)
• System setup (Diego Ugarte)
• non-CUDA GPU support (James Schloss)
• Differentiable simulations (me)
• Particle-mesh Ewald summation
• API will change some more, sorry
• At some point there will be a software paper

Ongoing and planned work - contributions welcome!

Introduction to PyMOL

• There is currently an advert out for a PhD student to join my group starting Oct 2024, see the MRC LMB website for more.

​

• An advert for a postdoc position will go out next week.

​

• Chat to me if interested.

Jobs

• Molly.jl contributors (incomplete list):
• Noé Blassel
• Sebastian Micluța-Câmpeanu
• Leandro Martínez
• Jaydev Singh Rao
• Ethan Meitz
• Diego Ugarte
• Pranay Venkatesh
• James Schloss
• Ehsan Irani
• Andrés Riedemann
• Maximilian Scheurer

​

• All Julia contributors and package authors, particularly Zygote, CUDA, StaticArrays, Chemfiles, ChainRules, AtomsBase and Unitful
• William Moses, Valentin Churavy and the Enzyme team for all the bug fixes

• OpenMM and Gromacs
• Sjors Scheres group, MRC-LMB
• David Jones group, UCL

Acknowledgements

• Find at https://github.com/noeblassel/SINEQSummerSchool2023/blob/main/notebooks/molly_average.ipynb
• You will also need the two data files (dipeptide_nowater.pdb and dipeptide_equil.pdb) to be in the same directory as the notebook.
• Recommend Julia 1.9 and Molly 0.18.
• To install relevant packages, press `]` from the Julia REPL and run `add IJulia Molly Bio3DView Zygote StaticArrays Unitful KernelDensity Measurements Plots`.
• You may need to press `]` from the Julia REPL and run `build IJulia` to get it to appear as a Jupyter kernel.

Molly exercise

Coding and implementing with Molly.jl

Joe Greener

MRC Laboratory of Molecular Biology

http://jgreener64.github.io

​

​

​

• Show further examples and features of Molly.jl.

​

• Introduce differentiable simulation and its implementation in Molly.

​

• Describe work done to improve an implicit solvent force field using differentiable simulations.

Outline

Barostat example

​

ACEpotentials.jl example

​

Force from AD example

​

Time correlation logger

• Current force summation kernel is simple: each thread calculates the force for one neighbouring pair and atomically adds it to the force array.
• Performance is actually okay, ~5x slower than the equivalent in OpenMM without serious optimisation.
• Enzyme is the only way to differentiate through Julia CUDA kernels like this.
• The next challenge: mature MD packages use clever neighbour ordering and reductions in much more complex kernels.

CUDA kernels

• The chain rule can be used to calculate gradients through a series of mathematical functions.

​

• This also applies to computer programs, provided that the elementary operations have known derivatives.

​

• If you write your program in an AD framework you can get the gradient of some output (loss) with respect to all parameters of the program.

​

• This is how neural networks are trained.

Automatic differentiation (AD)

https://en.wikipedia.org/wiki/Automatic_differentiation

• Neural networks are black box models, making them hard to interpret, and often don’t use the known structure of a system.

​

• Recently the field of deep learning is broadening to include the concept of differentiable programming, where gradients are calculated through arbitrary scientific models.

​

• Due to advances in hardware (GPUs) and software (PyTorch, Tensorflow, Jax, etc.).

Differentiable programming

Differentiable molecular simulation

Automatic differentiation over thousands of steps

More in “Automatic differentiation in machine learning: a survey”, Baydin et al. arXiv 2015

and “Neural Ordinary Differential Equations”, Chen et al. NeurIPS 2018

​

The variety of possible loss functions

Radius of gyration

Radial distribution function

Protein-ligand binding

Protein-protein

interaction

Phase change

Flexibility

Fit to experimental data

Supramolecular geometry

• Some packages do exist for differentiable simulations:
• Jax MD (in Jax)
• TorchMD (in PyTorch)
• DMFF (in Jax)

​

• These are all promising but are limited in some way by scope or performance.

​

• Molly can do differentiable simulations, including on proteins, and is being actively used for research in this area.

Software for differentiable simulations

• Zygote.jl and Enzyme.jl are used for AD. Gradients match finite differencing for reverse and forward mode AD.
• Zygote.jl is okay for general-purpose programming, but is limited by the lack support for mutation. This can lead to poor performance and memory usage.
• Enzyme.jl does AD at the LLVM level and supports mutation. Works well with “static-like” Julia code, support for the rest of Julia is growing.
• ChainRules.jl is useful for adding fast derivatives and working round bugs.
• Can we switch to using only Enzyme.jl long-term?

Foldit1 (PDB ID 6MRR) simulated with Molly.jl in the a99SB-ILDN force field with explicit solvent (not shown).

Differentiability in Molly.jl

• Run DMS on small proteins in Molly:
• Alanine dipeptide in water (2,917 atoms):�~25 ms per step with gradient on GPU,�~14 hours for 1 ns.
• Trp-cage with GBSA implicit solvent (284 atoms):�~12 ms per step with gradient on GPU,�~17 hours for 5 ns.

​

• Achievable to improve all-atom implicit solvent force fields with this iteration of the software.

​

• With further optimisation explicit solvent force fields will be in reach for improvement.

Differentiable simulation in Molly.jl

Sample gradients after 1 ns (10⁶ steps), loss is RMSD to starting structure

atom_N_σ

0.01064

​

atom_H_mass 0.0007533

​

inter_CO_coulomb_const 7.475e-7

​

inter_LJ_weight_14 0.001836

​

inter_PT_C/N/CT/C_k_1 -2.017e-5

• During training, the need to store values for reverse-mode automatic differentiation means that the memory required scales with the number of steps.

​

• This is not a problem when using the learned potential for simulations.

​

• Potential solutions:
• Gradient checkpointing.
• Invertible simulations.

GPU memory requirements

From https://github.com/cybertronai/gradient-checkpointing

From https://en.wikipedia.org/wiki/Automatic_differentiation

• Automatic differentiation gives exact gradients but with respect to the numerical integration.

​

• Some functional forms of force fields, e.g. hard sphere interactions, will give exploding gradients when used with standard integrators.

​

• Gradient clipping and using a stochastic thermostat seem to prevent gradients from exploding or vanishing.

​

• Fortunately there is lots of prior work on stabilising gradients through deep RNNs.

Exploding gradients

From Ingraham et al. 2019

• Long-range electrostatics with particle-mesh Ewald: difficult to implement, let alone differentiably. Currently using reaction field.

​

• Bond and angle constraints. A smaller time step should be used if not constraining bonds.

​

• Stochastic simulations, e.g. using certain thermostats during training or Langevin dynamics.

​

• Neighbour lists: not required to be differentiable since output is binary.

Algorithmic challenges

Differentiable examples

• Most of the atoms in a biomolecular system are solvent, but we are largely interested in the solute. Implicit solvent replaces solvent molecules with a continuous medium.

​

• Generalised Born models are a commonly used approximation with intra-solute screening determined by spheres on each atom.

​

• Regularly used across a range of biomolecular simulations.

Implicit solvent

Advantages

• Faster due to fewer atoms.
• Much faster exploration of conformational space due to lack of friction (can be tuned).
• Can fold small proteins in GPU-days with enhanced sampling.

Disadvantages

• Slower due to implicit solvent calculation.
• Slower due to larger non-bonded distance cutoff (no PME).
• Fundamentally limited due to lack of interaction with individual solvent molecules.
• Overcompacts disordered proteins.

Implicit solvent

• Improving implicit solvent models for disordered proteins and protein aggregation would be especially beneficial.

​

• Implicit solvent parameters are usually tuned separate to the rest of the force field, in contrast to recent advances in force fields using explicit solvent.

​

• Aim to improve radius of gyration and secondary structure content of IDPs while not compromising effectiveness on folded proteins.

​

• Fine tune from an existing force field, a99SB-disp and GB-Neck2.

Learning a better implicit solvent force field

Parameters changed

​

Common atom types only

​

Lennard-Jones σ/ε: 32

​

LJ/Coulomb 1-4 weighting: 2

​

Proper torsion ks: 36

​

GB-Neck2 element-dependent and general parameters: 27

​

Total: 97

• Training set is 11 small peptides/proteins of varying structure and disorder content.
• Loss value is KL divergence to residue-residue distance distributions from 2 μs explicit solvent simulations.
• Each epoch of training takes 1-2 days on a GPU cluster:
• Run 5 ns simulations for each protein to get gradients.
• Combine gradients and update force field.

Training set

• The trained force field, GB99dms, shows a better match to experiment for the radius of gyration and secondary structure of disordered proteins.

​

• Performance remains okay on folded proteins.

​

• See the paper at bioRxiv: “Differentiable simulation to develop molecular dynamics force fields for disordered proteins”.

Learning a better implicit solvent force field

• The next step is to apply this learning technique to improve all-atom force fields. There has been a lot of work in improving parameters to describe disordered proteins and aggregation.

DMS for all-atom force fields

From Mu et al. 2021

• There is currently an advert out for a PhD student to join my group starting Oct 2024, see the MRC LMB website for more.

​

• An advert for a postdoc position will go out next week.

​

• Chat to me if interested.

Jobs

• Molly.jl contributors (incomplete list):
• Noé Blassel
• Sebastian Micluța-Câmpeanu
• Leandro Martínez
• Jaydev Singh Rao
• Ethan Meitz
• Diego Ugarte
• Pranay Venkatesh
• James Schloss
• Ehsan Irani
• Andrés Riedemann
• Maximilian Scheurer

​

• All Julia contributors and package authors, particularly Zygote, CUDA, StaticArrays, Chemfiles, ChainRules, AtomsBase and Unitful
• William Moses, Valentin Churavy and the Enzyme team for all the bug fixes

• OpenMM and Gromacs
• Sjors Scheres group, MRC-LMB
• David Jones group, UCL

Acknowledgements