From 1f32acea16dd688afe9251f31b014be1e11ca839 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Thu, 16 Jan 2025 17:10:38 +0100
Subject: [PATCH 1/8] Add Graph class and tests for Graph and Collector +
 Dataloader refactoring

---
 pina/graph.py           | 324 +++++++++++++++++++++++++++-------------
 tests/test_collector.py | 125 ++++++++++++++++
 tests/test_graph.py     | 143 ++++++++++++++++++
 3 files changed, 491 insertions(+), 101 deletions(-)
 create mode 100644 tests/test_collector.py
 create mode 100644 tests/test_graph.py

diff --git a/pina/graph.py b/pina/graph.py
index bde5bbf50..22d70825f 100644
--- a/pina/graph.py
+++ b/pina/graph.py
@@ -1,118 +1,240 @@
-""" Module for Loss class """
+from logging import warning
 
-import logging
-from torch_geometric.nn import MessagePassing, InstanceNorm, radius_graph
-from torch_geometric.data import Data
 import torch
+from . import LabelTensor
+from torch_geometric.data import Data
+from torch_geometric.utils import to_undirected
+
 
 class Graph:
     """
-    PINA Graph managing the PyG Data class.
+    Class for the graph construction.
     """
-    def __init__(self, data):
-        self.data = data
-            
-    @staticmethod
-    def _build_triangulation(**kwargs):
-        logging.debug("Creating graph with triangulation mode.")
-
-        # check for mandatory arguments
-        if "nodes_coordinates" not in kwargs:
-            raise ValueError("Nodes coordinates must be provided in the kwargs.")
-        if "nodes_data" not in kwargs:
-            raise ValueError("Nodes data must be provided in the kwargs.")
-        if "triangles" not in kwargs:
-            raise ValueError("Triangles must be provided in the kwargs.")
-
-        nodes_coordinates = kwargs["nodes_coordinates"]
-        nodes_data = kwargs["nodes_data"]
-        triangles = kwargs["triangles"]
-
-
-
-        def less_first(a, b):
-            return [a, b] if a < b else [b, a]
-
-        list_of_edges = []
-
-        for triangle in triangles:
-            for e1, e2 in [[0, 1], [1, 2], [2, 0]]:
-                list_of_edges.append(less_first(triangle[e1],triangle[e2]))
-
-        array_of_edges = torch.unique(torch.Tensor(list_of_edges), dim=0) # remove duplicates
-        array_of_edges = array_of_edges.t().contiguous()
-        print(array_of_edges)
-
-        # list_of_lengths = []
 
-        # for p1,p2 in array_of_edges:
-        #     x1, y1 = tri.points[p1]
-        #     x2, y2 = tri.points[p2]
-        #     list_of_lengths.append((x1-x2)**2 + (y1-y2)**2)
-
-        # array_of_lengths = np.sqrt(np.array(list_of_lengths))
+    def __init__(self,
+                 x,
+                 pos,
+                 edge_index,
+                 edge_attr=None,
+                 build_edge_attr=False,
+                 undirected=False,
+                 additional_params=None):
+        """
+        Constructor for the Graph class.
+        :param x: The node features.
+        :type x: torch.Tensor or list[torch.Tensor]
+        :param pos: The node positions.
+        :type pos: torch.Tensor or list[torch.Tensor]
+        :param edge_index: The edge index.
+        :type edge_index: torch.Tensor or list[torch.Tensor]
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor or list[torch.Tensor]
+        :param build_edge_attr: Whether to build the edge attributes.
+        :type build_edge_attr: bool
+        :param undirected: Whether to build an undirected graph.
+        :type undirected: bool
+        :param additional_params: Additional parameters.
+        :type additional_params: dict
+        """
+        self.data = []
+        x, pos, edge_index = Graph._check_input_consistency(x, pos, edge_index)
+
+        # Check input dimension consistency and store the number of graphs
+        data_len = self._check_len_consistency(x, pos)
+
+        # Initialize additional_parameters (if present)
+        if additional_params is not None:
+            if not isinstance(additional_params, dict):
+                raise TypeError("additional_params must be a dictionary.")
+            for param, val in additional_params.items():
+                # Check if the values are tensors or lists of tensors
+                if isinstance(val, torch.Tensor):
+                    # If the tensor is 3D, we split it into a list of 2D tensors
+                    # In this case there must be a additional parameter for each
+                    # node
+                    if val.ndim == 3:
+                        additional_params[param] = [val[i] for i in
+                                                    range(val.shape[0])]
+                    # If the tensor is 2D, we replicate it for each node
+                    elif val.ndim == 2:
+                        additional_params[param] = [val] * data_len
+                    # If the tensor is 1D, each graph has a scalar values as
+                    # additional parameter
+                    if val.ndim == 1:
+                        if len(val) == data_len:
+                            additional_params[param] = [val[i] for i in
+                                                        range(len(val))]
+                        else:
+                            additional_params[param] = [val for _ in
+                                                        range(data_len)]
+                elif not isinstance(val, list):
+                    raise TypeError("additional_params values must be tensors "
+                                    "or lists of tensors.")
+        else:
+            additional_params = {}
+
+        # Make the graphs undirected
+        if undirected:
+            if isinstance(edge_index, list):
+                edge_index = [to_undirected(e) for e in edge_index]
+            else:
+                edge_index = to_undirected(edge_index)
+
+        if build_edge_attr:
+            if edge_attr is not None:
+                warning("Edge attributes are provided, build_edge_attr is set "
+                        "to True. The provided edge attributes will be ignored.")
+            edge_attr = self._build_edge_attr(pos, edge_index)
+
+        # Prepare internal lists to create a graph list (same positions but
+        # different node features)
+        if isinstance(x, list) and isinstance(pos,
+                                              (torch.Tensor, LabelTensor)):
+            # Replicate the positions, edge_index and edge_attr
+            pos, edge_index = [pos] * data_len, [edge_index] * data_len
+            if edge_attr is not None:
+                edge_attr = [edge_attr] * data_len
+        # Prepare internal lists to create a list containing a single graph
+        elif isinstance(x, (torch.Tensor, LabelTensor)) and isinstance(pos, (
+                torch.Tensor, LabelTensor)):
+            # Encapsulate the input tensors into lists
+            x, pos, edge_index = [x], [pos], [edge_index]
+            if isinstance(edge_attr, torch.Tensor):
+                edge_attr = [edge_attr]
+        # Prepare internal lists to create a list of graphs (same node features
+        # but different positions)
+        elif (isinstance(x, (torch.Tensor, LabelTensor))
+              and isinstance(pos, list)):
+            # Replicate the node features
+            x = [x] * data_len
+        elif not isinstance(x, list) and not isinstance(pos, list):
+            raise TypeError("x and pos must be lists or tensors.")
+
+        # Perform the graph construction
+        self._build_graph_list(x, pos, edge_index, edge_attr, additional_params)
+
+    def _build_graph_list(self, x, pos, edge_index, edge_attr,
+                          additional_params):
+        for i, (x_, pos_, edge_index_) in enumerate(zip(x, pos, edge_index)):
+            if isinstance(x_, LabelTensor):
+                x_ = x_.tensor
+            add_params_local = {k: v[i] for k, v in additional_params.items()}
+            if edge_attr is not None:
+
+                self.data.append(Data(x=x_, pos=pos_, edge_index=edge_index_,
+                                      edge_attr=edge_attr[i],
+                                      **add_params_local))
+            else:
+                self.data.append(Data(x=x_, pos=pos_, edge_index=edge_index_,
+                                      **add_params_local))
 
-        # return array_of_edges, array_of_lengths
+    @staticmethod
+    def _build_edge_attr(pos, edge_index):
+        if isinstance(pos, torch.Tensor):
+            pos = [pos]
+            edge_index = [edge_index]
+        distance = [pos_[edge_index_[0]] - pos_[edge_index_[1]] ** 2 for
+                    pos_, edge_index_ in zip(pos, edge_index)]
+        return distance
 
