Skip to content

Commit d2d6c97

Browse files
authored
Fix hf_pt_spectrum script (#447)
1 parent e60e7a7 commit d2d6c97

File tree

2 files changed

+46
-181
lines changed

2 files changed

+46
-181
lines changed

FirstAnalysis/hf_analysis_utils.py

Lines changed: 28 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
author: 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

Comments
 (0)