xnn.gnn.models.mace.RealAgnosticDensityInteractionBlock#

class xnn.gnn.models.mace.RealAgnosticDensityInteractionBlock(node_attrs_irreps, node_feats_irreps, edge_attrs_irreps, edge_feats_irreps, target_irreps, hidden_irreps, avg_num_neighbors, radial_MLP)[source]#

Bases: RealAgnosticInteractionBlock

Non-residual interaction with learned density normalization.

Identical to RealAgnosticInteractionBlock except that the aggregated message is divided by 1 + rho_i – a learned, per-node neighbor density rho_i = sum_j tanh(d(e_ij)^2) built from the radial edge features – instead of the global avg_num_neighbors constant. This is the interaction of the MACE-MP “density” foundation generation (0b2 / 0b3 / MPA-0 / OMAT-0 / MATPES).

Parameters:

avg_num_neighbors (float)

forward(node_attrs, node_feats, edge_attrs, edge_feats, edge_index)[source]#

Compute one density-normalized update (non-residual variant).

Parameters:

  • node_attrs (torch.Tensor) – Per-node one-hot element attributes.

  • node_feats (torch.Tensor) – Incoming node features, shape (num_nodes, node_feats_irreps.dim).

  • edge_attrs (torch.Tensor) – Edge spherical-harmonic attributes.

  • edge_feats (torch.Tensor) – Scalar radial edge features feeding the radial MLP.

  • edge_index (torch.Tensor) – Edge index of shape (2, num_edges) ([senders, receivers]).

Returns:

The reshaped message features and None (no separate self-connection for the non-residual block).

Return type:

tuple of (torch.Tensor, None)