Skip to content

Commit

Permalink
Revert formatting of unrelated files
Browse files Browse the repository at this point in the history
  • Loading branch information
@vkucera
vkucera committed Apr 10, 2024
1 parent 6e129b2 commit e5f95ae
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 42 deletions.
108 changes: 81 additions & 27 deletions PWGHF/D2H/Macros/compute_fraction_cutvar.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,23 @@ def main(config):
cfg = json.load(fil)

hist_rawy, hist_effp, hist_effnp = ([] for _ in range(3))
for filename_rawy, filename_eff in zip(cfg["rawyields"]["inputfiles"], cfg["efficiencies"]["inputfiles"]):
infile_rawy = ROOT.TFile.Open(os.path.join(cfg["rawyields"]["inputdir"], filename_rawy))
for filename_rawy, filename_eff in zip(
cfg["rawyields"]["inputfiles"], cfg["efficiencies"]["inputfiles"]
):
infile_rawy = ROOT.TFile.Open(
os.path.join(cfg["rawyields"]["inputdir"], filename_rawy)
)
hist_rawy.append(infile_rawy.Get(cfg["rawyields"]["histoname"]))
hist_rawy[-1].SetDirectory(0)
infile_rawy.Close()

infile_eff = ROOT.TFile.Open(os.path.join(cfg["efficiencies"]["inputdir"], filename_eff))
infile_eff = ROOT.TFile.Open(
os.path.join(cfg["efficiencies"]["inputdir"], filename_eff)
)
hist_effp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["prompt"]))
hist_effnp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["nonprompt"]))
hist_effnp.append(
infile_eff.Get(cfg["efficiencies"]["histonames"]["nonprompt"])
)
hist_effp[-1].SetDirectory(0)
hist_effnp[-1].SetDirectory(0)
infile_eff.Close()
Expand All @@ -63,14 +71,20 @@ def main(config):
)
set_object_style(hist_covariance)

output = ROOT.TFile(os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate")
output = ROOT.TFile(
os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate"
)
n_sets = len(hist_rawy)
for ipt in range(hist_rawy[0].GetNbinsX()):
pt_min = hist_rawy[0].GetXaxis().GetBinLowEdge(ipt + 1)
pt_max = hist_rawy[0].GetXaxis().GetBinUpEdge(ipt + 1)

rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = (np.zeros(n_sets) for _ in range(6))
for iset, (hrawy, heffp, heffnp) in enumerate(zip(hist_rawy, hist_effp, hist_effnp)):
rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = (
np.zeros(n_sets) for _ in range(6)
)
for iset, (hrawy, heffp, heffnp) in enumerate(
zip(hist_rawy, hist_effp, hist_effnp)
):
rawy.itemset(iset, hrawy.GetBinContent(ipt + 1))
effp.itemset(iset, heffp.GetBinContent(ipt + 1))
effnp.itemset(iset, heffnp.GetBinContent(ipt + 1))
Expand All @@ -81,26 +95,40 @@ def main(config):
minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp)
minimiser.minimise_system(cfg["minimisation"]["correlated"])

hist_corry_prompt.SetBinContent(ipt + 1, minimiser.get_prompt_yield_and_error()[0])
hist_corry_prompt.SetBinError(ipt + 1, minimiser.get_prompt_yield_and_error()[1])
hist_corry_nonprompt.SetBinContent(ipt + 1, minimiser.get_nonprompt_yield_and_error()[0])
hist_corry_nonprompt.SetBinError(ipt + 1, minimiser.get_nonprompt_yield_and_error()[1])
hist_corry_prompt.SetBinContent(
ipt + 1, minimiser.get_prompt_yield_and_error()[0]
)
hist_corry_prompt.SetBinError(
ipt + 1, minimiser.get_prompt_yield_and_error()[1]
)
hist_corry_nonprompt.SetBinContent(
ipt + 1, minimiser.get_nonprompt_yield_and_error()[0]
)
hist_corry_nonprompt.SetBinError(
ipt + 1, minimiser.get_nonprompt_yield_and_error()[1]
)
hist_covariance.SetBinContent(ipt + 1, minimiser.get_prompt_nonprompt_cov())
hist_covariance.SetBinError(ipt + 1, 0)

