Implement the ANI2x model architecture

[wiederm]

Description

The ANI2x provides an important architecture, sitting in expressiveness between the already implemented SchNET and PaiNN architecture, and has been widely adopted. This PR will re-implement the architecture using the previously adopted split in representation and interaction.

Todos

Notable points that this PR has either accomplished or will accomplish.

• [x] Implement ani representation
• [x] implement ensemble and interaction layers

Status

• [x] Ready to go

[wiederm]

Currently, the models in modelforge have the following modular abstraction:

• representation module: takes as input cartesian coordinates and outputs a model-specific representation of the molecular system. E.g., the SchNET representation module outputs radial symmetry features. The representation module is initialized with parameters to control the dimensions of the output tensor.
• interaction module: takes in the model specific representation of the molecular system and outputs the per-atom properties (scalar and/or vector). This is the model specific, trainable architecture.
• property pooling: takes in the per-atom properties and generates per-molecule properties. This includes a linear layer and then aggregation of the properties.

To translate this to the ANI architecture: The representation module takes cartesian coordinates and calculates radial and angular symmetry functions. The radial symmetry functions that we have currently implemented computes the terms of the following equation G_{m}^{R} = \sum_{j!=i}^{N} exp(-eta(|R_{ij} - R_s|)^2)) for a molecule m. R_{ij} is the pariwise distance, R_{s} the position of the peak of the gaussian and eta is a parameter that is currently calculated from the width of the gaussian. Applying a cutoff function to the terms results in the same functional form as used in torchani (eta is calculated slightly different though).

The angular symmetry function needs to be implemented and added to the package.

[wiederm]

Note: to reproduce the original ANI results, it is necessary to scale some of the used constants. Since many of these constants don't have units attached it is not apparent how they change with the unit system, so this is not a trivial task unfortunately.

I will document the changes here for now:

• Radial eta: (used to control the width of the radial symmetry function) is multiplied by 100 to reproduce the output of the Torchani implementation when using nanometer instead of Angstrom.

[wiederm]

The current implementation can reproduce the output of the radial and angular symmetry functions of the ANI model in the current torchani implementation (given the constants used there).

We have tests both for the radial as well as the angular symmetry functions here: https://github.com/choderalab/modelforge/blob/f6642b0513170d17bb4d94c71f7e624ff94a48a7/modelforge/tests/test_ani.py

While the symmetry functions of ANI are relatively straightforward, what is difficult is to get the indices in the right order to perform the calculations. One key function was taken directly from the torchani repository: https://github.com/aiqm/torchani/blob/17204c6dccf6210753bc8c0ca4c92278b60719c9/torchani/aev.py#L230 It is responsible for enumerating all triplets of atoms (ijk) for which the distances (i,j) and (j,k) are below a certain threshold.

The next step is to set up the per-element MLP that consumes the radial and angular symmetry vectors and calculate the per-atom energies in ANI.

[wiederm]

I have made some ad-hoc changes to the current structure during this PR that I want to outline here. Some of these should be spun off into new PRs to be refined later.

• I changed the name from GaussianRBF to RadialSymmetryFunction, which is more in line with its close cousin, AngularSymmetryFunction. Right now, we are only supporting Gaussian RBF, but we also want to support other functional forms. This will be addressed in a future PR.
• I changed the dimensions of the atomic_species to (nr_of_atoms,). Making this change also required adopting the EmbeddingModule. I made this change because I don't think there is ever a reason why this has more dimensions. If we want to embed multiple properties, we can pass them to the embedding module as well, but they should be stored under a different name.
• I added a only_unique_pairs property to the Neighbor and Pairlist. The SchNET family of NNPs wants a complete (redundant) pair list, while ANI wants a non-redundant pair list.
• Previously, the term 'cutoff' was ambiguous. It wasn't always clear if we meant the cutoff module or the actual cutoff distance. I resolved this by using cutoff_module throughout to indicate the module.

[wiederm]

I am not sure about these test failures. These are all passing locally. Will investigate.

[wiederm]

Indeed, switching from torchani 2.2 (which I am using locally) to 2.2.4 (which gh is installing) results in these tests to fail.

[codecov-commenter]

Codecov Report

Merging #70 (28a1e11) into main (228bc41) will increase coverage by 2.42%. The diff coverage is 96.63%.

Additional details and impacted files

[wiederm]

Test are now passing :-) ready for review