133 lines
3.4 KiB
Python
133 lines
3.4 KiB
Python
import pytest
|
|
import torch
|
|
from pina.model.block.message_passing import EquivariantGraphNeuralOperatorBlock
|
|
|
|
# Data for testing. Shapes: (time, nodes, features)
|
|
x = torch.rand(5, 10, 4)
|
|
pos = torch.rand(5, 10, 3)
|
|
vel = torch.rand(5, 10, 3)
|
|
|
|
# Edge index and attributes
|
|
edge_idx = torch.randint(0, 10, (2, 20))
|
|
edge_attributes = torch.randn(20, 2)
|
|
|
|
|
|
@pytest.mark.parametrize("node_feature_dim", [1, 3])
|
|
@pytest.mark.parametrize("edge_feature_dim", [0, 2])
|
|
@pytest.mark.parametrize("pos_dim", [2, 3])
|
|
@pytest.mark.parametrize("modes", [1, 5])
|
|
def test_constructor(node_feature_dim, edge_feature_dim, pos_dim, modes):
|
|
|
|
EquivariantGraphNeuralOperatorBlock(
|
|
node_feature_dim=node_feature_dim,
|
|
edge_feature_dim=edge_feature_dim,
|
|
pos_dim=pos_dim,
|
|
modes=modes,
|
|
)
|
|
|
|
# Should fail if modes is negative
|
|
with pytest.raises(AssertionError):
|
|
EquivariantGraphNeuralOperatorBlock(
|
|
node_feature_dim=node_feature_dim,
|
|
edge_feature_dim=edge_feature_dim,
|
|
pos_dim=pos_dim,
|
|
modes=-1,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("modes", [1, 5])
|
|
def test_forward(modes):
|
|
|
|
model = EquivariantGraphNeuralOperatorBlock(
|
|
node_feature_dim=x.shape[2],
|
|
edge_feature_dim=edge_attributes.shape[1],
|
|
pos_dim=pos.shape[2],
|
|
modes=modes,
|
|
)
|
|
|
|
output_ = model(
|
|
x=x,
|
|
pos=pos,
|
|
vel=vel,
|
|
edge_index=edge_idx,
|
|
edge_attr=edge_attributes,
|
|
)
|
|
|
|
# Checks on output shapes
|
|
assert output_[0].shape == x.shape
|
|
assert output_[1].shape == pos.shape
|
|
assert output_[2].shape == vel.shape
|
|
|
|
|
|
@pytest.mark.parametrize("modes", [1, 5])
|
|
def test_backward(modes):
|
|
|
|
model = EquivariantGraphNeuralOperatorBlock(
|
|
node_feature_dim=x.shape[2],
|
|
edge_feature_dim=edge_attributes.shape[1],
|
|
pos_dim=pos.shape[2],
|
|
modes=modes,
|
|
)
|
|
|
|
output_ = model(
|
|
x=x.requires_grad_(),
|
|
pos=pos.requires_grad_(),
|
|
vel=vel.requires_grad_(),
|
|
edge_index=edge_idx,
|
|
edge_attr=edge_attributes.requires_grad_(),
|
|
)
|
|
|
|
# Checks on gradients
|
|
loss = sum(torch.mean(output_[i]) for i in range(len(output_)))
|
|
loss.backward()
|
|
assert x.grad.shape == x.shape
|
|
assert pos.grad.shape == pos.shape
|
|
assert vel.grad.shape == vel.shape
|
|
|
|
|
|
@pytest.mark.parametrize("modes", [1, 5])
|
|
def test_equivariance(modes):
|
|
|
|
# Random rotation
|
|
rotation = torch.linalg.qr(torch.rand(pos.shape[2], pos.shape[2])).Q
|
|
if torch.det(rotation) < 0:
|
|
rotation[:, 0] *= -1
|
|
|
|
# Random translation
|
|
translation = torch.rand(1, pos.shape[2])
|
|
|
|
model = EquivariantGraphNeuralOperatorBlock(
|
|
node_feature_dim=x.shape[2],
|
|
edge_feature_dim=edge_attributes.shape[1],
|
|
pos_dim=pos.shape[2],
|
|
modes=modes,
|
|
).eval()
|
|
|
|
# Transform inputs (no translation for velocity)
|
|
pos_rot = pos @ rotation.T + translation
|
|
vel_rot = vel @ rotation.T
|
|
|
|
# Get model outputs
|
|
out1 = model(
|
|
x=x,
|
|
pos=pos,
|
|
vel=vel,
|
|
edge_index=edge_idx,
|
|
edge_attr=edge_attributes,
|
|
)
|
|
out2 = model(
|
|
x=x,
|
|
pos=pos_rot,
|
|
vel=vel_rot,
|
|
edge_index=edge_idx,
|
|
edge_attr=edge_attributes,
|
|
)
|
|
|
|
# Unpack outputs
|
|
h1, pos1, vel1 = out1
|
|
h2, pos2, vel2 = out2
|
|
|
|
assert torch.allclose(pos2, pos1 @ rotation.T + translation, atol=1e-5)
|
|
assert torch.allclose(vel2, vel1 @ rotation.T, atol=1e-5)
|
|
assert torch.allclose(h1, h2, atol=1e-5)
|