Fitting a SMIRNOFF Force Field with PyTorch - OpenFF 2026 Virtual Workshop

Hi everyone. My name is Dr. Ashley Mitchell. I'm Open Force Fields uh science communicator. Um if you've seen a few of these before, you might have known me by a different name. Uh but it's still me. Um today we're going to talk about how to fit a force field with the uh new fitting stack that um Open FF is currently uh trying to implement. So our previous fitting stack uses a piece of software called force balance which is bespoke optimization software for molecular dynamics force fields and we're we're transitioning towards using PyTorch. So, PyTorch is a uh more generic generalpurpose machine learning uh package. Uh and the SMEI and Descent packages implement PyTorch optimization for uh for for molecular dynamics force fields. Uh so we're just going to demonstrate a really basic workflow using these new tools to fit and validate a Schmoff force field. So the tools we'll be using today are PyTorch which are just Python bindings to the torch C++ machine learning library SME and descent which is a differentiable computation platform force fields via PyTorch. So PyTorch provides this really nice abstraction called a tensor. You can think of it as kind of like a numpy array except instead of only running on your CPU, you can also put it on your GPU or on some sort of heterogeneous compute. And then uh the tensor also keeps track of all the operations you've performed on it so that a gradient can be automatically computed. And this is really the magic of everything that we're showing here. We'll be able to write an arbitrary loss function and then compute the gradient automatically. So we can write anything we want in our loss function. We don't have to do any maths ourselves. The computer will figure out how to modify the force field to improve it by uh relative to that loss function. Descent uh provides fitting machinery around SMEI which is the actual differentiable computation engine which basically just implements molecular dynamics force fields uh through tenses instead of like by hand through C++ and then for this demonstration I'll get some initial structures of QC archive but you can get initials uh you can get your fitting data from anywhere you And we'll use the smoff force field format which is our native force field format for uh producing force fields that don't rely on atom types. So instead of using atom types we directly infer chemistry using smirks patterns and smarts patterns. This notebook is accessible online. You can go to this GitHub URL and I'll put it in the description as well or you can follow this QR code. uh and there's instructions there on how to install all all of these software and run this notebook. So the plan today is that I've gathered some train and test data sets from QC archive. Uh if you want to see how I've done that uh it's all on the GitHub force uh GitHub repository. Uh so I won't show it in the in the video. Then we're going to modify an old schernnoff force field so that it can be fitted efficiently using the new stack. We're going to benchmark that force field against some alkanes and then demonstrate how to transform the data set that we gathered from QC archive into a format that can be ingested by the PyTorch SME descent stack. Then we'll configure and perform an optimization of the force field and evaluate that final force field against the original benchmarks. So in this section here that I've hidden, we're just gathering data from QCRO. So to do a force field fitting, you will need to get data from somewhere. This is an alternate. This is not an alternative to bespoke fit. Um uh you can get that that data from anywhere. Um as long as you can get it into the format we'll demonstrate in a sec. Uh and you can set up uh benchmarking however you like. These aren't really a part of the flow. So uh we're not demonstrating them. If you are looking for an alternative to to bespoke fit uh well bespoke fit is still available. Uh we've moved it into application mode rather than library mode. So we're not going to be updating its uh dependencies uh or its requirements. Uh but but the software is still available and you can go uh install it in the usual way. and the old environments will continue to work. Uh, and there's a new uh automated force field fitting regime being produced by the coal lab called Presto. Uh, they're collaborators of ours. Uh, their their lab worked substantially on bespoke fit and that's sort of the spiritual successor to bespoke pit uses a completely different um bidding methodology. a lot more AI and machine learning involved. Uh but it seems to work so pretty well so far. So have a Google that. Okay. So to get started, we're going to choose an initial force field. And just so we can demonstrate a really big improvement, we're going to be using the SMOF9 Frost version 1.1.0 force field. So this is one of the firstoff uh force fields that we produced uh or that was produced even before open force field was a thing. Um and and this is what all of our force fields are based on sort of uh the ancestor. So uh we're going to make two transformations to this just to make it easier to train. The first is that we're teaching we're ditching its toolkit AM1BC parameter handler and instead computing charges with Naggle. So Naggle is our uh graph neural network for molecules um library uh and we're using a neural network designed to compute charges uh from uh from molecules. So this uh this neural network is called ash GC which is a um a neural network that's designed to reproduce the AM1 BCC charges that we would have produced um previously. Um Ashg provides a bunch of advantages over computing AM1 BCC charges directly. First it's way faster. Uh so ash GC is actually fast enough to compute charges for an entire protein in uh in a matter of seconds less than a minute. Uh whereas for AM1 BCC that was completely intractable. It would have taken years and it would have used all the memory in the universe. Uh and secondly, ashgc is completely unaware of uh the uh the coordinates that you have attached to your molecule. So we can sort of guarantee that it's confirmer independent because it's not aware of your confidence. So uh so that's that first transformation. Instead of using a semi-impirical QM method to calculate charges, we calculate them with a message passing rap network. Uh secondly, we're going to create non-constrained variant. So non-constrained variant was never published with 99 frost, but they're really important for force field fitting uh because constraints are not differentiable. So uh to create an unconstrained variant, we simply dregister the constraints uh parameter handler and then save that. If you wanted to start with say Sage 2.3, you wouldn't have to do either of these transformations because it already works with uh Nagalashgc and it already has an unconstrained version published. So uh that would be much simpler. Okay. So next we're going to do an initial benchmark against a small group of of alkanes just to see uh whether there's room for improvement of this force field. Uh we're using some software called Yams. So Yams is our internal benchmarking tool that we're developing at the moment. Uh you're welcome to use it but it's not a priority of ours to make it super accessible to every user. it's more important to us that we can rapidly modify it. Um, so yeah, use it at your own risk, I guess. So, yams lets us just modify uh optimize uh a set of uh molecules from a benchmark uh against a force field that we pass it uh and then we can plot metrics from that. So here we have um delta delta E calculation by count. So uh this is the difference in the computed difference of energies between the uh quantum calculation which was the input data for the benchmark and the molecular mechanics calculation which was uh the parameters. And so what we're looking for here is a very sharp peak above zero. And we can see that this is sort of spread out on both sides and it has some uh more distant peaks as well. Uh this is just a simple cumulative density function of RMSD. So this is the root mean squared deviation between the quantum chemical minimum structures and the molecular mechanics minimum structures. uh a perfect force field that perfectly reproduced the structures of the quantum mechanics uh simulation would have a perfectly vertical line where the entire cumulative distribution function is focused at uh RSD of zero and then this is a torsion fingerprint deviation you still using a CDF so same idea but instead this is looking at the torsions and so we can see that there is some room for improvement uh on the torsions and again a perfect force field would look like a vertical line here. Um the more leftwards uh the the line moves the better. So uh those were one set of lipid benchmarks. We can look at a uh so so those were optimizations and these are actual torsions. And so we're now looking at this torsion here on this branched alkan. And we can see that our initial force field is repitulating the shape reasonably well of the quantum chemical baseline which is this black. And we can see the blue is the force field. But there's definitely some room for improvement. So uh this torsion is really what we're going to be focusing on in seeing what we improve. Okay. So we we know our target. We've got a force field that we want to improve. We've got some training data that we want to improve it against. And we know what we're looking for uh to to evaluate an improved force field. So the next thing to do is to get that uh data set into a format that SME and descent understand. So uh for that we're going to use the hugging face data set format. So hugging face uh the hugging face format is an open source format that's really designed for uh training on massive data. So it's really suit well suited here. The really big advantage is that can it can be backed by file rather than memory. So you can use this really painlessly for larger than memory uh data sets. We're not going to use it for a larger than memory data set today. Uh but having it backed by a file is the only way to do it. And it also gives us some safety if if everything crashes unexpectedly at some point or the container that you're running this calculation uh gets gets uh preempted or something then uh then you just have to save a file rather than the contents of memory. So uh the hugging face data set format is a lot like p a pandas data frame. So it's basically a dictionary that goes from strings which are column titles to lists. Each element in the list is an entry at that uh in that table column and then the rows are the elements of the list that all have the same index. Uh so there's four columns in a SME data set. There's smiles which is a mapped explicit hydrogen smiles string identifying the molecule. Uh so that not that's a very dense string format that not only uh uniquely defines the actual molecule, it also defines the ordering of the atoms. And so then the uh the numbers that go into the coordinates and forces uh tenses uh can um can use that atom ordering so that uh there's no ambiguity of where all the atoms go. The energies is uh a list of uh energies uh each row being a different molecule. Um and the the the each element of the energies list has a number of energies equal to the number of confidence that we have data for for the molecule in that row. then uh the coordinates of the of each single point calculation. But for these calculations, we'll we'll use the optimized single points for optimizations and the end of the uh optimization for uh torsion drives. Uh and and so each row will have a shape of the number of compliments. So that same number of energies and then the number of atoms and then there's three uh an xyz coordinate for each atom and the forces are similar but they they give you forces on each atom um at the energy minimum or wherever you are in the uh in the calculation for that confirmer um uh in in the same format. So uh n by3 numbers for each conf. So I've stored the QC archive data in pickle files. They're available at the GitHub repository. And so we just need to load them in and then formulate the data set. And so what this looks like is just going through and mapping each of the columns we've just talked about to whatever arcane identifier QC archive uses to uh represent that same information. And we can see that involves uh gener sometimes generating new smiles that we have to make sure we've got the orders of atoms right for um it involves some changing uh some coordinate changing. Uh we were only provided with the SCF gradient. So we have to multiply that by minus one to get us into forces which we do down here. And so, uh, we basically just build up this, uh, dictionary of of lists, uh, one by one as we go through each atom. And then we iterate over them to collapse all of the, um, the molecules that are then represented on multiple rows into a single row. That looks something like this. And then we create a data set from that. And for a data set of uh 2,105 confirm that takes about a second. Okay. So we now have our force field that we want to start with. We've benchmarked it and we've prepared our data set in a format that we that we like. So the next step is to parameterize. So that means going through all those smile strings and and mapping them to fully parameterized uh systems that are ready to compute a molecular mechanics energy from. So uh so this is this is very familiar if you've used force field before. uh you just convert each molecule into a topology uh by iterating over all the smiles and then in parallel we compute interchanges from them from the force field that we're using. So once we have interchanges, the interchange uses the force field to parameterize the molecule and then we have a map from okay this parameter in the molecule whether that be a torsion or an angle or a bond or whatever has this value and those values are tied back to the original force field through the interchange. And so uh these interchanges can then be converted to tenses. And so this is where SMEI starts to come into the picture. So uh because from now on we're basically going to be using uh PyTorch tensors, we need to define a device. So you could imagine changing this to CUDA if you wanted to compute on a on a CUDA GPU instead of a CPU. Um I'm demonstrating with a CPU because everyone has one. My machine's actually AMD. So, uh, getting it working in a way that would be, uh, usable by everyone else is a bit of a pain, but you're welcome to investigate what that looks like for you. Uh, and then we call this SME convert interchange for uh, function to convert them all into tensor topologies and tensor force fields. So a tensor force field is a tensor, one of these uh magical um uh numpy like buffed up numpy arrays and a tensor topology maps a parameter that needs a force field value to the actual force field in the tensor. And so the reason this exists is so that uh we can just optimize the tensor force field and we still retain uh our knowledge about which of these parameters is tied to each of the molecules that we're doing our optimization over. We just do a quick sanity check and then we create a uh dictionary from each smile string to the tensor topology. Then we make sure that all the tenses are on the device we want and we're ready to fit. So first we need to decide what we're going to fit. Uh and that looks like these parameter config options. So what we're saying here is that we want to fit the bonds angles and proper torsions. All the other parameters in in the force field will not be fit. We're not fitting any attributes. So here parameters and attributes are taking on slightly unusual terminology. They're terminology from the XML format that SMROF is built around. Um so occasionally you'll want to optimize something that's actually an represented as an attribute in that force field rather than as an uh in that XML format rather than as a parameter. And they are treated differently by descent. Uh this is quite rare. Uh but you might need it for um Leonard Jones things and stuff like that. Uh if you wanted to optimize the cutoff um or something like that. Pretty pretty obscure things in force field fitting but uh theoretically could be useful. And so for each parameter we want to fit uh we decide which columns within that parameter we want to fit. So uh for bonds bonds are represented by a hookian term. So there's obviously two parameters in a hooks in hooks law. There's the force constant of the bond and then the uh equilibrium length of the bond. And so we say we want to fit both the force constant and the length and we define their relative scale. Uh so most optimizers work better if every number they're trying to optimize is about the same size. So we just scale everything so they're the same size. Uh and then we set limits. So it doesn't make sense for the force constant or the length to be less than zero. So we set limits of uh the lower bound is zero and the upper bound is none. Angles look very similar because they're also hooking. And then the proper torsions look a bit different because we're only looking at scaling the force constant. Um uh so so we could also uh optimize the phase and the um potentially the number of of terms but that's uh that's an integer which makes it a bit more difficult. Uh so that's the magic. Uh we we decide which parameters we want to fill. And so if you're optimizing your force field yourself, you might just want to stick to the proper torsions. You might want to bring in other parameters that we haven't optimized here. Uh you might want to optimize the phase along with everything else. Um but this is where you configure that. And then once we have those parameters and attributes and the tensor tops by smiles or actually the tensor force field, we can produce this trainable object. And this trainable object is just a wrapper around the tensor force field that allows us to get out a tensor that only includes the scaled parameters that we actually want to modify. Uh okay. And then we're going to perform the actual optimization. So this is a huge amount of code. Well, it's not a huge amount of code. It's like less than 100 lines. It's but it's conceptually quite dense because this code forms the entire optimization. So we first set up some configuration. Uh we set the number of epochs. So an epoch is a basically this is going to be a loop of things that we do and an epoch is a single iteration of that loop. though. Um whatever your uh whatever process you you perform during your optimization um the number of times you do that entire process is is the number of epochs. Um our loss function that we're defining uh it can be parameterized by the learning and force weights. Uh the loss function is completely arbitrary. You can choose whatever you want. This is just a demonstration that um we could have some of that configuration outside the the loop. Uh then we set a learning rate. This hasn't been optimized. Might be completely wrong. Uh but this is just how uh how quickly the optimizer pushes the force field uh down the the ramp basically. Uh so you can get better bidding in fewer epochs by optimizing this. It doesn't like making it too big or too small can both be bad. Um this the script I'm going to show you supports uh processing data sets that are larger than memory by doing it in batches. uh because we're using the uh hugging face data format, we can use this train this data loader class that lets us do that all automatically and we can just define the size of the batches. So you'd want to tune this so that each batch fits in memory but that it's as large as possible so that uh uh you get the advantages of um of cing and stuff like that. So we're just going to set it to the length of the entire data set so that there's only one batch so that we don't have to deal with that complexity. Uh we set a directory uh to to write some logs out to. So we're using PyTorch which is uh an open-source machine learning algorithm. So its competitor from Meta is called TensorFlow. And then Meta also puts out a uh piece of software to monitor TensorFlow optimizations called Tensor Board. And because Torch is open source, there's a a compatibility layer called TensorBoard X. And so we're just defining the log directory that we're going to write our TensorBoard uh files out to so that we can read it. And then we just say okay what do we want to actually train? We want to train the data set that we produced earlier. Uh okay and then we we set up the data loader. So the data loader we're just doing things like uh maybe we need to microoptimize the memory. So we'll pin that. uh we need a function that uh that collates data from the data set, but that's not super important because we're using hugging face. So we're just using a an identity function. Uh and we set the batch size here as well. And so then we can iterate over the batch the train over the data loader to get our batches out. Then we take our trainable which remember is that wrapper around the tensor force field and the configuration of the optimizer that lets us pull out scaled uh values that we want to optimize and we convert it to those values and then put the put it on the device that we want to run our calculations on. So for this calculation will be CPU but maybe for yours it's GPU or CUDA or something else. And then finally we set up the optimizer. So, we're using Atom, which is like the optimizer everyone says to use. It's a general purpose adaptive learning rate optimizer that you can throw anything at it and it'll do okay on. Uh so, uh that that's a that's an easy choice. Uh so, all we need to give Adam once we've set this stuff up is the gradient uh with respect to the loss function and the values of the loss function. and it will and the current values of the parameters and it will push those parameters towards uh minimizing the loss function. Very cool little bit of software this this optimizer. Okay. So then we set up tensorport x we set up a progress monitor and then we start iterating through our epochs. So first thing we do is we reinitialize all our intermediate values so that uh so that we don't uh carry over any uh information from the previous epoch by accident. So we zero out our loss and our gradient and we uh get new force a new force field that we can train. Then we uh uh compute our loss function in batches. So this is where the real magic happens. So first we just go through and get the batch that we want to compute and put it on the device we want. So this looks like uh put put the values on the device uh for each sample within the batch. We're iterating over the data loader. So that'll be the batch. Uh so then we get the length of the batch because the batches are not necessarily all the same size. The last one might be smaller for instance. And then we predict the energies and forces uh for that batch. So this is the first bit of magic is descent as this autoiffiable molecular mechanics energy computation method. And so it takes in a batch of uh of of data which uh uh which gives us uh all the all the data we want to train against takes in a force field and those topologies that we created earlier. So, so we now have uh what we want to train the force field against where the force field is currently up to in its journey and uh a mapping from the molecules in the batch to uh to to definitions for which parts of the force field it should use. Then we just have uh a mean as a as an additional parameter that you probably don't need to change. uh and then it spits out both the reference energies and the computed energy and the reference force and the predicted force from all of that. So these reference energies come from the batch come from the data set that we're training against and these predicted forces and energies come from uh the topologies and the force fields and remember these predicted force fields have been computed completely uh autoiffiably. So uh as long as we use them to compute the loss function, we will retain the dependence that they have on the force field, which is where the magic happens. So that's the first bit of magic. Here's the second bit of magic. This doesn't look like magic because it's literally just a very boring um Python math like a couple of lines, right? So this works exactly the same way as a numpy array. These uh tenses are getting broadcast against each other and then um squared individually and then we sum them and divide them by the true batch size. We do that the same for the force and this computes an L2 loss for both the energy and the force. And remember because these are tenses this is still being differentiated. So we still know what these values uh what their different what their derivative is with respect to this original force field. And then we combine them with those weights that we used earlier into a total loss function. And then uh the the last bit of absolute magic is that we auto differentiate them with autograd. And so autograd takes in that loss function that we computed and then the trainable parameters the parameters that we wanted to train against and it gives us a gradient. So this gradient is another tensor that tells uh tells the optimizer which direction to push the force field in in that very high dimensional space uh to improve the force field with respect to that batch that that loss function. So you can write like any math function you can imagine in here as long as it depends on this force field through differentiable PyTorch functions which is basically all of the PyTorch maths functions. This line of code will produce a gradient that tells you how those that loss function differs with the force field. So you can put regularization in here, you can come up with some obscure scheme for incorporating I don't know anything you want. Um if you can express it as a number that gets smaller as the force field gets better, you put that number into this batch loss uh variable and you'll get a derivative out of it. So okay. So then um that's the gradient for the for the current batch. The gradient for all batches combined is just the sum of all the gradients for the batches. So uh the first time we go through a batch we set the the the epoch gradient as opposed to the batch gradient to the batch gradient. And then each other time we just add the batch gradient to the epoch gradient. Uh we do this in this weird obscure way instead of like collecting all of the batch gradients and then summing them at the end because we just want to keep the memory as free as possible. Uh and this detach function um tells it that we no longer need to track where that this derivative comes from. So again, we're never going to differentiate the gradient with respect to the force field, right? We've already computed the gradient. So we detach it so that it it knows not to keep a track of that. Um which obviously gets more and more complicated for it as it incorporates more and more information and by this stage it's incorporated quite a lot of information. Okay. So then uh we uh we compute uh the the total losses for the epoch both in terms of energy and force so that we can write it out to tensorboard. uh we're we're doing detaches there as well because we don't need to keep track of the gradient anymore. And then we write those results out to TensorBoard. And finally, we perform the actual optimization step. So here we're saying the gradient of the trainable parameters is that gradient we just computed. Optimizer take a step. You already know where to find the gradient because we gave you the trainable parameters earlier. and then uh optimizer zero out your gradient so that in the next day beat everything's clean. Okay, you can see I ran this for about an hour uh before recording this uh and we got a nice progress bar uh for all 5,000 epochs uh and our loss function is down to 20. We can actually see if that's any any improvement by looking at uh tensorboard which we can do in the notebook. You probably wouldn't necessarily do this uh for a longer running uh fit but but you can. And so from here we can see the total loss function has uh improved over the whole course of the fit. So this fit could probably go for longer and we'd still see improvement. If we break that down into energy and forces, we see that the energy actually worsens for the first couple of frames of the fit. And that's because the forces are dragging it down. And then uh once the energy has sort of reached over this pump and gotten into a new regime, uh both the forces and the energy improve together. And that looks like in aggregate a a smooth improvement of the overall loss function even though initially there was a worsening of the energy. And so you can imagine that these dynamics might be different if we modified the weights of the losses of of the different contributors to the loss function. Okay, so we've now done our optimization. We just need to get uh our force field back out. So uh the trainable still holds uh the original force field. So we just give it the optimized parameters and it will give us a tensor force field and then we just uh convert that tensor force field to shownoff force field. Uh this write shernoff uh function may eventually make its way into uh a published release of something but it's not there yet. Uh so uh it's it's in the uh GitHub repository, so you'll be able to use it if you download this notebook. Um uh but yeah, might have some bugs. We're not sure. Uh it's probably okay. Then uh we can look at this line by line. So here is each modified line next to its original line. You can see these are formatted a little bit differently, but basically we're looking at uh the lengths, you know, changed a little bit. The force constants have changed a little bit more. Units are still the same. So, this looks like a force field that's gone through an optimization. The numbers are all different. Um, we can see that the parameters that have actually changed are parameters that are relevant to alkanes, right? All of these are talking about carbons and hydrogens. Um, and we've got bonds, angles, and proper torsions, which is what we asked for it to optimize in the proper torsions. Periodicities haven't changed, phases haven't changed, just the force constants have changed. That's all what we asked for. So, it looks like we got what we wanted. So, let's benchmark it and see if we've improved things or if the loss function was lying to us. So I'm passing in both the original force field and the optimized one. Um we've done all of this with the unconstrained force field because optimizations work much better if force fields are differentiable. Uh and we use the same mole store and torsion store that we had before. Uh give it uh with the same molecules and we give it the new force field. So we can see that uh this orange um series is the optimized force field and it's moved this blue block much closer to the rest of it. Apart from that hasn't changed much. Uh here the orange is generally more to the left. So we can infer that this is a generally improved force field even though we might have some worsening of some molecules around here. Uh and then the torsions have not improved. That's interesting. Uh so there's a little bit of um stochasticity to these fits. Uh when I've run this previously, it has improved. Uh but I guess not on this one. So let's look at the actual torsion we were interested in. We see that this one torsion has improved dramatically. So uh it's now almost perfectly reproducing the the torsion profile. And that's probably because we were cheating. uh we uh have we've really overfitted this force field to alkanes, right? So uh uh this this torsion looks really good, but perhaps some other force fields uh some other molecules might perform much worse. Um but that's that's a win for us. We wanted to uh improve our force field on this one torsion and we've accomplished that. And you can imagine that the ways that we'd make it more resistant to this sort of overfitting would be add a lot more data uh think about regularization uh um uh and and fit against generally just uh just more data and uh and and use a a a really good force field as a starting point rather than the oldest one we can possibly think of. Okay, so thanks for sitting through that. What we've done is we've grabbed some data off QC archive behind the scenes and then we've processed that into a format that our new PyTorchbased fitting stack can understand. We've uh uh prepared a force field so that it can be really efficiently applied to all of our molecules and then we've applied it to all of our molecules. uh and then we've configured and uh and optimized the the force field with respect to those molecules and we've seen that our benchmarks improve. Um so all this code is freely available. Um again I'll go back up to the top and you can uh look at the description or go to these links uh uh at for open force field we we do uh open science forever. So these are all licensed under MIT very permissive licenses. Uh so uh please have a go um and please feel free to raise issues either at the repository for this notebook or at any of our other um GitHub repositories if there's any dramas any anything that uh anything that you think could be improved uh or anything that went unexpectedly wrong. We'd love to hear from you. Um, so I'm Ashley Mitchell and thanks for watching.