Atomic Cluster Expansion Force Fields in OpenMM

[davkovacs]

I am trying to implement the Atomic Cluster Expansion (ACE) force fields into OpenMM using this library. For details of the method see https://doi.org/10.1021/acs.jctc.1c00647

I have the following functions implemented in C:

double energy(char*, double* X, int* Z, double* cell, int* pbc, int Nat);

where X : position of atoms (flattened) Z: Atomic number of atoms cell: simulation cell vectors pbc: whether to use periodic boundary conditions Nat: Total number of atoms

void forces(char*, double *F, double* X, int* Z, double* cell, int* pbc, int Nat);

where the arguments are analogous to the energy function, F is just an empty array that is pre-allocated for the results.

I have implemented Python wrappers for these functions, so I would need to get the positions of atoms, the atomic numbers, cell vectors all as numpy arrays.

I would really appreciate if someone with experience in OpenMM data structures could help me a little.

When I am implementing MLPotential , MLPotentialImpl and MLPotentialImplFactory I don't really see how I can obtain the above data needed and how I should return the energies and forces. I made very little progress so far, you can find it in the following repo: https://github.com/ACEsuit/ACEinterfaces/tree/OpenMM/OpenMM

[raimis]

I think, this repo won't be useful for your case.

You want to implement a plugin for a custom force (http://docs.openmm.org/latest/developerguide/03_writing_plugins.html#creating-new-forces). If your code is running on CPU only, you should look at OpenMM-PLUMED (https://github.com/openmm/openmm-plumed) as an example.

[peastman]

Also see the example plugin, which can make a good starting point for new plugins.

[davkovacs]

Thank you, my code runs on CPU only. This example plugin seems very useful.

@peastman could you perhaps help a little more where in the example plugin I need to make the modification so that it can use the two C functions I mentioned above.

I reckon I need to modify this file: https://github.com/openmm/openmmexampleplugin/blob/master/platforms/reference/src/ReferenceExampleKernels.cpp But I am still a little confused. In my ML model there are no bonds (everything is bonded to everything). Also I see no reference to cells and periodic boundary conditions here.

I have three functions I want to slot in to the appropriate places:

1.) ace_init(str *path) which initializes Julia and reads in the model inside Julia from a file.

2.) energy which evaluates the model given Positions, Atomic Numbers, Cell, Periodic Boundary conditions and Number of atoms.

3.) forces which returns the forces given the same arguments.

Sorry if I am asking too many obvious questions, this is my first time trying to contribute to a C++ code, the C functions we wrote are actually just wrapper to our actual code that is written in Julia. And as you may have relized from my questions, I also do not yet have any experience with OpenMM, but I have collaborators who are very keen to try the ACE models through OpenMM.

[peastman]

If you haven't already, you should start by reading through the developer guide. It describes the whole architecture for plugins. You're correct that ReferenceCalcExampleForceKernel is where the calculations get done. Notice that it has methods called initialize(), which is where any initialization can get done, and execute(), which is where forces and energies get calculated. They're both calculated in the same method, since it usually is more efficient to compute them both at once, but it has arguments called includeForces and includeEnergy that tell you which ones are actually needed. If either of those is false, you're free to skip the corresponding calculation.

ReferenceKernels.cpp implements the kernels for all the standard forces. You can find lots of examples in there for how to do various things. For example, you asked about periodic boundary conditions. That file defines a function

static Vec3* extractBoxVectors(ContextImpl& context) {
    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
    return data->periodicBoxVectors;
}

It takes the context argument that was passed to execute() and returns an array of three Vec3 objects containing the current periodic box vectors.