MLIPs

mlcg.nn currently implements the SchNet graph neural network, as well as several utility classes for computing distance expansions and cutoffs. The nn subpackage also contains several useful classes for extracting other properties from energy predictions or aggregating the predictions from several different model types.

SchNet Utilities

These classes are used to define a SchNet graph neural network. For “typical” SchNet models, users may find the class StandardSchNet to be helpful in getting started quickly.

class mlcg.nn.StandardSchNet(rbf_layer, cutoff, output_hidden_layer_widths, hidden_channels=128, embedding_size=100, num_filters=128, num_interactions=3, activation=Tanh(), max_num_neighbors=1000, aggr='add', nls_distance_method='torch')[source]

Small wrapper class for mlcg.nn.SchNet to simplify the definition of the SchNet model through an input file. The upper distance cutoff attribute in is set by default to match the upper cutoff value in the cutoff function.

Parameters:

  • rbf_layer (Module) – radial basis function used to project the distances \(r_{ij}\).

  • cutoff (Module) – smooth cutoff function to supply to the CFConv

  • output_hidden_layer_widths (List[int]) – List giving the number of hidden nodes of each hidden layer of the MLP used to predict the target property from the learned representation.

  • hidden_channels (int) – dimension of the learned representation, i.e. dimension of the embeding projection, convolution layers, and interaction block.

  • embedding_size (int) – dimension of the input embeddings (should be larger than AtomicData.atom_types.max()+1).

  • num_filters (int) – number of nodes of the networks used to filter the projected distances

  • num_interactions (int) – number of interaction blocks

  • activation (Module) – activation function

  • max_num_neighbors (int) – The maximum number of neighbors to return for each atom in data. If the number of actual neighbors is greater than max_num_neighbors, returned neighbors are picked randomly.

  • aggr (str) – Aggregation scheme for continuous filter output. For all options, see here for more options.

class mlcg.nn.schnet.SchNet(embedding_layer, interaction_blocks, rbf_layer, output_network, self_interaction=False, max_num_neighbors=1000, nls_distance_method='torch')[source]

PyTorch Geometric implementation of SchNet Code adapted from [PT_geom_schnet] which is based on the architecture described in [Schnet] .

Parameters:

  • embedding_layer (Module) – Initial embedding layer that transforms atoms/coarse grain bead types into embedded features

  • interaction_blocks (list of torch.nn.Module or torch.nn.Sequential) – Sequential interaction blocks of the model, where each interaction block applies

  • rbf_layer (Module) – The set of radial basis functions that expands pairwise distances between atoms/CG beads.

  • output_network (Module) – Output neural network that predicts scalar energies from SchNet features. This network should transform (num_examples * num_atoms, hidden_channels) to (num_examples * num atoms, 1).

  • upper_distance_cutoff – Upper distance cutoff used for making neighbor lists.

  • self_interaction (bool) – If True, self interactions/distancess are calculated. But it never had a function due to a bug in the implementation (see static method neighbor_list). Should be kept False. This option shall not be deleted for compatibility.

  • max_num_neighbors (int) – Maximum number of neighbors to return for a given node/atom when constructing the molecular graph during forward passes. This attribute is passed to the torch_cluster radius_graph routine keyword max_num_neighbors, which normally defaults to 32. Users should set this to higher values if they are using higher upper distance cutoffs and expect more than 32 neighbors per node/atom.

  • nls_distance_method (str) – Method for computing a neighbor list. Supported values are torch, nvalchemi_naive, nvalchemi_cell and custom.

class mlcg.nn.schnet.InteractionBlock(cfconv_layer, hidden_channels=128, activation=Tanh())[source]

Interaction blocks for SchNet. Consists of atomwise transformations of embedded features that are continuously convolved with filters generated from radial basis function-expanded pairwise distances.

Parameters:

  • cfconv_layer (Module) – Continuous filter convolution layer for convolutions of radial basis function-expanded distances with embedded features

  • hidden_channels (int) – Hidden dimension of embedded features

  • activation (Module) – Activation function applied to linear layer outputs

class mlcg.nn.schnet.CFConv(filter_network, cutoff, in_channels=128, out_channels=128, num_filters=128, aggr='add')[source]

Continuous filter convolutions for SchNet.

Parameters:

  • filter_net – Neural network for generating filters from expanded pairwise distances

  • cutoff (Module) – Cutoff envelope to apply to the output of the filter generating network.

  • in_channels (int) – Hidden input dimensions

  • out_channels (int) – Hidden output dimensions

  • num_filters (int) – Number of filters

  • aggr (str) – Aggregation scheme for continuous filter output. For all options, see here <https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html?highlight=MessagePassing#the-messagepassing-base-class>.

