pydocstyle on example

Author: niklasmei

examples/pretraining_maskpred_example.py

@@ -1,3 +1,5 @@
+"""Minimal example for use of maskpred pretraining."""
+
 from graphnet.models.gnn.pretraining_maskpred import mask_pred_frame
 from graphnet.models import Model
 from torch_geometric.data import Data
@@ -15,34 +17,43 @@
 
 
 class simple_model(Model):
+    """Just for a dummy model."""
+
     def __init__(
         self,
-    ):
+    ) -> None:
+        """Construct."""
         super().__init__()
         self.net = torch.nn.Sequential(
             torch.nn.Linear(4, 10), torch.nn.SELU(), torch.nn.Linear(10, 5)
         )
 
     def forward(self, data: Data):
+        """Forward pass."""
         x = self.net(data.x)
         x_rep = scatter(src=x, index=data.batch, dim=0, reduce="max")
         return x, x_rep
 
 
 class simple_target_gen(Model):
+    """Just for a dummy charge target."""
+
     def __init__(
         self,
-    ):
+    ) -> None:
+        """Construct."""
         super().__init__()
 
     def forward(self, data: Data):
+        """Forward pass."""
         target = torch.sum(
             scatter(src=data.x, index=data.batch, dim=0, reduce="max"), dim=1
         )
         return target.view(-1, 1)
 
 
-def test():
+def test() -> None:
+    """Function that just evaluates the model to test it and has a save example commented in the end."""
     graph_definition = KNNGraph(
         detector=Prometheus(),
         node_definition=NodesAsPulses(),
@@ -68,21 +79,6 @@ def test():
         data = batch
         break
 
-    # Standard Parameters
-    #              encoder: Model,
-    #              encoder_out_dim: int = None,
-    #              masked_ratio: float = 0.25,
-    #              masked_feat: List[int] = [0,1,2,3,4],
-    #              learned_masking_value: bool = True,
-    #              hlc_pos: int = None,
-    #              mask_pred_net: Model = None,
-    #              default_hidden_dim: int = 1000,
-    #              default_nb_linear: int = 5,
-    #              final_loss: str = 'mse',
-    #              add_charge_pred: bool = False,
-    #              need_charge_rep: bool = False,
-    #              custom_charge_target: Tensor = None,
-
     dummy_model = simple_model()
     dummy_target = simple_target_gen()