folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add egno (#602)

Browse Source 
Co-authored-by: GiovanniCanali <giovanni.canali98@yahoo.it>
This commit is contained in:
avisquid
2025-10-03 14:37:56 -04:00
committed by GitHub
parent b5e4d13663
commit 2108c76d14
11 changed files with 885 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pina/model/__init__.py
View File

@@ -14,6 +14,7 @@ __all__ = [
"Spline",
"GraphNeuralOperator",
"PirateNet",
"EquivariantGraphNeuralOperator",
]
from .feed_forward import FeedForward, ResidualFeedForward
@@ -26,3 +27,4 @@ from .low_rank_neural_operator import LowRankNeuralOperator
from .spline import Spline
from .graph_neural_operator import GraphNeuralOperator
from .pirate_network import PirateNet
from .equivariant_graph_neural_operator import EquivariantGraphNeuralOperator

4
pina/model/block/message_passing/__init__.py
View File

@@ -5,9 +5,13 @@ __all__ = [
"DeepTensorNetworkBlock",
"EnEquivariantNetworkBlock",
"RadialFieldNetworkBlock",
"EquivariantGraphNeuralOperatorBlock",
]
from .interaction_network_block import InteractionNetworkBlock
from .deep_tensor_network_block import DeepTensorNetworkBlock
from .en_equivariant_network_block import EnEquivariantNetworkBlock
from .radial_field_network_block import RadialFieldNetworkBlock
from .equivariant_graph_neural_operator_block import (
EquivariantGraphNeuralOperatorBlock,
)

62
pina/model/block/message_passing/en_equivariant_network_block.py
View File

@@ -3,7 +3,7 @@
import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import degree
from ....utils import check_positive_integer
from ....utils import check_positive_integer, check_consistency
from ....model import FeedForward
@@ -27,6 +27,12 @@ class EnEquivariantNetworkBlock(MessagePassing):
positions are updated by adding the incoming messages divided by the
degree of the recipient node.
When velocity features are used, node velocities are passed through a small
MLP to compute updates, which are then combined with the aggregated position
messages. The node positions are updated both by the normalized position
messages and by the updated velocities, ensuring equivariance while
incorporating dynamic information.
.. seealso::
**Original reference** Satorras, V. G., Hoogeboom, E., Welling, M.
@@ -40,6 +46,7 @@ class EnEquivariantNetworkBlock(MessagePassing):
node_feature_dim,
edge_feature_dim,
pos_dim,
use_velocity=False,
hidden_dim=64,
n_message_layers=2,
n_update_layers=2,
@@ -54,6 +61,8 @@ class EnEquivariantNetworkBlock(MessagePassing):
:param int node_feature_dim: The dimension of the node features.
:param int edge_feature_dim: The dimension of the edge features.
:param int pos_dim: The dimension of the position features.
:param bool use_velocity: Whether to use velocity features in the
message passing. Default is False.
:param int hidden_dim: The dimension of the hidden features.
Default is 64.
:param int n_message_layers: The number of layers in the message
@@ -80,6 +89,7 @@ class EnEquivariantNetworkBlock(MessagePassing):
:raises AssertionError: If `hidden_dim` is not a positive integer.
:raises AssertionError: If `n_message_layers` is not a positive integer.
:raises AssertionError: If `n_update_layers` is not a positive integer.
:raises AssertionError: If `use_velocity` is not a boolean.
"""
super().__init__(aggr=aggr, node_dim=node_dim, flow=flow)
@@ -90,6 +100,10 @@ class EnEquivariantNetworkBlock(MessagePassing):
check_positive_integer(hidden_dim, strict=True)
check_positive_integer(n_message_layers, strict=True)
check_positive_integer(n_update_layers, strict=True)
check_consistency(use_velocity, bool)
# Initialization
self.use_velocity = use_velocity
# Layer for computing the message
self.message_net = FeedForward(
@@ -119,7 +133,17 @@ class EnEquivariantNetworkBlock(MessagePassing):
func=activation,
)
def forward(self, x, pos, edge_index, edge_attr=None):
# If velocity is used, instantiate layer for velocity updates
if self.use_velocity:
self.update_vel_net = FeedForward(
input_dimensions=node_feature_dim,
output_dimensions=1,
inner_size=hidden_dim,
n_layers=n_update_layers,
func=activation,
)
def forward(self, x, pos, edge_index, edge_attr=None, vel=None):
"""
Forward pass of the block, triggering the message-passing routine.
@@ -130,11 +154,19 @@ class EnEquivariantNetworkBlock(MessagePassing):
:param torch.Tensor edge_index: The edge indices.
:param edge_attr: The edge attributes. Default is None.
:type edge_attr: torch.Tensor | LabelTensor
:param vel: The velocity of the nodes. Default is None.
:type vel: torch.Tensor | LabelTensor
:return: The updated node features and node positions.
:rtype: tuple(torch.Tensor, torch.Tensor)
:raises: ValueError: If ``use_velocity`` is True and ``vel`` is None.
"""
if self.use_velocity and vel is None:
raise ValueError(
"Velocity features are enabled, but no velocity is passed."
)
return self.propagate(
edge_index=edge_index, x=x, pos=pos, edge_attr=edge_attr
edge_index=edge_index, x=x, pos=pos, edge_attr=edge_attr, vel=vel
)
def message(self, x_i, x_j, pos_i, pos_j, edge_attr):
@@ -202,10 +234,9 @@ class EnEquivariantNetworkBlock(MessagePassing):
return agg_message, agg_m_ij
def update(self, aggregated_inputs, x, pos, edge_index):
def update(self, aggregated_inputs, x, pos, edge_index, vel):
"""
Update the node features and the node coordinates with the received
messages.
Update node features, positions, and optionally velocities.
:param tuple(torch.Tensor) aggregated_inputs: The messages to be passed.
:param x: The node features.
@@ -213,17 +244,26 @@ class EnEquivariantNetworkBlock(MessagePassing):
:param pos: The euclidean coordinates of the nodes.
:type pos: torch.Tensor | LabelTensor
:param torch.Tensor edge_index: The edge indices.
:param vel: The velocity of the nodes.
:type vel: torch.Tensor | LabelTensor
:return: The updated node features and node positions.
:rtype: tuple(torch.Tensor, torch.Tensor)
:rtype: tuple(torch.Tensor, torch.Tensor) |
tuple(torch.Tensor, torch.Tensor, torch.Tensor)
"""
# aggregated_inputs is tuple (agg_message, agg_m_ij)
agg_message, agg_m_ij = aggregated_inputs
# Degree for normalization of position updates
c = degree(edge_index[1], pos.shape[0]).unsqueeze(-1).clamp(min=1)
# If velocity is used, update it and use it to update positions
if self.use_velocity:
vel = self.update_vel_net(x) * vel
# Update node features with aggregated m_ij
x = self.update_feat_net(torch.cat((x, agg_m_ij), dim=-1))
# Degree for normalization of position updates
c = degree(edge_index[1], pos.shape[0]).unsqueeze(-1).clamp(min=1)
pos = pos + agg_message / c
# Update positions with aggregated messages m_ij and velocities
pos = pos + agg_message / c + (vel if self.use_velocity else 0)
return x, pos
return (x, pos, vel) if self.use_velocity else (x, pos)

188
pina/model/block/message_passing/equivariant_graph_neural_operator_block.py Normal file
View File

@@ -0,0 +1,188 @@
"""Module for the Equivariant Graph Neural Operator block."""
import torch
from ....utils import check_positive_integer
from .en_equivariant_network_block import EnEquivariantNetworkBlock
class EquivariantGraphNeuralOperatorBlock(torch.nn.Module):
"""
A single block of the Equivariant Graph Neural Operator (EGNO).
This block combines a temporal convolution with an equivariant graph neural
network (EGNN) layer. It preserves equivariance while modeling complex
interactions between nodes in a graph over time.
.. seealso::
**Original reference**
Xu, M., Han, J., Lou, A., Kossaifi, J., Ramanathan, A., Azizzadenesheli,
K., Leskovec, J., Ermon, S., Anandkumar, A. (2024).
*Equivariant Graph Neural Operator for Modeling 3D Dynamics*
DOI: `arXiv preprint arXiv:2401.11037.
<https://arxiv.org/abs/2401.11037>`_
"""
def __init__(
self,
node_feature_dim,
edge_feature_dim,
pos_dim,
modes,
hidden_dim=64,
n_message_layers=2,
n_update_layers=2,
activation=torch.nn.SiLU,
aggr="add",
node_dim=-2,
flow="source_to_target",
):
"""
Initialization of the :class:`EquivariantGraphNeuralOperatorBlock`
class.
:param int node_feature_dim: The dimension of the node features.
:param int edge_feature_dim: The dimension of the edge features.
:param int pos_dim: The dimension of the position features.
:param int modes: The number of Fourier modes to use in the temporal
convolution.
:param int hidden_dim: The dimension of the hidden features.
Default is 64.
:param int n_message_layers: The number of layers in the message
network. Default is 2.
:param int n_update_layers: The number of layers in the update network.
Default is 2.
:param torch.nn.Module activation: The activation function.
Default is :class:`torch.nn.SiLU`.
:param str aggr: The aggregation scheme to use for message passing.
Available options are "add", "mean", "min", "max", "mul".
See :class:`torch_geometric.nn.MessagePassing` for more details.
Default is "add".
:param int node_dim: The axis along which to propagate. Default is -2.
:param str flow: The direction of message passing. Available options
are "source_to_target" and "target_to_source".
The "source_to_target" flow means that messages are sent from
the source node to the target node, while the "target_to_source"
flow means that messages are sent from the target node to the
source node. See :class:`torch_geometric.nn.MessagePassing` for more
details. Default is "source_to_target".
:raises AssertionError: If ``modes`` is not a positive integer.
"""
super().__init__()
# Check consistency
check_positive_integer(modes, strict=True)
# Initialization
self.modes = modes
# Temporal convolution weights - real and imaginary parts
self.weight_scalar_r = torch.nn.Parameter(
torch.rand(node_feature_dim, node_feature_dim, modes)
)
self.weight_scalar_i = torch.nn.Parameter(
torch.rand(node_feature_dim, node_feature_dim, modes)
)
self.weight_vector_r = torch.nn.Parameter(torch.rand(2, 2, modes) * 0.1)
self.weight_vector_i = torch.nn.Parameter(torch.rand(2, 2, modes) * 0.1)
# EGNN block
self.egnn = EnEquivariantNetworkBlock(
node_feature_dim=node_feature_dim,
edge_feature_dim=edge_feature_dim,
pos_dim=pos_dim,
use_velocity=True,
hidden_dim=hidden_dim,
n_message_layers=n_message_layers,
n_update_layers=n_update_layers,
activation=activation,
aggr=aggr,
node_dim=node_dim,
flow=flow,
)
def forward(self, x, pos, vel, edge_index, edge_attr=None):
"""
Forward pass of the Equivariant Graph Neural Operator block.
:param x: The node feature tensor of shape
``[time_steps, num_nodes, node_feature_dim]``.
:type x: torch.Tensor | LabelTensor
:param pos: The node position tensor (Euclidean coordinates) of shape
``[time_steps, num_nodes, pos_dim]``.
:type pos: torch.Tensor | LabelTensor
:param vel: The node velocity tensor of shape
``[time_steps, num_nodes, pos_dim]``.
:type vel: torch.Tensor | LabelTensor
:param edge_index: The edge connectivity of shape ``[2, num_edges]``.
:type edge_index: torch.Tensor
:param edge_attr: The edge feature tensor of shape
``[time_steps, num_edges, edge_feature_dim]``. Default is None.
:type edge_attr: torch.Tensor | LabelTensor, optional
:return: The updated node features, positions, and velocities, each with
the same shape as the inputs.
:rtype: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
"""
# Prepare features
center = pos.mean(dim=1, keepdim=True)
vector = torch.stack((pos - center, vel), dim=-1)
# Compute temporal convolution
x = x + self._convolution(
x, "mni, iom -> mno", self.weight_scalar_r, self.weight_scalar_i
)
vector = vector + self._convolution(
vector,
"mndi, iom -> mndo",
self.weight_vector_r,
self.weight_vector_i,
)
# Split position and velocity
pos, vel = vector.unbind(dim=-1)
pos = pos + center
# Reshape to (time * nodes, feature) for egnn
x = x.reshape(-1, x.shape[-1])
pos = pos.reshape(-1, pos.shape[-1])
vel = vel.reshape(-1, vel.shape[-1])
if edge_attr is not None:
edge_attr = edge_attr.reshape(-1, edge_attr.shape[-1])
x, pos, vel = self.egnn(
x=x,
pos=pos,
edge_index=edge_index,
edge_attr=edge_attr,
vel=vel,
)
# Reshape back to (time, nodes, feature)
x = x.reshape(center.shape[0], -1, x.shape[-1])
pos = pos.reshape(center.shape[0], -1, pos.shape[-1])
vel = vel.reshape(center.shape[0], -1, vel.shape[-1])
return x, pos, vel
def _convolution(self, x, einsum_idx, real, img):
"""
Compute the temporal convolution.
:param torch.Tensor x: The input features.
:param str einsum_idx: The indices for the einsum operation.
:param torch.Tensor real: The real part of the convolution weights.
:param torch.Tensor img: The imaginary part of the convolution weights.
:return: The convolved features.
:rtype: torch.Tensor
"""
# Number of modes to use
modes = min(self.modes, (x.shape[0] // 2) + 1)
# Build complex weights
weights = torch.complex(real[..., :modes], img[..., :modes])
# Convolution in Fourier space
fourier = torch.fft.rfftn(x, dim=[0])[:modes]
out = torch.einsum(einsum_idx, fourier, weights)
return torch.fft.irfftn(out, s=x.shape[0], dim=0)

219
pina/model/equivariant_graph_neural_operator.py Normal file
View File

@@ -0,0 +1,219 @@
"""Module for the Equivariant Graph Neural Operator model."""
import torch
from ..utils import check_positive_integer
from .block.message_passing import EquivariantGraphNeuralOperatorBlock
class EquivariantGraphNeuralOperator(torch.nn.Module):
"""
Equivariant Graph Neural Operator (EGNO) for modeling 3D dynamics.
EGNO is a graph-based neural operator that preserves equivariance with
respect to 3D transformations while modeling temporal and spatial
interactions between nodes. It combines:
1. Temporal convolution in the Fourier domain to capture long-range
temporal dependencies efficiently.
2. Equivariant Graph Neural Network (EGNN) layers to model interactions
between nodes while respecting geometric symmetries.
This design allows EGNO to learn complex spatiotemporal dynamics of
physical systems, molecules, or particles while enforcing physically
meaningful constraints.
.. seealso::
**Original reference**
Xu, M., Han, J., Lou, A., Kossaifi, J., Ramanathan, A., Azizzadenesheli,
K., Leskovec, J., Ermon, S., Anandkumar, A. (2024).
*Equivariant Graph Neural Operator for Modeling 3D Dynamics*
DOI: `arXiv preprint arXiv:2401.11037.
<https://arxiv.org/abs/2401.11037>`_
"""
def __init__(
self,
n_egno_layers,
node_feature_dim,
edge_feature_dim,
pos_dim,
modes,
time_steps=2,
hidden_dim=64,
time_emb_dim=16,
max_time_idx=10000,
n_message_layers=2,
n_update_layers=2,
activation=torch.nn.SiLU,
aggr="add",
node_dim=-2,
flow="source_to_target",
):
"""
Initialization of the :class:`EquivariantGraphNeuralOperator` class.
:param int n_egno_layers: The number of EGNO layers.
:param int node_feature_dim: The dimension of the node features in each
EGNO layer.
:param int edge_feature_dim: The dimension of the edge features in each
EGNO layer.
:param int pos_dim: The dimension of the position features in each
EGNO layer.
:param int modes: The number of Fourier modes to use in the temporal
convolution.
:param int time_steps: The number of time steps to consider in the
temporal convolution. Default is 2.
:param int hidden_dim: The dimension of the hidden features in each EGNO
layer. Default is 64.
:param int time_emb_dim: The dimension of the sinusoidal time
embeddings. Default is 16.
:param int max_time_idx: The maximum time index for the sinusoidal
embeddings. Default is 10000.
:param int n_message_layers: The number of layers in the message
network of each EGNO layer. Default is 2.
:param int n_update_layers: The number of layers in the update network
of each EGNO layer. Default is 2.
:param torch.nn.Module activation: The activation function.
Default is :class:`torch.nn.SiLU`.
:param str aggr: The aggregation scheme to use for message passing.
Available options are "add", "mean", "min", "max", "mul".
See :class:`torch_geometric.nn.MessagePassing` for more details.
Default is "add".
:param int node_dim: The axis along which to propagate. Default is -2.
:param str flow: The direction of message passing. Available options
are "source_to_target" and "target_to_source".
The "source_to_target" flow means that messages are sent from
the source node to the target node, while the "target_to_source"
flow means that messages are sent from the target node to the
source node. See :class:`torch_geometric.nn.MessagePassing` for more
details. Default is "source_to_target".
:raises AssertionError: If ``n_egno_layers`` is not a positive integer.
:raises AssertionError: If ``time_emb_dim`` is not a positive integer.
:raises AssertionError: If ``max_time_idx`` is not a positive integer.
:raises AssertionError: If ``time_steps`` is not a positive integer.
"""
super().__init__()
# Check consistency
check_positive_integer(n_egno_layers, strict=True)
check_positive_integer(time_emb_dim, strict=True)
check_positive_integer(max_time_idx, strict=True)
check_positive_integer(time_steps, strict=True)
# Initialize parameters
self.time_steps = time_steps
self.time_emb_dim = time_emb_dim
self.max_time_idx = max_time_idx
# Initialize EGNO layers
self.egno_layers = torch.nn.ModuleList()
for _ in range(n_egno_layers):
self.egno_layers.append(
EquivariantGraphNeuralOperatorBlock(
node_feature_dim=node_feature_dim,
edge_feature_dim=edge_feature_dim,
pos_dim=pos_dim,
modes=modes,
hidden_dim=hidden_dim,
n_message_layers=n_message_layers,
n_update_layers=n_update_layers,
activation=activation,
aggr=aggr,
node_dim=node_dim,
flow=flow,
)
)
# Linear layer to adjust the scalar feature dimension
self.linear = torch.nn.Linear(
node_feature_dim + time_emb_dim, node_feature_dim
)
def forward(self, graph):
"""
Forward pass of the :class:`EquivariantGraphNeuralOperator` class.
:param graph: The input graph object with the following attributes:
- 'x': Node features, shape ``[num_nodes, node_feature_dim]``.
- 'pos': Node positions, shape ``[num_nodes, pos_dim]``.
- 'vel': Node velocities, shape ``[num_nodes, pos_dim]``.
- 'edge_index': Graph connectivity, shape ``[2, num_edges]``.
- 'edge_attr': Edge attrs, shape ``[num_edges, edge_feature_dim]``.
:type graph: Data | Graph
:return: The output graph object with updated node features,
positions, and velocities. The output graph adds to 'x', 'pos',
'vel', and 'edge_attr' the time dimension, resulting in shapes:
- 'x': ``[time_steps, num_nodes, node_feature_dim]``
- 'pos': ``[time_steps, num_nodes, pos_dim]``
- 'vel': ``[time_steps, num_nodes, pos_dim]``
- 'edge_attr': ``[time_steps, num_edges, edge_feature_dim]``
:rtype: Data | Graph
:raises ValueError: If the input graph does not have a 'vel' attribute.
"""
# Check that the graph has the required attributes
if "vel" not in graph:
raise ValueError("The input graph must have a 'vel' attribute.")
# Compute the temporal embedding
emb = self._embedding(torch.arange(self.time_steps)).to(graph.x.device)
emb = emb.unsqueeze(1).repeat(1, graph.x.shape[0], 1)
# Expand dimensions
x = graph.x.unsqueeze(0).repeat(self.time_steps, 1, 1)
x = self.linear(torch.cat((x, emb), dim=-1))
pos = graph.pos.unsqueeze(0).repeat(self.time_steps, 1, 1)
vel = graph.vel.unsqueeze(0).repeat(self.time_steps, 1, 1)
# Manage edge index
offset = torch.arange(self.time_steps).reshape(-1, 1)
offset = offset.to(graph.x.device) * graph.x.shape[0]
src = graph.edge_index[0].unsqueeze(0) + offset
dst = graph.edge_index[1].unsqueeze(0) + offset
edge_index = torch.stack([src, dst], dim=0).reshape(2, -1)
# Manage edge attributes
if graph.edge_attr is not None:
edge_attr = graph.edge_attr.unsqueeze(0)
edge_attr = edge_attr.repeat(self.time_steps, 1, 1)
else:
edge_attr = None
# Iteratively apply EGNO layers
for layer in self.egno_layers:
x, pos, vel = layer(
x=x,
pos=pos,
vel=vel,
edge_index=edge_index,
edge_attr=edge_attr,
)
# Build new graph
new_graph = graph.clone()
new_graph.x, new_graph.pos, new_graph.vel = x, pos, vel
if edge_attr is not None:
new_graph.edge_attr = edge_attr
return new_graph
def _embedding(self, time):
"""
Generate sinusoidal temporal embeddings.
:param torch.Tensor time: The time instances.
:return: The sinusoidal embedding tensor.
:rtype: torch.Tensor
"""
# Compute the sinusoidal embeddings
half_dim = self.time_emb_dim // 2
logs = torch.log(torch.as_tensor(self.max_time_idx)) / (half_dim - 1)
freqs = torch.exp(-torch.arange(half_dim) * logs)
args = torch.as_tensor(time)[:, None] * freqs[None, :]
emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
# Apply padding if the embedding dimension is odd
if self.time_emb_dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1), mode="constant")
return emb