The first tutorial of OpenMM-Torch with NNPOps

[raimis]

• Write a tutorial
• Add an automated test
• Add a link to README.md

Solves #59

Try on Colab: https://colab.research.google.com/github/raimis/openmm-torch/blob/example/tutorials/openmm-torch-nnpops.ipynb

[raimis]

The first tutorial of OpenMM-Torch is done.

There are two outstanding issue:

• An official NNPOps package is still not available: Conda-forge package of NNPOps NNPOps#26
• Context.setVelocitiesToTemperature is broken (Setting velocities trigger an exception #61) and we need an OpenMM release with Fix GIL for OpenMM::Context::setVelocitiesToTemperature openmm#3424.

Anyway, an imperfect tutorial is better than no tutorial.

[peastman]

This is really nice!

[mikemhenry]

Does this work now that #61 is closed?

[peastman]

It's fixed in the main branch, but the fix hasn't gotten into a release yet.

[jchodera]

Would be great to get feedback from @dominicrufa too!

[jchodera]

This is really great! Just had some minor comments.

[jchodera]

If we use mamba instead of conda here, it can speed up installs significantly (presuming we also used mambaforge above).

@mikemhenry : Are we ready to drop the -c mmh yet?

[mikemhenry]

Not yet, if we can't drop it by the time we want to merge the tutorial, we can make an issue to drop it so I won't forget once we get NNPOps on conda-forge

[jchodera]

What about refactoring the process of removing MM forces and creating/adding the ML potential into a small function folks can copy-and-paste into their own code?

For example, this code could be bundled together:

def replace_mm_with_nnp(system, topology, NNP):
    """
    Create an instance of the specified machine learning potential and replace all MM forces with the ML potential.

    """
    # Remove MM forces
    while system.getNumForces() > 0:
        system.removeForce(0)

    # Assert no forces or constraints remain
    assert system.getNumConstraints() == 0 
    assert system.getNumForces() == 0

    # Get atomic numbers
    atomic_numbers = [ atom.element.atomic_number for atom in topology.atoms() ]

    # Create an instance of the model
    nnp = NNP(atomic_numbers)
   
    # Save the NNP to a file and load it with OpenMM-Torch
    # TODO: Use a temporary file instead
    pt.jit.script(nnp).save('model.pt')
    force = TorchForce('model.pt')

    # Add the NNP to the system
    system.addForce(force)
    assert ala2.system.getNumForces() == 1 

[raimis]

As commented below, this a low-level tutorial. All the convenience functionality should be in OpenMM-ML.

[jchodera]

How does this work for atoms that should not have their potential computed by the NNP? Would we use atomic numbers of 0 for those atoms? For example, if we wanted to only compute ligand energies for a solvated protein:ligand system?

[peastman]

I think that's beyond the scope of this first tutorial. Mixing ML with MM could be a more advanced tutorial. Or we could just refer people to OpenMM-ML, which has code to automate it.

[jchodera]

It seems weird not to even mention how it would work, since this is a major use case. If folks wanted to just simulate only ML, why would they use OpenMM?

I agree a separate tutorial would be useful!

What about periodic ML systems? Do we support those? Or just non-periodic systems?

Maybe we could have three tutorials:

• Vacuum ML
• Periodic ML
• Hybrid ML/MM periodic

[peastman]

Periodic systems just take a minor change to the model. See https://github.com/openmm/openmm-torch/blob/master/README.md.

[raimis]

The intention of this tutorial is to demonstrate how to connect 3 dots: OpenMM-Torch, TorchANI, and NNPOps. All the advanced stuff should be done at a high-level with OpenMM-ML.

[raimis]

No more comments, so I merge the tutorial. We could revise it when there will be more tutorials in place and we will have a better idea.

[jchodera]

@dominicrufa: Can you provide some feedback too?

[mikemhenry]

@raimis nnpops is now on conda-forge 🎉 mamba info nnpops So you can drop the -c -mmh now