ViSNet¶

class mlip.models.visnet.models.Visnet(config: ~mlip.models.visnet.config.VisnetConfig, dataset_info: ~mlip.data.dataset_info.DatasetInfo, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)¶

The ViSNet model flax module. It is derived from the MLIPNetwork class.

References

Yusong Wang, Tong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, and Tie-Yan Liu. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nature Communications, 15(1), January 2024. ISSN: 2041-1723. URL: https://dx.doi.org/10.1038/s41467-023-43720-2.

config¶

Hyperparameters / configuration for the ViSNet model, see VisnetConfig.

Type:: mlip.models.visnet.config.VisnetConfig

dataset_info¶

Hyperparameters dictated by the dataset (e.g., cutoff radius or average number of neighbors).

Type:: mlip.data.dataset_info.DatasetInfo

__call__(edge_vectors: Array, node_species: Array, senders: Array, receivers: Array) → Array¶: Compute node-wise energy summands. This function must be overridden by the implementation of MLIPNetwork.

class mlip.models.visnet.config.VisnetConfig(*, num_layers: Annotated[int, Gt(gt=0)] = 4, num_channels: Annotated[int, Gt(gt=0)] = 256, l_max: Annotated[int, Ge(ge=0)] = 2, num_heads: Annotated[int, Gt(gt=0)] = 8, num_rbf: Annotated[int, Gt(gt=0)] = 32, trainable_rbf: bool = False, activation: Activation = Activation.SILU, attn_activation: Activation = Activation.SILU, vecnorm_type: VecNormType = VecNormType.NONE, atomic_energies: str | dict[int, float] | None = None, num_species: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = None)¶

Hyperparameters for the ViSNet model.

num_layers¶

Number of ViSNet layers. Default is 2.

Type:: int

num_channels¶

The number of channels. Default is 256.

Type:: int

l_max¶

Highest harmonic order included in the Spherical Harmonics series. Default is 2.

Type:: int

num_heads¶

Number of heads in the attention block. Default is 8.

Type:: int

num_rbf¶

Number of basis functions used in the embedding block. Default is 32.

Type:: int

trainable_rbf¶

Whether to add learnable weights to each of the radial embedding basis functions. Default is False.

Type:: bool

activation¶

Activation function for the output block. Options are “silu” (default), “ssp” (which is shifted softplus), “tanh”, “sigmoid”, and “swish”.

Type:: mlip.models.options.Activation

attn_activation¶

Activation function for the attention block. Options are “silu” (default), “ssp” (which is shifted softplus), “tanh”, “sigmoid”, and “swish”.

Type:: mlip.models.options.Activation

vecnorm_type¶

The type of the vector norm. The options are “none” (default), “max_min”, and “rms”.

Type:: mlip.models.options.VecNormType

atomic_energies¶

How to treat the atomic energies. If set to None (default) or the string "average", then the average atomic energies stored in the dataset info are used. It can also be set to the string "zero" which means not to use any atomic energies in the model. Lastly, one can also pass an atomic energies dictionary via this parameter different from the one in the dataset info, that is used.

Type:: str | dict[int, float] | None

num_species¶

The number of elements (atomic species descriptors) allowed. If None (default), infer the value from the atomic energies map in the dataset info.

Type:: int | None