add doc strings

deeprank2/neuralnets/cnn/model3d.py

@@ -23,7 +23,16 @@
 # ----------------------------------------------------------------------
 
 
-class CnnRegression(nn.Module):  # noqa: D101
+class CnnRegression(nn.Module):
+    """Convolutional Neural Network architecture for regression.
+
+    This type of network is used to predict a single scalar value of a continuous variable.
+
+    Args:
+        num_features: Number of features in the input data.
+        box_shape: Shape of the input data.
+    """
+
     def __init__(self, num_features: int, box_shape: tuple[int]):
         super().__init__()
 
@@ -76,7 +85,16 @@ def forward(self, data):
 # ----------------------------------------------------------------------
 
 
-class CnnClassification(nn.Module):  # noqa: D101
+class CnnClassification(nn.Module):
+    """Convolutional Neural Network architecture for binary classification.
+
+    This type of network is used to predict the class of an input data point.
+
+    Args:
+        num_features: Number of features in the input data.
+        box_shape: Shape of the input data.
+    """
+
     def __init__(self, num_features, box_shape):
         super().__init__()
 

deeprank2/neuralnets/gnn/alignmentnet.py

@@ -6,7 +6,19 @@
 __author__ = "Daniel-Tobias Rademaker"
 
 
-class GNNLayer(nn.Module):  # noqa: D101
+class GNNLayer(nn.Module):
+    """Custom-defined layer of a Graph Neural Network.
+
+    Args:
+        nmb_edge_projection: Number of features in the edge projection.
+        nmb_hidden_attr: Number of features in the hidden attributes.
+        nmb_output_features: Number of output features.
+        message_vector_length: Length of the message vector.
+        nmb_mlp_neurons: Number of neurons in the MLP.
+        act_fn: Activation function. Defaults to nn.SiLU().
+        is_last_layer: Whether this is the last layer of the GNN. Defaults to True.
+    """
+
     def __init__(
         self,
         nmb_edge_projection,
@@ -104,7 +116,26 @@ def output(self, hidden_features, get_attention=True):
         return output
 
 
-class SuperGNN(nn.Module):  # noqa: D101
+class SuperGNN(nn.Module):
+    """SuperGNN is a class that defines multiple GNN layers.
+
+    In particular, the `preproc_edge_mlp` and `preproc_node_mlp` are meant to
+    preprocess the edge and node attributes, respectively.
+
+    The `modlist` is a list of GNNLayer objects.
+
+    Args:
+        nm_edge_attr: Number of edge features.
+        nmb_node_attr: Number of node features.
+        nmb_hidden_attr: Number of hidden features.
+        nmb_mlp_neurons: Number of neurons in the MLP.
+        nmb_edge_projection: Number of edge projections.
+        nmb_gnn_layers: Number of GNN layers.
+        nmb_output_features: Number of output features.
+        message_vector_length: Length of the message vector.
+        act_fn: Activation function. Defaults to nn.SiLU().
+    """
+
     def __init__(
         self,
         nmb_edge_attr,
@@ -172,7 +203,26 @@ def run_through_network(self, edges, edge_attr, node_attr, with_output_attention
         return self.modlist[-1].output(node_attr, True)  # (boolean-positional-value-in-call)
 
 
-class AlignmentGNN(SuperGNN):  # noqa: D101
+class AlignmentGNN(SuperGNN):
+    """Graph Neural Network.
+
+    It applies different layers to the nodes and edges of a graph (`preproc_edge_mlp` and `preproc_node_mlp`),
+    and then applies multiple GNN layers (`modlist`).
+    It can be used for both regression and classification tasks.
+
+    Args:
+        nm_edge_attr: Number of edge features.
+        nmb_node_attr: Number of node features.
+        nmb_output_features: Number of output features.
+        nmb_hidden_attr: Number of hidden features.
+        message_vector_length: Length of the message vector.
+        nmb_mlp_neurons: Number of neurons in the MLP.
+        nmb_gnn_layers: Number of GNN layers.
+        nmb_edge_projection: Number of edge projections.
+        act_fn: Activation function. Defaults to nn.SiLU().
+
+    """
+
     def __init__(
         self,
         nmb_edge_attr,

deeprank2/neuralnets/gnn/foutnet.py

@@ -13,8 +13,7 @@
 class FoutLayer(nn.Module):
     """FoutLayer.
 
-    This layer is described by eq. (1) of
-    Protein Interface Predition using Graph Convolutional Network
+    This layer is described by eq. (1) of Protein Interface Predition using Graph Convolutional Network
     by Alex Fout et al. NIPS 2018.
 
     Args:
@@ -70,7 +69,18 @@ def __repr__(self):
         return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
 
 
-class FoutNet(nn.Module):  # noqa: D101
+class FoutNet(nn.Module):
+    """FoutNet.
+
+    Architecture based on the FoutLayer. It also uses community pooling to reduce the number of nodes.
+    It can be used for both regression and classification tasks.
+
+    Args:
+        input_shape: Size of each input sample.
+        output_shape: Size of each output sample. Defaults to 1.
+        input_shape_edge: Size of each input edge. Defaults to None.
+    """
+
     def __init__(
         self,
         input_shape,

deeprank2/neuralnets/gnn/ginet.py

@@ -10,7 +10,16 @@
 # ruff: noqa: ANN001, ANN201
 
 
-class GINetConvLayer(nn.Module):  # noqa: D101
+class GINetConvLayer(nn.Module):
+    """GiNet convolutional layer for graph neural networks.
+
+    Args:
+        in_channels: Number of input features.
+        out_channels: Number of output features.
+        number_edge_features: Number of edge features. Defaults to 1.
+        bias: If set to :obj:`False`, the layer will not learn an additive bias. Defaults to False.
+    """
+
     def __init__(self, in_channels, out_channels, number_edge_features=1, bias=False):
         super().__init__()
 
@@ -54,10 +63,19 @@ def __repr__(self):
         return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
 
 
-class GINet(nn.Module):  # noqa: D101
-    # input_shape -> number of node input features
-    # output_shape -> number of output value per graph
-    # input_shape_edge -> number of edge input features
+class GINet(nn.Module):
+    """GINet.
+
+    A graph neural network architecture based on the GiNet convolutional layer.
+    It uses community pooling to reduce the number of nodes. It can be used for both
+    regression and classification tasks.
+
+    Args:
+        input_shape: Number of input features.
+        output_shape: Number of output value per graph. Defaults to 1.
+        input_shape_edge: Number of edge input features. Defaults to 1.
+    """
+
     def __init__(self, input_shape, output_shape=1, input_shape_edge=1):
         super().__init__()
         self.conv1 = GINetConvLayer(input_shape, 16, input_shape_edge)

deeprank2/neuralnets/gnn/ginet_nocluster.py

@@ -7,7 +7,16 @@
 # ruff: noqa: ANN001, ANN201
 
 
-class GINetConvLayer(nn.Module):  # noqa: D101
+class GINetConvLayer(nn.Module):
+    """GiNet convolutional layer for graph neural networks.
+
+    Args:
+        in_channels: Number of input features.
+        out_channels: Number of output features.
+        number_edge_features: Number of edge features. Defaults to 1.
+        bias: If set to :obj:`False`, the layer will not learn an additive bias. Defaults to False.
+    """
+
     def __init__(self, in_channels, out_channels, number_edge_features=1, bias=False):
         super().__init__()
 
@@ -51,10 +60,18 @@ def __repr__(self):
         return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
 
 
-class GINet(nn.Module):  # noqa: D101
-    # input_shape -> number of node input features
-    # output_shape -> number of output value per graph
-    # input_shape_edge -> number of edge input features
+class GINet(nn.Module):
+    """GINet.
+
+    A graph neural network architecture based on the GiNet convolutional layer.
+    It can be used for both regression and classification tasks.
+
+    Args:
+        input_shape: Number of input features.
+        output_shape: Number of output value per graph. Defaults to 1.
+        input_shape_edge: Number of edge input features. Defaults to 1.
+    """
+
     def __init__(self, input_shape, output_shape=1, input_shape_edge=1):
         super().__init__()
         self.conv1 = GINetConvLayer(input_shape, 16, input_shape_edge)

deeprank2/neuralnets/gnn/naive_gnn.py

@@ -7,7 +7,14 @@
 # ruff: noqa: ANN001, ANN201
 
 
-class NaiveConvolutionalLayer(nn.Module):  # noqa: D101
+class NaiveConvolutionalLayer(nn.Module):
+    """Naive convolutional layer for graph neural networks.
+
+    Args:
+        count_node_features: Number of node features.
+        count_edge_features: Number of edge features.
+    """
+
     def __init__(self, count_node_features, count_edge_features):
         super().__init__()
         message_size = 32
@@ -31,15 +38,20 @@ def forward(self, node_features, edge_node_indices, edge_features):
         return self._node_mlp(node_input)
 
 
-class NaiveNetwork(nn.Module):  # noqa: D101
-    def __init__(self, input_shape: int, output_shape: int, input_shape_edge: int):
-        """NaiveNetwork.
+class NaiveNetwork(nn.Module):
+    """NaiveNetwork.
 
-        Args:
-            input_shape: Number of node input features.
-            output_shape: Number of output value per graph.
-            input_shape_edge: Number of edge input features.
-        """
+    Implementation of a naive graph neural network suited for classification tasks.
+    It uses two naive convolutional layers and a MLP to predict the output.
+    It can be used for both regression and classification tasks.
+
+    Args:
+        input_shape: Number of node input features.
+        output_shape: Number of output value per graph.
+        input_shape_edge: Number of edge input features.
+    """
+
+    def __init__(self, input_shape: int, output_shape: int, input_shape_edge: int):
         super().__init__()
         self._external1 = NaiveConvolutionalLayer(input_shape, input_shape_edge)
         self._external2 = NaiveConvolutionalLayer(input_shape, input_shape_edge)

deeprank2/neuralnets/gnn/sgat.py

@@ -87,7 +87,19 @@ def __repr__(self):
         return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
 
 
-class SGAT(nn.Module):  # noqa:D101
+class SGAT(nn.Module):
+    """SGAT.
+
+    Implementation of a simple graph attention network.
+    It uses two graph attention layers and a MLP to predict the output.
+    It can be used for both regression and classification tasks.
+
+    Args:
+        input_shape: Size of each input sample.
+        output_shape: Size of each output sample. Defaults to 1.
+        input_shape_edge: Size of each input edge. Defaults to None.
+    """
+
     def __init__(
         self,
         input_shape,