Allegro

The Allegro model as proposed in Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics.

Supervised

The base allegro_model as proposed in the original paper. It has the ability to output graph scalars, 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 = 2

  The number of Allegro layers.

• 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 features in the model

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• mlp_latent_dimensions: List[int] = [128]

  The MLP dimensions.

• latent_mlp_latent_dimensions: List[int] = [1024, 1024, 1024]

  The latent MLP dimensions.

• env_embed_multiplicity: int = 32

  The environment embedding multiplicity.

• two_body_latent_mlp_latent_dimensions: List[int] = [128, 256, 512, 1024]

  The two body interaction MLP dimensions.

• num_bessel: int = 8

  The highest order of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel basis has trainable weights

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• embed_initial_edge: bool = True

  Whether to embed the initial edge features.

• latent_resnet: bool = True

  Whether to use a latent resnet.

• 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_node_vector.

• y_graph_scalars_loss_config: Optional[Dict] = None

  The loss config for the y_graph_scalars.

Supervised mean and variance

The mean_var_allegro_model which predicts a mean and variance for the y_graph_scalars during training and inference. It has the same comfig as the allegro_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 = 2

  The number of Allegro layers.

• 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 features in the model

• mlp_irreps: str = "16x0e"

  The output MLP irreps.

• mlp_latent_dimensions: List[int] = [128]

  The MLP dimensions.

• latent_mlp_latent_dimensions: List[int] = [1024, 1024, 1024]

  The latent MLP dimensions.

• env_embed_multiplicity: int = 32

  The environment embedding multiplicity.

• two_body_latent_mlp_latent_dimensions: List[int] = [128, 256, 512, 1024]

  The two body interaction MLP dimensions.

• num_bessel: int = 8

  The highest order of bessel functions to use.

• bessel_basis_trainable: bool = True

  Whether the bessel basis has trainable weights

• num_polynomial_cutoff: int = 6

  The cut-off polynomial envelope power.

• avg_num_neighbours: Optional[float] = None

  The average number of neighbours in the dataset.

• embed_initial_edge: bool = True

  Whether to embed the initial edge features.

• latent_resnet: bool = True

  Whether to use a latent resnet.

• 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_node_vector.