canv_rawy, histos_rawy, leg_r = minimiser.plot_result(f"_pt{pt_min:.0f}_{pt_max:.0f}")
canv_rawy, histos_rawy, leg_r = minimiser.plot_result(
f"_pt{pt_min:.0f}_{pt_max:.0f}"
)
output.cd()
canv_rawy.Write()
for _, hist in histos_rawy.items():
hist.Write()

canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies(f"_pt{pt_min:.0f}_{pt_max:.0f}")
canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies(
f"_pt{pt_min:.0f}_{pt_max:.0f}"
)
output.cd()
canv_eff.Write()
for _, hist in histos_eff.items():
hist.Write()

canv_frac, histos_frac, leg_f = minimiser.plot_fractions(f"_pt{pt_min:.0f}_{pt_max:.0f}")
canv_frac, histos_frac, leg_f = minimiser.plot_fractions(
f"_pt{pt_min:.0f}_{pt_max:.0f}"
)
output.cd()
canv_frac.Write()
for _, hist in histos_frac.items():
Expand All @@ -114,21 +142,47 @@ def main(config):
output_name_rawy_pdf = f"Distr_{cfg['output']['file'].replace('.root', '.pdf')}"
output_name_eff_pdf = f"Eff_{cfg['output']['file'].replace('.root', '.pdf')}"
output_name_frac_pdf = f"Frac_{cfg['output']['file'].replace('.root', '.pdf')}"
output_name_covmat_pdf = f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}"
output_name_covmat_pdf = (
f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}"
)
if ipt == 0:
canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}[")
canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}[")
canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}[")
canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}[")
canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}")
canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}")
canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}")
canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}")
canv_rawy.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}["
)
canv_eff.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}["
)
canv_frac.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}["
)
canv_cov.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}["
)
canv_rawy.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}"
)
canv_eff.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}"
)
canv_frac.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}"
)
canv_cov.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}"
)
if ipt == hist_rawy[0].GetNbinsX() - 1:
canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}]")
canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}]")
canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}]")
canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}]")
canv_rawy.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}]"
)
canv_eff.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}]"
)
canv_frac.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}]"
)
canv_cov.SaveAs(
f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}]"
)

