Most straightforward way to distil model for prototyping

[JSLJ23]

Hi Nequip developers,

I was hoping to get some help on building the model in the most straightforward and distilled way whereby I can prototype training and inference in a protein ligand molecular dynamics simulation with OpenMM's torchforce plugin. What would be the most straightforward way to define the model which can take in atomic_numbers or atomic_symbols and output energies and forces?

One good example is how torchANI supports such an interface whereby the model can be simply imported as:

import torchani

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torchani.models.ANI2x(periodic_table_index=True).to(device)

coordinates = torch.tensor([[[0.03192167, 0.00638559, 0.01301679],
                             [-0.83140486, 0.39370209, -0.26395324],
                             [-0.66518241, -0.84461308, 0.20759389],
                             [0.45554739, 0.54289633, 0.81170881],
                             [0.66091919, -0.16799635, -0.91037834]]],
                           requires_grad=True, device=device)

species = torch.tensor([[6, 1, 1, 1, 1]], device=device)

energy = model((species, coordinates)).energies
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative

Does Nequip support such a similar API or what would the most friction-less way to go about achieving something like this?

Thank you.

[Linux-cpp-lisp]

There exists direct support for NequIP / Allegro / nequip models with OpenMM, please see #288.

In general, the simple way to use a trained nequip framework model from Python is either to use the ASE calculator (see various GitHub discussions/issues) or to call it directly with the AtomicDataDict (really just Dict[str, torch.Tensor]) containing your input structure in the relevant predefined keys.