xnn.gnn.models.nequip.InteractionBlock#

class xnn.gnn.models.nequip.InteractionBlock(irreps_in, irreps_out, irreps_node_attr, irreps_edge_attr, n_radial, invariant_layers=2, invariant_neurons=64, avg_num_neighbors=None, use_sc=True, nonlinearity_scalars={'e': 'silu'})[source]#

Bases: Module

The NequIP equivariant convolution (upstream nequip.nn.InteractionBlock).

Messages are the tensor product of the neighbour node features with the edge spherical harmonics, weighted per edge by a radial MLP acting on the invariant radial embedding:

m_ij  = TP(linear_1(h_j), Y(r_ij); w = fc(radial_ij))
h_i'  = linear_2( sum_j m_ij / sqrt(avg_num_neighbors) ) + sc(h_i, z_i)

where sc is the element-dependent self-connection (e3nn.o3.FullyConnectedTensorProduct with the one-hot species). Member names match upstream (linear_1/fc/tp/linear_2/ sc) so state dicts transplant directly.

Parameters:

  • irreps_in (e3nn.o3.Irreps) – Irreps of the input node features.

  • irreps_out (e3nn.o3.Irreps) – Irreps of the output node features (the gate’s input irreps).

  • irreps_node_attr (e3nn.o3.Irreps) – Irreps of the one-hot species node attributes.

  • irreps_edge_attr (e3nn.o3.Irreps) – Irreps of the spherical-harmonic edge attributes.

  • n_radial (int) – Width of the invariant radial edge embedding feeding fc.

  • invariant_layers (int, optional) – Hidden layers of the radial MLP, by default 2 (upstream example.yaml; the bare upstream module defaults to 1).

  • invariant_neurons (int, optional) – Hidden width of the radial MLP, by default 64.

  • avg_num_neighbors (float or None, optional) – Divide the aggregated message by sqrt(avg_num_neighbors); None (default) disables the normalization.

  • use_sc (bool, optional) – Include the self-connection, by default True.

  • nonlinearity_scalars (dict, optional) – Upstream-style {"e": <name>} choice of the radial-MLP nonlinearity, by default {"e": "silu"}.

forward(x, node_attrs, edge_index, edge_sh, edge_radial)[source]#

Compute one convolution update.

Parameters:

  • x (torch.Tensor) – Node features of shape (N, irreps_in.dim).

  • node_attrs (torch.Tensor) – One-hot species node attributes of shape (N, n_species).

  • edge_index (torch.Tensor) – Edge index of shape (2, E); features are gathered from row 0 (the neighbour) and scattered onto row 1 (the centre).

  • edge_sh (torch.Tensor) – Spherical-harmonic edge attributes, shape (E, irreps_sh.dim).

  • edge_radial (torch.Tensor) – Invariant radial edge embedding, shape (E, n_radial).

Returns:

Updated node features of shape (N, irreps_out.dim).

Return type:

torch.Tensor