-        return Data(
-            x=nodes_data,
-            pos=nodes_coordinates.T,
-            
-            edge_index=array_of_edges,
-        )
+    @staticmethod
+    def _check_len_consistency(x, pos):
+        if isinstance(x, list) and isinstance(pos, list):
+            if len(x) != len(pos):
+                raise ValueError("x and pos must have the same length.")
+            return max(len(x), len(pos))
+        elif isinstance(x, list) and not isinstance(pos, list):
+            return len(x)
+        elif not isinstance(x, list) and isinstance(pos, list):
+            return len(pos)
+        else:
+            return 1
 
     @staticmethod
-    def _build_radius(**kwargs):
-        logging.debug("Creating graph with radius mode.")
-
-        # check for mandatory arguments
-        if "nodes_coordinates" not in kwargs:
-            raise ValueError("Nodes coordinates must be provided in the kwargs.")
-        if "nodes_data" not in kwargs:
-            raise ValueError("Nodes data must be provided in the kwargs.")
-        if "radius" not in kwargs:
-            raise ValueError("Radius must be provided in the kwargs.")
-
-        nodes_coordinates = kwargs["nodes_coordinates"]
-        nodes_data = kwargs["nodes_data"]
-        radius = kwargs["radius"]
-
-        edges_data = kwargs.get("edge_data", None) 
-        loop = kwargs.get("loop", False)
-        batch = kwargs.get("batch", None)
-
-        logging.debug(f"radius: {radius}, loop: {loop}, "
-                        f"batch: {batch}")
-
-        edge_index = radius_graph(
-            x=nodes_coordinates.tensor,
-            r=radius,
-            loop=loop,
-            batch=batch,
-        )
-
-        logging.debug(f"edge_index computed")
-        return Data(
-            x=nodes_data.tensor,
-            pos=nodes_coordinates.tensor,
-            edge_index=edge_index,
-            edge_attr=edges_data,
-        )
+    def _check_input_consistency(x, pos, edge_index=None):
+        # If x is a 3D tensor, we split it into a list of 2D tensors
+        if isinstance(x, torch.Tensor) and x.ndim == 3:
+            x = [x[i] for i in range(x.shape[0])]
+
+        # If pos is a 3D tensor, we split it into a list of 2D tensors
+        if isinstance(pos, torch.Tensor) and pos.ndim == 3:
+            pos = [pos[i] for i in range(pos.shape[0])]
+
+        # If edge_index is a 3D tensor, we split it into a list of 2D tensors
+        if isinstance(edge_index, torch.Tensor) and edge_index.ndim == 3:
+            edge_index = [edge_index[i] for i in range(edge_index.shape[0])]
+        return x, pos, edge_index
+
+
+class RadiusGraph(Graph):
+    def __init__(self,
+                 x,
+                 pos,
+                 r,
+                 build_edge_attr=False,
+                 undirected=False,
+                 additional_params=None, ):
+        x, pos, edge_index = Graph._check_input_consistency(x, pos)
+
+        if isinstance(pos, (torch.Tensor, LabelTensor)):
+            edge_index = RadiusGraph._radius_graph(pos, r)
+        else:
+            edge_index = [RadiusGraph._radius_graph(p, r) for p in pos]
+
+        super().__init__(x=x, pos=pos, edge_index=edge_index,
+                         build_edge_attr=build_edge_attr,
+                         undirected=undirected,
+                         additional_params=additional_params)
 
     @staticmethod
-    def build(mode, **kwargs):
+    def _radius_graph(points, r):
         """
-        Constructor for the `Graph` class.
+        Implementation of the radius graph construction.
+        :param points: The input points.
+        :type points: torch.Tensor
+        :param r: The radius.
+        :type r: float
+        :return: The edge index.
+        :rtype: torch.Tensor
         """
-        if mode == "radius":
-            graph = Graph._build_radius(**kwargs)
-        elif mode == "triangulation":
-            graph = Graph._build_triangulation(**kwargs)
+        dist = torch.cdist(points, points, p=2)
+        edge_index = torch.nonzero(dist <= r, as_tuple=False).t()
+        return edge_index
+
+
+class KNNGraph(Graph):
+    def __init__(self,
+                 x,
+                 pos,
+                 k,
+                 build_edge_attr=False,
+                 undirected=False,
+                 additional_params=None,
+                 ):
+        x, pos, edge_index = Graph._check_input_consistency(x, pos)
+        if isinstance(pos, (torch.Tensor, LabelTensor)):
+            edge_index = KNNGraph._knn_graph(pos, k)
         else:
-            raise ValueError(f"Mode {mode} not recognized")
-        
-        return Graph(graph)
-
+            edge_index = [KNNGraph._knn_graph(p, k) for p in pos]
+        super().__init__(x=x, pos=pos, edge_index=edge_index,
+                         build_edge_attr=build_edge_attr,
+                         undirected=undirected,
+                         additional_params=additional_params)
 
