reformatting again

Author: niklasmei

examples/pretraining_maskpred_example.py

@@ -9,34 +9,39 @@
 
 import torch
 
-from graphnet.models.detector.prometheus  import  Prometheus
-from graphnet.models.graphs  import  KNNGraph
-from graphnet.models.graphs.nodes  import  NodesAsPulses
+from graphnet.models.detector.prometheus import Prometheus
+from graphnet.models.graphs import KNNGraph
+from graphnet.models.graphs.nodes import NodesAsPulses
 
 
 class simple_model(Model):
-    def __init__(self,
-                 ):
+    def __init__(
+        self,
+    ):
         super().__init__()
-        self.net = torch.nn.Sequential(torch.nn.Linear(4,10),
-                                       torch.nn.SELU(),
-                                       torch.nn.Linear(10,5))
+        self.net = torch.nn.Sequential(
+            torch.nn.Linear(4, 10), torch.nn.SELU(), torch.nn.Linear(10, 5)
+        )
 
-    def forward(self, data:Data):
+    def forward(self, data: Data):
         x = self.net(data.x)
-        x_rep = scatter(src=x, index=data.batch, dim = 0, reduce='max')
+        x_rep = scatter(src=x, index=data.batch, dim=0, reduce="max")
         return x, x_rep
-    
+
+
 class simple_target_gen(Model):
-    def __init__(self,
-                 ):
+    def __init__(
+        self,
+    ):
         super().__init__()
 
-    def forward(self, data:Data):
-        target = torch.sum(scatter(src=data.x, index=data.batch, dim = 0, reduce='max'), dim=1)
-        return target.view(-1,1)
+    def forward(self, data: Data):
+        target = torch.sum(
+            scatter(src=data.x, index=data.batch, dim=0, reduce="max"), dim=1
+        )
+        return target.view(-1, 1)
+
 
-    
 def test():
     graph_definition = KNNGraph(
         detector=Prometheus(),
@@ -50,7 +55,7 @@ def test():
         truth_table="mc_truth",
         features=["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t"],
         truth=["injection_energy", "injection_zenith"],
-        data_representation = graph_definition,
+        data_representation=graph_definition,
     )
 
     dataloader = DataLoader(
@@ -63,15 +68,15 @@ def test():
         data = batch
         break
 
-    #Standard Parameters
+    # 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_hidden_dim: int = 1000,
     #              default_nb_linear: int = 5,
     #              final_loss: str = 'mse',
     #              add_charge_pred: bool = False,
@@ -81,23 +86,23 @@ def test():
     dummy_model = simple_model()
     dummy_target = simple_target_gen()
 
-    model = mask_pred_frame(encoder=dummy_model,
-                            encoder_out_dim=5,
-                            masked_feat = [0,1],
-                            learned_masking_value=True,
-                            final_loss = 'cosine',
-                            add_charge_pred = True,
-                            need_charge_rep = False,
-                            custom_charge_target = dummy_target,
-                            )
-
-    
+    model = mask_pred_frame(
+        encoder=dummy_model,
+        encoder_out_dim=5,
+        masked_feat=[0, 1],
+        learned_masking_value=True,
+        final_loss="cosine",
+        add_charge_pred=True,
+        need_charge_rep=False,
+        custom_charge_target=dummy_target,
+    )
+
     out = model(data)
     print(out)
 
-    #for saving
-    #model.save_pretrained_model('some/path')
+    # for saving
+    # model.save_pretrained_model('some/path')
 
 
 if __name__ == "__main__":
-    test()
+    test()

src/graphnet/models/gnn/pretraining_maskpred.py

@@ -256,9 +256,12 @@ def forward(self, data: Union[Data, List[Data]]) -> List[Tensor]:
         nodes = rep[~mask.bool()]
         btch = data.batch[~mask.bool()]
 
-        loss = scatter(src=self.loss_func(nodes,target,return_elements=True), 
-                       index = btch, reduce="mean", 
-                       dim=0).view(-1,1)
+        loss = scatter(
+            src=self.loss_func(nodes, target, return_elements=True),
+            index=btch,
+            reduce="mean",
+            dim=0,
+        ).view(-1, 1)
 
         if self.add_charge_pred:
             if self.custom_charge_target is None:

src/graphnet/training/loss_functions.py

@@ -544,6 +544,7 @@ def __init__(self, vmfs_factor: float = 0.05) -> None:
             prediction_keys=[[0, 1, 2], [0, 1, 2, 3]],
         )
 
+
 class NegCosLoss(LossFunction):
     """Negative Cosine error loss."""