diff --git a/pina/model/__init__.py b/pina/model/__init__.py index 0db0acd7..c75f9b65 100644 --- a/pina/model/__init__.py +++ b/pina/model/__init__.py @@ -10,7 +10,7 @@ "AveragingNeuralOperator", "LowRankNeuralOperator", "Spline", - "GNO" + "GraphNeuralOperator" ] from .feed_forward import FeedForward, ResidualFeedForward @@ -21,4 +21,4 @@ from .avno import AveragingNeuralOperator from .lno import LowRankNeuralOperator from .spline import Spline -from .gno import GNO \ No newline at end of file +from .gno import GraphNeuralOperator \ No newline at end of file diff --git a/pina/model/gno.py b/pina/model/gno.py index 3e9af8ae..9a8b8887 100644 --- a/pina/model/gno.py +++ b/pina/model/gno.py @@ -1,6 +1,6 @@ import torch from torch.nn import Tanh -from .layers import GraphIntegralLayer +from .layers import GNOBlock from .base_no import KernelNeuralOperator @@ -46,7 +46,7 @@ def __init__( internal_func = Tanh if shared_weights: - self.layers = GraphIntegralLayer( + self.layers = GNOBlock( width=width, edges_features=edge_features, n_layers=internal_n_layers, @@ -58,7 +58,7 @@ def __init__( self.forward = self.forward_shared else: self.layers = torch.nn.ModuleList( - [GraphIntegralLayer( + [GNOBlock( width=width, edges_features=edge_features, n_layers=internal_n_layers, @@ -101,7 +101,7 @@ def forward_shared(self, x, edge_index, edge_attr): return x -class GNO(KernelNeuralOperator): +class GraphNeuralOperator(KernelNeuralOperator): """ TODO add docstring """ diff --git a/pina/model/layers/__init__.py b/pina/model/layers/__init__.py index 50827dc4..3e3e7168 100644 --- a/pina/model/layers/__init__.py +++ b/pina/model/layers/__init__.py @@ -15,7 +15,7 @@ "AVNOBlock", "LowRankBlock", "RBFBlock", - "GraphIntegralLayer" + "GNOBlock" ] from .convolution_2d import ContinuousConvBlock @@ -32,4 +32,4 @@ from .avno_layer import AVNOBlock from .lowrank_layer import LowRankBlock from .rbf_layer import RBFBlock -from .graph_integral_kernel import GraphIntegralLayer +from .gno_block import GNOBlock diff --git a/pina/model/layers/graph_integral_kernel.py b/pina/model/layers/gno_block.py similarity index 96% rename from pina/model/layers/graph_integral_kernel.py rename to pina/model/layers/gno_block.py index 70d172ca..34929fe8 100644 --- a/pina/model/layers/graph_integral_kernel.py +++ b/pina/model/layers/gno_block.py @@ -2,10 +2,11 @@ from torch_geometric.nn import MessagePassing -class GraphIntegralLayer(MessagePassing): +class GNOBlock(MessagePassing): """ TODO: Add documentation """ + def __init__( self, width, @@ -27,7 +28,7 @@ def __init__( :type n_layers: int """ from pina.model import FeedForward - super(GraphIntegralLayer, self).__init__(aggr='mean') + super(GNOBlock, self).__init__(aggr='mean') self.width = width if layers is None and inner_size is None: inner_size = width diff --git a/tests/test_model/test_gno.py b/tests/test_model/test_gno.py index e07d8453..8fb10d8e 100644 --- a/tests/test_model/test_gno.py +++ b/tests/test_model/test_gno.py @@ -1,7 +1,7 @@ import pytest import torch from pina.graph import KNNGraph -from pina.model import GNO +from pina.model import GraphNeuralOperator from torch_geometric.data import Batch x = [torch.rand(100, 6) for _ in range(10)] @@ -10,7 +10,6 @@ input_ = Batch.from_data_list(graph.data) - @pytest.mark.parametrize( "shared_weights", [ @@ -21,29 +20,29 @@ def test_constructor(shared_weights): lifting_operator = torch.nn.Linear(6, 16) projection_operator = torch.nn.Linear(16, 3) - GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - internal_layers=[16, 16], - shared_weights=shared_weights) + GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + internal_layers=[16, 16], + shared_weights=shared_weights) - GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - inner_size=16, - internal_n_layers=10, - shared_weights=shared_weights) + GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + inner_size=16, + internal_n_layers=10, + shared_weights=shared_weights) int_func = torch.nn.Softplus ext_func = torch.nn.ReLU - GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - internal_n_layers=10, - shared_weights=shared_weights, - internal_func=int_func, - external_func=ext_func) + GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + internal_n_layers=10, + shared_weights=shared_weights, + internal_func=int_func, + external_func=ext_func) @pytest.mark.parametrize( @@ -56,14 +55,15 @@ def test_constructor(shared_weights): def test_forward_1(shared_weights): lifting_operator = torch.nn.Linear(6, 16) projection_operator = torch.nn.Linear(16, 3) - model = GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - internal_layers=[16, 16], - shared_weights=shared_weights) + model = GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + internal_layers=[16, 16], + shared_weights=shared_weights) output_ = model(input_) assert output_.shape == torch.Size([1000, 3]) + @pytest.mark.parametrize( "shared_weights", [ @@ -74,15 +74,16 @@ def test_forward_1(shared_weights): def test_forward_2(shared_weights): lifting_operator = torch.nn.Linear(6, 16) projection_operator = torch.nn.Linear(16, 3) - model = GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - inner_size=32, - internal_n_layers=2, - shared_weights=shared_weights) + model = GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + inner_size=32, + internal_n_layers=2, + shared_weights=shared_weights) output_ = model(input_) assert output_.shape == torch.Size([1000, 3]) + @pytest.mark.parametrize( "shared_weights", [ @@ -93,17 +94,18 @@ def test_forward_2(shared_weights): def test_backward(shared_weights): lifting_operator = torch.nn.Linear(6, 16) projection_operator = torch.nn.Linear(16, 3) - model = GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - internal_layers=[16, 16], - shared_weights=shared_weights) + model = GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + internal_layers=[16, 16], + shared_weights=shared_weights) input_.x.requires_grad = True output_ = model(input_) l = torch.mean(output_) l.backward() assert input_.x.grad.shape == torch.Size([1000, 6]) + @pytest.mark.parametrize( "shared_weights", [ @@ -114,14 +116,14 @@ def test_backward(shared_weights): def test_backward_2(shared_weights): lifting_operator = torch.nn.Linear(6, 16) projection_operator = torch.nn.Linear(16, 3) - model = GNO(lifting_operator=lifting_operator, - projection_operator=projection_operator, - edge_features=3, - inner_size=32, - internal_n_layers=2, - shared_weights=shared_weights) + model = GraphNeuralOperator(lifting_operator=lifting_operator, + projection_operator=projection_operator, + edge_features=3, + inner_size=32, + internal_n_layers=2, + shared_weights=shared_weights) input_.x.requires_grad = True output_ = model(input_) l = torch.mean(output_) l.backward() - assert input_.x.grad.shape == torch.Size([1000, 6]) \ No newline at end of file + assert input_.x.grad.shape == torch.Size([1000, 6])