Update execute.py

modules/execute.py

@@ -5,12 +5,16 @@ import pickle
 from generator import Generator, Parameters, SurrogateDataset
 from reconstruction import Reconstruction
 from surrogate import Surrogate
-from optimizer import Optimizer
+from optimizerA import Optimizer
 from matplotlib import pyplot as plt
 import time
 import torch
 import numpy as np
 import os
+import json
+import sys
+
+os.nice(18)
 
 def reset_weights(m):
     '''
@@ -24,7 +28,7 @@ def reset_weights(m):
 if __name__ == "__main__":
     multiprocessing.set_start_method('spawn')
 
-    outpath = 'run2_EM_homo_simple3'
+    outpath = 'FF02'
     
     os.system('mkdir  '+outpath)
     outpath += '/'
@@ -35,26 +39,30 @@ if __name__ == "__main__":
 
     divide = 10 if test else 1
 
+    num_layers = 3
+
     #so far these are only layer thicknesses
-    start_pars = {'thickness_absorber_0': .1,
-                  #'thickness_absorber_1': .1, 
-                  #'thickness_absorber_2': .1, 
-                  #'thickness_absorber_3': .1, 
-                  #'thickness_absorber_4': .1, 
+    start_pars = {'thickness_absorber_0': 20,
+                  'thickness_absorber_1': 17, 
+                  'thickness_absorber_2': 35, 
                   #'thickness_absorber_5': .1, 
                   #'thickness_absorber_6': .1, 
                   #'thickness_absorber_7': .1, 
                   #'thickness_absorber_8': .1, 
-                  'thickness_scintillator_0': 0.5, 
-                  #'thickness_scintillator_1': 0.5,
-                  #'thickness_scintillator_2': 0.5,
-                  #'thickness_scintillator_3': 0.5,
-                  #'thickness_scintillator_4': 0.5,
+                  'thickness_scintillator_0': 5, 
+                  'thickness_scintillator_1': 12,
+                  'thickness_scintillator_2': 3,
                   #'thickness_scintillator_5': 0.5,
                   #'thickness_scintillator_6': 0.5,
                   #'thickness_scintillator_7': 0.5,
                   #'thickness_scintillator_8': 0.5,
-                  }
+                  'material_abs_0' : 0.33,
+                  'material_abs_1' : 0.33,
+                  'material_abs_2' : 0.33, # Randomly chosen values see what happens, ideally we could pass an array here, this for testing
+                  'material_scint_0' : 0.33,
+                  'material_scint_1' : 0.66,
+                  'material_scint_2' : 0.33, 
+                 }
     box = np.array(len(start_pars)*[1.5]) #will be adjusted to some extent by optimizer
     #start parameters from run 8
     #start_pars = {'thickness_absorber_0': 0.7642903, 'thickness_absorber_1': 10.469371, 'thickness_scintillator_0': 30.585306, 'thickness_scintillator_1': 22.256506}
@@ -62,9 +70,9 @@ if __name__ == "__main__":
     
     gen = Generator(box, 
                     Parameters(parameters = start_pars),
-                    n_vars = 30,
-                    n_events_per_var = 400//divide,
-                    particles=[['gamma',0.22]])#,['pi+',2.11]])#0)
+                    n_vars = 25,
+                    n_events_per_var = 200//divide,
+                    particles=[['gamma',0.22,'pi+',2.11]])
     
     
     # reconstruction
@@ -75,16 +83,19 @@ if __name__ == "__main__":
     surrogate_model = Surrogate(gen.n_parameters, 
                            gen.n_target_parameters + gen.n_context_parameters, #context
                            gen.n_target_parameters,  #reco
-                           n_time_steps=100)
+                           n_time_steps=50)
     
     optimizer = Optimizer(gen, surrogate_model, reco_model, gen.parameters,
-                          constraints= {'length': 25}) #one meter
-    
-    n_epochs_pre = 30//divide
-    n_epochs_main = 100//divide
+                          constraints= {'length': 200, 'upper' : 1, 'lower': 0,'diff' : 0.66, 'cost' : 50000}) # length of detector, upper limit, lower limit and max diff of material parameters
+    import sys
+
+    n_epochs_pre = 25//divide
+    n_epochs_main = 40//divide
     
-    parameters = gen.parameters
+    parameters = gen.parameters 
 
+    
+    
     #if saved evolution exists, load it
     import os.path
     if False and os.path.isfile('evolution.npy'):
@@ -106,6 +117,7 @@ if __name__ == "__main__":
              'stds': std}
         )
         evolution.append([ps, losses, reco_losses])
+
         with open(outpath+'evolution.pkl', 'wb') as f:
             pickle.dump(evolution, f)
         # human readable
@@ -116,6 +128,7 @@ if __name__ == "__main__":
         # now plot all entries of the dictionary ps in a plot and save it, make a new point for each of the items in the evolution list and use the dictionary keys as legend
         # 'transpose' dictionary
         pltdict = {k: [] for k in ps_dict.keys()}
+            
         for i in range(len(evolution)):
             for k in pltdict.keys():
                 pltdict[k] += [evolution[i][0][k]]
@@ -126,7 +139,7 @@ if __name__ == "__main__":
             if 'absorber' in k:
                 style = '--'
             plt.plot([i for i in range(len(pltdict[k]))], np.where(np.array(pltdict[k]) < 0., 0., np.array(pltdict[k])), label = k, linestyle=style)
-        
+            
         plt.legend()
         # create a second y-axis for the loss values. Make the line black and thicker than the others; make the loss y axis log scale
         plt.twinx()
@@ -137,10 +150,20 @@ if __name__ == "__main__":
         
         plt.legend()
         plt.xlabel('step')
-        plt.savefig(outpath+'evolution.png')
+        plt.savefig(outpath+'evolution_continuous.png')
         plt.close()
 
+        for k in pltdict.keys():
+            if 'material' in k:
+                style = '-'
+                if 'abs' in k:
+                    style = '--'
+                plt.plot([i for i in range(len(pltdict[k]))], np.where(np.array(pltdict[k]) < 0., 0., np.array(pltdict[k])*10), label = k, linestyle=style)
 
+        plt.legend()
+        plt.xlabel('step')
+        plt.savefig(outpath+'evolution_material.png')
+        plt.close()
 
 
     #save_evolution(parameters, 1., 1.)
@@ -201,6 +224,7 @@ if __name__ == "__main__":
         #these are un-normalised quantities
         surr_out, reco_out, true_in = surrogate_model.apply_model_in_batches(surrogate_dataset, batch_size=512)
         make_check_plot(surr_out, reco_out, true_in, evolution, outname = "pretrain")
+    print(parameters)
     pre_train()
     
 
@@ -288,8 +312,3 @@ if __name__ == "__main__":
             time.sleep(2)
         save_evolution(updated_detector_parameters, o_loss, reco_loss, reco_model, gen.means, gen.stds) #save the optimised parameters here, not the ones for the next box generation
         #exit()
-
-    
-    
-
-