fix doc model part 1

pina/model/graph_neural_operator.py

@@ -1,5 +1,5 @@
 """
-Module for the Graph Neural Operator and Graph Neural Kernel.
+Module for the Graph Neural Operator model class.
 """
 
 import torch
@@ -10,7 +10,18 @@ from .kernel_neural_operator import KernelNeuralOperator
 
 class GraphNeuralKernel(torch.nn.Module):
     """
-    TODO add docstring
+    Graph Neural Kernel model class.
+
+    This class implements the Graph Neural Kernel network.
+
+    .. seealso::
+
+        **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K.,
+        Liu, B., Bhattacharya, K., Stuart, A., Anandkumar, A. (2020).
+        *Neural Operator: Graph Kernel Network for Partial Differential
+        Equations*.
+        DOI: `arXiv preprint arXiv:2003.03485.
+        <https://arxiv.org/abs/2003.03485>`_
     """
 
     def __init__(
@@ -26,28 +37,24 @@ class GraphNeuralKernel(torch.nn.Module):
         shared_weights=False,
     ):
         """
-        The Graph Neural Kernel constructor.
+        Initialization of the :class:`GraphNeuralKernel` class.
 
-        :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.
+        :param int width: The width of the kernel.
+        :param int edge_features: The number of edge features.
+        :param int n_layers: The number of kernel layers. Default is ``2``.
+        :param int internal_n_layers: The number of layers of the neural network
+            inside each kernel layer. Default is ``0``.
+        :param internal_layers: The number of neurons for each layer of the
+            neural network inside each kernel layer. Default is ``None``.
+        :type internal_layers: list[int] | tuple[int]
+        :param torch.nn.Module internal_func: The activation function used
+            inside each kernel layer. If ``None``, it uses the
+            :class:`torch.nn.Tanh`. activation. Default is ``None``.
+        :param torch.nn.Module external_func: The activation function applied to
+            the output of the each kernel layer. If ``None``, it uses the
+            :class:`torch.nn.Tanh`. activation. Default is ``None``.
+        :param bool shared_weights: If ``True``, the weights of each kernel
+            layer are shared. Default is ``False``.
         """
         super().__init__()
         if external_func is None:
@@ -85,11 +92,33 @@ class GraphNeuralKernel(torch.nn.Module):
             self._forward_func = self._forward_unshared
 
     def _forward_unshared(self, x, edge_index, edge_attr):
+        """
+        Forward pass for the Graph Neural Kernel with unshared weights.
+
+        :param x: The input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor edge_index: The edge index.
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor | LabelTensor
+        :return: The output tensor.
+        :rtype: 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):
+        """
+        Forward pass for the Graph Neural Kernel with shared weights.
+
+        :param x: The input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor edge_index: The edge index.
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor | LabelTensor
+        :return: The output tensor.
+        :rtype: torch.Tensor
+        """
         for _ in range(self.n_layers):
             x = self.layers(x, edge_index, edge_attr)
         return x
@@ -98,19 +127,34 @@ class GraphNeuralKernel(torch.nn.Module):
         """
         The forward pass of the Graph Neural Kernel.
 
-        :param x: The input batch.
-        :type x: torch.Tensor
-        :param edge_index: The edge index.
-        :type edge_index: torch.Tensor
+        :param x: The input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor edge_index: The edge index.
         :param edge_attr: The edge attributes.
-        :type edge_attr: torch.Tensor
+        :type edge_attr: torch.Tensor | LabelTensor
+        :return: The output tensor.
+        :rtype: torch.Tensor
         """
         return self._forward_func(x, edge_index, edge_attr)
 
 
 class GraphNeuralOperator(KernelNeuralOperator):
     """
-    TODO add docstring
+    Graph Neural Operator model class.
+
+    The Graph Neural Operator is a general architecture for learning operators,
+    which map functions to functions. It can be trained both with Supervised
+    and Physics-Informed learning strategies. The Graph Neural Operator performs
+    graph convolution by means of a Graph Neural Kernel.
+
+    .. seealso::
+
+        **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K.,
+        Liu, B., Bhattacharya, K., Stuart, A., Anandkumar, A. (2020).
+        *Neural Operator: Graph Kernel Network for Partial Differential
+        Equations*.
+        DOI: `arXiv preprint arXiv:2003.03485.
+        <https://arxiv.org/abs/2003.03485>`_
     """
 
     def __init__(
@@ -127,34 +171,29 @@ class GraphNeuralOperator(KernelNeuralOperator):
         shared_weights=True,
     ):
         """
-        The Graph Neural Operator constructor.
+        Initialization of the :class:`GraphNeuralOperator` class.
 
-        :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
+        param torch.nn.Module lifting_operator: The lifting neural network
+            mapping the input to its hidden dimension.
+        :param torch.nn.Module projection_operator: The projection neural
+            network mapping the hidden representation to the output function.
+        :param int edge_features: The number of edge features.
+        :param int n_layers: The number of kernel layers. Default is ``10``.
+        :param int internal_n_layers: The number of layers of the neural network
+            inside each kernel layer. Default is ``0``.
+        :param int inner_size: The size of the hidden layers of the neural
+            network inside each kernel layer. Default is ``None``.
+        :param internal_layers: The number of neurons for each layer of the
+            neural network inside each kernel layer. Default is ``None``.
+        :type internal_layers: list[int] | tuple[int]
+        :param torch.nn.Module internal_func: The activation function used
+            inside each kernel layer. If ``None``, it uses the
+            :class:`torch.nn.Tanh`. activation. Default is ``None``.
+        :param torch.nn.Module external_func: The activation function applied to
+            the output of the each kernel layer. If ``None``, it uses the
+            :class:`torch.nn.Tanh`. activation. Default is ``None``.
+        :param bool shared_weights: If ``True``, the weights of each kernel
+            layer are shared. Default is ``False``.
         """
 
         if internal_func is None:
@@ -182,8 +221,9 @@ class GraphNeuralOperator(KernelNeuralOperator):
         """
         The forward pass of the Graph Neural Operator.
 
-        :param x: The input batch.
-        :type x: torch_geometric.data.Batch
+        :param torch_geometric.data.Batch x: The input graph.
+        :return: The output tensor.
+        :rtype: torch.Tensor
         """
         x, edge_index, edge_attr = x.x, x.edge_index, x.edge_attr
         x = self.lifting_operator(x)