To solve the Initial Structure to Relaxed Structure (IS2RS) task, we use a pre-trained S2EF model to iteratively run a structure optimization process in order to arrive at a relaxed structure. It is also possible to train a model that can directly predict relaxed structures, but most approaches to this task use an iterative process.

To make IS2RS predictions, the following steps can be followed:

  1. Define or load a configuration (config) that specifies the details of the task, such as the model to be used, the optimization algorithm, the dataset to be used for training, and the trainer object. The config may also include information about the relaxation dataset.
  2. Create a ForcesTrainer object, which is responsible for coordinating the training process.
  3. Train an S2EF model or load an existing S2EF checkpoint. The S2EF model is used to predict the energy and forces of a structure, which are used in the structure optimization process.
  4. Run relaxations by calling a method or function that initiates the structure optimization process. The ForcesTrainer object will handle the details of iteratively updating the structure based on the S2EF model's predictions and the optimization algorithm.

GNN:

Each atom is represented as a node with its features computed using a simple torch.nn.Embedding layer on the atomic number.

All pairs of atoms with a defined cutoff radius (=6A) are assumed to have edges between them, with their features computed as the concatenation of

  1. a Gaussian expansion of the distance between the atoms, and the

  2. source

  3. target node features.

use the GaussianSmearing layer (reproduced below) from the PyTorch Geometric library for computing distance features:

class GaussianSmearing(torch.nn.Module):
    def __init__(self, start=0.0, stop=5.0, num_gaussians=50):
        super(GaussianSmearing, self).__init__()
        offset = torch.linspace(start, stop, num_gaussians)
        self.coeff = -0.5 / (offset[1] - offset[0]).item()**2
        self.register_buffer('offset', offset)

    def forward(self, dist):
        dist = dist.view(-1, 1) - self.offset.view(1, -1)
        return torch.exp(self.coeff * torch.pow(dist, 2))

Message passing

message-passing scheme needed to predict system energy and forces.

Given the node and edge features →

  1. sum up edge features for all edges eij connecting node i to its neighbors j
  2. pass the resultant vector through a fully-connected layer to project it down to a scalar