mlcolvar.core.nn.graph.radial.RadialEmbeddingBlock

class mlcolvar.core.nn.graph.radial.RadialEmbeddingBlock(cutoff: float, long_range_cutoff: float = -1.0, n_bases: int = 8, n_polynomials: int = 6, basis_type: str = 'bessel')

Bases: Module

Radial embedding block [1]

References

[1]

Gasteiger, J.; Groß, J.; Günnemann, S. Directional Message Passing

for Molecular Graphs; ICLR 2020.

__init__(cutoff: float, long_range_cutoff: float = -1.0, n_bases: int = 8, n_polynomials: int = 6, basis_type: str = 'bessel') → None

Initializes a radial embedding block

Parameters:

• cutoff (float) – Cutoff radius.

• long_range_cutoff (float) – Long range cutoff for interaction between subsystem atoms, not used if negative, by default -1.0

• n_bases (int, optional) – Size of the basis set, by default 8

• n_polynomials (int, optional) – Order of the polynomial for the polynomial cutoff, by default 6

• basis_type (str, optional) – Type fo the basis function, by default ‘bessel’

Raises:

RuntimeError – _description_

Methods

__init__(cutoff[, long_range_cutoff, ...])	Initializes a radial embedding block
forward(edge_lengths[, edge_masks_lr])	The forward pass of RadialEmbeddingBlock

forward(edge_lengths: Tensor, edge_masks_lr: Tensor = None) → Tensor

The forward pass of RadialEmbeddingBlock

Parameters:

• edge_lengths (torch.Tensor (shape: [n_edges, 1])) – Lengths of edges.

• edge_masks_lr (torch.Tensor (shape: [1, n_edges])) – Mask for long range edges.

Returns:

edge_embedding – The radial edge embedding.

Return type:

torch.Tensor (shape: [n_edges, n_bases])

Attributes

T_destination
call_super_init
dump_patches
training