-    def __repr__(self):
-        return f"Graph(data={self.data})"
\ No newline at end of file
+    @staticmethod
+    def _knn_graph(points, k):
+        """
+        Implementation of the k-nearest neighbors graph construction.
+        :param points: The input points.
+        :type points: torch.Tensor
+        :param k: The number of nearest neighbors.
+        :type k: int
+        :return: The edge index.
+        :rtype: torch.Tensor
+        """
+        dist = torch.cdist(points, points, p=2)
+        knn_indices = torch.topk(dist, k=k + 1, largest=False).indices[:, 1:]
+        row = torch.arange(points.size(0)).repeat_interleave(k)
+        col = knn_indices.flatten()
+        edge_index = torch.stack([row, col], dim=0)
+        return edge_index
diff --git a/tests/test_collector.py b/tests/test_collector.py
new file mode 100644
index 000000000..f25a01af4
--- /dev/null
+++ b/tests/test_collector.py
@@ -0,0 +1,125 @@
+import torch
+import pytest
+from pina import Condition, LabelTensor, Graph
+from pina.condition import InputOutputPointsCondition, DomainEquationCondition
+from pina.graph import RadiusGraph
+from pina.problem import AbstractProblem, SpatialProblem
+from pina.domain import CartesianDomain
+from pina.equation.equation import Equation
+from pina.equation.equation_factory import FixedValue
+from pina.operators import laplacian
+
+def test_supervised_tensor_collector():
+    class SupervisedProblem(AbstractProblem):
+        output_variables = None
+        conditions = {
+            'data1' : Condition(input_points=torch.rand((10,2)),
+                                output_points=torch.rand((10,2))),
+            'data2' : Condition(input_points=torch.rand((20,2)),
+                                output_points=torch.rand((20,2))),
+            'data3' : Condition(input_points=torch.rand((30,2)),
+                                output_points=torch.rand((30,2))),
+        }
+    problem = SupervisedProblem()
+    collector = problem.collector
+    for v in collector.conditions_name.values():
+        assert v in problem.conditions.keys()
+    assert all(collector._is_conditions_ready.values())
+
+def test_pinn_collector():
+    def laplace_equation(input_, output_):
+        force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+                      torch.sin(input_.extract(['y']) * torch.pi))
+        delta_u = laplacian(output_.extract(['u']), input_)
+        return delta_u - force_term
+
+    my_laplace = Equation(laplace_equation)
+    in_ = LabelTensor(torch.tensor([[0., 1.]], requires_grad=True), ['x', 'y'])
+    out_ = LabelTensor(torch.tensor([[0.]], requires_grad=True), ['u'])
+    class Poisson(SpatialProblem):
+        output_variables = ['u']
+        spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
+
+        conditions = {
+            'gamma1':
+                Condition(domain=CartesianDomain({
+                    'x': [0, 1],
+                    'y': 1
+                }),
+                    equation=FixedValue(0.0)),
+            'gamma2':
+                Condition(domain=CartesianDomain({
+                    'x': [0, 1],
+                    'y': 0
+                }),
+                    equation=FixedValue(0.0)),
+            'gamma3':
+                Condition(domain=CartesianDomain({
+                    'x': 1,
+                    'y': [0, 1]
+                }),
+                    equation=FixedValue(0.0)),
+            'gamma4':
+                Condition(domain=CartesianDomain({
+                    'x': 0,
+                    'y': [0, 1]
+                }),
+                    equation=FixedValue(0.0)),
+            'D':
+                Condition(domain=CartesianDomain({
+                    'x': [0, 1],
+                    'y': [0, 1]
+                }),
+                    equation=my_laplace),
+            'data':
+                Condition(input_points=in_, output_points=out_)
+        }
+
+        def poisson_sol(self, pts):
+            return -(torch.sin(pts.extract(['x']) * torch.pi) *
+                     torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi**2)
+
+        truth_solution = poisson_sol
+
+    problem = Poisson()
+    collector = problem.collector
+    for k,v in problem.conditions.items():
+        if isinstance(v, InputOutputPointsCondition):
+            assert collector._is_conditions_ready[k] == True
+            assert list(collector.data_collections[k].keys()) == ['input_points', 'output_points']
+        else:
+            assert collector._is_conditions_ready[k] == False
+            assert collector.data_collections[k] == {}
+
+    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+    problem.discretise_domain(10, 'grid', locations=boundaries)
+    problem.discretise_domain(10, 'grid', locations='D')
+    assert all(collector._is_conditions_ready.values())
+    for k,v in problem.conditions.items():
+        if isinstance(v, DomainEquationCondition):
+            assert list(collector.data_collections[k].keys()) == ['input_points', 'equation']
+
+
+def test_supervised_graph_collector():
+    pos = torch.rand((100,3))
+    x = [torch.rand((100,3)) for _ in range(10)]
+    graph_list_1 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
+    out_1 = torch.rand((10,100,3))
+    pos = torch.rand((50,3))
+    x = [torch.rand((50,3)) for _ in range(10)]
+    graph_list_2 = RadiusGraph(pos=pos, x=x, build_edge_attr=True, r=.4)
+    out_2 = torch.rand((10,50,3))
+    class SupervisedProblem(AbstractProblem):
+        output_variables = None
+        conditions = {
+            'data1' : Condition(input_points=graph_list_1,
+                                output_points=out_1),
+            'data2' : Condition(input_points=graph_list_2,
+                                output_points=out_2),
+        }
+
+    problem = SupervisedProblem()
+    collector = problem.collector
+    assert all(collector._is_conditions_ready.values())
+    for v in collector.conditions_name.values():
+        assert v in problem.conditions.keys()
diff --git a/tests/test_graph.py b/tests/test_graph.py
new file mode 100644
index 000000000..6886dd923
--- /dev/null
+++ b/tests/test_graph.py
@@ -0,0 +1,143 @@
+import pytest
+import torch
+from pina import Graph
+from pina.graph import RadiusGraph, KNNGraph
+
+
+@pytest.mark.parametrize(
+    "x, pos",
+    [
+        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
+        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
+        (torch.rand(3,10,2), torch.rand(3,10,3)),
+        (torch.rand(3,10,2), torch.rand(3,10,3)),
+    ]
+)
+def test_build_multiple_graph_multiple_val(x, pos):
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k = 3)
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+
+def test_build_single_graph_multiple_val():
+    x = torch.rand(10, 2)
+    pos = torch.rand(10, 3)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 1
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert len(graph.data) == 1
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+    x = torch.rand(10, 2)
+    pos = torch.rand(10, 3)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    assert len(graph.data) == 1
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    data = graph.data
+    assert len(graph.data) == 1
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+
+@pytest.mark.parametrize(
+    "pos",
+    [
+        ([torch.rand(10, 3) for _ in range(3)]),
+        ([torch.rand(10, 3) for _ in range(3)]),
+        (torch.rand(3, 10, 3)),
+        (torch.rand(3, 10, 3))
+    ]
+)
+def test_build_single_graph_single_val(pos):
+    x = torch.rand(10, 2)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+    x = torch.rand(10, 2)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=False, k=3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+def test_additional_parameters_1():
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+    additional_parameters = {'y': torch.ones(3)}
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
+                  additional_params=additional_parameters)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(hasattr(d, 'y') for d in data)
+    assert all(d_.y == 1 for d_ in data)
+
+@pytest.mark.parametrize(
+    "additional_parameters",
+    [
+        ({'y': torch.rand(3,10,1)}),
+        ({'y': [torch.rand(10,1) for _ in range(3)]}),
+    ]
+)
+def test_additional_parameters_2(additional_parameters):
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
+                  additional_params=additional_parameters)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(hasattr(d, 'y') for d in data)
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))

From 05faaaacc2c9c394416386bd9ce8c5be5c28429f Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 4 Feb 2025 18:11:06 +0100
Subject: [PATCH 2/8] Implement Graph Neural Operator #231

---
 pina/model/__init__.py                     |   2 +
 pina/model/gno.py                          | 173 +++++++++++++++++++++
 pina/model/layers/__init__.py              |   2 +
 pina/model/layers/graph_integral_kernel.py |  82 ++++++++++
 4 files changed, 259 insertions(+)
 create mode 100644 pina/model/gno.py
 create mode 100644 pina/model/layers/graph_integral_kernel.py

diff --git a/pina/model/__init__.py b/pina/model/__init__.py
index 3224d0af3..0db0acd71 100644
--- a/pina/model/__init__.py
+++ b/pina/model/__init__.py
@@ -10,6 +10,7 @@
     "AveragingNeuralOperator",
     "LowRankNeuralOperator",
     "Spline",
+    "GNO"
 ]
 
 from .feed_forward import FeedForward, ResidualFeedForward
@@ -20,3 +21,4 @@
 from .avno import AveragingNeuralOperator
 from .lno import LowRankNeuralOperator
 from .spline import Spline