PaiNN Utilities

These classes are used to define a PaiNN graph neural network. For “typical” PaiNN models, users may find the class StandardPaiNN to be helpful in getting started quickly.

class mlcg.nn.StandardPaiNN(rbf_layer, cutoff, output_hidden_layer_widths, hidden_channels=128, embedding_size=100, num_interactions=3, activation=Tanh(), max_num_neighbors=1000, aggr='add', epsilon=1e-08, nls_distance_method='torch')[source]

Small wrapper class for PaiNN to simplify the definition of the PaiNN model through an input file. The upper distance cutoff attribute in is set by default to match the upper cutoff value in the cutoff function.

Parameters:

  • rbf_layer (Module) – Radial basis function used to project the distances \(r_{ij}\).

  • cutoff (Module) – Smooth cutoff function to supply to the PaiNNInteraction.

  • output_hidden_layer_widths (List[int]) – List giving the number of hidden nodes of each hidden layer of the MLP used to predict the target property from the learned scalar representation.

  • hidden_channels (int) – Dimension of the learned representation, i.e. dimension of the embedding projection, convolution layers, and interaction block.

  • embedding_size (int) – Dimension of the input embeddings (should be larger than AtomicData.atom_types.max()+1).

  • num_interactions (int) – Number of interaction blocks.

  • activation (Module) – Activation function.

  • max_num_neighbors (int) – The maximum number of neighbors to return for each atom in data. If the number of actual neighbors is greater than max_num_neighbors, returned neighbors are picked randomly.

  • aggr (str) –

    Aggregation scheme for continuous filter output. For all options, see here for more options.

  • epsilon (float) – Stability constant added in norm to prevent numerical instabilities.

class mlcg.nn.PaiNN(embedding_layer, interaction_blocks, mixing_blocks, rbf_layer, output_network, max_num_neighbors, nls_distance_method='torch')[source]

Implementation of PaiNN - polarizable interaction neural network Code adapted from https://schnetpack.readthedocs.io/en/latest/api/generated/representation.PaiNN.html which is based on the architecture described in http://proceedings.mlr.press/v139/schutt21a.html

Parameters:

  • embedding_layer (Module) – Initial embedding layer that transforms atoms/coarse grain bead types into embedded features

  • interaction_blocks (Union[PaiNNInteraction, List[PaiNNInteraction]]) – list of PaiNNInteraction or single PaiNNInteraction block. Sequential interaction blocks of the model, where each interaction block applies.

  • mixing_blocks (Union[PaiNNMixing, List[PaiNNMixing]]) – List of PaiNNMixing or single PaiNNMixing block. Sequential mixing blocks of the model, where each mixing applies after each interaction block.

  • rbf_layer (Module) – The set of radial basis functions that expands pairwise distances between atoms/CG beads.

  • output_network (Module) – Output neural network that predicts scalar energies from scalar PaiNN features. This network should transform (num_examples * num_atoms, hidden_channels) to (num_examples * num atoms, 1).

  • max_num_neighbors (int) – Maximum number of neighbors to return for a given node/atom when constructing the molecular graph during forward passes. This attribute is passed to the torch_cluster radius_graph routine keyword max_num_neighbors, which normally defaults to 32. Users should set this to higher values if they are using higher upper distance cutoffs and expect more than 32 neighbors per node/atom.

  • nls_distance_method (str) – Method for computing a neighbor list. Supported values are torch, nvalchemi_naive, nvalchemi_cell and custom.

So3krates Utilities

These classes are used to define a so3krates graph neural network. For “typical” so3krates models, users may find the class StandardPaiNN to be helpful in getting started quickly.

class mlcg.nn.So3krates(embedding_layer, rbf_layer, interaction_blocks, layer_norm, output_network, degrees, max_num_neighbors=1000, normalize_sph=True, nls_distance_method='torch')[source]

Main SO3krates model for equivariant message passing neural networks.

Parameters:

  • embedding_layer (nn.Module) – Embedding layer for atom types.

  • rbf_layer (nn.Module) – Radial basis function layer to project distances.

  • interaction_blocks (List[So3kratesLayer]) – List of SO3krates layers.

  • layer_norm (nn.Module) – Layer normalization module.

  • output_network (nn.Module) – Output network for final energy prediction.

  • degrees (List[int]) – List of spherical harmonic degrees.

  • max_num_neighbors (int) – Maximum number of neighbors per atom.

  • normalize_sph (bool) – Whether to normalize spherical harmonics, default is True.

  • nls_distance_method (str) – Method for computing a neighbor list. Supported values are torch, nvalchemi_naive, nvalchemi_cell and custom.

