might work

modules/execute.py

@@ -46,7 +46,7 @@ if __name__ == "__main__":
                   #'thickness_absorber_6': .1, 
                   #'thickness_absorber_7': .1, 
                   #'thickness_absorber_8': .1, 
-                  'thickness_scintillator_0': 10.5, 
+                  'thickness_scintillator_0': 3.5, 
                   #'thickness_scintillator_1': 0.5,
                   #'thickness_scintillator_2': 0.5,
                   #'thickness_scintillator_3': 0.5,
@@ -190,25 +190,24 @@ if __name__ == "__main__":
     
         surrogate_dataset = SurrogateDataset(ds, reco_result.detach().cpu().numpy())#important to detach here
         print('surr pre-training 0')
-        surrogate_model.train_model(surrogate_dataset, batch_size=256, n_epochs= n_epochs_pre//2, lr=0.03)
+        surrogate_model.train_model(surrogate_dataset, batch_size=128, n_epochs= n_epochs_pre//2, lr=0.01)
         print('surr pre-training 1')
-        surrogate_model.train_model(surrogate_dataset, batch_size=256, n_epochs= n_epochs_pre, lr=0.01)
+        surrogate_model.train_model(surrogate_dataset, batch_size=128, n_epochs= n_epochs_pre, lr=0.01)
         print('surr pre-training 2')
-        surrogate_model.train_model(surrogate_dataset, batch_size=512, n_epochs= n_epochs_pre, lr=0.001)
+        surrogate_model.train_model(surrogate_dataset, batch_size=128, n_epochs= n_epochs_pre, lr=0.001)
         print('surr pre-training 3')
-        surrogate_model.train_model(surrogate_dataset, batch_size=1024, n_epochs= n_epochs_pre, lr=0.001)
+        surrogate_model.train_model(surrogate_dataset, batch_size=512, n_epochs= n_epochs_pre, lr=0.001)
         print('surr pre-training 4')
-        surrogate_model.train_model(surrogate_dataset, batch_size=1024, n_epochs= n_epochs_pre, lr=0.0003)
+        surrogate_model.train_model(surrogate_dataset, batch_size=2048, n_epochs= n_epochs_pre, lr=0.0003)
         print('surr pre-training 5')
-        surrogate_model.train_model(surrogate_dataset, batch_size=1024, n_epochs= n_epochs_pre, lr=1e-4)
+        surrogate_model.train_model(surrogate_dataset, batch_size=2048, n_epochs= n_epochs_pre, lr=1e-4)
         print('surr pre-training 6')
-        surrogate_model.train_model(surrogate_dataset, batch_size=1024, n_epochs= n_epochs_pre, lr=3e-5)
+        surrogate_model.train_model(surrogate_dataset, batch_size=2048, n_epochs= n_epochs_pre, lr=3e-5)
         #these are un-normalised quantities
-        surr_out, reco_out, true_in = surrogate_model.apply_model_in_batches(surrogate_dataset, batch_size=1024)
+        surr_out, reco_out, true_in = surrogate_model.apply_model_in_batches(surrogate_dataset, batch_size=2048)
         make_check_plot(surr_out, reco_out, true_in, evolution, outname = "pretrain")
     pre_train()
-    exit() #DEBUG
-    
+    #exit() #DEBUG
 
     best_surrogate_loss = 1e10
     def update_best_surrogate_loss(loss):

modules/one_dim_flow_surrogate.py

@@ -311,11 +311,12 @@ class Surrogate(torch.nn.Module):
 
         self.n_detector_parameters = n_detector_parameters
         self.model = PiecewiseLinear(n_conditions = n_detector_parameters + n_true_plus_context_inputs, #true context is energy etc,
-                                     xmin = -10., xmax = 10., #here the reco stuff is normalised so that should (TM) work
+                                     xmin = -5., xmax = 5., #here the reco stuff is normalised so that should (TM) work
+                                     nb=265
                                      )
         
 
-        self.optimizer = torch.optim.Adam(self.parameters(), lr=0.0001)
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.0001)
         self.n_reco_parameters = n_out_parameters # called n_reco_parameters for consistency with the original code
         self.device = torch.device('cuda')
         self.n_time_steps = None #compatibility
@@ -377,6 +378,8 @@ class Surrogate(torch.nn.Module):
                             )
                 #print(derivatives)
 
+                #print('reco',reco_step_inputs.mean(), reco_step_inputs.std())
+
                 loss = ( 0.5 * (model_out**2 + math.log(2*pi)) - derivatives.log() ).mean()
                 #print(loss)