Update reconstruction.py

modules/reconstruction.py

@@ -15,6 +15,13 @@ class Reconstruction(torch.nn.Module):
         super(Reconstruction, self).__init__()
         self.n_parameters = n_detector_parameters
 
+        self.preprocess = torch.nn.Sequential(
+            torch.nn.Linear(n_detector_parameters, 100),
+            torch.nn.ELU(),
+            torch.nn.Linear(100,100),
+            torch.nn.ELU(),
+            torch.nn.Linear(100, n_input_parameters)
+        )
         # take into account that 
         self.layers = torch.nn.Sequential(
             #torch.nn.BatchNorm1d(n_detector_parameters + n_input_parameters),
@@ -35,8 +42,10 @@ class Reconstruction(torch.nn.Module):
 
     def forward(self, detector_parameters, x):
         # concatenate the detector parameters and the input
-        x = torch.cat([detector_parameters, x], dim=1)
-        return self.layers(x)
+        scaling_factors = self.preprocess(detector_parameters)
+        x_scaled = torch.multiply(x, scaling_factors)
+        concatenated_input = torch.cat([detector_parameters, x_scaled], dim=1)
+        return self.layers(concatenated_input)
     
     def loss(self, y_pred, y):
         #to make it more stable filter nans or infs based on the predicted first from y, then from pedicted