class mlcg.nn.StandardSo3krates(rbf_layer, cutoff, output_hidden_layer_widths, hidden_channels=132, embedding_size=100, num_interactions=6, degrees=[1, 2, 3], n_heads=4, activation=SiLU(), max_num_neighbors=1000, normalize_sph=True, nls_distance_method='torch')[source]

Standard configuration for the SO3krates model.

Provides a convenient interface for constructing a SO3krates model with typical settings, including embedding, interaction, and output modules.

Parameters:

  • rbf_layer (nn.Module) – Radial basis function layer.

  • cutoff (nn.Module) – Cutoff function for edge weights.

  • output_hidden_layer_widths (List[int]) – Widths for hidden layers in output MLP.

  • hidden_channels (int) – Number of hidden channels.

  • embedding_size (int) – Size of atom type embedding.

  • num_interactions (int) – Number of SO3krates layers.

  • degrees (List[int]) – List of spherical harmonic degrees.

  • n_heads (int) – Number of attention heads.

  • activation (nn.Module) – Activation function.

  • max_num_neighbors (int) – Maximum number of neighbors per atom.

  • normalize_sph (bool) – Whether to normalize spherical harmonics, default is True.

MACE Utilities

These classes are used to define a MACE graph neural network, for which the base implementation is required and available here. For “typical” MACE models, users may find the class StandardMACE to be helpful in getting started quickly.

class mlcg.nn.StandardMACE(r_max, num_bessel, num_polynomial_cutoff, max_ell, interaction_cls, interaction_cls_first, num_interactions, hidden_irreps, MLP_irreps, avg_num_neighbors, atomic_numbers, correlation, gate, max_num_neighbors=1000, pair_repulsion=False, distance_transform='None', radial_MLP=None, radial_type='bessel', cueq_config=None, use_cueq=False, nls_distance_method='torch')[source]

Standard configuration of the MACE model.

This class provides a convenient interface for constructing a MACE model with typical settings and block choices, including embedding, interaction, and readout modules.

Parameters:

  • r_max (float) – Cutoff radius for neighbor list.

  • num_bessel (int) – Number of Bessel functions for radial basis.

  • num_polynomial_cutoff (int) – Number of polynomial cutoff functions.

  • max_ell (int) – Maximum angular momentum for spherical harmonics.

  • interaction_cls (str) – Class name for interaction blocks.

  • interaction_cls_first (str) – Class name for the first interaction block.

  • num_interactions (int) – Number of interaction blocks.

  • hidden_irreps (str) – Irreducible representations for hidden features. For example if only a scalar representation with 128 channels is used can be “128x0e”. If also a vector representation is used can be “128x0e + 128x1o”.

  • MLP_irreps (str) – Irreducible representations for MLP layers.

  • avg_num_neighbors (float) – Average number of neighbors per atom used for normalization and numerical stability.

  • atomic_numbers (List[int]) – List of atomic numbers in the system.

  • correlation (Union[int, List[int]]) – Correlation order(s) for product blocks.

  • gate (Optional[Callable]) – Activation function for non-linearities.

  • max_num_neighbors (int, optional) – Maximum number of neighbors per atom.

  • pair_repulsion (bool, optional) – Whether to use pairwise repulsion.

  • distance_transform (str, optional) – Distance transformation type.

  • radial_MLP (Optional[List[int]], optional) – Radial MLP architecture.

  • radial_type (Optional[str], optional) – Radial basis type.

  • cueq_config (Optional[Dict[str, Any]], optional) – cuEquivariance configuration.

  • use_cueq (Optional[bool], optional) – Whether to use cuEquivariance acceleration.

class mlcg.nn.mace.MACE(atomic_numbers, node_embedding, radial_embedding, spherical_harmonics, interactions, products, readouts, r_max, max_num_neighbors, pair_repulsion_fn=None, nls_distance_method='torch')[source]

Implementation of MACE neural network model from https://github.com/ACEsuit/mace

Parameters:

  • atomic_numbers (torch.Tensor) – Tensor of atomic numbers present in the system.

  • node_embedding (torch.nn.Module) – Module for embedding node (atom) attributes.

  • radial_embedding (torch.nn.Module) – Module for embedding radial (distance) features.

  • spherical_harmonics (torch.nn.Module) – Module for computing spherical harmonics of edge vectors.

  • interactions (List[torch.nn.Module]) – List of interaction blocks.

  • products (List[torch.nn.Module]) – List of product basis blocks.

  • readouts (List[torch.nn.Module]) – List of readout blocks.

  • r_max (float) – Cutoff radius for neighbor list.

  • max_num_neighbors (int) – Maximum number of neighbors per atom.

  • pair_repulsion_fn (torch.nn.Module, optional) – Optional pairwise repulsion energy function.

  • nls_distance_method (str) – Method for computing a neighbor list. Supported values are torch, nvalchemi_naive, nvalchemi_cell and custom.

