source/SR_event_generator/sr_generator.py

0001 import numpy as np
0002 import pandas as pd
0003 import yaml
0004 import os
0005 import ROOT
0006 import matplotlib.pyplot as plt
0007 import time
0008 import argparse
0009 #import pyhepmc_ng as hep
0010 import pyhepmc as hep
0011
0012 # =============================================================
0013 class sr_generator:
0014     # ---------------------------------------
0015     def __init__(self,config):
0016
0017         with open(config,'r') as stream:
0018             self.config = yaml.safe_load(stream)
0019
0020         print('********************************************************')
0021         print('Synchrotron-radiation event generator from Synrad+ table')
0022         print('authors: Reynier Cruz-Torres*, Benjamen Sterwerf**')
0023         print('* Lawrence Berkeley National Laboratory')
0024         print('** University of California, Berkeley')
0025         print('')
0026         if 'input_single_photons' in self.config:
0027             self.input_single_photons = self.config['input_single_photons']
0028             print('Loading single photons from:',self.input_single_photons)
0029         else:
0030             print("Provide 'input_single_photons' in config file")
0031             exit()
0032
0033         if 'n_events' in self.config:
0034             self.n_events = self.config['n_events']
0035             print("Number of events to run:",self.n_events)
0036         else:
0037             print("Number of events not provided in config file. Defaulting to 10")
0038             self.n_events = 10
0039
0040         if 'integration_window' in self.config:
0041             self.integration_window = (float)(self.config['integration_window'])
0042             print("Time integration window selected:",self.integration_window)
0043         else:
0044             print("Time integration window not provided in config file. Defaulting to 100 ns")
0045             self.integration_window = (float)(1e-07)
0046
0047         if 'seed' in self.config:
0048             self.seed = self.config['seed']
0049         else:
0050             self.seed = 0
0051
0052         if self.seed == 0:
0053             print("The generator seed won't be constrained")
0054         else:
0055             print("The generator seed will be:",self.seed)
0056             ROOT.gRandom.SetSeed(self.seed)
0057
0058         self.outpath = 'output_plots'
0059         if not os.path.exists(self.outpath):
0060             os.makedirs(self.outpath)
0061
0062         self.pid = 22 # 22 = photon, can be changed for testing
0063         self.status_000 = 4 # try 1, 2, or 4 (see page 31 in https://arxiv.org/pdf/1912.08005.pdf)
0064         self.status_xyz = 1
0065
0066         if self.seed > 0:
0067             self.outfname = 'SR_out_int_window_{}ns_nevents_{}_pid_{}_status_{}_{}_seed_{}.hepmc'.format(self.integration_window*1e+09,
0068             self.n_events,self.pid,self.status_000,self.status_xyz,self.seed)
0069         else:
0070             self.outfname = 'SR_out_int_window_{}ns_nevents_{}_pid_{}_status_{}_{}.hepmc'.format(self.integration_window*1e+09,
0071             self.n_events,self.pid,self.status_000,self.status_xyz)
0072
0073         print('Generated events will be saved in',self.outfname)
0074         print('********************************************************')
0075
0076         self.load_single_photons()
0077
0078     # ---------------------------------------
0079     def load_single_photons(self):
0080         print('Loading single photons')
0081
0082         self.df = pd.read_csv(self.input_single_photons)
0083         self.df = self.df.sort_values('NormFact')
0084
0085         n_entries = len(self.df)
0086         self.h1_df = ROOT.TH1D('h1_df',';entry;W [#gamma/sec]',n_entries,0,n_entries)
0087         for i in range(n_entries):
0088             self.h1_df.SetBinContent(i+1,self.df['NormFact'].iloc[i])
0089
0090     # ---------------------------------------
0091     def generate_an_event(self):
0092         event = []
0093         integrated_so_far = 0.
0094         while integrated_so_far < self.integration_window:
0095             x = self.h1_df.FindBin(self.h1_df.GetRandom())
0096             if x >= 1800000: continue
0097             photon = self.df.iloc[x]
0098             integrated_so_far += 1./photon['NormFact']
0099             event.append(photon)
0100         return event
0101
0102     # ---------------------------------------
0103     def generate(self):
0104         print('Generating SR events')
0105
0106         events = []
0107         hep_events = []
0108         photons_per_event = []
0109         z_dist = []
0110         rho_dist = []
0111
0112         for i in range(self.n_events):
0113             event = self.generate_an_event()
0114
0115             # ---------------------------------------------------
0116             # Save to hepmc format
0117             # implemented following the example from:
0118             # https://github.com/scikit-hep/pyhepmc/blob/master/tests/test_basic.py
0119
0120             evt = hep.GenEvent(hep.Units.GEV, hep.Units.MM)
0121             evt.event_number = i+1
0122             particles_out = []
0123             particles_in = []
0124             vertices = []
0125
0126             # loop over each photon in the event
0127             for g in range(len(event)):
0128
0129                 x = event[g]['x']
0130                 y = event[g]['y']
0131                 z = event[g]['z']
0132
0133                 z_dist.append(z)
0134                 rho_dist.append(np.sqrt(x*x+y*y))
0135
0136                 px = event[g]['px']
0137                 py = event[g]['py']
0138                 pz = event[g]['pz']
0139                 E = np.sqrt(px**2 + py**2 + pz**2)
0140
0141                 pinx = E*x/np.sqrt(x*x+y*y+z*z)
0142                 piny = E*y/np.sqrt(x*x+y*y+z*z)
0143                 pinz = E*z/np.sqrt(x*x+y*y+z*z)
0144
0145                 # Particles going into the vertex
0146                 pin = hep.GenParticle((pinx,piny,pinz,E),self.pid,self.status_000)
0147                 pin.generated_mass = 0.
0148                 evt.add_particle(pin)
0149                 particles_in.append(pin)
0150
0151                 # Particles coming out of the vertex
0152                 pout = hep.GenParticle((px,py,pz,E),self.pid,self.status_xyz)
0153                 pout.generated_mass = 0.
0154                 evt.add_particle(pout)
0155                 particles_out.append(pout)
0156
0157                 # make sure vertex is not optimized away by WriterAscii
0158                 v1 = hep.GenVertex((x,y,z,0.))
0159                 v1.add_particle_in(pin)
0160                 v1.add_particle_out(pout)
0161                 evt.add_vertex(v1)
0162                 vertices.append(v1)
0163
0164             # -------------------
0165             #evt.weights = [1.0]
0166             if i == 0:
0167                 evt.run_info = hep.GenRunInfo()
0168                 #evt.run_info.weight_names = ["0"]
0169             hep_events.append(evt)
0170             # ---------------------------------------------------
0171
0172             photons_per_event.append(len(event))
0173
0174             if i < 6:
0175                 events.append(event)
0176
0177         with hep.io.WriterAscii(self.outfname) as f:
0178             for e in hep_events:
0179                 f.write_event(e)
0180
0181         # --------------------------
0182         # Make some plots
0183         self.plot_histo(photons_per_event,'# photons per event','Nphotons_per_event.png')
0184         self.plot_2d_scatter(events,'x','y','x [mm]','y [mm]',(-40,40),(-40,40),'events_x_v_y.png')
0185         self.plot_2d_scatter(events,'z','x','z [mm]','x [mm]',(-5000,3000),(-40,40),'events_z_v_x.png')
0186         self.plot_2d_scatter(events,'z','y','z [mm]','y [mm]',(-5000,3000),(-40,40),'events_z_v_y.png')
0187
0188     # ---------------------------------------
0189     def plot_histo(self,hist,xlabel,fname):
0190         plt.figure()
0191         plt.hist(hist)
0192         plt.xlabel(xlabel)
0193         plt.ylabel('Counts')
0194         plt.tight_layout()
0195         plt.savefig(os.path.join(self.outpath,fname),dpi=400)
0196
0197     # ---------------------------------------
0198     def plot_2d_scatter(self,events,xl,yl,labelx,labely,xlim,ylim,fname):
0199         plt.figure(figsize=(13,8))
0200
0201         for i in range(6):
0202             plt.subplot(2,3,i+1)
0203
0204             x, y = [], []
0205
0206             for j in range(len(events[i])):
0207                 x.append(events[i][j][xl])
0208                 y.append(events[i][j][yl])
0209
0210             plt.scatter(x,y,marker='o',alpha=0.2)
0211
0212             plt.xlim(xlim[0],xlim[1])
0213             plt.ylim(ylim[0],ylim[1])
0214             plt.xlabel(labelx)
0215             plt.ylabel(labely)
0216
0217             plt.text(-8,0,r'{} $\gamma$'.format(len(events[i])),fontsize=15)
0218
0219         plt.tight_layout()
0220         plt.savefig(os.path.join(self.outpath,fname),dpi=400)
0221
0222 # =============================================================
0223 if __name__=='__main__':
0224     t0 = time.time()
0225
0226     parser = argparse.ArgumentParser(description='Generating synchrotron radiation events')
0227     parser.add_argument('-c', '--configFile',
0228                         help='Path of config file for analysis',
0229                         action='store', type=str,
0230                         default='config.yaml')
0231     args = parser.parse_args()
0232
0233     generator = sr_generator(args.configFile)
0234     generator.generate()
0235
0236     print('Overall running time:',np.round((time.time()-t0)/60.,2),'min')