Graph transformer neural network force field for prediction of atomic forces and energies in molecular dynamic simulations

1. A computational method for simulating motion of elements within a multi-element system using a graph transformer neural network (GTFF), the method comprising:

converting a molecular dynamics snapshot of the elements within the multi-element system into a graph with atoms as nodes of the graph;

defining a matrix such that each column of the matrix represents a node in the graph;

defining a distance matrix according to a set of relative positions of each of the atoms;

iterating through the GTFF using an attention mechanism, operating on the matrix and augmented by incorporating the distance matrix, to pass hidden state from a current layer of the GTFF to a next layer of the GTFF;

performing a combination over the columns of the matrix to produce a scalar molecular energy;

making a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to compute a prediction of force acting on each atom; and

returning the prediction of the force acting on each atom.

2. The computational method of claim 1 , wherein the combination over the columns of the matrix is implemented as a sum.

3. The computational method of claim 1 , wherein the combination over the columns of the matrix is implemented as a mean.

4. The computational method of claim 1 , wherein the combination over the columns of the matrix is implemented using a neural network.

5. The computational method of claim 1 , further comprising including, in the matrix as nodes of the graph, other molecular objects in addition to the atoms, other molecular objects including one or more of bonds or quantum state of the atoms.

6. The computational method of claim 5 , further comprising defining distances of the other molecular objects to other objects in the graph such that the distance matrix remains well-defined for each element.

7. The computational method of claim 1 , further comprising training the GTFF as a deep neural network using a loss function formed to compare the force predictions to ground truth forces on a labeled dataset, such that network weights are optimized with respect to the loss function using gradient descent.

8. The computational method of claim 1 , wherein the computational method is integrated into molecular dynamics (MD) software.

9. The computational method of claim 1 , further comprising utilizing softmax in the attention mechanism for forward propagation of the hidden state through the GTFF.

10. A computational system for simulating motion of elements within a multi-element system using a graph transformer neural network (GTFF), the system comprising:

a memory storing instructions of a GTFF algorithm of molecular dynamics (MD) software; and

a processor programmed to execute the instructions to perform operations including to

convert a molecular dynamics snapshot of the elements within the multi-element system into a graph with atoms as nodes of the graph;

define a matrix such that each column of the matrix represents a node in the graph;

define a distance matrix according to a set of relative positions of each of the atoms;

iterate through the GTFF using an attention mechanism, operating on the matrix and augmented by incorporating the distance matrix, to pass hidden state from a current layer of the GTFF to a next layer of the GTFF;

perform a combination over the columns of the matrix to produce a scalar molecular energy;

make a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to compute a prediction of force acting on each atom; and

return the prediction of the force acting on each atom.

11. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to perform the combination over the columns of the matrix as a sum.

12. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to perform the combination over the columns of the matrix as a mean.

13. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to perform the combination over the columns of the matrix using a neural network.

14. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to include, in the matrix as nodes of the graph, other molecular objects in addition to the atoms, the other molecular objects including one or more of bonds or quantum state of the atoms.

15. The computational system of claim 14 , wherein the processor is further programmed to execute the instructions to define distances of the other molecular objects to other objects in the graph such that the distance matrix remains well-defined for each element.

16. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to train the GTFF as a deep neural network using a loss function formed to compare the force predictions to ground truth forces on a labeled dataset, such that network weights are optimized with respect to the loss function using gradient descent.

17. The computational system of claim 10 , wherein the processor is further programmed to execute the instructions to utilize softmax in the attention mechanism for forward propagation of the hidden state through the GTFF.

18. A non-transitory computer-readable medium comprising instructions for simulating motion of elements within a multi-element system using a graph transformer neural network (GTFF) that, when executed by a processor, cause the processor to:

convert a molecular dynamics snapshot of the elements within the multi-element system into a graph with atoms as nodes of the graph;

define a matrix such that each column of the matrix represents a node in the graph;

define a distance matrix according to a set of relative positions of each of the atoms;

iterate through the GTFF using an attention mechanism, operating on the matrix and augmented by incorporating the distance matrix, to pass hidden state from a current layer of the GTFF to a next layer of the GTFF;

perform a combination over the columns of the matrix to produce a scalar molecular energy;

make a backward pass through the GTFF, iteratively calculating derivatives at each of the layers of the GTFF to compute a prediction of force acting on each atom; and

return the prediction of the force acting on each atom.

19. The medium of claim 18 , further comprising instructions that, when executed by a processor, cause the processor to:

include, in the matrix as nodes of the graph, other molecular objects in addition to the atoms, other molecular objects including one or more of bonds or quantum state of the atoms; and

define distances of the other molecular objects to other objects in the graph such that the distance matrix remains well-defined for each element.

20. The medium of claim 18 , further comprising instructions that, when executed by a processor, cause the processor to train the GTFF as a deep neural network using a loss function formed to compare the force predictions to ground truth forces on a labeled dataset, such that network weights are optimized with respect to the loss function using gradient descent.