Allegro Utilities

These classes are used to define an Allegro graph neural network, for which the base implementation is required and available here. For “typical” MACE models, users may find the class StandardAllegro to be helpful in getting started quickly.

class mlcg.nn.StandardAllegro(r_max=4.0, embedding_size=100, l_max=2, parity=False, radial_chemical_embed=None, radial_chemical_embed_dim=None, per_edge_type_cutoff=None, scalar_embed_mlp_hidden_layers_depth=1, scalar_embed_mlp_hidden_layers_width=64, scalar_embed_mlp_nonlinearity='silu', num_layers=2, num_scalar_features=64, num_tensor_features=16, allegro_mlp_hidden_layers_depth=1, allegro_mlp_hidden_layers_width=64, allegro_mlp_nonlinearity='silu', tp_path_channel_coupling=True, readout_mlp_hidden_layers_depth=1, readout_mlp_hidden_layers_width=32, readout_mlp_nonlinearity='silu', avg_num_neighbors=10.0, weight_individual_irreps=True, per_type_energy_scales=None, per_type_energy_shifts=None, per_type_energy_scales_trainable=False, per_type_energy_shifts_trainable=False, pair_potential=None, forward_normalize=True, use_cueq=False, nls_distance_method='torch')[source]

Standard implementation of the Allegro model with configurable parameters.

This class provides a higher-level interface to the Allegro model, allowing users to specify model parameters directly rather than constructing each component manually.

Parameters:

  • r_max (float, default=4.0) – Cutoff radius for neighbor identification.

  • embedding_size (int, default=100) – Dimension of the atom embedding space.

  • l_max (int, default=2) – Maximum spherical harmonic degree to use.

  • parity (bool, default=False) – Whether to use parity-equivariant features.

  • radial_chemical_embed (Optional[dict], default=None) – Configuration for the radial chemical embedding.

  • radial_chemical_embed_dim (Optional[int], default=None) – Dimension of the radial chemical embedding.

  • per_edge_type_cutoff (Optional[Dict], default=None) – Cutoff parameters for each edge type.

  • scalar_embed_mlp_hidden_layers_depth (int, default=1) – Depth of the scalar embedding MLP.

  • scalar_embed_mlp_hidden_layers_width (int, default=64) – Width of the scalar embedding MLP.

  • scalar_embed_mlp_nonlinearity (Optional[str], default="silu") – Nonlinearity to use in the scalar embedding MLP.

  • num_layers (int, default=2) – Number of message-passing layers.

  • num_scalar_features (int, default=64) – Number of scalar features.

  • num_tensor_features (int, default=16) – Number of tensor features.

  • allegro_mlp_hidden_layers_depth (int, default=1) – Depth of the Allegro MLP.

  • allegro_mlp_hidden_layers_width (int, default=64) – Width of the Allegro MLP.

  • allegro_mlp_nonlinearity (Optional[str], default="silu") – Nonlinearity to use in the Allegro MLP.

  • tp_path_channel_coupling (bool, default=True) – Whether to use channel coupling in tensor product paths.

  • readout_mlp_hidden_layers_depth (int, default=1) – Depth of the readout MLP.

  • readout_mlp_hidden_layers_width (int, default=32) – Width of the readout MLP.

  • readout_mlp_nonlinearity (Optional[str], default="silu") – Nonlinearity to use in the readout MLP.

  • avg_num_neighbors (Optional[float], default=10.0) – Average number of neighbors per atom for normalization.

  • weight_individual_irreps (bool, default=True) – Whether to use separate weights for each irreducible representation.

  • per_type_energy_scales (Optional[Union[float, Sequence[float]]], default=None) – Scaling factors for energies of each atom type.

  • per_type_energy_shifts (Optional[Union[float, Sequence[float]]], default=None) – Shift values for energies of each atom type.

  • per_type_energy_scales_trainable (Optional[bool], default=False) – Whether the energy scales are trainable parameters.

  • per_type_energy_shifts_trainable (Optional[bool], default=False) – Whether the energy shifts are trainable parameters.

  • pair_potential (Optional[Dict], default=None) – Optional pair potential to add to the energy.

  • forward_normalize (bool, default=True) – Whether to normalize weights in the forward pass.