66author: Fabrizio Grosa <[email protected] >, CERN 77"""
88
9- import numpy as np
9+ import numpy as np # pylint: disable=import-error
1010
1111
1212# pylint: disable=too-many-arguments
@@ -46,12 +46,7 @@ def compute_crosssection(
4646 - crosssec_unc: cross-section statistical uncertainty
4747 """
4848
49- crosssection = (
50- rawy
51- * frac
52- * sigma_mb
53- / (2 * delta_pt * delta_y * eff_times_acc * n_events * br )
54- )
49+ crosssection = rawy * frac * sigma_mb / (2 * delta_pt * delta_y * eff_times_acc * n_events * br )
5550 if method_frac == "Nb" :
5651 crosssec_unc = rawy_unc / (rawy * frac ) * crosssection
5752 else :
@@ -68,7 +63,7 @@ def compute_fraction_fc(
6863 cross_sec_fd ,
6964 raa_prompt = 1.0 ,
7065 raa_fd = 1.0 ,
71- ):
66+ ) -> "tuple[list[float], list[float]]" :
7267 """
7368 Method to get fraction of prompt / FD fraction with fc method
7469
@@ -89,16 +84,13 @@ def compute_fraction_fc(
8984 - frac_fd: list of fraction of non-prompt D (central, min, max)
9085 """
9186
92- if not isinstance (cross_sec_prompt , list ) and isinstance (cross_sec_prompt , float ):
93- cross_sec_prompt = [cross_sec_prompt ]
94- if not isinstance (cross_sec_fd , list ) and isinstance (cross_sec_fd , float ):
95- cross_sec_fd = [cross_sec_fd ]
96- if not isinstance (raa_prompt , list ) and isinstance (raa_prompt , float ):
97- raa_prompt = [raa_prompt ]
98- if not isinstance (raa_fd , list ) and isinstance (raa_fd , float ):
99- raa_fd = [raa_fd ]
87+ cross_sec_prompt_l = cross_sec_prompt if isinstance (cross_sec_prompt , list ) else [cross_sec_prompt ]
88+ cross_sec_fd_l = cross_sec_fd if isinstance (cross_sec_fd , list ) else [cross_sec_fd ]
89+ raa_prompt_l = raa_prompt if isinstance (raa_prompt , list ) else [raa_prompt ]
90+ raa_fd_l = raa_fd if isinstance (raa_fd , list ) else [raa_fd ]
10091
101- frac_prompt , frac_fd = [], []
92+ frac_prompt : list [float ] = []
93+ frac_fd : list [float ] = []
10294 if acc_eff_prompt == 0 :
10395 frac_fd_cent = 1.0
10496 frac_prompt_cent = 0.0
@@ -112,40 +104,14 @@ def compute_fraction_fc(
112104 frac_fd = [frac_fd_cent , frac_fd_cent , frac_fd_cent ]
113105 return frac_prompt , frac_fd
114106
115- for i_sigma , (sigma_p , sigma_f ) in enumerate (zip (cross_sec_prompt , cross_sec_fd )):
116- for i_raa , (raa_p , raa_f ) in enumerate (zip (raa_prompt , raa_fd )):
107+ for i_sigma , (sigma_p , sigma_f ) in enumerate (zip (cross_sec_prompt_l , cross_sec_fd_l )):
108+ for i_raa , (raa_p , raa_f ) in enumerate (zip (raa_prompt_l , raa_fd_l )):
117109 if i_sigma == 0 and i_raa == 0 :
118- frac_prompt_cent = 1.0 / (
119- 1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p
120- )
121- frac_fd_cent = 1.0 / (
122- 1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f
123- )
110+ frac_prompt_cent = 1.0 / (1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p )
111+ frac_fd_cent = 1.0 / (1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f )
124112 else :
125- frac_prompt .append (
126- 1.0
127- / (
128- 1
129- + acc_eff_fd
130- / acc_eff_prompt
131- * sigma_f
132- / sigma_p
133- * raa_f
134- / raa_p
135- )
136- )
137- frac_fd .append (
138- 1.0
139- / (
140- 1
141- + acc_eff_prompt
142- / acc_eff_fd
143- * sigma_p
144- / sigma_f
145- * raa_p
146- / raa_f
147- )
148- )
113+ frac_prompt .append (1.0 / (1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p ))
114+ frac_fd .append (1.0 / (1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f ))
149115
150116 if frac_prompt and frac_fd :
151117 frac_prompt .sort ()
@@ -172,7 +138,7 @@ def compute_fraction_nb(
172138 sigma_mb ,
173139 raa_ratio = 1.0 ,
174140 taa = 1.0 ,
175- ):
141+ ) -> "list[float]" :
176142 """
177143 Method to get fraction of prompt / FD fraction with Nb method
178144
@@ -196,102 +162,39 @@ def compute_fraction_nb(
196162 - frac: list of fraction of prompt (non-prompt) D (central, min, max)
197163 """
198164
199- if not isinstance (crosssection , list ) and isinstance (crosssection , float ):
200- crosssection = [crosssection ]
201-
202- if not isinstance (raa_ratio , list ) and isinstance (raa_ratio , float ):
203- raa_ratio = [raa_ratio ]
165+ crosssection_l = crosssection if isinstance (crosssection , list ) else [crosssection ]
166+ raa_ratio_l = raa_ratio if isinstance (raa_ratio , list ) else [raa_ratio ]
204167
205- frac = []
206- for i_sigma , sigma in enumerate (crosssection ):
207- for i_raa_ratio , raa_rat in enumerate (raa_ratio ):
168+ frac : list [ float ] = []
169+ for i_sigma , sigma in enumerate (crosssection_l ):
170+ for i_raa_ratio , raa_rat in enumerate (raa_ratio_l ):
208171 raa_other = 1.0
209172 if i_sigma == 0 and i_raa_ratio == 0 :
210173 if raa_rat == 1.0 and taa == 1.0 : # pp
211- frac_cent = (
212- 1
213- - sigma
214- * delta_pt
215- * delta_y
216- * acc_eff_other
217- * br
218- * n_events
219- * 2
220- / rawy
221- / sigma_mb
222- )
174+ frac_cent = 1 - sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 / rawy / sigma_mb
223175 else : # p-Pb or Pb-Pb: iterative evaluation of Raa needed
224176 delta_raa = 1.0
225177 while delta_raa > 1.0e-3 :
226178 raw_fd = (
227- taa
228- * raa_rat
229- * raa_other
230- * sigma
231- * delta_pt
232- * delta_y
233- * acc_eff_other
234- * br
235- * n_events
236- * 2
179+ taa * raa_rat * raa_other * sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2
237180 )
238181 frac_cent = 1 - raw_fd / rawy
239182 raa_other_old = raa_other
240- raa_other = (
241- frac_cent
242- * rawy
243- * sigma_mb
244- / 2
245- / acc_eff_same
246- / delta_pt
247- / delta_y
248- / br
249- / n_events
250- )
183+ raa_other = frac_cent * rawy * sigma_mb / 2 / acc_eff_same / delta_pt / delta_y / br / n_events
251184 delta_raa = abs ((raa_other - raa_other_old ) / raa_other )
252185 else :
253186 if raa_rat == 1.0 and taa == 1.0 : # pp
254- frac .append (
255- 1
256- - sigma
257- * delta_pt
258- * delta_y
259- * acc_eff_other
260- * br
261- * n_events
262- * 2
263- / rawy
264- / sigma_mb
265- )
187+ frac .append (1 - sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 / rawy / sigma_mb )
266188 else : # p-Pb or Pb-Pb: iterative evaluation of Raa needed
267189 delta_raa = 1.0
268190 frac_tmp = 1.0
269191 while delta_raa > 1.0e-3 :
270192 raw_fd = (
271- taa
272- * raa_rat
273- * raa_other
274- * sigma
275- * delta_pt
276- * delta_y
277- * acc_eff_other
278- * br
279- * n_events
280- * 2
193+ taa * raa_rat * raa_other * sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2
281194 )
282195 frac_tmp = 1 - raw_fd / rawy
283196 raa_other_old = raa_other
284- raa_other = (
285- frac_tmp
286- * rawy
287- * sigma_mb
288- / 2
289- / acc_eff_same
290- / delta_pt
291- / delta_y
292- / br
293- / n_events
294- )
197+ raa_other = frac_tmp * rawy * sigma_mb / 2 / acc_eff_same / delta_pt / delta_y / br / n_events
295198 delta_raa = abs ((raa_other - raa_other_old ) / raa_other )
296199 frac .append (frac_tmp )
297200
@@ -323,8 +226,6 @@ def get_hist_binlimits(histo):
323226 n_limits = histo .GetNbinsX () + 1
324227 low_edge = histo .GetBinLowEdge (1 )
325228 bin_width = histo .GetBinWidth (1 )
326- bin_limits = np .array (
327- [low_edge + i_bin * bin_width for i_bin in range (n_limits )], "d"
328- )
229+ bin_limits = np .array ([low_edge + i_bin * bin_width for i_bin in range (n_limits )], "d" )
329230
330231 return bin_limits
0 commit comments