DimeNet Model

[JustinBakerMath]

Introduces the key components of the DimeNet model to the HydraGNN library by modifying the PyG implementation.

The core addition is the DIMEStack which is inherits from Base to provide convolutional layers on request using the _get_conv() method.

Notice: Because the PyG implementations use the glorot initialization, the construction of the hidden dimension must be larger than 1 to avoid division by zero. Only the input_dim and output_dim are provided to _get_conv(), without constraint on the dimension of the convolution. Therefore, DIMEStack will use the input_dim if larger than 1 or the output_dim otherwise. If neither is larger than one, then the assertion will fail.

The DIMEStack also computes the necessary convolutional arguments, being
rbf: (tensor) the radial basis function output.
sbf: (tensor) the spherical basis function output.
i: The message passing target.
j: The message passing source.
idx_kj: A portion of the bond triple.
idx_ji: The second component of the bond triple.

Notice: DimeNet uses the triplet() method to compute the bond angles between triplets of nodes. This information is computed on every call but there is room for computational speed ups by computing once and then updating iteratively.

The stack modifies the Embedding and InteractionBlock and preserves theOutputBlock of the baseline DimeNet implementation. For the Embedding layer, the baseline atomic feature embedding is removed as embedding is handled by HydraGNN.

The baseline InteractionBlock does not find the optimal path when using einsum so this call is segmented into two equivalent steps.

Six hyperparameters are introduced.

• num_before_skip/num_after_skip control layers before and after the attention mechanism in the OutputBlock.
• num_bilinear determines the number of bilinear layers used in the InteractionBlock
• num_radial/envelope_exponent determine the radial basis functions used.
• num_spherical determines the number of spherical basis functions used

These arguments are checked to exist or set to none, and verified not to be none before use by DimeNet.

Other updates include documentation and testing.

[JustinBakerMath]

• First, with regard to formatting. I would like to follow the pre-commit instructions, but I encounter known compatibility issues between click and black. If an old version of black is being used, is it possible to get a click version pinned in the pre-commit configuration.
• Second, the current errors 1 and 2 indicate import issues inside of torch_geometric with the sympy dependency. Trouble shooting this error is harder, as I am relying on the torch_geometric install in the CI to build the dependencies. The most brute-force option is to re-write the SphericalBasisLayer to be independent of sympy. The opposite end seems to be to include sympy in one of the requirements-*.txt.
• I've included a screenshot of my local pytest output.

[JustinBakerMath]

Closed in favor of DimeNet++ integration and request to include test results for non-unit test datasets.