+from .gno import GNO
\ No newline at end of file
diff --git a/pina/model/gno.py b/pina/model/gno.py
new file mode 100644
index 000000000..991ca397b
--- /dev/null
+++ b/pina/model/gno.py
@@ -0,0 +1,173 @@
+import torch
+from torch.nn import Tanh
+from .layers import GraphIntegralLayer
+from .base_no import KernelNeuralOperator
+
+
+class GraphNeuralKernel(torch.nn.Module):
+    """
+    TODO add docstring
+    """
+
+    def __init__(
+            self,
+            width,
+            edge_features,
+            n_layers=2,
+            internal_n_layers=0,
+            internal_layers=None,
+            internal_func=None,
+            external_func=None,
+            shared_weights=False
+    ):
+        """
+        The Graph Neural Kernel constructor.
+
+        :param width: The width of the kernel.
+        :type width: int
+        :param edge_features: The number of edge features.
+        :type edge_features: int
+        :param n_layers: The number of kernel layers.
+        :type n_layers: int
+        :param internal_n_layers: The number of layers the FF Neural Network internal to each Kernel Layer.
+        :type internal_n_layers: int
+        :param internal_layers: Number of neurons of hidden layers(s) in the FF Neural Network inside for each Kernel Layer.
+        :type internal_layers: list | tuple
+        :param internal_func: The activation function used inside the computation of the representation of the edge features in the Graph Integral Layer.
+        :param external_func: The activation function applied to the output of the Graph Integral Layer.
+        :type external_func: torch.nn.Module
+        :param shared_weights: If ``True`` the weights of the Graph Integral Layers are shared.
+        """
+        super().__init__()
+        if external_func is None:
+            external_func = Tanh
+        if internal_func is None:
+            internal_func = Tanh
+
+        if shared_weights:
+            self.layers = GraphIntegralLayer(
+                width=width,
+                edges_features=edge_features,
+                n_layers=internal_n_layers,
+                layers=internal_layers,
+                internal_func=internal_func,
+                external_func=external_func)
+            self.n_layers = n_layers
+            self.forward = self.forward_shared
+        else:
+            self.layers = torch.nn.ModuleList(
+                [GraphIntegralLayer(
+                    width=width,
+                    edges_features=edge_features,
+                    n_layers=internal_n_layers,
+                    layers=internal_layers,
+                    internal_func=internal_func,
+                    external_func=external_func
+                )
+                    for _ in range(n_layers)]
+            )
+
+    def forward(self, x, edge_index, edge_attr):
+        """
+        The forward pass of the Graph Neural Kernel used when the weights are not shared.
+
+        :param x: The input batch.
+        :type x: torch.Tensor
+        :param edge_index: The edge index.
+        :type edge_index: torch.Tensor
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor
+        """
+        for layer in self.layers:
+            x = layer(x, edge_index, edge_attr)
+        return x
+
+    def forward_shared(self, x, edge_index, edge_attr):
+        """
+        The forward pass of the Graph Neural Kernel used when the weights are shared.
+
+        :param x: The input batch.
+        :type x: torch.Tensor
+        :param edge_index: The edge index.
+        :type edge_index: torch.Tensor
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor
+        """
+        for _ in range(self.n_layers):
+            x = self.layers(x, edge_index, edge_attr)
+        return x
+
+
+class GNO(KernelNeuralOperator):
+    """
+    TODO add docstring
+    """
+
+    def __init__(
+            self,
+            lifting_operator,
+            projection_operator,
+            edge_features,
+            n_layers=10,
+            internal_n_layers=0,
+            inner_size=None,
+            internal_layers=None,
+            internal_func=None,
+            external_func=None,
+            shared_weights=True
+    ):
+        """
+        The Graph Neural Operator constructor.
+
+        :param lifting_operator: The lifting operator mapping the node features to its hidden dimension.
+        :type lifting_operator: torch.nn.Module
+        :param projection_operator: The projection operator mapping the hidden representation of the nodes features to the output function.
+        :type projection_operator: torch.nn.Module
+        :param edge_features: Number of edge features.
+        :type edge_features: int
+        :param n_layers: The number of kernel layers.
+        :type n_layers: int
+        :param internal_n_layers: The number of layers the Feed Forward Neural Network internal to each Kernel Layer.
+        :type internal_n_layers: int
+        :param internal_layers: Number of neurons of hidden layers(s) in the FF Neural Network inside for each Kernel Layer.
+        :type internal_layers: list | tuple
+        :param internal_func: The activation function used inside the computation of the representation of the edge features in the Graph Integral Layer.
+        :type internal_func: torch.nn.Module
+        :param external_func: The activation function applied to the output of the Graph Integral Kernel.
+        :type external_func: torch.nn.Module
+        :param shared_weights: If ``True`` the weights of the Graph Integral Layers are shared.
+        :type shared_weights: bool
+        """
+
+        if internal_func is None:
+            internal_func = Tanh
+        if external_func is None:
+            external_func = Tanh
+
+        super().__init__(
+            lifting_operator=lifting_operator,
+            integral_kernels=GraphNeuralKernel(
+                width=lifting_operator.out_features,
+                edge_features=edge_features,
+                internal_n_layers=internal_n_layers,
+                internal_layers=internal_layers,
+                external_func=external_func,
+                internal_func=internal_func,
+                n_layers=n_layers,
+                shared_weights=shared_weights
+            ),
+            projection_operator=projection_operator
+        )
+
+    def forward(self, x):
+        """
+        The forward pass of the Graph Neural Operator.
+
+        :param x: The input batch.
+        :type x: torch_geometric.data.Batch
+        """
+        x, edge_index, edge_attr = x.x, x.edge_index, x.edge_attr
+        x = self.lifting_operator(x)
+        x = self.integral_kernels(x, edge_index, edge_attr)
+        x = self.projection_operator(x)
+        return x
diff --git a/pina/model/layers/__init__.py b/pina/model/layers/__init__.py
index 5108522c5..50827dc4b 100644
--- a/pina/model/layers/__init__.py
+++ b/pina/model/layers/__init__.py
@@ -15,6 +15,7 @@
     "AVNOBlock",
     "LowRankBlock",
     "RBFBlock",
+    "GraphIntegralLayer"
 ]
 
 from .convolution_2d import ContinuousConvBlock
@@ -31,3 +32,4 @@
 from .avno_layer import AVNOBlock
 from .lowrank_layer import LowRankBlock
 from .rbf_layer import RBFBlock
+from .graph_integral_kernel import GraphIntegralLayer
diff --git a/pina/model/layers/graph_integral_kernel.py b/pina/model/layers/graph_integral_kernel.py
new file mode 100644
index 000000000..713b0d78d
--- /dev/null
+++ b/pina/model/layers/graph_integral_kernel.py
@@ -0,0 +1,82 @@
+import torch
+from torch_geometric.nn import MessagePassing
+
+
+class GraphIntegralLayer(MessagePassing):
+    """
+    TODO: Add documentation
+    """
+    def __init__(
+            self,
+            width,
+            edges_features,
+            n_layers=0,
+            layers=None,
+            internal_func=None,
+            external_func=None
+    ):
+        """
+        Initialize the Graph Integral Layer, inheriting from the MessagePassing class of PyTorch Geometric.
+
+        :param width: The width of the hidden representation of the nodes features
+        :type width: int
+        :param edges_features: The number of edge features.
+        :type edges_features: int
+        :param n_layers: The number of layers in the Feed Forward Neural Network used to compute the representation of the edges features.
+        :type n_layers: int
+        """
+        from pina.model import FeedForward
+        super(GraphIntegralLayer, self).__init__(aggr='mean')
+        self.width = width
+        self.dense = FeedForward(input_dimensions=edges_features,
+                                 output_dimensions=width ** 2,
+                                 n_layers=n_layers,
+                                 layers=layers,
+                                 func=internal_func)
+        self.W = torch.nn.Linear(width, width)
+        self.func = external_func()
+
+    def message(self, x_j, edge_attr):
+        """
+        This function computes the message passed between the nodes of the graph. Overwrite the default message function defined in the MessagePassing class.
+
+        :param x_j: The node features of the neighboring.
+        :type x_j: torch.Tensor
+        :param edge_attr: The edge features.
+        :type edge_attr: torch.Tensor
+        :return: The message passed between the nodes of the graph.
+        :rtype: torch.Tensor
+        """
+        x = self.dense(edge_attr).view(-1, self.width, self.width)
+        return torch.einsum('bij,bj->bi', x, x_j)
+
+    def update(self, aggr_out, x):
+        """
+        This function updates the node features of the graph. Overwrite the default update function defined in the MessagePassing class.
+
+        :param aggr_out: The aggregated messages.
+        :type aggr_out: torch.Tensor
+        :param x: The node features.
+        :type x: torch.Tensor
+        :return: The updated node features.
+        :rtype: torch.Tensor
+        """
+        aggr_out = aggr_out + self.W(x)
+        return aggr_out
+
+    def forward(self, x, edge_index, edge_attr):
+        """
+        The forward pass of the Graph Integral Layer.
+
+        :param x: Node features.
+        :type x: torch.Tensor
+        :param edge_index: Edge index.
+        :type edge_index: torch.Tensor
+        :param edge_attr: Edge features.
+        :type edge_attr: torch.Tensor
+        :return: Output of a single iteration over the Graph Integral Layer.
+        :rtype: torch.Tensor
+        """
+        return self.func(
+            self.propagate(edge_index, x=x, edge_attr=edge_attr)
+        )

From 6eba0cf3b699c3edb7fae713e4ba72563a469c74 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 4 Feb 2025 19:47:10 +0100
Subject: [PATCH 3/8] Implement PinaGraphDataset

---
 pina/data/dataset.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/pina/data/dataset.py b/pina/data/dataset.py
index 8b5f998f9..4eeb20e02 100644
--- a/pina/data/dataset.py
+++ b/pina/data/dataset.py
@@ -93,8 +93,7 @@ def __getitem__(self, idx):
 
 
 class PinaGraphDataset(PinaDataset):
