def test_project_3d_to_2d(self):
hist_3d = _get_hist(3)
# populate overflow bins to make sure that they are treated as expected
fill_hist(hist_3d, np.random.uniform(-1, 0, (100, 3)))
fill_hist(hist_3d, np.random.uniform(1, 2, (100, 3)))
val3d, err3d = hu.get_array(hist_3d), hu.get_array(hist_3d, errors=True)
hist_xy = hu.project(hist_3d, 'xy')
val, err = hu.get_array(hist_xy), hu.get_array(hist_xy, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=2))
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=2)))
hist_yz = hu.project(hist_3d, 'yz')
val, err = hu.get_array(hist_yz), hu.get_array(hist_yz, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=0))
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=0)))
hist_zx = hu.project(hist_3d, 'zx')
val, err = hu.get_array(hist_zx), hu.get_array(hist_zx, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=1).T)
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=1)).T)
hist_yx = hu.project(hist_3d, 'yx')
val, err = hu.get_array(hist_yx), hu.get_array(hist_yx, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=2).T)
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=2)).T)
def test_project_THnD_integer_index(self):
"""Check that using integer indices works as well"""
hist_3d = _get_hist(3)
proj_lab = hu.project(hist_3d, 'xy')
proj_ind = hu.project(hist_3d, [0, 1])
val_lab, err_lab = hu.get_array(proj_lab), hu.get_array(proj_lab, errors=True)
val_ind, err_ind = hu.get_array(proj_ind), hu.get_array(proj_ind, errors=True)
npt.assert_equal(val_lab, val_ind)
npt.assert_equal(err_lab, err_ind)
def test_project_5d_to_4d(self):
"""Test that also returning into a THn works"""
hist5d = _get_hist(5)
val5d, err5d = hu.get_array(hist5d), hu.get_array(hist5d, errors=True)
hist_0234 = hu.project(hist5d, [0, 2, 3, 4])
val, err = hu.get_array(hist_0234), hu.get_array(hist_0234, errors=True)
npt.assert_allclose(val, np.sum(val5d, axis=1))
npt.assert_allclose(err, np.sqrt(np.sum(err5d**2, axis=1)))
hist_0314 = hu.project(hist5d, [0, 3, 1, 4])
val, err = hu.get_array(hist_0314), hu.get_array(hist_0314, errors=True)
npt.assert_allclose(val, np.swapaxes(np.sum(val5d, axis=2), 1, 2))
npt.assert_allclose(err, np.swapaxes(np.sqrt(np.sum(err5d**2, axis=2)), 1, 2))
def get_corr_map(corr_f, variable):
"""
Get the correction map for one file (either chi1 or chi2) for a given
variable
"""
gen_h = corr_f.Get('fold_costh_phi_JpsiPt_JpsiRap_gen_HX')
reco_h = corr_f.Get('fold_costh_phi_JpsiPt_JpsiRap_reco_HX')
# project the histograms onto the variable and rebin them into the analysis
# binning
proj_dir = 0 if variable == 'costh' else 1
gen_h = rebin_1d_binning(project(gen_h, proj_dir),
ANALYSIS_BINNING[variable])
reco_h = rebin_1d_binning(project(reco_h, proj_dir),
ANALYSIS_BINNING[variable])
return divide(reco_h, gen_h)
def test_project_4d_to_3d(self):
# Fill a test histogram
hist4d = _get_hist(4)
val4d, err4d = hu.get_array(hist4d), hu.get_array(hist4d, errors=True)
hist_123 = hu.project(hist4d, [1, 2, 3])
self.assertTrue(isinstance(hist_123, r.TH3))
val, err = hu.get_array(hist_123), hu.get_array(hist_123, errors=True)
npt.assert_allclose(val, np.sum(val4d, axis=0))
npt.assert_allclose(err, np.sqrt(np.sum(err4d**2, axis=0)))
hist_021 = hu.project(hist4d, [0, 2, 1])
self.assertTrue(isinstance(hist_021, r.TH3))
val, err = hu.get_array(hist_021), hu.get_array(hist_021, errors=True)
npt.assert_allclose(val, np.swapaxes(np.sum(val4d, axis=3), 1, 2))
npt.assert_allclose(err, np.swapaxes(np.sqrt(np.sum(err4d**2, axis=3)), 1, 2))
def get_correction_map(cmfile, use_pt=True, use_acc=False):
"""
Get the correction map
"""
map_name = 'fold_costh_phi_JpsiPt_JpsiRap_{}_HX'
reco = 'acc' if use_acc else 'reco'
if use_pt:
gen_dist = project(cmfile.Get(map_name.format('gen')), [0, 1, 2])
reco_dist = project(cmfile.Get(map_name.format(reco)), [0, 1, 2])
else:
gen_dist = project(cmfile.Get(map_name.format('gen')), [1, 0])
reco_dist = project(cmfile.Get(map_name.format(reco)), [1, 0])
accmap = divide(reco_dist, gen_dist)
return accmap
def test_project_2d_to_1d(self):
hist2d = _get_hist(2)
# populate overflow bins to make sure that they are treated as expected
fill_hist(hist2d, np.random.uniform(-1, 0, (100, 2)))
fill_hist(hist2d, np.random.uniform(1, 2, (100, 2)))
val2d, err2d = hu.get_array(hist2d), hu.get_array(hist2d, errors=True)
hist_x = hu.project(hist2d, 'x')
val, err = hu.get_array(hist_x), hu.get_array(hist_x, errors=True)
npt.assert_equal(val, np.sum(val2d, axis=1))
npt.assert_equal(err, np.sqrt(np.sum(err2d**2, axis=1)))
hist_y = hu.project(hist2d, 'y')
val, err = hu.get_array(hist_y), hu.get_array(hist_y, errors=True)
npt.assert_equal(val, np.sum(val2d, axis=0))
npt.assert_equal(err, np.sqrt(np.sum(err2d**2, axis=0)))
def get_acc_mask(cmfile, use_pt, min_acc):
"""
Define an acceptance (only) map mask to exclude events with very low
acceptance values
"""
logging.info('Masking all bins of correction map with acceptance < {}'
.format(min_acc))
if use_pt:
gen_dist = project(cmfile.Get(MAP_NAME.format('gen')), [0, 1, 2])
acc_dist = project(cmfile.Get(MAP_NAME.format('acc')), [0, 1, 2])
else:
gen_dist = project(cmfile.Get(MAP_NAME.format('gen')), [1, 0])
acc_dist = project(cmfile.Get(MAP_NAME.format('acc')), [1, 0])
accmap = divide(acc_dist, gen_dist)
mask = get_array(accmap) < min_acc
return mask
def test_project_4d_to_1d(self):
hist4d = _get_hist(4)
val4d, err4d = hu.get_array(hist4d), hu.get_array(hist4d, errors=True)
hist_0 = hu.project(hist4d, 0)
self.assertTrue(isinstance(hist_0, r.TH1))
val, err = hu.get_array(hist_0), hu.get_array(hist_0, errors=True)
npt.assert_allclose(val, np.sum(val4d, axis=(1, 2, 3)))
npt.assert_allclose(err, np.sqrt(np.sum(err4d**2, axis=(1,2,3))))
def main(args):
"""Main"""
frame, scen = deduce_scenario(args.genlevelfile)
genfile = r.TFile.Open(args.genlevelfile)
genhist = project(genfile.Get(GENNAME), [1, 0])
func = w_costh_phi(set_vals={'norm': 1e5})
chi2, ndf = fit_to_dist(genhist, func, args.verbose)
fit_res = fit_dict(func, chi2, ndf)
update_fit_file(args.outfile, fit_res, frame, scen)
def get_correction_map(cmfile, use_pt=False, acc_only=False, min_acc=None):
"""
Get the correction map from the correction map file
"""
# TODO: Potentially rebin this thing
reco = 'acc' if acc_only else 'reco'
if use_pt:
gen_dist = project(cmfile.Get(MAP_NAME.format('gen')), [0, 1, 2])
reco_dist = project(cmfile.Get(MAP_NAME.format(reco)), [0, 1, 2])
else:
gen_dist = project(cmfile.Get(MAP_NAME.format('gen')), [1, 0])
reco_dist = project(cmfile.Get(MAP_NAME.format(reco)), [1, 0])
accmap = divide(reco_dist, gen_dist)
if min_acc is not None:
acc_mask = get_acc_mask(cmfile, use_pt, min_acc)
else:
acc_mask = None
return AcceptanceCorrectionProvider(accmap, mask=acc_mask)
def test_project_4d_to_2d(self):
hist4d = _get_hist(4)
val4d, err4d = hu.get_array(hist4d), hu.get_array(hist4d, errors=True)
hist_01 = hu.project(hist4d, [0, 1])
self.assertTrue(isinstance(hist_01, r.TH2))
val, err = hu.get_array(hist_01), hu.get_array(hist_01, errors=True)
# For some reason the sum does not give the expected shape here,
# However transposing helps
# But the projected histogram has the expected shape and contens
# TODO: investigate why
npt.assert_allclose(val, np.sum(val4d, axis=(2, 3)).T)
npt.assert_allclose(err, np.sqrt(np.sum(err4d**2, axis=(2, 3))).T)