def get_corrmap(cfile, corr_frame, corr_var=None, mask_acc=False, acc_only=False, **kwargs):
# acc map only
if not mask_acc or acc_only:
if corr_var is None:
# ACCMAP only
if acc_only:
acc_map = cfile.Get('acc_map_costh_phi_{}'.format(corr_frame))
else:
acc_map = cfile.Get('acc_map_costh_phi_{}_eff'.format(corr_frame))
else:
# ACCMAP only
if acc_only:
acc_map = cfile.Get('acc_map_costh_phi_{}_{}'.format(corr_var, corr_frame))
else:
acc_map = cfile.Get('acc_map_costh_phi_{}_{}_eff'.format(corr_var, corr_frame))
return AcceptanceCorrectionProvider(acc_map, **kwargs)
if corr_var is None:
acc_map = get_array(cfile.Get('acc_map_costh_phi_{}'.format(corr_frame)))
acc_eff_name = 'acc_map_costh_phi_{}_eff'.format(corr_frame)
else:
acc_map = get_array(cfile.Get('acc_map_costh_phi_{}_{}'.format(corr_var, corr_frame)))
acc_eff_name = 'acc_map_costh_phi_{}_{}_eff'.format(corr_var, corr_frame)
acc_mask = acc_map < kwargs.pop('min_acc')
return AcceptanceCorrectionProvider(cfile.Get(acc_eff_name), mask=acc_mask, **kwargs)
def test_masking_init(self):
"""
Test if masking also works when finding zero bins and using precision
limit or acceptance limit
"""
# Mask the precision bins first, to avoid division by zero
min_prec = 0.143
arr, err = get_array(self.acc_map), get_array(self.acc_map,
errors=True)
masked_bins = (err / arr > min_prec).astype(bool)
# mask the min acceptance bins
min_acc = 0.00425
masked_bins |= (arr < min_acc).astype(bool)
mask_bin_vals = np.random.uniform(0, 1, (10, 2))
idcs = set()
for v in mask_bin_vals:
idxx = self.acc_map.GetXaxis().FindBin(v[0])
idxy = self.acc_map.GetYaxis().FindBin(v[1])
idcs.add((idxx, idxy))
self.acc_map.SetBinContent(idxx, idxy, 0)
# add the bins that are masked due to being zero
for x, y in idcs:
masked_bins[x - 1, y - 1] = True # ROOT binning starts at 1
corr_prov = AcceptanceCorrectionProvider(self.acc_map,
mask_prec=min_prec)
npt.assert_equal((corr_prov.corr_map == -1).astype(bool), masked_bins)
def get_costh_phi_in_bins(hist_3d):
"""Get all costh phi histograms in each bin of the 3rd variable"""
arr = get_array(hist_3d)
binning = np.array([get_binning(hist_3d, 'X'), get_binning(hist_3d, 'Y')])
err = get_array(hist_3d, errors=True)
return [from_array(arr[:,:,i], binning, errors=err[:,:,i])
for i in xrange(arr.shape[2])]
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_get_array_w_errors_no_overflow(self):
hist = _get_hist(1)
arr = hu.get_array(hist, errors=True)
self.assertEqual(arr.shape[0], hist.GetNbinsX())
self._comp_array_hist(arr, hist, error=True)
hist = _get_hist(2)
arr = hu.get_array(hist, errors=True)
self.assertEqual(arr.shape[0], hist.GetNbinsX())
self.assertEqual(arr.shape[1], hist.GetNbinsY())
self._comp_array_hist(arr, hist, error=True)
hist = _get_hist(3)
arr = hu.get_array(hist, errors=True)
self.assertEqual(arr.shape[0], hist.GetNbinsX())
self.assertEqual(arr.shape[1], hist.GetNbinsY())
self.assertEqual(arr.shape[2], hist.GetNbinsZ())
self._comp_array_hist(arr, hist, error=True)
hist = _get_hist(4)
arr = hu.get_array(hist, errors=True)
self.assertEqual(arr.shape[0], hist.GetAxis(0).GetNbins())
self.assertEqual(arr.shape[1], hist.GetAxis(1).GetNbins())
self.assertEqual(arr.shape[2], hist.GetAxis(2).GetNbins())
self.assertEqual(arr.shape[3], hist.GetAxis(3).GetNbins())
self._comp_array_hist(arr, hist, error=True)
def test_uniform_rebinning(self):
hist = _get_hist(1)
targ_bin = np.linspace(0, 1, 6)
# since hist has 10 bin, rebin(hist, 5) should return the same histogram
# as the one when we rebin it according to a custom binning
npt.assert_allclose(hu.get_array(hu.rebin_1d_binning(hist, targ_bin)),
hu.get_array(hu.rebin(hist, [(0, 5)])))
def test_rel_uncer(self):
for i in xrange(1, 4):
hist = _get_hist(i)
uncer_hist = hu.uncer_hist(hist)
vals = hu.get_array(hist)
exp_uncers = np.zeros_like(vals)
np.divide(1, np.sqrt(vals), out=exp_uncers, where=vals!=0)
npt.assert_allclose(exp_uncers, hu.get_array(uncer_hist))
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 get_pt_bin(amap, pt_val):
"""
Get the pt bin costh-phi map from the passed (3d) acceptance map
"""
pt_bin = find_bin(get_binning(amap, 2), np.array([pt_val]))[0]
val, err = get_array(amap), get_array(amap, errors=True)
ctp_binning = np.array([get_binning(amap, i) for i in [0, 1]])
return from_array(val[:, :, pt_bin], ctp_binning, errors=err[:, :, pt_bin])
def shift_by_median(ppd, use_val=None):
"""
Shift the ppd by the median to center it around 0
"""
if use_val is None:
return ppd
else:
med = use_val
binning = get_binning(ppd)
return from_array(get_array(ppd), binning - med, errors=get_array(ppd, errors=True))
def test_non_full_coverage(self):
"""
Check that if the target binning does not cover the full range of the
original histogram things still work
"""
hist = _get_hist(1)
targ_bin = np.linspace(0.1, 0.6, 6)
targ_vals = hu.get_array(hist)[1:6]
npt.assert_allclose(hu.get_array(hu.rebin_1d_binning(hist, targ_bin)),
targ_vals)
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_non_uniform_binning(self):
hist = _get_hist(1)
targ_bin = np.array([0, 0.1, 0.2, 0.4, 0.6, 0.8, 1.0])
# manually rebin this histogram here, since we now what we want
vals = hu.get_array(hist)
targ_vals = np.array([vals[0], vals[1], vals[2] + vals[3],
vals[4] + vals[5], vals[6] + vals[7],
vals[8] + vals[9]])
npt.assert_allclose(hu.get_array(hu.rebin_1d_binning(hist, targ_bin)),
targ_vals)
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)
def test_masking_low_prec_init(self):
"""
Test if the masking in init works for bins with too low precision
"""
# sprt(1e4 / (20*10)) / (1e4 / 20 * 10), average poisson uncertainty
min_prec = 0.143 # tuned to be slightly below the expected uncertainty
arr, err = get_array(self.acc_map), get_array(self.acc_map,
errors=True)
masked_bins = (err / arr > min_prec).astype(bool)
corr_prov = AcceptanceCorrectionProvider(self.acc_map,
mask_prec=min_prec)
# the correction map has to have the same bins masked as the acceptance map
npt.assert_equal((corr_prov.corr_map == -1).astype(bool), masked_bins)
def test_get_array_3dim(self):
hist = _get_hist(3)
arr = hu.get_array(hist)
self.assertEqual(arr.shape[0], hist.GetNbinsX())
self.assertEqual(arr.shape[1], hist.GetNbinsY())
self.assertEqual(arr.shape[2], hist.GetNbinsZ())
self._comp_array_hist(arr, hist)
def get_contour(hist):
"""
Get the outer contour of all filled points in the histogram
"""
vals = get_array(hist) > 0
xbinning, ybinning = get_binning(hist, 0), get_binning(hist, 1)
xvals = 0.5 * (xbinning[:-1] + xbinning[1:])
yvals = 0.5 * (ybinning[:-1] + ybinning[1:])
filled = []
for ix, xv in enumerate(xvals):
for iy, yv in enumerate(yvals):
if vals[ix, iy]:
filled.append([xv, yv])
filled = np.array(filled)
hull = ConvexHull(filled)
# Append the first point again at the end to "close" the contour
xcont = filled[hull.vertices, 0]
xcont = np.append(xcont, np.array([xcont[0]]))
ycont = filled[hull.vertices, 1]
ycont = np.append(ycont, np.array(ycont[0]))
return r.TGraph(len(hull.vertices) + 1, xcont, ycont)
def test_masking_min_acc_init(self):
"""Test if masking works for bins with too low acceptance (i.e. content)"""
min_acc = 0.00475 # (slightly below average occupancy)
arr = get_array(self.acc_map)
masked_bins = (arr < min_acc).astype(bool)
corr_prov = AcceptanceCorrectionProvider(self.acc_map, min_acc=min_acc)
npt.assert_equal((corr_prov.corr_map == -1).astype(bool), masked_bins)
def to_bw_hist(hist):
"""Fill all filled bins with value 1 and all empty ones with 0"""
arr = get_array(hist)
# TODO: generalize and put into hist_utils
binning = np.array([get_binning(hist, 0), get_binning(hist, 1)])
arr = arr != 0
return from_array(arr, binning)
def get_combined_ppd(inputfiles, var):
"""
Get the combined ppd from all inputfiles
"""
ppds = [get_scaled_ppd(f, var) for f in inputfiles]
# PPDs all have the same binning
ppd_binning = get_binning(ppds[0])
ppd_vals = np.array([get_array(p) for p in ppds])
ppd_errs = np.array([get_array(p, errors=True) for p in ppds])
# Get the maximum value in each gin and its uncertainty
max_idx = np.argmax(ppd_vals, axis=0)
# Necessary for 2d indexing. There might be an easier way for this
idx = np.arange(0, len(ppd_vals[0]))
max_ppd = ppd_vals[max_idx, idx]
max_err = ppd_errs[max_idx, idx]
return from_array(max_ppd, ppd_binning, errors=max_err)
def test_get_array_4dim(self):
"""Test if THns also work"""
hist = _get_hist(4)
arr = hu.get_array(hist)
self.assertEqual(arr.shape[0], hist.GetAxis(0).GetNbins())
self.assertEqual(arr.shape[1], hist.GetAxis(1).GetNbins())
self.assertEqual(arr.shape[2], hist.GetAxis(2).GetNbins())
self.assertEqual(arr.shape[3], hist.GetAxis(3).GetNbins())
self._comp_array_hist(arr, hist)
def get_combined_ppd_2d(inputfiles, var1, var2):
"""
Get the combined 2d ppd from all inputfiles
"""
ppds = [get_scaled_ppd_2d(f, var1, var2, 100, 100) for f in inputfiles]
ppd_binning = np.array([get_binning(ppds[0], 0), get_binning(ppds[0], 1)])
ppd_vals = np.array([get_array(p) for p in ppds])
# TODO: at some point find out how argmax works in multiple dimensions
return from_array(np.max(ppd_vals, axis=0), ppd_binning)
def _test_from_array_nd_w_overflow(self, n_dim):
hist = _get_hist(n_dim)
arr = hu.get_array(hist, overflow=True)
axes = 'X'
if n_dim == 2:
axes = 'XY'
if n_dim == 3:
axes = 'XYZ'
binning = np.array([hu.get_binning(hist, ax) for ax in axes])
arr_hist = hu.from_array(arr, binning)
npt.assert_equal(hu.get_array(arr_hist, overflow=True), arr)
npt.assert_equal(hu.get_binning(arr_hist, 'X'), hu.get_binning(hist, 'X'))
if n_dim > 1:
npt.assert_equal(hu.get_binning(arr_hist, 'Y'), hu.get_binning(hist, 'Y'))
if n_dim > 2:
npt.assert_equal(hu.get_binning(arr_hist, 'Z'), hu.get_binning(hist, 'Z'))
err = hu.get_array(hist, errors=True, overflow=True)
arr_err_hist = hu.from_array(arr, binning, errors=err)
npt.assert_equal(hu.get_array(arr_err_hist, overflow=True), arr)
npt.assert_equal(hu.get_array(arr_err_hist, overflow=True, errors=True), err)
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_3d_to_1d(self):
hist3d = _get_hist(3)
# populate overflow bins to make sure that they are treated as expected
fill_hist(hist3d, np.random.uniform(-1, 0, (100, 3)))
fill_hist(hist3d, np.random.uniform(1, 2, (100, 3)))
val3d, err3d = hu.get_array(hist3d), hu.get_array(hist3d, errors=True)
hist_x = hu.project(hist3d, 'x')
val, err = hu.get_array(hist_x), hu.get_array(hist_x, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=(1,2)))
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=(1,2))))
hist_y = hu.project(hist3d, 'y')
val, err = hu.get_array(hist_y), hu.get_array(hist_y, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=(0,2)))
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=(0,2))))
hist_z = hu.project(hist3d, 'z')
val, err = hu.get_array(hist_z), hu.get_array(hist_z, errors=True)
npt.assert_equal(val, np.sum(val3d, axis=(0,1)))
npt.assert_equal(err, np.sqrt(np.sum(err3d**2, axis=(0,1))))
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 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 make_overlay_plot(pt_map, pt_data, **kwargs):
"""
Plot the coverage of the pt_data onto the
"""
amap_x, amap_y = get_binning(pt_map, 0), get_binning(pt_map, 1)
if np.min(amap_x) == 0:
costh = pt_data.costh_HX_fold.abs()
else:
costh = pt_data.costh_HX_fold
data_dist = hist2d(costh,
pt_data.phi_HX_fold,
x_hist_sett=(len(amap_x) - 1, amap_x),
y_hist_sett=(len(amap_y) - 1, amap_y))
coverage = get_array(data_dist) > 0
cov_graph = get_mask_graph(amap_x, amap_y, coverage)
can = mkplot(pt_map, **kwargs)
mkplot(cov_graph,
can=can,
drawOpt='sameE5',
attr=[{
'color': r.kRed,
'fillalpha': (r.kRed, 0),
'marker': 1
}])
mkplot([
r.TLine(v, np.min(amap_y), v, np.max(amap_y))
for v in [-0.625, -0.45, 0.45, 0.625]
],
attr=[{
'color': 12,
'line': 7,
'width': 2
}],
can=can,
drawOpt='same')
return can
def get_coverage_contour(hist, coverage=0.683):
"""
Get the contour from the passed histogram that surpasses the specified coerage
"""
vals = get_array(hist)
sum_vals = np.sum(vals)
def _coverage(level):
"""Calculate the coverage corresponding to the passed level"""
return np.sum(vals * (vals >= level)) / sum_vals
# do some pre-processing to start from a slightly better bracket for the
# secant method
dec_cov = np.array(
[_coverage(0.05 * i * np.max(vals)) for i in xrange(21)])
q_bin = find_bin(dec_cov, np.array([coverage]))
search_brack = [
q_bin * 0.05 * np.max(vals), (q_bin + 1) * 0.05 * np.max(vals)
]
cov_level = root_scalar(lambda x: _coverage(x) - coverage,
bracket=search_brack)
filled = vals >= cov_level.root
x_vals, y_vals = get_binning(hist, 'X'), get_binning(hist, 'Y')
# get the bin centers
x_vals = 0.5 * (x_vals[1:] + x_vals[:-1])
y_vals = 0.5 * (y_vals[1:] + y_vals[:-1])
filled_coords = []
for ix, xv in enumerate(x_vals):
for iy, yv in enumerate(y_vals):
if filled[ix, iy]:
filled_coords.append([xv, yv])
return contour_as_tgraph(np.array(filled_coords))
def __init__(self, acc_map, min_acc=0, mask_prec=None, mask=None):
"""
Args:
acc_map (TH2D, TH3D or THnD): costh-phi map or costh-phi-var map
obtained by applying all cuts and selections (and possibly
efficiency weightings). For each bin 1 / (bin content) will be
the weight for the acceptance correction
min_acc (float, optional): Mask all bins with an acceptance below
this value (default = 0)
mask_prec (float, optional): If not None, mask all bins for which
the relative error is larger than the passed value
mask (np.array, optional): Array with the same dimensions as the
acceptance map. All bins containing a non False value will be
masked. Overrides the min_acc and mask_prec argument (i.e. they
will be ignored) but still respects zero bin masking
"""
self.hist = acc_map
logging.debug('Using acceptance map \'{}\''.format(
self.hist.GetName()))
# Corrections are 1 / acceptance map
acc_values = get_array(self.hist)
if mask is not None:
if mask.shape != acc_values.shape:
logging.error('mask and acceptance map need to have the same '
'dimensions. mask: {}, map: {}'.format(
mask.shape, acc_values.shape))
if min_acc != 0:
logging.info(
'Ignoring min_acc={} because a mask is used'.format(
min_acc))
if mask_prec is not None:
logging.info(
'Ignoring mask_prec={} because a mask is used'.format(
mask_prec))
# mask the values without acceptance in the acceptance map
# this will also make them return -1 for the correction map
logging.debug('Masking {} bins according to the mask'.format(
np.sum(mask)))
empty_mask = (acc_values == 0)
logging.debug('Masking {} empty bins'.format(np.sum(empty_mask)))
masked_vals = empty_mask | mask
else:
if min_acc < 0 or min_acc > 1:
logging.warning('The minimum acceptance should be a value '
'between 0 and 1, but is {}'.format(min_acc))
masked_vals = (acc_values <= min_acc).astype(bool)
logging.debug('Minimum acceptance = {}: Masking {} bins'.format(
min_acc, np.sum(masked_vals)))
if mask_prec is not None:
if isinstance(mask_prec, float):
acc_errs = get_array(self.hist, errors=True)
rel_uncer = np.zeros_like(acc_errs)
np.divide(acc_errs,
acc_values,
out=rel_uncer,
where=acc_values != 0)
mask_uncer = (rel_uncer > mask_prec).astype(bool)
logging.debug(
'Minimum precision = {}: Masking {} bins'.format(
mask_prec, np.sum(mask_uncer)))
masked_vals |= mask_uncer
else:
logging.error(
'mask_prec has to be a float value. Not using'
' it to mask bins with too low precision.')
acc_values = ~masked_vals * acc_values + -1 * masked_vals
logging.debug('{} of {} bins are masked in the correction map'.format(
np.sum(masked_vals), acc_values.size))
self.corr_map = 1.0 / acc_values
self.var_binnings = []
self.ndim = self.corr_map.ndim
for i in xrange(self.ndim):
self.var_binnings.append(get_binning(acc_map, i))
