NequIP

The NequIP model as proposed in E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials.

Supervised

The base nequip_model as proposed in the original paper. It has the ability to output graph scalars, a graph vector, node scalars, and a node vector. The model has the following config

• num_node_feats: int

  The number of node features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_edge_feats: int

  The number of edge features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_layers: int = 4

  The number of layers of the model.

• max_ell: int = 2

  The maximum SO(3) irreps dimension to use in the model (during interactions).

• parity: bool = True

  Whether to use parity odd irreps.

• num_features: int = 32

  The number of features of the hidden embeddings.

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• num_bessel: int = 8

  The number of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel function weights are trainable.

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• self_connection: bool = True

  Whether to use self connections

• resnet: bool = True

  Whether to use a resnet.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• scaling_mean: float = 0.0

  The scaling mean of the model output. This is usually computed as the mean of (molecule_energy - molecule_self_energy) / molecule_num_atoms

• scaling_std: float = 1.0

  The scaling std of the model output. This is usually computed as the std of (molecule_energy - molecule_self_energy)

• compute_forces: bool = False

  Whether to compute forces as the gradient of the y_graph_scalars and use those as the y_node_vector output.

• y_node_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_node_scalars.

• y_node_vector_loss_config: Optional[Dict] = None

  The loss config for the y_edge_scalars.

• y_graph_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_graph_scalars.

• y_graph_vector_loss_config: Optional[Dict] = None

  The loss config for the y_graph_vector.

Supervised mean and variance

The mean_var_nequip_model which predicts a mean and variance for the y_graph_scalars during training and inference. It has the same comfig as the nequip_model but without the ability to specify a y_graph_scalars_loss_config which is hardcoded to be the torch.nn.GaussianNLLLoss. The model config is as follows

• num_node_feats: int

  The number of node features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_edge_feats: int

  The number of edge features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_layers: int = 4

  The number of layers of the model.

• max_ell: int = 2

  The maximum SO(3) irreps dimension to use in the model (during interactions).

• parity: bool = True

  Whether to use parity odd irreps.

• num_features: int = 32

  The number of features of the hidden embeddings.

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• num_bessel: int = 8

  The number of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel function weights are trainable.

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• self_connection: bool = True

  Whether to use self connections

• resnet: bool = True

  Whether to use a resnet.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• scaling_mean: float = 0.0

  The scaling mean of the model output. This is usually computed as the mean of (molecule_energy - molecule_self_energy) / molecule_num_atoms

• scaling_std: float = 1.0

  The scaling std of the model output. This is usually computed as the std of (molecule_energy - molecule_self_energy)

• compute_forces: bool = False

  Whether to compute forces as the gradient of the y_graph_scalars and use those as the y_node_vector output.

• y_node_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_node_scalars.

• y_node_vector_loss_config: Optional[Dict] = None

  The loss config for the y_edge_scalars.

• y_graph_vector_loss_config: Optional[Dict] = None

  The loss config for the y_graph_vector.

Supervised feature scale-shift

The ssf_nequip_model which adds scale and shift for the hidden embeddings in the MACE model as proposed in Scaling & Shifting Your Features: A New Baseline for Efficient Model Tuning. This is useful for transfer learning. The config is the same as the nequip_model. The model config is as follows

• num_node_feats: int

  The number of node features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_edge_feats: int

  The number of edge features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_layers: int = 4

  The number of layers of the model.

• max_ell: int = 2

  The maximum SO(3) irreps dimension to use in the model (during interactions).

• parity: bool = True

  Whether to use parity odd irreps.

• num_features: int = 32

  The number of features of the hidden embeddings.

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• num_bessel: int = 8

  The number of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel function weights are trainable.

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• self_connection: bool = True

  Whether to use self connections

• resnet: bool = True

  Whether to use a resnet.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• scaling_mean: float = 0.0

  The scaling mean of the model output. This is usually computed as the mean of (molecule_energy - molecule_self_energy) / molecule_num_atoms

• scaling_std: float = 1.0

  The scaling std of the model output. This is usually computed as the std of (molecule_energy - molecule_self_energy)

• compute_forces: bool = False

  Whether to compute forces as the gradient of the y_graph_scalars and use those as the y_node_vector output.

• y_node_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_node_scalars.

• y_node_vector_loss_config: Optional[Dict] = None

  The loss config for the y_edge_scalars.

• y_graph_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_graph_scalars.

• y_graph_vector_loss_config: Optional[Dict] = None

  The loss config for the y_graph_vector.

Supervised adapters

The adapter_nequip_model which adds adapters at the beginning and end of the message passing backbone of the NequIP model as proposed in AdapterGNN. This is useful for transfer learning. The config is the same as the nequip_model but with the added kwargs for ratio_adapter_down and initial_s. The model config is as follows

• num_node_feats: int

  The number of node features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_edge_feats: int

  The number of edge features. Must be equal to the initial node feature dimension (sum of one-hot encoded features).

• num_layers: int = 4

  The number of layers of the model.

• max_ell: int = 2

  The maximum SO(3) irreps dimension to use in the model (during interactions).

• parity: bool = True

  Whether to use parity odd irreps.

• num_features: int = 32

  The number of features of the hidden embeddings.

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• num_bessel: int = 8

  The number of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel function weights are trainable.

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• self_connection: bool = True

  Whether to use self connections

• resnet: bool = True

  Whether to use a resnet.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• scaling_mean: float = 0.0

  The scaling mean of the model output. This is usually computed as the mean of (molecule_energy - molecule_self_energy) / molecule_num_atoms

• scaling_std: float = 1.0

  The scaling std of the model output. This is usually computed as the std of (molecule_energy - molecule_self_energy)

• compute_forces: bool = False

  Whether to compute forces as the gradient of the y_graph_scalars and use those as the y_node_vector output.

• y_node_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_node_scalars.

• y_node_vector_loss_config: Optional[Dict] = None

  The loss config for the y_edge_scalars.

• y_graph_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_graph_scalars.

• y_graph_vector_loss_config: Optional[Dict] = None

  The loss config for the y_graph_vector.