-
Notifications
You must be signed in to change notification settings - Fork 0
/
run_simple_energy_loss_parallel.py
148 lines (115 loc) · 5.33 KB
/
run_simple_energy_loss_parallel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter, AutoMinorLocator
import matplotlib.ticker as mtick
import scipy #.optimize
import scipy.interpolate
import scipy.integrate
from temperature_profile import brick_profile, Bjorken_hydro_profile
from parton_emission_rates import energy_loss_rates
from solver_euler import parton_evolution_solver_euler
from solver_rk import parton_evolution_solver_rk
import time
from joblib import Parallel, delayed
import signal
class TimeoutException(Exception): # Custom exception class
pass
def timeout_handler(signum, frame): # Custom signal handler
raise TimeoutException
# Change the behavior of SIGALRM
signal.signal(signal.SIGALRM, timeout_handler)
hbarc=0.1973
# Initial conditions
def P_g_tau0(p):
p0=1.75
return np.power(p0*p0+p*p,-5.)
def simulate(ii, params, p_min=1, p_max=20, num_p_bins=20):
print(f'Working on design {ii+1}')
start_time = time.time()
signal.alarm(600) # time out after 10 minutes
try:
params=params.flatten()
param_dict={
'T0_in_GeV':0.3,
'tau0':0.2,
'T_final_in_GeV':0.15,
'alpha_s':params[0],
'N_f':0,
'mD_factor':1.0,
'exponent_inel':params[1],
'exponent_el':params[2],
'RAA_pT_binnings':np.linspace(p_min, p_max, num_p_bins),
'scale_inel':params[3],
'scale_el':params[4]
}
T0_in_GeV=param_dict['T0_in_GeV']
tau0=param_dict['tau0']
T_final_in_GeV=param_dict['T_final_in_GeV']
alpha_s=param_dict['alpha_s']
N_f=param_dict['N_f']
mD_factor=param_dict['mD_factor']
exponent_inel=param_dict['exponent_inel']
exponent_el=param_dict['exponent_el']
RAA_pT_binnings=param_dict['RAA_pT_binnings']
# The new parameters
scale_inel=param_dict['scale_inel']
scale_el=param_dict['scale_el']
#################################################
############## Temperature profile ##############
##################################################
#T_profile=brick_profile(T0_in_GeV=T0_in_GeV)
T_profile=Bjorken_hydro_profile(T0_in_GeV=T0_in_GeV, tau0=tau0)
######################################################
############## Parton energy loss rates ##############
######################################################
K_factor_fct_inel=lambda T, scale_inel=scale_inel, exponent_inel=exponent_inel: (1.+np.power(T/scale_inel,exponent_inel))
K_factor_fct_elastic=lambda T, scale_el=scale_el, exponent_el=exponent_el: (1.+np.power(T/scale_el,exponent_el))
energy_loss_rate=energy_loss_rates(alpha_s = alpha_s, N_f=N_f, mD_factor=mD_factor, K_factor_fct_inel=K_factor_fct_inel, K_factor_fct_elastic=K_factor_fct_elastic)
#######################################################
############## Parton energy loss solver ##############
#######################################################
# Initialize and use the solver
num_p_solver=num_p_bins
pmin_solver=p_min
pmax_solver=p_max
#parton_evolution_solver=parton_evolution_solver_euler(initial_condition_fct=P_g_tau0, tau0=tau0, T_profile=T_profile, energy_loss_rate=energy_loss_rate, num_p=num_p_solver, pmin=pmin_solver, pmax=pmax_solver)
#P_final_fct=parton_evolution_solver.evolve_to_min_temperature(dtau=dtau_adaptive, T_min_in_GeV=T_final_in_GeV, use_adaptive_timestep=True)
parton_evolution_solver=parton_evolution_solver_rk(initial_condition_fct=P_g_tau0, tau0=tau0, T_profile=T_profile, energy_loss_rate=energy_loss_rate, num_p=num_p_solver, pmin=pmin_solver, pmax=pmax_solver)
P_final_fct=parton_evolution_solver.evolve_to_min_temperature(T_min_in_GeV=T_final_in_GeV)
# Compute some "RAA"-equivalent
P_initial=P_g_tau0(RAA_pT_binnings)
P_final=P_final_fct(RAA_pT_binnings)
result=P_final/P_initial
end_time = time.time()
delta_t = end_time - start_time
if delta_t > 60 :
print(f'For model parameters {params} takes {delta_t} S')
return result
except TimeoutException: # what to do if it takes longer than 10 minutes
print(f'For model parameters {params} takes more than 10 mins')
return [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
else:
# Reset the alarm
signal.alarm(0)
def run_simulation(design_matrix, p_min=1, p_max=20, num_p_bins=20):
"""Simulate energy loss of a jet for each design point
parameters
----------
see definition of "param_dict"
p_min, p_max, num_p_bins : Fix the momentem range and number of bins for R_AA calculations
Returns
-------
numpy array
R_AA binned according to the momentem range speceifed at the input
"""
par = True # Run in parallel
# Evaluate simulation at validation points
simulate_run = lambda ii, params: simulate(ii, params, p_min=p_min, p_max=p_max, num_p_bins=num_p_bins)
if par:
cores = 4
observations = Parallel(n_jobs=cores)(delayed(simulate_run)(ii,params) for ii, params in enumerate(design_matrix))
else:
observations = [simulate_run(ii,params) for ii, params in enumerate(design_matrix)]
observations=np.array(observations)
return observations