def Eval(self, wcp=None):
""" Set a WC point and evaluate """
if isinstance(WCPoint, dict):
wcp = WCPoint(wcp)
elif isinstance(wcp, str):
values = wcp.replace(" ", "").split(',')
wcp = WCPoint(values, names=self.GetWCnames())
elif isinstance(wcp, list):
wcp = WCPoint(wcp, names=self.GetWCnames())
if wcp is None:
if self._WCPoint is None:
self._WCPoint = WCPoint(names=self._wcnames)
else:
self._WCPoint = wcp
self.EvalInSelfPoint()
def EvalPointError(self, pt, val = 0.0):
""" Evaluate the error fit at a particular WC phase space point """
if not isinstance(pt, WCPoint):
wc_name = pt
pt = WCPoint()
pt.SetStrength(wc_name, val)
i = 0
v,x1,x2,x3,x4,c = 0., 0., 0., 0., 0., 0.
n1,n2,n3,n4 = '', '', '', ''
err_pair,idx_pair = (0.,0.),(0.,0.)
for i in range(self.ErrSize()):
c = self.err_coeffs[i]
err_pair = self.err_pairs[i]
idx_pair = self.pairs[err_pair[0]]; n1 = self.names[idx_pair[0]]; n2 = self.names[idx_pair[1]]
idx_pair = self.pairs[err_pair[1]]; n3 = self.names[idx_pair[0]]; n4 = self.names[idx_pair[1]]
x1 = 1.0 if n1 == kSMstr else pt.GetStrength(n1)
x2 = 1.0 if n2 == kSMstr else pt.GetStrength(n2)
x3 = 1.0 if n3 == kSMstr else pt.GetStrength(n3)
x4 = 1.0 if n4 == kSMstr else pt.GetStrength(n4)
v += x1*x2*x3*x4*c;
return sqrt(v);
def test_wcfit():
chk_str = ''
unit_chk = True
all_chks, units = [0] * 2
tolerance = 1e-4
# The structure constants
s00 = 1.0
s10 = 1.5
s11 = 1.25
pts = []
vals = [-1.0, 1.25, 0.5, 2.5, 4]
for x in vals:
y = s00 * 1.0 + s10 * x + s11 * x * x
pts.append(WCPoint(f'EFTrwgt0_ctG_{x}', y))
chk_x = 1.5
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_pt = WCPoint(f'EFTrwgt0_ctG_{chk_x}', 0.0)
print('Running unit tests for WCFit class')
all_chks = 0
units = 0
fit_base = WCFit(pts, 'base')
unit_chk = (abs(fit_base.EvalPoint(chk_pt) - chk_y) < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 1 ---')
print('chk_x : ', chk_x)
print('chk_y : ', chk_y)
print('EvalPoint: ', fit_base.EvalPoint(chk_pt))
print('test: ', chk_str)
print('--------------\n')
fit_new = WCFit()
fit_new.SetTag('new')
fit_new.AddFit(fit_base)
unit_chk = (abs(fit_base.EvalPoint(chk_pt) - chk_y) < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 2 ---')
print('chk_x : ', chk_x)
print('chk_y : ', chk_y)
print('EvalPoint: ', fit_new.EvalPoint(chk_pt))
print('test: ', chk_str)
print('--------------\n')
fit_new.AddFit(fit_base) #CAREFUL b/c WCFit is mutable
unit_chk = (abs(fit_new.EvalPoint(chk_pt) - 2 * chk_y) < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 3 ---')
print('chk_x : ', chk_x)
print('chk_y : ', 2 * chk_y)
print('EvalPoint: ', fit_new.EvalPoint(chk_pt))
print('test: ', chk_str)
print('--------------\n')
#fit_base = WCFit(pts,'base') #redefine b/c WCFit is mutable
unit_chk = (abs(fit_base.EvalPoint(chk_pt) - chk_y) < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 4 ---')
print('chk_x : ', chk_x)
print('chk_y : ', chk_y)
print('EvalPoint: ', fit_base.EvalPoint(chk_pt))
print('test: ', chk_str)
print('--------------\n')
print(f'Passed Checks: {all_chks}/{units}')
return (all_chks == units)
def test_histeft():
chk_str = ''
unit_chk = True
all_chks, units = [0] * 2
result, expected, diff, tolerance = [0] * 4
tolerance = 1e-4
wc_names = ['sm', 'ctG', 'ctZ']
# The structure constants
s00 = 1.0
s10 = 1.5
s11 = 1.25
sconst = [s00, s10, s11]
# A dummy WC name to use
wc_name = 'ctG'
pts = []
vals = [-1.0, 1.25, 0.5, 2.5, 4]
for x in vals:
y = s00 * 1.0 + s10 * x + s11 * x * x
pts.append(WCPoint(f'EFTrwgt0_{wc_name}_{x}', y))
fit_1 = WCFit(pts, 'f1')
fit_2 = WCFit()
fit_2.SetTag('f2')
fit_2.AddFit(fit_1)
fit_2.AddFit(fit_1)
chk_x = 1.5
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_vals = {wc_name: chk_x, 'ctZ': 0.0}
chk_pt = WCPoint(f'EFTrwgt0_{wc_name}_{chk_x}', 0.0)
print('Running unit tests for HistEFT class')
all_chks = 0
units = 0
h_base = HistEFT("h_base", wc_names[1::], hist.Cat("sample", "sample"),
hist.Bin("n", "", 1, 0, 1))
val = ak.Array([0.5])
eftval = ak.Array([0.0002579])
sconst = sconst + sconst
h_base.fill(n=val,
sample='test',
weight=np.ones_like(val),
eft_coeff=[ak.Array(sconst)])
expected = 1.0
result = list(h_base.values().values())[0][0]
unit_chk = (abs(result - expected) < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 1 ---')
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
###########################
h_base.set_wilson_coefficients(**chk_vals)
expected = fit_1.EvalPoint(chk_pt)
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 2 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
###########################
chk_x = 0.75
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_vals = {wc_name: chk_x, 'ctZ': 0.0}
chk_pt.SetStrength(wc_name, chk_x)
h_base.set_wilson_coefficients(**chk_vals)
expected = fit_1.EvalPoint(chk_pt)
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 3 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
###########################
h_base.fill(n=val,
sample='test',
weight=np.ones_like(val),
eft_coeff=[ak.Array(sconst) * 2])
h_base.set_wilson_coefficients(**chk_vals)
# First make sure the original WCFits weren't messed with
expected = chk_y + 2 * chk_y
result = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
diff = abs(expected - result)
tolerance = 1e-10
unit_chk = (diff < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 4 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('fit_1 + fit_2: ', result)
print('difference : ', diff)
print('tolerance : ', tolerance)
print('test: ', chk_str)
print('--------------\n')
# Now check that the TH1EFT actually worked
expected = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 5 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
###########################
h_new = h_base.copy()
chk_x = 0.975
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_vals = {wc_name: chk_x, 'ctZ': 0.0}
chk_pt.SetStrength(wc_name, chk_x)
h_new.set_wilson_coefficients(**chk_vals)
# First check that h_new has the right value
expected = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
result = list(h_new.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 6 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
chk_x = 0.75 # Needs to be w/e chk_x was before UNIT 6
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_vals = {wc_name: chk_x, 'ctZ': 0.0}
chk_pt.SetStrength(wc_name, chk_x)
# Next check that the h_base was unaffected when we scaled h_new
expected = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 7 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
# Check HistEFT.add()
expected = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(
chk_pt) #fits for h_base
expected += fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(
chk_pt) #fits for h_new
h_base.add(h_new)
h_base.set_wilson_coefficients(**chk_vals) #evaluate h_base at chk_pt
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 8 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
# Check HistEFT.add() reweight
chk_x = 0.75 # Needs to be w/e chk_x was before UNIT 6
chk_y = s00 * 1.0 + s10 * chk_x + s11 * chk_x * chk_x
chk_vals = {wc_name: chk_x, 'ctZ': 0.0}
chk_pt.SetStrength(wc_name, chk_x)
expected = fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
expected += fit_1.EvalPoint(chk_pt) + fit_2.EvalPoint(chk_pt)
h_base.set_wilson_coefficients(**chk_vals)
result = list(h_base.values().values())[0][0]
unit_chk = abs(result - expected) < tolerance
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 9 ---')
print('chk_x : ', chk_pt.GetStrength(wc_name))
print('expected : ', expected)
print('GetBinContent: ', result)
print('test: ', chk_str)
print('--------------\n')
###########################
print(f'Passed Checks: {all_chks}/{units}')
return (all_chks == units)
def test_stats():
chk_str = ''
unit_chk = True
all_chks, units = [0] * 2
result, expected, diff, tolerance = [0] * 4
tolerance = 0.001
# Basically the SM 'strength'
x0 = 1.0
# Dummy WC names to use (needs to match dimension of pt
wc_names = ['sm', 'ctG', 'ctZ']
# The structure constants, need to match dimension of pt
svals = [
1.15, # (00)
1.35,
1.25, # (10) (11)
0.25,
0.75,
1.00, # (20) (21) (22)
]
# Make sure there are enough pts to fully determine the fit!
pts = [
[x0, -1.00, 0.00],
[x0, -0.50, 0.25],
[x0, 0.00, 0.35],
[x0, 0.25, 0.05],
[x0, 0.50, -0.05],
[x0, 0.75, 0.25],
[x0, 1.00, -0.35],
]
wc_pts = []
idx = 0
for pt in pts:
y = fval(pt, svals)
s = f'EFTrwgt{idx}'
for i in range(1, len(pt)): # NOTE: pt better not be size 0!!
wc_str = wc_names[i]
s += f'_{wc_str}_{pt[i]}'
#print(s,y)
wc_pts.append(WCPoint(s, y))
idx += 1
fit_1 = WCFit(wc_pts, 'f1')
fit_2 = WCFit()
fit_2.SetTag('f2')
nevents = 5000
for i in range(nevents):
fit_2.AddFit(fit_1)
###########################
print('Running unit tests for stats unc.')
all_chks = 0
units = 0
# Needs to be the same size as wc_names
chk_x = [x0, 1.2, 0.4]
chk_y = 0.0
chk_e = 0.0
for i in range(nevents):
v = fval(chk_x, svals)
chk_y += v
chk_e += v * v
chk_e = chk_e**.5
chk_wcstr = 'EFTrwgt0'
sidx = 0
for i in range(len(wc_names)):
if i: # Need to skip first entry since that's the SM 'strength'
chk_wcstr += f'_{wc_names[i]}_{chk_x[i]}'
for j in range(i + 1):
v = svals[sidx]
#print(f'{i}{j}: {v}')
sidx = sidx + 1
print()
chk_pt = WCPoint(chk_wcstr, 0.0)
###########################
# Basic check for proper adding of quadratic structure constants
# Note: We expect the diff to grow with increased number of events due to the numeric precison
expected = chk_y
result = fit_2.EvalPoint(chk_pt)
diff = abs(expected - result)
tolerance = 1e-4
unit_chk = (diff < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 1 ---')
print('evts : ', nevents)
print('chk_wcstr: ', chk_wcstr)
print('expected : ', expected)
print('result : ', result)
print('diff : ', diff)
print('tolerance: ', tolerance)
print('test: ', chk_str)
print('--------------\n')
# Check the error calculation
# Note: We expect the diff to grow with increased number of events due to the numeric precison
expected = chk_e
result = fit_2.EvalPointError(chk_pt)
diff = abs(expected - result)
tolerance = 1e-05 * (10 * nevents)**.5
unit_chk = (diff < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 2 ---')
print('evts : ', nevents)
print('chk_wcstr: ', chk_wcstr)
print('expected : ', expected)
print('result : ', result)
print('diff : ', diff)
print('tolerance: ', tolerance)
print('test: ', chk_str)
print('--------------\n')
# Now do the percent error
# Note: The diff here also appears to grow apparently due to numeric precison, but much more slowly (it is still kind of concerning)
expected = chk_e / chk_y
result = fit_2.EvalPointError(chk_pt) / fit_2.EvalPoint(chk_pt)
diff = abs(expected - result)
tolerance = 1e-04
unit_chk = (diff < tolerance)
all_chks += unit_chk
units += 1
chk_str = 'Passed' if unit_chk else 'Failed'
print('--- UNIT 3 ---')
print('evts : ', nevents)
print('chk_wcstr: ', chk_wcstr)
print('expected : ', expected)
print('result : ', result)
print('diff : ', diff)
print('tolerance: ', tolerance)
print('test: ', chk_str)
print('--------------\n')
###########################
print(f'Passed Checks: {all_chks}/{units}')
return (all_chks == units)
def SetSMpoint(self):
""" Set SM WC point and evaluate """
wc = WCPoint(names=self._wcnames)
wc.SetSMPoint()
self.Eval(wc)
class HistEFT(coffea.hist.Hist):
def __init__(self, label, wcnames, *axes, **kwargs):
""" Initialize """
if isinstance(wcnames, str) and ',' in wcnames:
wcnames = wcnames.replace(' ', '').split(',')
n = len(wcnames) if isinstance(wcnames, list) else wcnames
self._wcnames = wcnames
self._nwc = n
self._ncoeffs = int(1 + 2 * n + n * (n - 1) / 2)
self.CreatePairs()
self._WCPoint = None
super().__init__(label, *axes, **kwargs)
self.EFTcoeffs = {}
self.EFTerrs = {}
self.WCFit = {}
def CreatePairs(self):
""" Create pairs... same as for WCFit class """
self.idpairs = []
self.errpairs = []
n = self._nwc
for f in range(n + 1):
for i in range(f + 1):
self.idpairs.append((f, i))
for j in range(len(self.idpairs) - 1):
self.errpairs.append([i, j])
self.errpairs = np.array(self.errpairs)
def GetErrCoeffs(self, coeffs):
""" Get all the w*w coefficients """
#return [coeffs[p[0]]*coeffs[p[1]] if (p[1] == p[0]) else 2*(coeffs[p[0]]*coeffs[p[1]]) for p in self.errpairs]
return np.where(
self.errpairs[:, 0] == self.errpairs[:, 1],
coeffs[self.errpairs[:, 0]] * coeffs[self.errpairs[:, 1]],
2 * coeffs[self.errpairs[:, 0]] * coeffs[self.errpairs[:, 1]])
def copy(self, content=True):
""" Copy """
out = HistEFT(self._label,
self._wcnames,
*self._axes,
dtype=self._dtype)
if self._sumw2 is not None: out._sumw2 = {}
if content:
out._sumw = copy.deepcopy(self._sumw)
out._sumw2 = copy.deepcopy(self._sumw2)
out.EFTcoeffs = copy.deepcopy(self.EFTcoeffs)
out.EFTerrs = copy.deepcopy(self.EFTerrs)
return out
def identity(self):
return self.copy(content=False)
def clear(self):
self._sumw = {}
self._sumw2 = None
self.EFTcoeffs = {}
self.EFTerrs = {}
self.WCFit = {}
def GetNcoeffs(self):
""" Number of coefficients """
return self._ncoeffs
def GetNcoeffsErr(self):
""" Number of w*w coefficients """
return int((self._ncoeffs + 1) * (self._ncoeffs) / 2)
def GetSparseKeys(self, **values):
""" Get tuple from values """
return tuple(d.index(values[d.name]) for d in self.sparse_axes())
def fill(self, EFTcoefficients, **values):
""" Fill histogram, incuding EFT fit coefficients """
if EFTcoefficients is None or len(EFTcoefficients) == 0:
super().fill(**values)
return
values_orig = values.copy()
weight = values.pop("weight", None)
sparse_key = tuple(d.index(values[d.name]) for d in self.sparse_axes())
if sparse_key not in self.EFTcoeffs:
self.EFTcoeffs[sparse_key] = []
self.EFTerrs[sparse_key] = []
for i in range(self.GetNcoeffs()):
self.EFTcoeffs[sparse_key].append(
np.zeros(shape=self._dense_shape, dtype=self._dtype))
for i in range(self.GetNcoeffsErr()):
self.EFTerrs[sparse_key].append(
np.zeros(shape=self._dense_shape, dtype=self._dtype))
errs = []
iCoeff, iErr = 0, 0
EFTcoefficients = np.asarray(EFTcoefficients)
if self.dense_dim() > 0:
dense_indices = tuple(
d.index(values[d.name]) for d in self._axes
if isinstance(d, coffea.hist.hist_tools.DenseAxis))
xy = np.atleast_1d(
np.ravel_multi_index(dense_indices, self._dense_shape))
if len(EFTcoefficients) > 0:
#EFTcoefficients = EFTcoefficients.regular()
errs = [self.GetErrCoeffs(x) for x in EFTcoefficients]
errs = np.asarray(errs)
for coef in np.transpose(EFTcoefficients):
#coef = coffea.util._ensure_flat(coef)
self.EFTcoeffs[sparse_key][iCoeff][:] += np.bincount(
xy,
weights=coef,
minlength=np.array(self._dense_shape).prod()).reshape(
self._dense_shape)
iCoeff += 1
# Calculate errs...
for err in np.transpose(errs):
self.EFTerrs[sparse_key][iErr][:] += np.bincount(
xy,
weights=err,
minlength=np.array(self._dense_shape).prod()).reshape(
self._dense_shape)
iErr += 1
else:
for coef in np.transpose(EFTcoefficients):
self.EFTcoeffs[sparse_key][iCoeff] += np.sum(coef)
# Calculate errs...
errs = np.asarray(errs)
for err in np.transpose(errs):
self.EFTerrs[sparse_key][iErr][:] += np.sum(err)
super().fill(**values_orig)
#######################################################################################
def SetWCFit(self, key=None):
if key == None:
for key in list(self._sumw.keys())[-1:]:
self.SetWCFit(key)
return
self.WCFit[key] = []
bins = np.transpose(np.asarray(self.EFTcoeffs[key])
) #np.array((self.EFTcoeffs[key])[:]).transpose()
errs = np.asarray((self.EFTerrs[key])[:]).transpose()
ibin = 0
for fitcoeff, fiterrs in zip(bins, errs):
self.WCFit[key].append(
WCFit(tag='%i' % ibin,
names=self._wcnames,
coeffs=fitcoeff,
errors=fiterrs))
#self.WCFit[key][-1].Dump()
ibin += 1
def add(self, other):
""" Add another histogram into this one, in-place """
#super().add(other)
if not self.compatible(other):
raise ValueError(
"Cannot add this histogram with histogram %r of dissimilar dimensions"
% other)
raxes = other.sparse_axes()
def add_dict(left, right):
for rkey in right.keys():
lkey = tuple(
self.axis(rax).index(rax[ridx])
for rax, ridx in zip(raxes, rkey))
if lkey in left:
left[lkey] += right[rkey]
else:
left[lkey] = copy.deepcopy(right[rkey])
if self._sumw2 is None and other._sumw2 is None: pass
elif self._sumw2 is None:
self._init_sumw2()
add_dict(self._sumw2, other._sumw2)
elif other._sumw2 is None:
add_dict(self._sumw2, other._sumw)
else:
add_dict(self._sumw2, other._sumw2)
add_dict(self._sumw, other._sumw)
add_dict(self.EFTcoeffs, other.EFTcoeffs)
add_dict(self.EFTerrs, other.EFTerrs)
return self
def DumpFits(self, key=''):
""" Display all the fit parameters for all bins """
if key == '':
for k in self.EFTcoeffs.keys():
self.DumpFits(k)
return
for fit in (len(self.WCFit[key])):
fit.Dump()
def ScaleFits(self, SF, key=''):
""" Scale all the fits by some amount """
if key == '':
for k in self.EFTcoeffs.keys():
self.ScaleFits(SF, k)
return
for fit in self.WCFit[key]:
fit.Scale(SF)
def __getitem__(self, keys):
""" Extended from parent class """
if not isinstance(keys, tuple): keys = (keys, )
if len(keys) > self.dim():
raise IndexError("Too many indices for this histogram")
elif len(keys) < self.dim():
if Ellipsis in keys:
idx = keys.index(Ellipsis)
slices = (slice(None), ) * (self.dim() - len(keys) + 1)
keys = keys[:idx] + slices + keys[idx + 1:]
else:
slices = (slice(None), ) * (self.dim() - len(keys))
keys += slices
sparse_idx, dense_idx, new_dims = [], [], []
for s, ax in zip(keys, self._axes):
if isinstance(ax, coffea.hist.hist_tools.SparseAxis):
sparse_idx.append(ax._ireduce(s))
new_dims.append(ax)
else:
islice = ax._ireduce(s)
dense_idx.append(islice)
new_dims.append(ax.reduced(islice))
dense_idx = tuple(dense_idx)
def dense_op(array):
return np.block(
coffea.hist.hist_tools.assemble_blocks(array, dense_idx))
out = HistEFT(self._label, self._wcnames, *new_dims, dtype=self._dtype)
if self._sumw2 is not None: out._init_sumw2()
for sparse_key in self._sumw:
if not all(k in idx for k, idx in zip(sparse_key, sparse_idx)):
continue
if sparse_key in out._sumw:
out._sumw[sparse_key] += dense_op(self._sumw[sparse_key])
if self._sumw2 is not None:
out._sumw2[sparse_key] += dense_op(self._sumw2[sparse_key])
else:
out._sumw[sparse_key] = dense_op(self._sumw[sparse_key]).copy()
if self._sumw2 is not None:
out._sumw2[sparse_key] = dense_op(
self._sumw2[sparse_key]).copy()
for sparse_key in self.EFTcoeffs:
if not all(k in idx for k, idx in zip(sparse_key, sparse_idx)):
continue
if sparse_key in out.EFTcoeffs:
for i in range(len(out.EFTcoeffs[sparse_key])):
out.EFTcoeffs[sparse_key][i] += dense_op(
self.EFTcoeffs[sparse_key][i])
out.EFTerrs[sparse_key][i] += dense_op(
self.EFTerrs[sparse_key][i])
else:
out.EFTcoeffs[sparse_key] = []
out.EFTerrs[sparse_key] = []
for i in range(self.GetNcoeffs()):
out.EFTcoeffs[sparse_key].append(
np.zeros(shape=self._dense_shape, dtype=self._dtype))
for i in range(self.GetNcoeffsErr()):
out.EFTerrs[sparse_key].append(
np.zeros(shape=self._dense_shape, dtype=self._dtype))
for i in range(len(self.EFTcoeffs[sparse_key])):
out.EFTcoeffs[sparse_key][i] += dense_op(
self.EFTcoeffs[sparse_key][i]).copy()
out.EFTerrs[sparse_key][i] += dense_op(
self.EFTerrs[sparse_key][i]).copy()
return out
def sum(self, *axes, **kwargs):
""" Integrates out a set of axes, producing a new histogram
Project() and integrate() depends on sum() and are heritated """
overflow = kwargs.pop('overflow', 'none')
axes = [self.axis(ax) for ax in axes]
reduced_dims = [ax for ax in self._axes if ax not in axes]
out = HistEFT(self._label,
self._wcnames,
*reduced_dims,
dtype=self._dtype)
if self._sumw2 is not None: out._init_sumw2()
sparse_drop = []
dense_slice = [slice(None)] * self.dense_dim()
dense_sum_dim = []
for axis in axes:
if isinstance(axis, coffea.hist.hist_tools.DenseAxis):
idense = self._idense(axis)
dense_sum_dim.append(idense)
dense_slice[idense] = overflow_behavior(overflow)
elif isinstance(axis, coffea.hist.hist_tools.SparseAxis):
isparse = self._isparse(axis)
sparse_drop.append(isparse)
dense_slice = tuple(dense_slice)
dense_sum_dim = tuple(dense_sum_dim)
def dense_op(array):
if len(dense_sum_dim) > 0:
return np.sum(array[dense_slice], axis=dense_sum_dim)
return array
for key in self._sumw.keys():
new_key = tuple(k for i, k in enumerate(key)
if i not in sparse_drop)
if new_key in out._sumw:
out._sumw[new_key] += dense_op(self._sumw[key])
if self._sumw2 is not None:
out._sumw2[new_key] += dense_op(self._sumw2[key])
else:
out._sumw[new_key] = dense_op(self._sumw[key]).copy()
if self._sumw2 is not None:
out._sumw2[new_key] = dense_op(self._sumw2[key]).copy()
for key in self.EFTcoeffs.keys():
new_key = tuple(k for i, k in enumerate(key)
if i not in sparse_drop)
if new_key in out.EFTcoeffs:
#out.EFTcoeffs[new_key] += dense_op(self.EFTcoeffs[key])
#out.EFTerrs [new_key] += dense_op(self.EFTerrs [key])
for i in range(len(self.EFTcoeffs[key])):
out.EFTcoeffs[new_key][i] += dense_op(
self.EFTcoeffs[key][i])
for i in range(len(self.EFTerrs[key])):
out.EFTerrs[new_key][i] += dense_op(self.EFTerrs[key][i])
else:
out.EFTcoeffs[new_key] = []
out.EFTerrs[new_key] = []
for i in range(len(self.EFTcoeffs[key])):
out.EFTcoeffs[new_key].append(
dense_op(self.EFTcoeffs[key][i]).copy())
for i in range(len(self.EFTerrs[key])):
out.EFTerrs[new_key].append(
dense_op(self.EFTerrs[key][i]).copy())
return out
def project(self, *axes, **kwargs):
""" Project histogram onto a subset of its axes
Same as in parent class """
overflow = kwargs.pop('overflow', 'none')
axes = [self.axis(ax) for ax in axes]
toremove = [ax for ax in self.axes() if ax not in axes]
return self.sum(*toremove, overflow=overflow)
def integrate(self, axis_name, int_range=slice(None), overflow='none'):
""" Integrates current histogram along one dimension
Same as in parent class """
axis = self.axis(axis_name)
full_slice = tuple(
slice(None) if ax != axis else int_range for ax in self._axes)
if isinstance(int_range, coffea.hist.hist_tools.Interval):
# Handle overflow intervals nicely
if int_range.nan(): overflow = 'justnan'
elif int_range.lo == -np.inf: overflow = 'under'
elif int_range.hi == np.inf: overflow = 'over'
return self[full_slice].sum(
axis.name, overflow=overflow) # slice may make new axis, use name
def remove(self, bins, axis):
""" Remove bins from a sparse axis
Same as in parent class """
axis = self.axis(axis)
if not isinstance(axis, coffea.hist.hist_tools.SparseAxis):
raise NotImplementedError(
"Hist.remove() only supports removing items from a sparse axis."
)
bins = [axis.index(binid) for binid in bins]
keep = [
binid.name for binid in self.identifiers(axis) if binid not in bins
]
full_slice = tuple(
slice(None) if ax != axis else keep for ax in self._axes)
return self[full_slice]
def group(self, old_axes, new_axis, mapping, overflow='none'):
""" Group a set of slices on old axes into a single new axis """
### WARNING: check that this function works properly... (TODO) --> Are the EFT coefficients properly grouped?
if not isinstance(new_axis, coffea.hist.hist_tools.SparseAxis):
raise TypeError(
"New axis must be a sparse axis. Note: Hist.group() signature has changed to group(old_axes, new_axis, ...)!"
)
if new_axis in self.axes() and self.axis(new_axis) is new_axis:
raise RuntimeError(
"new_axis is already in the list of axes. Note: Hist.group() signature has changed to group(old_axes, new_axis, ...)!"
)
if not isinstance(old_axes, tuple): old_axes = (old_axes, )
old_axes = [self.axis(ax) for ax in old_axes]
old_indices = [i for i, ax in enumerate(self._axes) if ax in old_axes]
new_dims = [new_axis] + [ax for ax in self._axes if ax not in old_axes]
out = HistEFT(self._label, self._wcnames, *new_dims, dtype=self._dtype)
if self._sumw2 is not None: out._init_sumw2()
for new_cat in mapping.keys():
the_slice = mapping[new_cat]
if not isinstance(the_slice, tuple): the_slice = (the_slice, )
if len(the_slice) != len(old_axes):
raise Exception(
"Slicing does not match number of axes being rebinned")
full_slice = [slice(None)] * self.dim()
for idx, s in zip(old_indices, the_slice):
full_slice[idx] = s
full_slice = tuple(full_slice)
reduced_hist = self[full_slice].sum(
*tuple(ax.name for ax in old_axes),
overflow=overflow) # slice may change old axis binning
new_idx = new_axis.index(new_cat)
for key in reduced_hist._sumw:
new_key = (new_idx, ) + key
out._sumw[new_key] = reduced_hist._sumw[key]
if self._sumw2 is not None:
out._sumw2[new_key] = reduced_hist._sumw2[key]
# Will this piece work??
out.EFTcoeffs = copy.deepcopy(reduced_hist.EFTcoeffs)
out.EFTerrs = copy.deepcopy(reduced_hist.EFTerrs)
return out
def rebin(self, old_axis, new_axis):
""" Rebin a dense axis """
old_axis = self.axis(old_axis)
if isinstance(new_axis, numbers.Integral):
new_axis = Bin(old_axis.name, old_axis.label,
old_axis.edges()[::new_axis])
new_dims = [ax if ax != old_axis else new_axis for ax in self._axes]
out = HistEFT(self._label, self._wcnames, *new_dims, dtype=self._dtype)
if self._sumw2 is not None: out._init_sumw2()
idense = self._idense(old_axis)
def view_ax(idx):
fullindex = [slice(None)] * self.dense_dim()
fullindex[idense] = idx
return tuple(fullindex)
binmap = [
new_axis.index(i) for i in old_axis.identifiers(overflow='allnan')
]
def dense_op(array):
anew = np.zeros(out._dense_shape, dtype=out._dtype)
for iold, inew in enumerate(binmap):
anew[view_ax(inew)] += array[view_ax(iold)]
return anew
for key in self._sumw:
out._sumw[key] = dense_op(self._sumw[key])
if self._sumw2 is not None:
out._sumw2[key] = dense_op(self._sumw2[key])
### TODO: check that this is working!
for key in self.EFTcoeffs.keys():
if key in out.EFTcoeffs:
for i in range(len(self.EFTcoeffs[key])):
out.EFTcoeffs[key][i] += dense_op(self.EFTcoeffs[key][i])
for i in range(len(self.EFTerrs[key])):
out.EFTerrs[key][i] += dense_op(self.EFTerrs[key][i])
else:
out.EFTcoeffs[key] = []
out.EFTerrs[key] = []
for i in range(len(self.EFTcoeffs[key])):
out.EFTcoeffs[key].append(
dense_op(self.EFTcoeffs[key][i]).copy())
for i in range(len(self.EFTerrs[key])):
out.EFTerrs[key].append(
dense_op(self.EFTerrs[key][i]).copy())
return out
###################################################################
### Evaluation
def Eval(self, wcp=None):
""" Set a WC point and evaluate """
if isinstance(WCPoint, dict):
wcp = WCPoint(wcp)
elif isinstance(wcp, str):
values = wcp.replace(" ", "").split(',')
wcp = WCPoint(values, names=self.GetWCnames())
elif isinstance(wcp, list):
wcp = WCPoint(wcp, names=self.GetWCnames())
if wcp is None:
if self._WCPoint is None:
self._WCPoint = WCPoint(names=self._wcnames)
else:
self._WCPoint = wcp
self.EvalInSelfPoint()
def EvalInSelfPoint(self):
""" Evaluate to self._WCPoint """
if len(self.WCFit.keys()) == 0: self.SetWCFit()
if not hasattr(self, '_sumw_orig'):
self._sumw_orig = self._sumw.copy()
self._sumw2_orig = self._sumw2.copy()
for key in self.WCFit.keys():
weights = np.array(
[wc.EvalPoint(self._WCPoint) for wc in self.WCFit[key]])
errors = np.array(
[wc.EvalPointError(self._WCPoint) for wc in self.WCFit[key]])
self._sumw[key] = self._sumw_orig[key] * weights
self._sumw2[key] = self._sumw2_orig[key] * errors
def SetSMpoint(self):
""" Set SM WC point and evaluate """
wc = WCPoint(names=self._wcnames)
wc.SetSMPoint()
self.Eval(wc)
def SetStrength(self, wc, val):
""" Set a WC strength and evaluate """
self._WCPoint.SetStrength(wc, val)
self.EvalInSelfPoint()
def GetWCnames(self):
return self._wcnames