output.cd()
hist_corry_prompt.Write()
Expand Down
60 changes: 45 additions & 15 deletions PWGHF/D2H/Macros/cut_variation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,16 +77,26 @@ def __check_input_consistency(self):
Helper method to check self consistency of inputs
"""

if len(self.eff_prompt) != self.n_sets or len(self.eff_nonprompt) != self.n_sets:
if (
len(self.eff_prompt) != self.n_sets
or len(self.eff_nonprompt) != self.n_sets
):
print("ERROR: number of raw yields and efficiencies not consistent! Exit")
sys.exit()

if len(self.unc_raw_yields) != self.n_sets:
print("ERROR: number of raw yields and raw-yield uncertainties not consistent! Exit")
print(
"ERROR: number of raw yields and raw-yield uncertainties not consistent! Exit"
)
sys.exit()

if len(self.unc_eff_prompt) != self.n_sets or len(self.unc_eff_nonprompt) != self.n_sets:
print("ERROR: number of raw yields and efficiency uncertainties not consistent! Exit")
if (
len(self.unc_eff_prompt) != self.n_sets
or len(self.unc_eff_nonprompt) != self.n_sets
):
print(
"ERROR: number of raw yields and efficiency uncertainties not consistent! Exit"
)
sys.exit()

def __initialise_objects(self):
Expand All @@ -108,7 +118,9 @@ def __initialise_objects(self):
self.unc_frac_prompt = np.zeros(shape=self.n_sets)
self.unc_frac_nonprompt = np.zeros(shape=self.n_sets)

for i_set, (rawy, effp, effnp) in enumerate(zip(self.raw_yields, self.eff_prompt, self.eff_nonprompt)):
for i_set, (rawy, effp, effnp) in enumerate(
zip(self.raw_yields, self.eff_prompt, self.eff_nonprompt)
):
self.m_rawy.itemset(i_set, rawy)
self.m_eff.itemset((i_set, 0), effp)
self.m_eff.itemset((i_set, 1), effnp)
Expand Down Expand Up @@ -140,7 +152,9 @@ def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100)
zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt)
):
for i_col, (rw_unc_col, effp_unc_col, effnp_unc_col) in enumerate(
zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt)
zip(
self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt
)
):
unc_row = np.sqrt(
rw_unc_row**2
Expand Down Expand Up @@ -175,12 +189,16 @@ def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100)
self.m_covariance = np.linalg.inv(np.linalg.cholesky(self.m_covariance))
self.m_covariance = self.m_covariance.T * self.m_covariance

self.m_corr_yields = self.m_covariance * (m_eff_tr * self.m_weights) * self.m_rawy
self.m_corr_yields = (
self.m_covariance * (m_eff_tr * self.m_weights) * self.m_rawy
)
self.m_res = self.m_eff * self.m_corr_yields - self.m_rawy

rel_delta = [
(self.m_corr_yields.item(0) - m_corr_yields_old.item(0)) / self.m_corr_yields.item(0),
(self.m_corr_yields.item(1) - m_corr_yields_old.item(1)) / self.m_corr_yields.item(1),
(self.m_corr_yields.item(0) - m_corr_yields_old.item(0))
/ self.m_corr_yields.item(0),
(self.m_corr_yields.item(1) - m_corr_yields_old.item(1))
/ self.m_corr_yields.item(1),
]

if rel_delta[0] < precision and rel_delta[1] < precision:
Expand All @@ -195,10 +213,18 @@ def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100)
for i_set, (effp, effnp) in enumerate(zip(self.eff_prompt, self.eff_nonprompt)):
rawyp = effp * self.m_corr_yields.item(0)
rawynp = effnp * self.m_corr_yields.item(1)
der_fp_p = (effp * (rawyp + rawynp) - effp**2 * self.m_corr_yields.item(0)) / (rawyp + rawynp) ** 2
der_fp_np = -effp * effnp * self.m_corr_yields.item(0) / (rawyp + rawynp) ** 2
der_fnp_np = (effnp * (rawyp + rawynp) - effnp**2 * self.m_corr_yields.item(1)) / (rawyp + rawynp) ** 2
der_fnp_p = -effp * effnp * self.m_corr_yields.item(1) / (rawyp + rawynp) ** 2
der_fp_p = (
effp * (rawyp + rawynp) - effp**2 * self.m_corr_yields.item(0)
) / (rawyp + rawynp) ** 2
der_fp_np = (
-effp * effnp * self.m_corr_yields.item(0) / (rawyp + rawynp) ** 2
)
der_fnp_np = (
effnp * (rawyp + rawynp) - effnp**2 * self.m_corr_yields.item(1)
) / (rawyp + rawynp) ** 2
der_fnp_p = (
-effp * effnp * self.m_corr_yields.item(1) / (rawyp + rawynp) ** 2
)

unc_fp = np.sqrt(
der_fp_p**2 * self.m_covariance.item(0, 0)
Expand Down Expand Up @@ -311,7 +337,9 @@ def plot_result(self, suffix=""):
rawy_nonprompt = self.m_corr_yields.item(1) * effnp
unc_rawy_nonprompt = np.sqrt(self.m_covariance.item(1, 1)) * effnp
unc_sum = np.sqrt(
unc_rawy_prompt**2 + unc_rawy_nonprompt**2 + 2 * self.m_covariance.item(1, 0) * effp * effnp
unc_rawy_prompt**2
+ unc_rawy_nonprompt**2
+ 2 * self.m_covariance.item(1, 0) * effp * effnp
)

hist_raw_yield_prompt.SetBinContent(i_bin + 1, rawy_prompt)
Expand All @@ -323,7 +351,9 @@ def plot_result(self, suffix=""):

set_object_style(hist_raw_yield)
set_object_style(hist_raw_yield_prompt, color=ROOT.kRed + 1, fillstyle=3145)
set_object_style(hist_raw_yield_nonprompt, color=ROOT.kAzure + 4, fillstyle=3154)
set_object_style(
hist_raw_yield_nonprompt, color=ROOT.kAzure + 4, fillstyle=3154
)
set_object_style(hist_raw_yield_sum, color=ROOT.kGreen + 2, fillstyle=0)

canvas = ROOT.TCanvas(f"cRawYieldVsCut{suffix}", "", 500, 500)
Expand Down

0 comments on commit e5f95ae

Please sign in to comment.