7bf7d34d0f
* Fix adaptive refinement (#571) --------- Co-authored-by: Dario Coscia <93731561+dario-coscia@users.noreply.github.com> * Remove collector * Fixes * Fixes * rm unnecessary comment * fix advection (#581) * Fix tutorial .html link (#580) * fix problem data collection for v0.1 (#584) * Message Passing Module (#516) * add deep tensor network block * add interaction network block * add radial field network block * add schnet block * add equivariant network block * fix + tests + doc files * fix egnn + equivariance/invariance tests Co-authored-by: Dario Coscia <dariocos99@gmail.com> --------- Co-authored-by: giovanni <giovanni.canali98@yahoo.it> Co-authored-by: AleDinve <giuseppealessio.d@student.unisi.it> * add type checker (#527) --------- Co-authored-by: Filippo Olivo <filippo@filippoolivo.com> Co-authored-by: Giovanni Canali <115086358+GiovanniCanali@users.noreply.github.com> Co-authored-by: giovanni <giovanni.canali98@yahoo.it> Co-authored-by: AleDinve <giuseppealessio.d@student.unisi.it>
85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
import pytest
|
|
import torch
|
|
from pina.model.block.message_passing import InteractionNetworkBlock
|
|
|
|
# Data for testing
|
|
x = torch.rand(10, 3)
|
|
edge_index = torch.randint(0, 10, (2, 20))
|
|
edge_attr = torch.randn(20, 2)
|
|
|
|
|
|
@pytest.mark.parametrize("node_feature_dim", [1, 3])
|
|
@pytest.mark.parametrize("edge_feature_dim", [0, 2])
|
|
def test_constructor(node_feature_dim, edge_feature_dim):
|
|
|
|
InteractionNetworkBlock(
|
|
node_feature_dim=node_feature_dim,
|
|
edge_feature_dim=edge_feature_dim,
|
|
hidden_dim=64,
|
|
n_message_layers=2,
|
|
n_update_layers=2,
|
|
)
|
|
|
|
# Should fail if node_feature_dim is negative
|
|
with pytest.raises(AssertionError):
|
|
InteractionNetworkBlock(node_feature_dim=-1)
|
|
|
|
# Should fail if edge_feature_dim is negative
|
|
with pytest.raises(AssertionError):
|
|
InteractionNetworkBlock(node_feature_dim=3, edge_feature_dim=-1)
|
|
|
|
# Should fail if hidden_dim is negative
|
|
with pytest.raises(AssertionError):
|
|
InteractionNetworkBlock(node_feature_dim=3, hidden_dim=-1)
|
|
|
|
# Should fail if n_message_layers is negative
|
|
with pytest.raises(AssertionError):
|
|
InteractionNetworkBlock(node_feature_dim=3, n_message_layers=-1)
|
|
|
|
# Should fail if n_update_layers is negative
|
|
with pytest.raises(AssertionError):
|
|
InteractionNetworkBlock(node_feature_dim=3, n_update_layers=-1)
|
|
|
|
|
|
@pytest.mark.parametrize("edge_feature_dim", [0, 2])
|
|
def test_forward(edge_feature_dim):
|
|
|
|
model = InteractionNetworkBlock(
|
|
node_feature_dim=x.shape[1],
|
|
edge_feature_dim=edge_feature_dim,
|
|
hidden_dim=64,
|
|
n_message_layers=2,
|
|
n_update_layers=2,
|
|
)
|
|
|
|
if edge_feature_dim == 0:
|
|
output_ = model(edge_index=edge_index, x=x)
|
|
else:
|
|
output_ = model(edge_index=edge_index, x=x, edge_attr=edge_attr)
|
|
assert output_.shape == x.shape
|
|
|
|
|
|
@pytest.mark.parametrize("edge_feature_dim", [0, 2])
|
|
def test_backward(edge_feature_dim):
|
|
|
|
model = InteractionNetworkBlock(
|
|
node_feature_dim=x.shape[1],
|
|
edge_feature_dim=edge_feature_dim,
|
|
hidden_dim=64,
|
|
n_message_layers=2,
|
|
n_update_layers=2,
|
|
)
|
|
|
|
if edge_feature_dim == 0:
|
|
output_ = model(edge_index=edge_index, x=x.requires_grad_())
|
|
else:
|
|
output_ = model(
|
|
edge_index=edge_index,
|
|
x=x.requires_grad_(),
|
|
edge_attr=edge_attr.requires_grad_(),
|
|
)
|
|
|
|
loss = torch.mean(output_)
|
|
loss.backward()
|
|
assert x.grad.shape == x.shape
|