-    pass
-'''
+
     def __init__(self, conditions_dict, max_conditions_lengths,
                  automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
@@ -113,7 +112,7 @@ def fetch_from_idx_list(self, idx):
             to_return_dict[condition] = {k: Batch.from_data_list([v[i]
                                                                   for i in cond_idx])
             if isinstance(v, list)
-            else v[cond_idx]
+            else v[cond_idx].reshape(-1, *v[cond_idx].shape[2:])
                                          for k, v in data.items()
                                          }
         return to_return_dict
@@ -132,5 +131,4 @@ def get_all_data(self):
         return self.fetch_from_idx_list(index)
 
     def __getitem__(self, idx):
-        return self._getitem_func(idx)
-'''
\ No newline at end of file
+        return self._getitem_func(idx)
\ No newline at end of file

From d591aee196e83be936c957bbaa4a7a1395d20f61 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 4 Feb 2025 20:37:47 +0100
Subject: [PATCH 4/8] Bug fix in GNO and implementation of tests

---
 pina/model/gno.py                          |   4 +
 pina/model/layers/graph_integral_kernel.py |   6 +-
 tests/test_model/test_gno.py               | 127 +++++++++++++++++++++
 3 files changed, 136 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_model/test_gno.py

diff --git a/pina/model/gno.py b/pina/model/gno.py
index 991ca397b..3e9af8ae2 100644
--- a/pina/model/gno.py
+++ b/pina/model/gno.py
@@ -16,6 +16,7 @@ def __init__(
             n_layers=2,
             internal_n_layers=0,
             internal_layers=None,
+            inner_size=None,
             internal_func=None,
             external_func=None,
             shared_weights=False
@@ -50,6 +51,7 @@ def __init__(
                 edges_features=edge_features,
                 n_layers=internal_n_layers,
                 layers=internal_layers,
+                inner_size=inner_size,
                 internal_func=internal_func,
                 external_func=external_func)
             self.n_layers = n_layers
@@ -61,6 +63,7 @@ def __init__(
                     edges_features=edge_features,
                     n_layers=internal_n_layers,
                     layers=internal_layers,
+                    inner_size=inner_size,
                     internal_func=internal_func,
                     external_func=external_func
                 )
@@ -150,6 +153,7 @@ def __init__(
                 width=lifting_operator.out_features,
                 edge_features=edge_features,
                 internal_n_layers=internal_n_layers,
+                inner_size=inner_size,
                 internal_layers=internal_layers,
                 external_func=external_func,
                 internal_func=internal_func,
diff --git a/pina/model/layers/graph_integral_kernel.py b/pina/model/layers/graph_integral_kernel.py
index 713b0d78d..70d172ca3 100644
--- a/pina/model/layers/graph_integral_kernel.py
+++ b/pina/model/layers/graph_integral_kernel.py
@@ -10,8 +10,9 @@ def __init__(
             self,
             width,
             edges_features,
-            n_layers=0,
+            n_layers=2,
             layers=None,
+            inner_size=None,
             internal_func=None,
             external_func=None
     ):
@@ -28,10 +29,13 @@ def __init__(
         from pina.model import FeedForward
         super(GraphIntegralLayer, self).__init__(aggr='mean')
         self.width = width
+        if layers is None and inner_size is None:
+            inner_size = width
         self.dense = FeedForward(input_dimensions=edges_features,
                                  output_dimensions=width ** 2,
                                  n_layers=n_layers,
                                  layers=layers,
+                                 inner_size=inner_size,
                                  func=internal_func)
         self.W = torch.nn.Linear(width, width)
         self.func = external_func()
diff --git a/tests/test_model/test_gno.py b/tests/test_model/test_gno.py
new file mode 100644
index 000000000..e07d8453c
--- /dev/null
+++ b/tests/test_model/test_gno.py
@@ -0,0 +1,127 @@
+import pytest
+import torch
+from pina.graph import KNNGraph
+from pina.model import GNO
+from torch_geometric.data import Batch
+
+x = [torch.rand(100, 6) for _ in range(10)]
+pos = [torch.rand(100, 3) for _ in range(10)]
+graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=6)
+input_ = Batch.from_data_list(graph.data)
+
+
+
+@pytest.mark.parametrize(
+    "shared_weights",
+    [
+        True,
+        False
+    ]
+)
+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)
+
+    GNO(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)
+
+
+@pytest.mark.parametrize(
+    "shared_weights",
+    [
+        True,
+        False
+    ]
+)
+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)
+    output_ = model(input_)
+    assert output_.shape == torch.Size([1000, 3])
+
+@pytest.mark.parametrize(
+    "shared_weights",
+    [
+        True,
+        False
+    ]
+)
+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)
+    output_ = model(input_)
+    assert output_.shape == torch.Size([1000, 3])
+
+@pytest.mark.parametrize(
+    "shared_weights",
+    [
+        True,
+        False
+    ]
+)
+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)
+    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",
+    [
+        True,
+        False
+    ]
+)
+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)
+    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

From dc87615ed1dc0f33aee0e8ff73aee00bf20454d3 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Tue, 4 Feb 2025 21:37:49 +0100
Subject: [PATCH 5/8] Add TemporalGraph class

---
 pina/graph.py       | 30 ++++++++++++++++++++++++++++++
 tests/test_graph.py | 13 ++++++++++++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/pina/graph.py b/pina/graph.py
index 22d70825f..7365bf0bb 100644
--- a/pina/graph.py
+++ b/pina/graph.py
@@ -1,6 +1,7 @@
 from logging import warning
 
 import torch
+
 from . import LabelTensor
 from torch_geometric.data import Data
 from torch_geometric.utils import to_undirected
@@ -238,3 +239,32 @@ def _knn_graph(points, k):
         col = knn_indices.flatten()
         edge_index = torch.stack([row, col], dim=0)
         return edge_index
