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))