def test_bytemaskedarray2b():
array = ak.from_iter(
[[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False)
mask = ak.layout.Index8(np.array([0, 1, 0, 0], dtype=np.int8))
maskedarray = ak.layout.ByteMaskedArray(mask, array, valid_when=False)
assert ak.to_list(maskedarray) == [
[0.0, 1.1, 2.2],
None,
[5.5],
[6.6, 7.7, 8.8, 9.9],
]
assert ak.to_list(maskedarray[ak.Array([[1, 1], None, [0], [0, -1]])]) == [
[1.1, 1.1],
None,
[5.5],
[6.6, 9.9],
]
def _getMatchIdxs(self, tree):
#no split here
truthMom = self._readArray(tree,"MergedSimCluster_boundaryEnergy")
truthEta = self._readArray(tree,"MergedSimCluster_impactPoint_eta")
truthPhi = self._readArray(tree,"MergedSimCluster_impactPoint_phi")
truthpos = ak1.concatenate([self._expand(truthEta),self._expand(truthPhi)],axis=-1)
impactEta = self._readArray(tree,"Track_HGCFront_eta")
impactPhi = self._readArray(tree,"Track_HGCFront_phi")
impactpos = ak1.concatenate([self._expand(impactEta),self._expand(impactPhi)],axis=-1)
trackPt = self._readArray(tree,"Track_pt")
trackVertEta = self._readArray(tree,"Track_eta")
trackMom = trackPt * np.cosh(trackVertEta)
#match by x,y, and momentum
finalidxs = []
for tpos, ipos, tmom, imom, ipt in zip(truthpos, impactpos, truthMom, trackMom, trackPt):
# create default
tpos, ipos, tmom, imom,ipt = tpos.to_numpy(), ipos.to_numpy(), tmom.to_numpy(), imom.to_numpy(), ipt.to_numpy()
tpos = np.expand_dims(tpos, axis=0) #one is truth
tmom = np.expand_dims(tmom, axis=0) #one is truth
ipos = np.expand_dims(ipos, axis=1)
imom = np.expand_dims(imom, axis=1)
ipt = np.expand_dims(ipt,axis=1)
#this is in cm.
posdiffsq = np.sum( (tpos[:,:,0:1]-ipos[:,:,0:1])**2 +deltaPhi(tpos[:,:,1:2],ipos[:,:,1:2])**2, axis=-1) # Trk x K
#this is in %
momdiff = 100.*np.abs(tmom - imom)/(imom+1e-3) #rel diff
#scale position by 100 (DeltaR)
totaldiff = np.sqrt(100.**2*posdiffsq + (momdiff*np.exp(-0.05*ipt))**2)#weight momentum difference less with higher momenta
closestSC = np.argmin(totaldiff, axis=1) # Trk
#more than 5 percent/1cm total difference
closestSC[totaldiff[np.arange(len(closestSC)),closestSC] > 5] = -1
finalidxs.append(closestSC)
return ak1.from_iter(finalidxs)
def test_sort_strings():
content1 = ak.from_iter(["one", "two", "three", "four", "five"],
highlevel=False)
assert ak.to_list(content1) == ["one", "two", "three", "four", "five"]
assert ak.to_list(content1.sort(0, True, False)) == [
"five",
"four",
"one",
"three",
"two",
]
assert ak.to_list(content1.sort(0, False, False)) == [
"two",
"three",
"one",
"four",
"five",
]
def test_rpad_recordarray():
array = ak.from_iter(
[{"x": [], "y": [2, 2]}, {"x": [1.1], "y": [1]}, {"x": [2.2, 2.2], "y": []}],
highlevel=False,
)
assert ak.to_list(array.rpad(5, 0)) == [
{"x": [], "y": [2, 2]},
{"x": [1.1], "y": [1]},
{"x": [2.2, 2.2], "y": []},
None,
None,
]
assert ak.to_list(array.rpad(2, 1)) == [
{"x": [None, None], "y": [2, 2]},
{"x": [1.1, None], "y": [1, None]},
{"x": [2.2, 2.2], "y": [None, None]},
]
def test_argmin_argmax():
array = ak.Array(
[
[
[np.datetime64("2022"), np.datetime64("2023"), np.datetime64("2025")],
[],
[np.datetime64("2027"), np.datetime64("2011")],
[np.datetime64("2013")],
],
[],
[[np.datetime64("2017"), np.datetime64("2019")], [np.datetime64("2023")]],
],
check_valid=True,
)
assert ak.argmin(array) == 4
assert ak.argmax(array) == 3
assert ak.to_list(ak.argmin(array, axis=0)) == [[1, 1, 0], [1], [0, 0], [0]]
assert ak.to_list(ak.argmax(array, axis=0)) == [[0, 0, 0], [1], [0, 0], [0]]
assert ak.to_list(ak.argmin(array, axis=1)) == [[3, 2, 0], [], [0, 0]]
assert ak.to_list(ak.argmax(array, axis=1)) == [[2, 0, 0], [], [1, 0]]
array = ak.from_iter(
[
[
[
np.datetime64("2021-01-20"),
np.datetime64("2021-01-10"),
np.datetime64("2021-01-30"),
]
],
[[]],
[None, None, None],
[
[
np.datetime64("2021-01-14"),
np.datetime64("2021-01-15"),
np.datetime64("2021-01-16"),
]
],
],
highlevel=False,
)
assert ak.to_list(array.argmin(axis=2)) == [[1], [None], [None, None, None], [0]]
def test_tag_particle_indices():
# will need an eventwise with Parents, Children, MCPID
components = {}
# 0 1 2 3 4 5 6 7 8 9 10
components["Children"] = [[[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [],
[], [], []]]
components["Parents"] = [[[], [], [1], [1], [], [2, 3], [3], [6], [6], [6],
[]]]
components["MCPID"] = [[1, -5, 5101, 10511, 5, 1, -5, -1, 7, 11, 12]]
tag_pids = ['hadrons']
hard_event_pids = [None]
expected = [[1, 4, 6]]
expected_creators = [[[], [], []]]
# 0 1 2 3 4 5 6 7 8 9 10
components["Children"] += [[[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [],
[], [], []]]
components["Parents"] += [[[], [], [1], [1], [], [2, 3], [3], [6], [6],
[6], []]]
components["MCPID"] += [[1, -6, 5101, 10511, 5, 5, -5, -1, 7, 11, 12]]
tag_pids.append([5])
hard_event_pids.append([10511, 5101])
expected.append([4, 5, 6])
expected_creators.append([[], [2, 3], [3]])
components = {k: ak.from_iter(v) for k, v in components.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**components)
for i in range(2):
eventWise.selected_event = i
tag_idx, create_idx = TrueTag.tag_particle_indices(
eventWise,
tag_pids=tag_pids[i],
hard_interaction_pids=hard_event_pids[i])
for idx, create in zip(tag_idx, create_idx):
assert idx in expected[i],\
f"Found {idx}, but tags should be {expected[i]}"
pos = expected[i].index(idx)
tst.assert_allclose(
sorted(create),
expected_creators[i][pos],
err_msg=f"Tag {idx} should be created by " +
f"{expected_creators[i][pos]} but " + f"{create} found")
def addUniqueIndices(self):
'''
Adds a '1' to exactly one hit representing
a truth object, per truth object. such that
number of truth objects = sum(unique)
and
object properties = truth_per_hit_property[unique]
This can be very helpful in loss functions etc.
'''
t_idx = self.truth['t_idx']
nplist=[]
for a in t_idx: #these are numpy arrays
a = a.to_numpy()
_,idx = np.unique(a,return_index=True)
un = np.zeros_like(a)
un[idx]=1
nplist.append(un)
akarr = ak1.from_iter(nplist)
self.truth['t_is_unique']=akarr
def test_indexedarray2b():
array = ak.from_iter(
[[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False)
index = ak.layout.Index64(np.array([0, -1, 2, 3], dtype=np.int64))
indexedarray = ak.layout.IndexedOptionArray64(index, array)
assert ak.to_list(indexedarray) == [
[0.0, 1.1, 2.2],
None,
[5.5],
[6.6, 7.7, 8.8, 9.9],
]
assert ak.to_list(indexedarray[ak.Array([[1, 1], None, [0],
[0, -1]])]) == [
[1.1, 1.1],
None,
[5.5],
[6.6, 9.9],
]
def test_numpyarray():
layout = ak.from_iter(
[0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], highlevel=False
)
numbatype = ak._connect._numba.arrayview.tonumbatype(layout.form)
assert ak_connect_numba_layout.typeof(layout).name == numbatype.name
lookup1 = ak_connect_numba_arrayview.Lookup(layout)
lookup2 = ak_connect_numba_arrayview.Lookup(layout.form)
numbatype.form_fill(0, layout, lookup2)
assert np.array_equal(lookup1.arrayptrs, lookup2.arrayptrs)
assert np.array_equal(lookup1.sharedptrs == -1, lookup2.sharedptrs == -1)
counter = [0]
def materialize():
counter[0] += 1
return layout
generator = ak.layout.ArrayGenerator(
materialize, form=layout.form, length=len(layout)
)
virtualarray = ak.layout.VirtualArray(generator)
lookup3 = ak_connect_numba_arrayview.Lookup(virtualarray)
assert len(lookup1.arrayptrs) + 3 == len(lookup3.arrayptrs)
array = ak.Array(virtualarray)
array.numba_type
assert counter[0] == 0
@numba.njit
def f1(x):
return x[5]
assert f1(array) == 5.5
assert counter[0] == 1
assert f1(array) == 5.5
assert counter[0] == 1
def test_bitmaskedarray2b():
array = ak.from_iter(
[[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False
)
mask = ak.layout.IndexU8(np.array([66], dtype=np.uint8))
maskedarray = ak.layout.BitMaskedArray(
mask, array, valid_when=False, length=4, lsb_order=True
) # lsb_order is irrelevant in this example
assert ak.to_list(maskedarray) == [
[0.0, 1.1, 2.2],
None,
[5.5],
[6.6, 7.7, 8.8, 9.9],
]
assert ak.to_list(maskedarray[ak.Array([[1, 1], None, [0], [0, -1]])]) == [
[1.1, 1.1],
None,
[5.5],
[6.6, 9.9],
]
def test_get_simplified_spec(
example_full_analysis_likelihood, example_simplified_analysis_likelihood
):
ylds = yields.Yields(
regions=['SRLMEM_mct2'],
samples=[
'diboson',
'multiboson',
'singletop',
'ttbar',
'tth',
'ttv',
'vh',
'wjets',
'zjets',
],
yields={
'SRLMEM_mct2': np.asarray(
[
[3.81835159e-01, 8.29630721e-02, 1.82488713e-01],
[0.00000000e00, 4.39047681e-03, 8.58322892e-05],
[1.19221557e01, 7.45065852e00, 4.70797363e00],
[1.12848942e00, 1.07669201e00, 7.45587698e-01],
[4.14071746e-02, 4.21658868e-02, 3.47387836e-02],
[5.44315116e-01, 2.52051808e-01, 2.30325982e-01],
[9.73768735e-02, 2.90478572e-01, 4.60765010e-01],
[1.42638837e00, 1.82756453e00, 7.67691098e-01],
[1.35629924e-01, 6.63743835e-02, 4.54261930e-02],
]
)
},
uncertainties={'SRLMEM_mct2': ak.from_iter([23.6, 25.1, 14.6])},
data={'SRLMEM_mct2': np.asarray([16.0, 11.0, 7.0])},
)
spec = simplified.get_simplified_spec(
example_full_analysis_likelihood, ylds, allowed_modifiers=[], prune_channels=[]
)
assert spec == example_simplified_analysis_likelihood
def test_sort_strings():
content1 = ak.from_iter(["one", "two", "three", "four", "five"],
highlevel=False)
assert ak.to_list(content1) == ["one", "two", "three", "four", "five"]
assert ak.to_list(ak.sort(content1, axis=0, ascending=True,
stable=False)) == [
"five",
"four",
"one",
"three",
"two",
]
assert ak.to_list(ak.sort(content1, axis=0, ascending=False,
stable=False)) == [
"two",
"three",
"one",
"four",
"five",
]
def test_ByteMaskedArray_setidentities():
content = ak.from_iter(
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9], [321]],
highlevel=False,
)
mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.layout.ByteMaskedArray(mask, content, valid_when=False)
assert ak.to_list(array) == [[0.0, 1.1, 2.2], [], None, None,
[6.6, 7.7, 8.8, 9.9]]
array.setidentities()
assert np.asarray(array.identities).tolist() == [[0], [1], [2], [3], [4]]
assert np.asarray(array.content.identities).tolist() == [
[0],
[1],
[2],
[3],
[4],
[-1],
]
assert ak.is_valid(ak.Array(array))
def append_tag_shower_widths(eventWise, silent=True):
"""
Calculate the widths of showers from the tag particles and
save the results in the eventWise.
Parameters
----------
eventWise : EventWise
object containting particle data
silent : bool
should progress be printed?
(Default value = True)
"""
if "TagIndex" not in eventWise.columns:
TrueTag.add_tag_particles(eventWise, silent=silent)
eventWise.selected_event = None
name = "TagWidth"
n_events = len(eventWise.TagIndex)
widths = list(getattr(eventWise, name, []))
start_point = len(widths)
if start_point >= n_events:
print("Finished")
return True
end_point = n_events
if not silent:
print(f" Will stop at {100*end_point/n_events}%")
for event_n in range(start_point, end_point):
if event_n % 10 == 0 and not silent:
print(f"{100*event_n/n_events}%", end='\r', flush=True)
if os.path.exists("stop"):
print(f"Completed event {event_n-1}")
break
eventWise.selected_event = event_n
tag_idxs = eventWise.TagIndex
widths_here = [decendants_width(eventWise, idx) for idx in tag_idxs]
widths.append(widths_here)
widths = ak.from_iter(widths)
eventWise.append(**{name: widths})
def test_ByteMaskedArray_localindex():
content = ak.from_iter(
[
[[0.0, 1.1, 2.2], [], [3.3, 4.4]],
[],
[[5.5]],
[[6.6, 7.7, 8.8, 9.9]],
[[], [10.0, 11.1, 12.2]],
],
highlevel=False,
)
mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.layout.ByteMaskedArray(mask, content, valid_when=False)
assert ak.to_list(array) == [
[[0.0, 1.1, 2.2], [], [3.3, 4.4]],
[],
None,
None,
[[], [10.0, 11.1, 12.2]],
]
assert ak.to_list(array.localindex(axis=0)) == [0, 1, 2, 3, 4]
assert ak.to_list(array.localindex(axis=-3)) == [0, 1, 2, 3, 4]
assert ak.to_list(array.localindex(axis=1)) == [[0, 1, 2], [], None, None,
[0, 1]]
assert ak.to_list(array.localindex(axis=-2)) == [[0, 1, 2], [], None, None,
[0, 1]]
assert ak.to_list(array.localindex(axis=2)) == [
[[0, 1, 2], [], [0, 1]],
[],
None,
None,
[[], [0, 1, 2]],
]
assert ak.to_list(array.localindex(axis=-1)) == [
[[0, 1, 2], [], [0, 1]],
[],
None,
None,
[[], [0, 1, 2]],
]
def test_filter_pileup():
with TempTestDir("tst") as dir_name:
# make a simple base with just one event
content = {}
content["Parents"] = [[[], [0]]]
content["Start_vertex_barcode"] = [[-1, 0]]
content["Vertex_barcode"] = [[-1]]
content["Z"] = [[0.]]
content["MCPID"] = [[1, 3]]
content["Event_n"] = [0]
content["Parents"] += [[[], [0], [0], [1, 2]]]
content["Start_vertex_barcode"] += [[-1, -2, -3, -4]]
content["Vertex_barcode"] += [[-3, -2, -1, -4]]
content["Z"] += [[0., 10., 20., 30.]]
content["MCPID"] += [[1, 23, 12, 11]]
content["Event_n"] += [1]
ak_content = {name: ak.from_iter(val) for name, val in content.items()}
simple_path = os.path.join(dir_name, "simple.parquet")
simple = Components.EventWise(simple_path)
simple.append(**ak_content)
# an empty filter should always come back empty
simple.selected_event = 0
returns = RemovePileup.filter_pileup(simple, [])
assert len(returns) == 0
simple.selected_event = 1
returns = RemovePileup.filter_pileup(simple, [])
assert len(returns) == 0
# otherwise it shoudl select indices that meet requirements
simple.selected_event = 0
returns = RemovePileup.filter_pileup(simple, [0, 1])
tst.assert_allclose(returns, [0, 1])
simple.selected_event = 1
returns = RemovePileup.filter_pileup(simple, [1, 2, 3])
tst.assert_allclose(returns, [1, 2])
returns = RemovePileup.filter_pileup(simple, [1, 2, 3],
vertex_resolution_mm=100)
tst.assert_allclose(returns, [1, 2, 3])
def test_IndexedArray_getitem():
content = ak.from_iter(
[0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], highlevel=False
)
index = ak.layout.Index64(np.array([3, 2, 2, 5, 0, 7], dtype=np.int64))
array = ak.Array(ak.layout.IndexedArray64(index, content), check_valid=True)
@numba.njit
def f1(x, i):
return x[i]
assert [f1(array, 0), f1(array, 1), f1(array, 2), f1(array, 3)] == [
3.3,
2.2,
2.2,
5.5,
]
@numba.njit
def f2(x, i1, i2):
return x[i1:i2]
assert ak.to_list(f2(array, 1, 5)) == [2.2, 2.2, 5.5, 0]
def sort_overlap(visibles, roots, eventWise):
visibles = ak.values_astype(visibles, int)
momentums = np.array((eventWise.Px[roots], eventWise.Py[roots],
eventWise.Rapidity[roots])).T
all_visibles = sorted(set(ak.flatten(visibles)))
new_visibles = [[] for _ in roots]
for i, vis in enumerate(all_visibles):
in_showers = [
n_s for n_s, shower in enumerate(visibles) if vis in shower
]
if len(in_showers) == 1:
new_visibles[in_showers[0]].append(vis)
continue
p = np.array(
[eventWise.Px[vis], eventWise.Py[vis], eventWise.Rapidity[vis]])
abs_p = np.sqrt(np.sum(p**2))
abs_m = np.sqrt(np.sum(momentums[in_showers]**2, axis=1))
cos_angle = np.sum(p * momentums[in_showers], axis=1) / (abs_p * abs_m)
new_visibles[in_showers[np.argmax(cos_angle)]].append(vis)
new_visibles = ak.from_iter(new_visibles)
assert len(set(ak.flatten(new_visibles))) == len(ak.flatten(new_visibles))
assert len(all_visibles) == len(ak.flatten(new_visibles))
return new_visibles
def test_append_tag_shower_widths():
# check a normal case
jetinputs_sourceidx = [0, 5, 6, 7]
bquakidx = [0, 1, 4, 10]
# 0 1 2 3 4 5 6 7 8 9 10
children = [[], [2, 3], [5], [6, 5], [], [], [7, 8, 9], [], [], [], []]
rapidity = [0, -1, -1, -1, -1, 1, 0, 0, -1, -1, -1]
phi = [0, -1, -1, -1, -1, 1, 0, 0, -1, -1, -1]
content = {
"JetInputs_SourceIdx": jetinputs_sourceidx,
"TagIndex": bquakidx,
"Children": children,
"Rapidity": rapidity,
"Phi": phi
}
content = {k: ak.from_iter([v]) for k, v in content.items()}
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**content)
AreaMeasures.append_tag_shower_widths(eventWise)
eventWise.selected_event = None
expected = [0., np.sqrt(2), 0., 0.]
tst.assert_allclose(eventWise.TagWidth[0], expected)
def test_sort_bytestrings():
array = ak.from_iter(
[b"one", b"two", b"three", b"two", b"two", b"one", b"three"],
highlevel=False)
assert ak.to_list(array) == [
b"one",
b"two",
b"three",
b"two",
b"two",
b"one",
b"three",
]
assert ak.to_list(array.sort(0, True, False)) == [
b"one",
b"one",
b"three",
b"three",
b"two",
b"two",
b"two",
]
def test_ByteMaskedArray():
content = ak.from_iter(
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False)
mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.layout.ByteMaskedArray(mask, content, valid_when=False)
assert ak.to_list(array.argsort(0, True, False)) == [[0, 0, 0], [],
[1, 1, 1, 0]]
assert ak.to_list(array.sort(0, True, False)) == [
[0.0, 1.1, 2.2],
[],
[6.6, 7.7, 8.8, 9.9],
]
assert ak.to_list(array.sort(0, False, False)) == [
[6.6, 7.7, 8.8],
[],
[0.0, 1.1, 2.2, 9.9],
]
assert ak.to_list(array.argsort(1, True, False)) == [
[0, 1, 2],
[],
None,
None,
[0, 1, 2, 3],
]
assert ak.to_list(array.sort(1, False, False)) == [
[2.2, 1.1, 0.0],
[],
None,
None,
[9.9, 8.8, 7.7, 6.6],
]
def test_minmax():
assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 1 + 5j
assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]])) == 3 + 3j
assert ak.min(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [
1 + 5j,
None,
3 + 3j,
]
assert ak.max(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [
2 + 4j,
None,
3 + 3j,
]
assert ak.argmin(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [0, None, 0]
assert ak.argmax(ak.from_iter([[1 + 5j, 2 + 4j], [], [3 + 3j]]),
axis=1).tolist() == [1, None, 0]
def test_add_softdrop_mask():
# the existing model is simple to test
params = {}
# event 0, empty
params['Jet_Label'] = [[]]
params['Jet_PT'] = [[]]
params['Jet_JoinDistance'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
# event 1
params['Jet_Label'] += [[[0, 1, 2]]]
params['Jet_PT'] += [[[60, 10, 40]]]
params['Jet_JoinDistance'] += [[[60, 10, 40]]]
params['Jet_Child1'] += [[[1, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0]]]
params = {name: ak.from_iter(params[name]) for name in params}
with TempTestDir("tst") as dir_name:
# try a normal mask
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.append_hyperparameters(Jet_DeltaR=0.5)
input_z_cut = 5.
input_beta = 3.
def fake_softdrop_existing(label, parent, pt, z_cut, beta, r0,
distance):
tst.assert_allclose(label, params["Jet_Label"][1, 0])
tst.assert_allclose(parent, params["Jet_Parent"][1, 0])
tst.assert_allclose(pt, params["Jet_PT"][1, 0])
tst.assert_allclose(beta, input_beta)
tst.assert_allclose(input_z_cut, z_cut)
tst.assert_allclose(r0, eventWise.Jet_DeltaR)
tst.assert_allclose(distance, params["Jet_JoinDistance"][1, 0])
return np.ones(1, dtype=bool)
with unittest.mock.patch('jet_tools.RemovePileup.softdrop_from_tree',
new=fake_softdrop_existing):
RemovePileup.add_softdrop_mask(eventWise, "Jet", input_z_cut,
input_beta)
tst.assert_allclose(eventWise.Jet_PSEx5p0Beta3p0R0p5ZCut, input_z_cut)
tst.assert_allclose(eventWise.Jet_PSEx5p0Beta3p0R0p5Beta, input_beta)
tst.assert_allclose(eventWise.Jet_PSEx5p0Beta3p0R0p5R0,
eventWise.Jet_DeltaR)
assert len(eventWise.Jet_PSEx5p0Beta3p0R0p5Mask) == 2
assert len(eventWise.Jet_PSEx5p0Beta3p0R0p5Mask[1]) == 1
assert len(eventWise.Jet_PSEx5p0Beta3p0R0p5Mask[1, 0]) == 1
# try with beta=0
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
eventWise.append_hyperparameters(Jet_DeltaR=0.5)
input_z_cut = 5.
input_beta = 0.
def fake_softdrop_existing(label, parent, pt, z_cut, beta, r0,
distance):
tst.assert_allclose(label, params["Jet_Label"][1, 0])
tst.assert_allclose(parent, params["Jet_Parent"][1, 0])
tst.assert_allclose(pt, params["Jet_PT"][1, 0])
tst.assert_allclose(beta, input_beta)
tst.assert_allclose(input_z_cut, z_cut)
assert distance is None
return np.ones(1, dtype=bool)
with unittest.mock.patch('jet_tools.RemovePileup.softdrop_from_tree',
new=fake_softdrop_existing):
RemovePileup.add_softdrop_mask(eventWise, "Jet", input_z_cut,
input_beta)
tst.assert_allclose(eventWise.Jet_PSEx5p0ZCut, input_z_cut)
tst.assert_allclose(eventWise.Jet_PSEx5p0Beta, input_beta)
assert len(eventWise.Jet_PSEx5p0Mask) == 2
assert len(eventWise.Jet_PSEx5p0Mask[1]) == 1
assert len(eventWise.Jet_PSEx5p0Mask[1, 0]) == 1
# testing with ca clustring will be more complex, as it requires almost all
# jet parameters to be present
# we can cheat a bit using add_all to get rapidity phi and pt
params = {}
# event 0, empty
params['Jet_Label'] = [[]]
params['Jet_Energy'] = [[]]
params['Jet_Px'] = [[]]
params['Jet_Py'] = [[]]
params['Jet_Pz'] = [[]]
params['Jet_Child1'] = [[]]
params['Jet_Parent'] = [[]]
# event 1
params['Jet_Label'] += [[[0, 1, 2]]]
params['Jet_Energy'] += [[[60., 10., 40.]]]
params['Jet_Px'] += [[[6., 1., 4.]]]
params['Jet_Py'] += [[[6., 1., 4.]]]
params['Jet_Pz'] += [[[6., 1., 4.]]]
params['Jet_Child1'] += [[[1, -1, -1]]]
params['Jet_Parent'] += [[[-1, 0, 0]]]
params = {name: ak.from_iter(params[name]) for name in params}
with TempTestDir("tst") as dir_name:
# try a normal mask
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
eventWise.append(**params)
Components.add_all(eventWise, "Jet_")
eventWise.append_hyperparameters(Jet_DeltaR=0.5)
input_z_cut = 5.
input_beta = 3.
def fake_softdrop_existing(label, parent, pt, z_cut, beta, r0,
distance):
assert np.sum(label != -1) == 3
# the leaves should make it into the label
assert 1 in label
assert 2 in label
assert len(label) == len(parent)
assert len(label) == len(pt)
assert len(label) == len(distance)
tst.assert_allclose(beta, input_beta)
tst.assert_allclose(input_z_cut, z_cut)
tst.assert_allclose(r0, eventWise.Jet_DeltaR)
return np.ones(1, dtype=bool)
with unittest.mock.patch('jet_tools.RemovePileup.softdrop_from_tree',
new=fake_softdrop_existing):
RemovePileup.add_softdrop_mask(eventWise,
"Jet",
input_z_cut,
input_beta,
existing=False)
tst.assert_allclose(eventWise.Jet_PSCA5p0Beta3p0R0p5ZCut, input_z_cut)
tst.assert_allclose(eventWise.Jet_PSCA5p0Beta3p0R0p5Beta, input_beta)
tst.assert_allclose(eventWise.Jet_PSCA5p0Beta3p0R0p5R0,
eventWise.Jet_DeltaR)
assert len(eventWise.Jet_PSCA5p0Beta3p0R0p5Mask) == 2
assert len(eventWise.Jet_PSCA5p0Beta3p0R0p5Mask[1]) == 1
assert len(eventWise.Jet_PSCA5p0Beta3p0R0p5Mask[1, 0]) == 3
def test_rpad_and_clip_indexed_array():
listoffsetarray = ak.from_iter(
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False)
backward = ak.from_iter(
[[6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], [], [0.0, 1.1, 2.2]],
highlevel=False)
index = ak.layout.Index64(np.array([4, 3, 2, 1, 0], dtype=np.int64))
indexedarray = ak.layout.IndexedArray64(index, listoffsetarray)
assert ak.to_list(indexedarray) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
[0.0, 1.1, 2.2],
]
assert ak.to_list(backward.rpad_and_clip(4, 1)) == ak.to_list(
indexedarray.rpad_and_clip(4, 1))
assert ak.to_list(indexedarray.rpad_and_clip(1,
0)) == [[6.6, 7.7, 8.8, 9.9]]
assert ak.to_list(indexedarray.rpad_and_clip(2,
0)) == [[6.6, 7.7, 8.8, 9.9],
[5.5]]
assert ak.to_list(indexedarray.rpad_and_clip(3, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
]
assert ak.to_list(indexedarray.rpad_and_clip(4, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
]
assert ak.to_list(indexedarray.rpad_and_clip(5, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
[0.0, 1.1, 2.2],
]
assert ak.to_list(indexedarray.rpad_and_clip(6, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
[0.0, 1.1, 2.2],
None,
]
assert ak.to_list(indexedarray.rpad_and_clip(7, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
[0.0, 1.1, 2.2],
None,
None,
]
assert ak.to_list(indexedarray.rpad_and_clip(8, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
[],
[0.0, 1.1, 2.2],
None,
None,
None,
]
assert ak.to_list(indexedarray.rpad_and_clip(1, 1)) == [
[6.6],
[5.5],
[3.3],
[None],
[0.0],
]
assert ak.to_list(indexedarray.rpad_and_clip(2, 1)) == [
[6.6, 7.7],
[5.5, None],
[3.3, 4.4],
[None, None],
[0.0, 1.1],
]
assert ak.to_list(indexedarray.rpad_and_clip(3, 1)) == [
[6.6, 7.7, 8.8],
[5.5, None, None],
[3.3, 4.4, None],
[None, None, None],
[0.0, 1.1, 2.2],
]
assert ak.to_list(indexedarray.rpad_and_clip(4, 1)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5, None, None, None],
[3.3, 4.4, None, None],
[None, None, None, None],
[0.0, 1.1, 2.2, None],
]
assert ak.to_list(indexedarray.rpad_and_clip(5, 1)) == [
[6.6, 7.7, 8.8, 9.9, None],
[5.5, None, None, None, None],
[3.3, 4.4, None, None, None],
[None, None, None, None, None],
[0.0, 1.1, 2.2, None, None],
]
def test_rpad_indexed_option_array():
listoffsetarray = ak.from_iter(
[[0.0, None, None], None, [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False,
)
backward = ak.from_iter(
[[6.6, 7.7, 8.8, 9.9], [5.5], [3.3, 4.4], None, [0.0, None, None]],
highlevel=False,
)
index = ak.layout.Index64(np.array([4, 3, 2, -1, 0], dtype=np.int64))
indexedarray = ak.layout.IndexedOptionArray64(index, listoffsetarray)
assert ak.to_list(indexedarray) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
]
assert ak.to_list(backward.rpad(4,
1)) == ak.to_list(indexedarray.rpad(4, 1))
assert ak.to_list(indexedarray.rpad(1, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
]
assert ak.to_list(indexedarray.rpad(1, 1)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
]
assert ak.to_list(indexedarray.rpad(3, 1)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5, None, None],
[3.3, 4.4, None],
None,
[0.0, None, None],
]
assert ak.to_list(indexedarray.rpad(4, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
]
assert ak.to_list(indexedarray.rpad(5, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
]
assert ak.to_list(indexedarray.rpad(6, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
None,
]
assert ak.to_list(indexedarray.rpad(7, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
None,
None,
]
assert ak.to_list(indexedarray.rpad(8, 0)) == [
[6.6, 7.7, 8.8, 9.9],
[5.5],
[3.3, 4.4],
None,
[0.0, None, None],
None,
None,
None,
]
assert ak.to_list(indexedarray.rpad_and_clip(1,
0)) == [[6.6, 7.7, 8.8, 9.9]]
assert ak.to_list(indexedarray.rpad_and_clip(1, 1)) == [
[6.6],
[5.5],
[3.3],
None,
[0.0],
]
def eval_model(model, dataset, config, outdir):
ibatch = 0
jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4)
for elem in tqdm(dataset, desc="Evaluating model"):
y_pred = model.predict(elem["X"], verbose=False)
np_outfile = "{}/pred_batch{}.npz".format(outdir, ibatch)
ygen = unpack_target(elem["ygen"],
config["dataset"]["num_output_classes"], config)
ycand = unpack_target(elem["ycand"],
config["dataset"]["num_output_classes"], config)
outs = {}
for key in y_pred.keys():
outs["gen_{}".format(key)] = ygen[key].numpy()
outs["cand_{}".format(key)] = ycand[key].numpy()
outs["pred_{}".format(key)] = y_pred[key]
jets_coll = {}
jets_const = {}
for typ in ["gen", "cand", "pred"]:
cls_id = np.argmax(outs["{}_cls".format(typ)], axis=-1)
valid = cls_id != 0
pt = awkward.from_iter(
[y[m][:, 0] for y, m in zip(outs["{}_pt".format(typ)], valid)])
eta = awkward.from_iter([
y[m][:, 0] for y, m in zip(outs["{}_eta".format(typ)], valid)
])
phi = np.arctan2(outs["{}_sin_phi".format(typ)],
outs["{}_cos_phi".format(typ)])
phi = awkward.from_iter([y[m][:, 0] for y, m in zip(phi, valid)])
e = awkward.from_iter([
y[m][:, 0]
for y, m in zip(outs["{}_energy".format(typ)], valid)
])
vec = vector.arr({"pt": pt, "eta": eta, "phi": phi, "e": e})
cluster = fastjet.ClusterSequence(vec.to_xyzt(), jetdef)
jets = cluster.inclusive_jets()
jet_constituents = cluster.constituent_index()
jets_coll[typ] = jets[jets.pt > 5.0]
jets_const[typ] = jet_constituents[jets.pt > 5.0]
for key in ["pt", "eta", "phi", "energy"]:
outs["jets_gen_{}".format(key)] = awkward.to_numpy(
awkward.flatten(getattr(jets_coll["gen"], key)))
outs["jets_cand_{}".format(key)] = awkward.to_numpy(
awkward.flatten(getattr(jets_coll["cand"], key)))
outs["jets_pred_{}".format(key)] = awkward.to_numpy(
awkward.flatten(getattr(jets_coll["pred"], key)))
# DeltaR match between genjets and PF/MLPF jets
cart = awkward.cartesian([jets_coll["gen"], jets_coll["pred"]],
nested=True)
jets_a, jets_b = awkward.unzip(cart)
drs = deltar(jets_a, jets_b)
match_gen_to_pred = [awkward.where(d < 0.1) for d in drs]
m0 = awkward.from_iter([m[0] for m in match_gen_to_pred])
m1 = awkward.from_iter([m[1] for m in match_gen_to_pred])
j1s = jets_coll["gen"][m0]
j2s = jets_coll["pred"][m1]
outs["jets_pt_gen_to_pred"] = np.stack([
awkward.to_numpy(awkward.flatten(j1s.pt)),
awkward.to_numpy(awkward.flatten(j2s.pt))
],
axis=-1)
cart = awkward.cartesian([jets_coll["gen"], jets_coll["cand"]],
nested=True)
jets_a, jets_b = awkward.unzip(cart)
drs = deltar(jets_a, jets_b)
match_gen_to_pred = [awkward.where(d < 0.1) for d in drs]
m0 = awkward.from_iter([m[0] for m in match_gen_to_pred])
m1 = awkward.from_iter([m[1] for m in match_gen_to_pred])
j1s = jets_coll["gen"][m0]
j2s = jets_coll["cand"][m1]
outs["jets_pt_gen_to_cand"] = np.stack([
awkward.to_numpy(awkward.flatten(j1s.pt)),
awkward.to_numpy(awkward.flatten(j2s.pt))
],
axis=-1)
np.savez(np_outfile, X=elem["X"], **outs)
ibatch += 1
def test_record():
array1 = ak.from_iter([{
"x": 1,
"y": 1.1
}, {
"x": 2,
"y": 2.2
}, {
"x": 3,
"y": 3.3
}],
highlevel=False)
assert ak.to_list(array1) == [
{
"x": 1,
"y": 1.1
},
{
"x": 2,
"y": 2.2
},
{
"x": 3,
"y": 3.3
},
]
array2 = array1.setitem_field(
"z", ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array2) == [
{
"x": 1,
"y": 1.1,
"z": []
},
{
"x": 2,
"y": 2.2,
"z": [1]
},
{
"x": 3,
"y": 3.3,
"z": [2, 2]
},
]
array3 = array1.setitem_field(
None, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [
{
"x": 1,
"y": 1.1,
"2": []
},
{
"x": 2,
"y": 2.2,
"2": [1]
},
{
"x": 3,
"y": 3.3,
"2": [2, 2]
},
]
array3 = array1.setitem_field(
0, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [
{
"x": 1,
"y": 1.1,
"0": []
},
{
"x": 2,
"y": 2.2,
"0": [1]
},
{
"x": 3,
"y": 3.3,
"0": [2, 2]
},
]
array1 = ak.from_iter([(1, 1.1), (2, 2.2), (3, 3.3)], highlevel=False)
assert ak.to_list(array1) == [(1, 1.1), (2, 2.2), (3, 3.3)]
array2 = array1.setitem_field(
"z", ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array2) == [
{
"0": 1,
"1": 1.1,
"z": []
},
{
"0": 2,
"1": 2.2,
"z": [1]
},
{
"0": 3,
"1": 3.3,
"z": [2, 2]
},
]
array3 = array1.setitem_field(
None, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2,
2])]
array3 = array1.setitem_field(
0, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [([], 1, 1.1), ([1], 2, 2.2), ([2,
2], 3, 3.3)]
array3 = array1.setitem_field(
1, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [(1, [], 1.1), (2, [1], 2.2), (3, [2,
2], 3.3)]
array3 = array1.setitem_field(
100, ak.from_iter([[], [1], [2, 2]], highlevel=False))
assert ak.to_list(array3) == [(1, 1.1, []), (2, 2.2, [1]), (3, 3.3, [2,
2])]
def test_add_bg_mass():
# will need TagIndex JetInputs_SourceIdx, DetectableTag_Leaves,
# Energy Px Py Pz
params = {}
# an empty event should have 0 mass
params["TagIndex"] = [ak.from_iter([3, 4])]
params["JetInputs_SourceIdx"] = [ak.from_iter([])]
params["DetectableTag_Leaves"] = [ak.from_iter([])]
params["Energy"] = [ak.from_iter([])]
params["Px"] = [ak.from_iter([])]
params["Py"] = [ak.from_iter([])]
params["Pz"] = [ak.from_iter([])]
# an event with no jet_inputs should also have no mass
params["TagIndex"] += [ak.from_iter([1, 2, 3, 4])]
params["JetInputs_SourceIdx"] += [ak.from_iter([])]
params["DetectableTag_Leaves"] += [ak.from_iter([[5, 6], [7, 8]])]
params["Energy"] += [ak.from_iter([0, 0, 3, 0, 0, 0, 1, 10, 9])]
params["Px"] += [ak.from_iter([0, 0, 1, 0, 0, 0, 1, 0, -3])]
params["Py"] += [ak.from_iter([0, 0, 1, 1, 0, 0, -1, 0, 3])]
params["Pz"] += [ak.from_iter([0, 0, 0, 0, 0, 0, 1, 1, 3])]
# a standard event can have mass
params["TagIndex"] += [ak.from_iter([1, 2, 3, 4])]
params["JetInputs_SourceIdx"] += [ak.from_iter([2, 6, 7])]
params["DetectableTag_Leaves"] += [ak.from_iter([[5, 6], [7, 8]])]
params["Energy"] += [ak.from_iter([0, 0, 3, 0, 0, 0, 1, 10, 9])]
params["Px"] += [ak.from_iter([0, 0, 1, 0, 0, 0, 1, 0, -3])]
params["Py"] += [ak.from_iter([0, 0, 1, 1, 0, 0, -1, 0, 3])]
params["Pz"] += [ak.from_iter([0, 0, 0, 0, 0, 0, 1, 1, 3])]
expected = [0, 0, np.sqrt(9 - 2)]
with TempTestDir("tst") as dir_name:
# this will raise a value error if given an empty eventWise
file_name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, file_name))
ew.append(**params)
# settign a selected index should not screw things up
ew.selected_event = 1
CompareClusters.add_bg_mass(ew)
ew.selected_event = None
tst.assert_allclose(ew.DetectableBG_Mass, expected)
def test_make_scale():
# this has two modes, ordinal and float
# either way it should return somehting sensible
# given an empty list it should not choke
empty = np.array([])
positions, scale_positions, scale_labels = CompareClusters.make_scale(
empty)
assert len(positions) == 0
assert len(scale_positions) == 0
assert len(scale_labels) == 0
# a list with just one number should return it
one = np.array([1])
positions, scale_positions, scale_labels = CompareClusters.make_scale(one)
assert len(positions) == 1
tst.assert_allclose(positions, one)
assert len(scale_positions) == 1
tst.assert_allclose(scale_positions, one)
assert len(scale_labels) == 1
assert scale_labels[0] == '1'
# ditto for one string
one = np.array(['frog'])
positions, scale_positions, scale_labels = CompareClusters.make_scale(one)
assert len(positions) == 1
assert len(scale_positions) == 1
assert len(scale_labels) == 1
assert scale_labels[0] == one[0]
# one nan or inf should return that, but it may also add 0 as a scale pos
one = np.array([np.nan])
positions, scale_positions, scale_labels = CompareClusters.make_scale(one)
assert len(positions) == 1
scale_loc = np.argmin(np.abs(scale_positions - positions[0]))
assert 'nan' in scale_labels[scale_loc] or 'None' in scale_labels[scale_loc]
one = np.array([-np.inf])
positions, scale_positions, scale_labels = CompareClusters.make_scale(one)
assert len(positions) == 1
scale_loc = np.argmin(np.abs(scale_positions - positions[0]))
assert 'inf' in scale_labels[scale_loc] and '-' in scale_labels[scale_loc]
one = np.array([np.inf])
positions, scale_positions, scale_labels = CompareClusters.make_scale(one)
assert len(positions) == 1
scale_loc = np.argmin(np.abs(scale_positions - positions[0]))
assert 'inf' in scale_labels[scale_loc] and '-' not in scale_labels[
scale_loc]
# now try some more complex arrangements
inputs = [[4, 3.2], [4, 3.2, np.nan], [4, 3.2, 0., np.nan],
[-4, 3.2, np.nan], [4, 3.2, np.inf], [4, 3.2, -np.inf],
[-4, 3.2, 0, np.nan, np.inf], [None, 'dog', 'cat'],
[
np.inf, 1, 3, 5, 3.4, 10, np.nan, 2.2, 7, 3, 5, 9, 2, 2,
-np.inf
]]
closest_contains = [['4', '3'], ['4', '3', 'nan'], ['4', '3', '0', 'nan'],
['-4', '3', 'nan'], ['4', '3', 'inf'], ['4', '3', '-'],
['-4', '3', '0', 'nan', 'inf'], ['none', 'dog', 'cat'],
[
'inf', '1', '3', '5', '3', '10', 'nan', '2', '7',
'3', '5', '9', '2', '2', '-'
]]
for inp, close in zip(inputs, closest_contains):
# make the test scale
inp = np.array(inp)
positions, scale_positions, scale_labels = CompareClusters.make_scale(
inp)
# check that for each of the values in the test scale
for i, pos in enumerate(positions):
s_position = np.argmin(np.abs(scale_positions - pos))
# the closest label contains the expected string
assert close[i] in scale_labels[s_position].lower(
), f"inputs={inp}, expected={close[i]}, found={scale_labels[s_position]}"
# finally check the behavior with tuples
inp = [None, ('knn', 5), ('knn', 4), ('distance', 0), ('distance', 2)]
# make the test scale
inp = ak.from_iter(inp)
positions, scale_positions, scale_labels = CompareClusters.make_scale(inp)
assert positions[1] > positions[2]
assert positions[4] > positions[3]
label0 = scale_labels[np.argmin(np.abs(scale_positions -
positions[0]))].lower()
assert 'none' in label0
label1 = scale_labels[np.argmin(np.abs(scale_positions -
positions[1]))].lower()
assert 'knn' in label1 and '5' in label1
label2 = scale_labels[np.argmin(np.abs(scale_positions -
positions[2]))].lower()
assert 'knn' in label2 and '4' in label2
label3 = scale_labels[np.argmin(np.abs(scale_positions -
positions[3]))].lower()
assert 'distance' in label3 and '0' in label3
label4 = scale_labels[np.argmin(np.abs(scale_positions -
positions[4]))].lower()
assert 'distance' in label4 and '2' in label4
def test_append_scores():
# set up as for filter jets
params = {}
# event 0
params['Jet_Label'] = [ak.from_iter([])]
params['Jet_Parent'] = [ak.from_iter([])]
params['Jet_Child1'] = [ak.from_iter([])]
params['Jet_PT'] = [ak.from_iter([])]
params['Energy'] = [ak.from_iter([])]
params['Px'] = [ak.from_iter([])]
params['Py'] = [ak.from_iter([])]
params['Pz'] = [ak.from_iter([])]
params['Children'] = [ak.from_iter([])]
params['Parents'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['JetInputs_SourceIdx'] = [ak.from_iter([])]
params['TagIndex'] = [ak.from_iter([])]
# event 1
params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3, 4, 5]])]
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1, 2, 2]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, 2, -1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [0.2, 0.1, 0., 0., .1]])]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.])
]
params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])]
params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])]
params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])]
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
# event 2
params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3]])]
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [21., 0.1, 0.]])]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.])
]
params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])]
params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])]
params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])]
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
params = {key: ak.from_iter(val) for key, val in params.items()}
with TempTestDir("tst") as dir_name:
# this will raise a value error if given an empty eventWise
file_name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, file_name))
ew.append(**params)
# mock JetQuality.quality_width_fracton and get_detectable_comparisons
with unittest.mock.patch('jet_tools.JetQuality.quality_width_fracton',
new=fake_quality_width_fraction):
with unittest.mock.patch(
'jet_tools.CompareClusters.get_detectable_comparisons',
new=fake_detectable_comparisons):
with unittest.mock.patch(
'jet_tools.TrueTag.add_detectable_fourvector',
new=fake_empty):
import jet_tools
CompareClusters.append_scores(ew)
tst.assert_allclose(ew.Jet_Bork, [np.inf, np.inf, np.inf])
tst.assert_allclose(ew.Jet_Quack[0], [0.5, np.nan])
tst.assert_allclose(ew.Jet_Quack[1], [0.5, -np.inf])
tst.assert_allclose(ew.Jet_Quack[2], [])
assert np.isnan(ew.Jet_AveBork)
tst.assert_allclose(ew.Jet_AveQuack, 0.5)
assert np.isnan(ew.Jet_QualityWidth)
tst.assert_allclose(ew.Jet_QualityFraction,
3 / Constants.dijet_mass)