+
+class TemporalGraph(Graph):
+    def __init__(
+            self,
+            x,
+            pos,
+            t,
+            edge_index=None,
+            edge_attr=None,
+            build_edge_attr=False,
+            undirected=False,
+            r=None
+    ):
+
+        x, pos, edge_index = self._check_input_consistency(x, pos, edge_index)
+        print(len(pos))
+        if edge_index is None:
+            edge_index = [RadiusGraph._radius_graph(p, r) for p in pos]
+        additional_params = {'t': t}
+        self._check_time_consistency(pos, t)
+        super().__init__(x=x, pos=pos, edge_index=edge_index, edge_attr=edge_attr,
+                         build_edge_attr=build_edge_attr,
+                         undirected=undirected,
+                         additional_params=additional_params)
+
+    @staticmethod
+    def _check_time_consistency(pos, times):
+        if len(pos) != len(times):
+            raise ValueError("pos and times must have the same length.")
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 6886dd923..e6ce88c00 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 from pina import Graph
-from pina.graph import RadiusGraph, KNNGraph
+from pina.graph import RadiusGraph, KNNGraph, TemporalGraph
 
 
 @pytest.mark.parametrize(
@@ -141,3 +141,14 @@ def test_additional_parameters_2(additional_parameters):
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
     assert all(hasattr(d, 'y') for d in data)
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+
+def test_temporal_graph():
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+    t = torch.rand(3)
+    graph = TemporalGraph(x=x, pos=pos, build_edge_attr=True, r=.3, t=t)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(hasattr(d, 't') for d in data)
+    assert all(d_.t == t_ for (d_, t_) in zip(data, t))

From d79017e5df9a7c1a4e93a193296e73695e29c2c6 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 5 Feb 2025 10:59:48 +0100
Subject: [PATCH 6/8] Refactor Graph class to support custom edge attribute
 logic

---
 pina/graph.py       | 173 ++++++++++++++++++++++++--------------------
 tests/test_graph.py |  41 ++++++++---
 2 files changed, 128 insertions(+), 86 deletions(-)

diff --git a/pina/graph.py b/pina/graph.py
index 7365bf0bb..d856ad036 100644
--- a/pina/graph.py
+++ b/pina/graph.py
@@ -5,6 +5,7 @@
 from . import LabelTensor
 from torch_geometric.data import Data
 from torch_geometric.utils import to_undirected
+import inspect
 
 
 class Graph:
@@ -12,14 +13,17 @@ class Graph:
     Class for the graph construction.
     """
 
-    def __init__(self,
-                 x,
-                 pos,
-                 edge_index,
-                 edge_attr=None,
-                 build_edge_attr=False,
-                 undirected=False,
-                 additional_params=None):
+    def __init__(
+            self,
+            x,
+            pos,
+            edge_index,
+            edge_attr=None,
+            build_edge_attr=False,
+            undirected=False,
+            custom_build_edge_attr=None,
+            additional_params=None
+    ):
         """
         Constructor for the Graph class.
         :param x: The node features.
@@ -34,45 +38,23 @@ def __init__(self,
         :type build_edge_attr: bool
         :param undirected: Whether to build an undirected graph.
         :type undirected: bool
+        :param custom_build_edge_attr: Custom function to build the edge
+        attributes.
+        :type custom_build_edge_attr: function
         :param additional_params: Additional parameters.
         :type additional_params: dict
         """
         self.data = []
-        x, pos, edge_index = Graph._check_input_consistency(x, pos, edge_index)
+        x, pos, edge_index = self._check_input_consistency(x, pos, edge_index)
 
         # Check input dimension consistency and store the number of graphs
         data_len = self._check_len_consistency(x, pos)
+        if inspect.isfunction(custom_build_edge_attr):
+            self._build_edge_attr = custom_build_edge_attr
 
-        # Initialize additional_parameters (if present)
-        if additional_params is not None:
-            if not isinstance(additional_params, dict):
-                raise TypeError("additional_params must be a dictionary.")
-            for param, val in additional_params.items():
-                # Check if the values are tensors or lists of tensors
-                if isinstance(val, torch.Tensor):
-                    # If the tensor is 3D, we split it into a list of 2D tensors
-                    # In this case there must be a additional parameter for each
-                    # node
-                    if val.ndim == 3:
-                        additional_params[param] = [val[i] for i in
-                                                    range(val.shape[0])]
-                    # If the tensor is 2D, we replicate it for each node
-                    elif val.ndim == 2:
-                        additional_params[param] = [val] * data_len
-                    # If the tensor is 1D, each graph has a scalar values as
-                    # additional parameter
-                    if val.ndim == 1:
-                        if len(val) == data_len:
-                            additional_params[param] = [val[i] for i in
-                                                        range(len(val))]
-                        else:
-                            additional_params[param] = [val for _ in
-                                                        range(data_len)]
-                elif not isinstance(val, list):
-                    raise TypeError("additional_params values must be tensors "
-                                    "or lists of tensors.")
-        else:
-            additional_params = {}
+        # Check consistency and initialize additional_parameters (if present)
+        additional_params = self._check_additional_params(additional_params,
+                                                          data_len)
 
         # Make the graphs undirected
         if undirected:
@@ -81,27 +63,17 @@ def __init__(self,
             else:
                 edge_index = to_undirected(edge_index)
 
-        if build_edge_attr:
-            if edge_attr is not None:
-                warning("Edge attributes are provided, build_edge_attr is set "
-                        "to True. The provided edge attributes will be ignored.")
-            edge_attr = self._build_edge_attr(pos, edge_index)
-
         # Prepare internal lists to create a graph list (same positions but
         # different node features)
         if isinstance(x, list) and isinstance(pos,
                                               (torch.Tensor, LabelTensor)):
             # Replicate the positions, edge_index and edge_attr
             pos, edge_index = [pos] * data_len, [edge_index] * data_len
-            if edge_attr is not None:
-                edge_attr = [edge_attr] * data_len
         # Prepare internal lists to create a list containing a single graph
         elif isinstance(x, (torch.Tensor, LabelTensor)) and isinstance(pos, (
                 torch.Tensor, LabelTensor)):
             # Encapsulate the input tensors into lists
             x, pos, edge_index = [x], [pos], [edge_index]
-            if isinstance(edge_attr, torch.Tensor):
-                edge_attr = [edge_attr]
         # Prepare internal lists to create a list of graphs (same node features
         # but different positions)
         elif (isinstance(x, (torch.Tensor, LabelTensor))
@@ -111,6 +83,10 @@ def __init__(self,
         elif not isinstance(x, list) and not isinstance(pos, list):
             raise TypeError("x and pos must be lists or tensors.")
 
+        # Build the edge attributes
+        edge_attr = self._check_and_build_edge_attr(edge_attr, build_edge_attr,
+                                                    data_len, edge_index, pos, x)
+
         # Perform the graph construction
         self._build_graph_list(x, pos, edge_index, edge_attr, additional_params)
 
@@ -130,12 +106,8 @@ def _build_graph_list(self, x, pos, edge_index, edge_attr,
                                       **add_params_local))
 
     @staticmethod
-    def _build_edge_attr(pos, edge_index):
-        if isinstance(pos, torch.Tensor):
-            pos = [pos]
-            edge_index = [edge_index]
-        distance = [pos_[edge_index_[0]] - pos_[edge_index_[1]] ** 2 for
-                    pos_, edge_index_ in zip(pos, edge_index)]
+    def _build_edge_attr(x, pos, edge_index):
+        distance = torch.abs(pos[edge_index[0]] - pos[edge_index[1]])
         return distance
 
     @staticmethod
@@ -166,15 +138,65 @@ def _check_input_consistency(x, pos, edge_index=None):
             edge_index = [edge_index[i] for i in range(edge_index.shape[0])]
         return x, pos, edge_index
 
+    @staticmethod
+    def _check_additional_params(additional_params, data_len):
+        if additional_params is not None:
+            if not isinstance(additional_params, dict):
+                raise TypeError("additional_params must be a dictionary.")
+            for param, val in additional_params.items():
+                # Check if the values are tensors or lists of tensors
+                if isinstance(val, torch.Tensor):
+                    # If the tensor is 3D, we split it into a list of 2D tensors
+                    # In this case there must be a additional parameter for each
+                    # node
+                    if val.ndim == 3:
+                        additional_params[param] = [val[i] for i in
+                                                    range(val.shape[0])]
+                    # If the tensor is 2D, we replicate it for each node
+                    elif val.ndim == 2:
+                        additional_params[param] = [val] * data_len
+                    # If the tensor is 1D, each graph has a scalar values as
+                    # additional parameter
+                    if val.ndim == 1:
+                        if len(val) == data_len:
+                            additional_params[param] = [val[i] for i in
+                                                        range(len(val))]
+                        else:
+                            additional_params[param] = [val for _ in
+                                                        range(data_len)]
+                elif not isinstance(val, list):
+                    raise TypeError("additional_params values must be tensors "
+                                    "or lists of tensors.")
+        else:
+            additional_params = {}
+        return additional_params
+
+    def _check_and_build_edge_attr(self, edge_attr, build_edge_attr, data_len,
+                                   edge_index, pos, x):
+        # Check if edge_attr is consistent with x and pos
+        if edge_attr is not None:
+            if build_edge_attr is True:
+                warning("edge_attr is not None. build_edge_attr will not be "
+                        "considered.")
+            if isinstance(edge_attr, list):
+                if len(edge_attr) != data_len:
+                    raise ValueError("edge_attr must have the same length as x "
+                                     "and pos.")
+            return [edge_attr] * data_len
+
+        if build_edge_attr:
+            return [self._build_edge_attr(x,pos_, edge_index_) for
+                    pos_, edge_index_ in zip(pos, edge_index)]
+
 
 class RadiusGraph(Graph):
-    def __init__(self,
-                 x,
-                 pos,
-                 r,
-                 build_edge_attr=False,
-                 undirected=False,
-                 additional_params=None, ):
+    def __init__(
+            self,
+            x,
+            pos,
+            r,
+            **kwargs
+    ):
         x, pos, edge_index = Graph._check_input_consistency(x, pos)
 
         if isinstance(pos, (torch.Tensor, LabelTensor)):
@@ -183,9 +205,7 @@ def __init__(self,
             edge_index = [RadiusGraph._radius_graph(p, r) for p in pos]
 
         super().__init__(x=x, pos=pos, edge_index=edge_index,
-                         build_edge_attr=build_edge_attr,
-                         undirected=undirected,
-                         additional_params=additional_params)
+                         **kwargs)
 
     @staticmethod
     def _radius_graph(points, r):
@@ -204,23 +224,20 @@ def _radius_graph(points, r):
 
 
 class KNNGraph(Graph):
-    def __init__(self,
-                 x,
-                 pos,
-                 k,
-                 build_edge_attr=False,
-                 undirected=False,
-                 additional_params=None,
-                 ):
+    def __init__(
+            self,
+            x,
+            pos,
+            k,
+            **kwargs
+    ):
         x, pos, edge_index = Graph._check_input_consistency(x, pos)
         if isinstance(pos, (torch.Tensor, LabelTensor)):
             edge_index = KNNGraph._knn_graph(pos, k)
         else:
             edge_index = [KNNGraph._knn_graph(p, k) for p in pos]
         super().__init__(x=x, pos=pos, edge_index=edge_index,
-                         build_edge_attr=build_edge_attr,
-                         undirected=undirected,
-                         additional_params=additional_params)
+                         **kwargs)
 
     @staticmethod
     def _knn_graph(points, k):
@@ -240,6 +257,7 @@ def _knn_graph(points, k):
         edge_index = torch.stack([row, col], dim=0)
         return edge_index
 
+
 class TemporalGraph(Graph):
     def __init__(
             self,
@@ -259,7 +277,8 @@ def __init__(
             edge_index = [RadiusGraph._radius_graph(p, r) for p in pos]
         additional_params = {'t': t}
         self._check_time_consistency(pos, t)
-        super().__init__(x=x, pos=pos, edge_index=edge_index, edge_attr=edge_attr,
+        super().__init__(x=x, pos=pos, edge_index=edge_index,
+                         edge_attr=edge_attr,
                          build_edge_attr=build_edge_attr,
                          undirected=undirected,
                          additional_params=additional_params)
diff --git a/tests/test_graph.py b/tests/test_graph.py
index e6ce88c00..5521be008 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -7,10 +7,12 @@
 @pytest.mark.parametrize(
     "x, pos",
     [
-        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
-        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
-        (torch.rand(3,10,2), torch.rand(3,10,3)),
-        (torch.rand(3,10,2), torch.rand(3,10,3)),
+        ([torch.rand(10, 2) for _ in range(3)],
+         [torch.rand(10, 3) for _ in range(3)]),
+        ([torch.rand(10, 2) for _ in range(3)],
+         [torch.rand(10, 3) for _ in range(3)]),
+        (torch.rand(3, 10, 2), torch.rand(3, 10, 3)),
+        (torch.rand(3, 10, 2), torch.rand(3, 10, 3)),
     ]
 )
 def test_build_multiple_graph_multiple_val(x, pos):
@@ -28,7 +30,7 @@ def test_build_multiple_graph_multiple_val(x, pos):
     assert all(d.edge_attr is not None for d in data)
     assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
 
-    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k = 3)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
     data = graph.data
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
     assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
@@ -112,36 +114,39 @@ def test_build_single_graph_single_val(pos):
     assert all(d.edge_attr is not None for d in data)
     assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
 
+
 def test_additional_parameters_1():
     x = torch.rand(3, 10, 2)
     pos = torch.rand(3, 10, 2)
     additional_parameters = {'y': torch.ones(3)}
     graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
-                  additional_params=additional_parameters)
+                        additional_params=additional_parameters)
     assert len(graph.data) == 3
     data = graph.data
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
     assert all(hasattr(d, 'y') for d in data)
     assert all(d_.y == 1 for d_ in data)
 
+
 @pytest.mark.parametrize(
     "additional_parameters",
     [
-        ({'y': torch.rand(3,10,1)}),
-        ({'y': [torch.rand(10,1) for _ in range(3)]}),
+        ({'y': torch.rand(3, 10, 1)}),
+        ({'y': [torch.rand(10, 1) for _ in range(3)]}),
     ]
 )
 def test_additional_parameters_2(additional_parameters):
     x = torch.rand(3, 10, 2)
     pos = torch.rand(3, 10, 2)
     graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
-                  additional_params=additional_parameters)
+                        additional_params=additional_parameters)
     assert len(graph.data) == 3
     data = graph.data
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
     assert all(hasattr(d, 'y') for d in data)
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
 
+
 def test_temporal_graph():
     x = torch.rand(3, 10, 2)
     pos = torch.rand(3, 10, 2)
@@ -152,3 +157,21 @@ def test_temporal_graph():
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
     assert all(hasattr(d, 't') for d in data)
     assert all(d_.t == t_ for (d_, t_) in zip(data, t))
+
+
+def test_custom_build_edge_attr_func():
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+
+    def build_edge_attr(x, pos, edge_index):
+        return torch.cat([pos[edge_index[0]], pos[edge_index[1]]], dim=-1)
+
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
+                        custom_build_edge_attr=build_edge_attr)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(hasattr(d, 'edge_attr') for d in data)
+    assert all(d.edge_attr.shape[1] == 4 for d in data)
+    assert all(torch.isclose(d.edge_attr,
+                             build_edge_attr(d.x, d.pos, d.edge_index)).all()
+               for d in data)

From a68b711bd2bd527c49ccf8a3c13b2bfa283510eb Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 5 Feb 2025 17:10:26 +0100
Subject: [PATCH 7/8] Refactor GNO model and enhance Graph class documentation
 and error handling. Remove TemporalGraph class

---
 pina/graph.py       | 109 +++++++++++++++++++++++---------------------
 tests/test_graph.py |  16 +------
 2 files changed, 59 insertions(+), 66 deletions(-)

diff --git a/pina/graph.py b/pina/graph.py
index d856ad036..8c167417c 100644
--- a/pina/graph.py
+++ b/pina/graph.py
@@ -25,25 +25,44 @@ def __init__(
             additional_params=None
     ):
         """
-        Constructor for the Graph class.
-        :param x: The node features.
+        Constructor for the Graph class. This object creates a list of PyTorch Geometric Data objects.
+        Based on the input of x and pos there could be the following cases:
+            1. 1 pos, 1 x: a single graph will be created
+            2. N pos, 1 x: N graphs will be created with the same node features
+            3. 1 pos, N x: N graphs will be created with the same nodes but different node features
+            4. N pos, N x: N graphs will be created
+
+        :param x: Node features. Can be a single 2D tensor of shape [num_nodes, num_node_features],
+                  or a 3D tensor of shape [n_graphs, num_nodes, num_node_features]
+                  or a list of such 2D tensors of shape [num_nodes, num_node_features].
         :type x: torch.Tensor or list[torch.Tensor]
-        :param pos: The node positions.
+        :param pos: Node coordinates. Can be a single 2D tensor of shape [num_nodes, num_coordinates],
+                    or a 3D tensor of shape [n_graphs, num_nodes, num_coordinates]
+                    or a list of such 2D tensors of shape [num_nodes, num_coordinates].
         :type pos: torch.Tensor or list[torch.Tensor]
-        :param edge_index: The edge index.
+        :param edge_index: The edge index defining connections between nodes.
+                           It should be a 2D tensor of shape [2, num_edges]
+                           or a 3D tensor of shape [n_graphs, 2, num_edges]
+                           or a list of such 2D tensors of shape [2, num_edges].
         :type edge_index: torch.Tensor or list[torch.Tensor]
-        :param edge_attr: The edge attributes.
-        :type edge_attr: torch.Tensor or list[torch.Tensor]
-        :param build_edge_attr: Whether to build the edge attributes.
-        :type build_edge_attr: bool
-        :param undirected: Whether to build an undirected graph.
-        :type undirected: bool
-        :param custom_build_edge_attr: Custom function to build the edge
-        attributes.
-        :type custom_build_edge_attr: function
-        :param additional_params: Additional parameters.
-        :type additional_params: dict
+        :param edge_attr: Edge features. If provided, should have the shape [num_edges, num_edge_features]
+                          or be a list of such tensors for multiple graphs.
+        :type edge_attr: torch.Tensor or list[torch.Tensor], optional
+        :param build_edge_attr: Whether to compute edge attributes during initialization.
+        :type build_edge_attr: bool, default=False
+        :param undirected: If True, converts the graph(s) into an undirected graph by adding reciprocal edges.
+        :type undirected: bool, default=False
+        :param custom_build_edge_attr: A user-defined function to generate edge attributes dynamically.
+                                       The function should take (x, pos, edge_index) as input and return a tensor
+                                       of shape [num_edges, num_edge_features].
+        :type custom_build_edge_attr: function or callable, optional
+        :param additional_params: Dictionary containing extra attributes to be added to each Data object.
+                                  Keys represent attribute names, and values should be tensors or lists of tensors.
+        :type additional_params: dict, optional
+
+        Note: if x, pos, and edge_index are both lists or 3D tensors, then len(x) == len(pos) == len(edge_index).
         """
+
         self.data = []
         x, pos, edge_index = self._check_input_consistency(x, pos, edge_index)
 
@@ -85,7 +104,8 @@ def __init__(
 
         # Build the edge attributes
         edge_attr = self._check_and_build_edge_attr(edge_attr, build_edge_attr,
-                                                    data_len, edge_index, pos, x)
+                                                    data_len, edge_index, pos,
+                                                    x)
 
         # Perform the graph construction
         self._build_graph_list(x, pos, edge_index, edge_attr, additional_params)
@@ -128,14 +148,32 @@ def _check_input_consistency(x, pos, edge_index=None):
         # If x is a 3D tensor, we split it into a list of 2D tensors
         if isinstance(x, torch.Tensor) and x.ndim == 3:
             x = [x[i] for i in range(x.shape[0])]
+        elif (not (isinstance(x, list) and all(t.ndim == 2 for t in x)) and
+              not (isinstance(x, torch.Tensor) and x.ndim == 2)):
+            raise TypeError("x must be either a list of 2D tensors or a 2D "
+                            "tensor or a 3D tensor")
 
         # If pos is a 3D tensor, we split it into a list of 2D tensors
         if isinstance(pos, torch.Tensor) and pos.ndim == 3:
             pos = [pos[i] for i in range(pos.shape[0])]
+        elif not (isinstance(pos, list) and all(
+                t.ndim == 2 for t in pos)) and not (
+                isinstance(pos, torch.Tensor) and pos.ndim == 2):
+            raise TypeError("pos must be either a list of 2D tensors or a 2D "
+                            "tensor or a 3D tensor")
 
         # If edge_index is a 3D tensor, we split it into a list of 2D tensors
-        if isinstance(edge_index, torch.Tensor) and edge_index.ndim == 3:
-            edge_index = [edge_index[i] for i in range(edge_index.shape[0])]
+        if edge_index is not None:
+            if isinstance(edge_index, torch.Tensor) and edge_index.ndim == 3:
+                edge_index = [edge_index[i] for i in range(edge_index.shape[0])]
+            elif not (isinstance(edge_index, list) and all(
+                    t.ndim == 2 for t in edge_index)) and not (
+                    isinstance(edge_index,
+                               torch.Tensor) and edge_index.ndim == 2):
+                raise TypeError(
+                    "edge_index must be either a list of 2D tensors or a 2D "
+                    "tensor or a 3D tensor")
+
         return x, pos, edge_index
 
     @staticmethod
@@ -180,12 +218,12 @@ def _check_and_build_edge_attr(self, edge_attr, build_edge_attr, data_len,
                         "considered.")
             if isinstance(edge_attr, list):
                 if len(edge_attr) != data_len:
-                    raise ValueError("edge_attr must have the same length as x "
+                    raise TypeError("edge_attr must have the same length as x "
                                      "and pos.")
             return [edge_attr] * data_len
 
         if build_edge_attr:
-            return [self._build_edge_attr(x,pos_, edge_index_) for
+            return [self._build_edge_attr(x, pos_, edge_index_) for
                     pos_, edge_index_ in zip(pos, edge_index)]
 
 
@@ -256,34 +294,3 @@ def _knn_graph(points, k):
         col = knn_indices.flatten()
         edge_index = torch.stack([row, col], dim=0)
         return edge_index
-
-
-class TemporalGraph(Graph):
-    def __init__(
-            self,
-            x,
-            pos,
-            t,
-            edge_index=None,
-            edge_attr=None,
-            build_edge_attr=False,
-            undirected=False,
-            r=None
-    ):
-
-        x, pos, edge_index = self._check_input_consistency(x, pos, edge_index)
-        print(len(pos))
-        if edge_index is None:
-            edge_index = [RadiusGraph._radius_graph(p, r) for p in pos]
-        additional_params = {'t': t}
-        self._check_time_consistency(pos, t)
-        super().__init__(x=x, pos=pos, edge_index=edge_index,
-                         edge_attr=edge_attr,
-                         build_edge_attr=build_edge_attr,
-                         undirected=undirected,
-                         additional_params=additional_params)
-
-    @staticmethod
-    def _check_time_consistency(pos, times):
-        if len(pos) != len(times):
-            raise ValueError("pos and times must have the same length.")
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 5521be008..660ec3428 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,7 +1,6 @@
 import pytest
 import torch
-from pina import Graph
-from pina.graph import RadiusGraph, KNNGraph, TemporalGraph
+from pina.graph import RadiusGraph, KNNGraph
 
 
 @pytest.mark.parametrize(
@@ -146,19 +145,6 @@ def test_additional_parameters_2(additional_parameters):
     assert all(hasattr(d, 'y') for d in data)
     assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
 
-
-def test_temporal_graph():
-    x = torch.rand(3, 10, 2)
-    pos = torch.rand(3, 10, 2)
-    t = torch.rand(3)
-    graph = TemporalGraph(x=x, pos=pos, build_edge_attr=True, r=.3, t=t)
-    assert len(graph.data) == 3
-    data = graph.data
-    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
-    assert all(hasattr(d, 't') for d in data)
-    assert all(d_.t == t_ for (d_, t_) in zip(data, t))
-
-
 def test_custom_build_edge_attr_func():
     x = torch.rand(3, 10, 2)
     pos = torch.rand(3, 10, 2)

From 9b7cdbfbcc9df3d2aa9f81f5f2aa6d7a9f1efaa8 Mon Sep 17 00:00:00 2001
From: FilippoOlivo <filippo@filippoolivo.com>
Date: Wed, 5 Feb 2025 17:23:46 +0100
Subject: [PATCH 8/8] Rename classes and modules for GNO

---
 pina/model/__init__.py                        |  4 +-
 pina/model/gno.py                             |  8 +-
 pina/model/layers/__init__.py                 |  4 +-
 ...{graph_integral_kernel.py => gno_block.py} |  5 +-
 tests/test_model/test_gno.py                  | 88 ++++++++++---------
 5 files changed, 56 insertions(+), 53 deletions(-)
 rename pina/model/layers/{graph_integral_kernel.py => gno_block.py} (96%)

diff --git a/pina/model/__init__.py b/pina/model/__init__.py
index 0db0acd71..c75f9b658 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 3e9af8ae2..9a8b88878 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 50827dc4b..3e3e71682 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 70d172ca3..34929fe89 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 e07d8453c..8fb10d8e5 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])