def test_NumpyArray():
a = ak._v2.contents.RegularArray(
v1_to_v2(ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout), 3)
assert ak.to_list(a[1]) == [
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29],
]
assert ak.to_list(a[1, -2]) == [20, 21, 22, 23, 24]
assert a[1, -2, 2] == 22
with pytest.raises(IndexError):
a[1, -2, 2, 0]
assert ak.to_list(a[1, -2, 2:]) == [22, 23, 24]
with pytest.raises(IndexError):
a[1, -2, 2:, 0]
with pytest.raises(IndexError):
a[1, -2, "hello"]
with pytest.raises(IndexError):
a[1, -2, ["hello", "there"]]
assert ak.to_list(a[1, -2, np.newaxis, 2]) == [22]
assert ak.to_list(a[1, -2, np.newaxis, np.newaxis, 2]) == [[22]]
assert ak.to_list(a[1, -2, ...]) == [20, 21, 22, 23, 24]
assert a[1, -2, ..., 2] == 22
with pytest.raises(IndexError):
a[1, -2, ..., 2, 2]
b = ak.layout.RegularArray(
ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3)
assert ak.to_list(b[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22]
assert ak.to_list(a[1, -2, [3, 1, 1, 2]]) == [23, 21, 21, 22]
with pytest.raises(IndexError):
a[1, -2, [3, 1, 1, 2], 2]
def set_genZ(events, genBranches, selection_options, debug):
if genBranches is None:
events["genZ_decayMode"] = ak.from_numpy(-1 * numpy.ones(len(events)))
else:
electron_idxs = abs(genBranches.pdgId) == 11
muon_idxs = abs(genBranches.pdgId) == 13
tau_idxs = abs(genBranches.pdgId) == 15
motherOfElectrons = genBranches.genPartIdxMother[electron_idxs]
motherOfMuons = genBranches.genPartIdxMother[muon_idxs]
motherOfTaus = genBranches.genPartIdxMother[tau_idxs]
ZToEleEvents = ak.sum(
(genBranches.pdgId[motherOfElectrons] == 23), axis=1) > 0
ZToMuEvents = ak.sum(
(genBranches.pdgId[motherOfMuons] == 23), axis=1) > 0
ZToTauEvents = ak.sum(
(genBranches.pdgId[motherOfTaus] == 23), axis=1) > 0
# W decay dudes
WToEEvents = ak.sum(
(abs(genBranches.pdgId[motherOfElectrons]) == 24), axis=1) > 0
WToMuEvents = ak.sum(
(abs(genBranches.pdgId[motherOfMuons]) == 24), axis=1) > 0
WToTauEvents = ak.sum(
(abs(genBranches.pdgId[motherOfTaus]) == 24), axis=1) > 0
events["genZ_decayMode"] = ak.from_numpy(
numpy.zeros(len(events))
) + 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents + 4 * WToEEvents + 5 * WToMuEvents + 6 * WToTauEvents
events["tau_motherID"] = ak.fill_none(
ak.firsts(genBranches.pdgId[motherOfTaus]), 0)
return events
def test_ArrayBuilder_append_2():
# issue #415
A = ak.from_numpy(np.array([0, 1, 2], dtype=np.float32))
B = ak.from_numpy(np.array([0, 1], dtype=np.float32))
builder = ak.ArrayBuilder()
with builder.list():
builder.append(A.tolist())
with builder.list():
builder.append(A.tolist())
with builder.list():
pass
with builder.list():
builder.append(B.tolist())
assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [],
[[0, 1]]]
assert str(ak.type(builder.snapshot())) == "4 * var * var * float64"
builder = ak.ArrayBuilder()
with builder.list():
builder.append(A)
with builder.list():
builder.append(A)
with builder.list():
pass
with builder.list():
builder.append(B)
assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [],
[[0, 1]]]
assert str(ak.type(builder.snapshot())) == "4 * var * var * float32"
@numba.njit
def f1(builder, A, B):
builder.begin_list()
builder.append(A)
builder.end_list()
builder.begin_list()
builder.append(A)
builder.end_list()
builder.begin_list()
builder.end_list()
builder.begin_list()
builder.append(B)
builder.end_list()
return builder
assert f1(ak.ArrayBuilder(), A, B).snapshot().tolist() == [
[[0, 1, 2]],
[[0, 1, 2]],
[],
[[0, 1]],
]
def test_fromnumpy():
assert ak.to_list(ak.from_numpy(np.array(["uno", "dos", "tres",
"quatro"]))) == [
"uno",
"dos",
"tres",
"quatro",
]
assert ak.to_list(
ak.from_numpy(np.array([["uno", "dos"],
["tres", "quatro"]]))) == [["uno", "dos"],
["tres", "quatro"]]
assert ak.to_list(
ak.from_numpy(np.array([["uno", "dos"], ["tres", "quatro"]]),
regulararray=True)) == [["uno", "dos"],
["tres", "quatro"]]
def test_localindex():
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)
assert ak.to_list(array.localindex(0)) == [0, 1, 2, 3, 4]
assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [], [0, 1], [0],
[0, 1, 2, 3]]
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,
)
assert ak.to_list(array.localindex(0)) == [0, 1, 2, 3]
assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [], [0], [0]]
assert ak.to_list(array.localindex(2)) == [
[[0, 1, 2], [], [0, 1]],
[],
[[0]],
[[0, 1, 2, 3]],
]
array = ak.from_numpy(np.arange(2 * 3 * 5).reshape(2, 3, 5),
regulararray=True,
highlevel=False)
assert ak.to_list(array.localindex(0)) == [0, 1]
assert ak.to_list(array.localindex(1)) == [[0, 1, 2], [0, 1, 2]]
assert ak.to_list(array.localindex(2)) == [
[[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]],
[[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]],
]
def to_awkward(
self,
featurization_key="last",
y_dtype="float32",
clear_after=False,
):
"""
Creates an awkward array out of the featurized systems
and the associated measurements.
Returns
-------
awkward array
Notes
-----
Awkward Array is a library for nested, variable-sized data,
including arbitrary-length lists, records, mixed types,
and missing data, using NumPy-like idioms.
Arrays are dynamically typed, but operations on
them are compiled and fast. Their behavior coincides
with NumPy when array dimensions are regular
and generalizes when they’re not.
"""
features = self.featurized_systems(key=featurization_key,
clear_after=clear_after)
# Features is a list of systems (s) and their features (f): [(s0f0, s0f1), (s1f0, s1f1)...]
# We are going to iterate over columns: (s0f0, s1f0... snf0), (s0f1, s1f1, ..., snf1)
# The result is that X will contain the an array of f0, then an array for f1... etc.
X = [ak.from_iter(subX) for subX in zip(*features)]
y = ak.from_numpy(self.measurements_as_array(dtype=y_dtype))
return X, y
def test_regular():
np_data = np.array([[1, 3, 5, 4, 2], [3, 7, 8, 2, 4], [2, 3, 1, 7, 7],
[5, 1, 9, 10, 2]])
ak_data = ak.from_numpy(np_data)
assert ak.ptp(ak_data, axis=1).tolist() == np.ptp(np_data, axis=1).tolist()
assert ak.ptp(ak_data, axis=0).tolist() == np.ptp(np_data, axis=0).tolist()
assert ak.ptp(ak_data) == np.ptp(np_data)
def test():
first = ak.from_numpy(np.array([1, 2, 3]))
deltas = ak.Array([[1, 2], [1, 2], [1, 2, 3]])
assert np.hstack(
(first[:,
np.newaxis], ak.fill_none(ak.pad_none(deltas, 3, axis=-1),
999))).tolist() == [[1, 1, 2, 999],
[2, 1, 2, 999],
[3, 1, 2, 3]]
def test():
parameters = ak.from_numpy(np.arange(3 * 3).reshape(-1, 3), regulararray=True)
array = ak.zip((parameters, parameters))
assert array.layout.purelist_isregular
assert ak.to_list(array) == [
[(0, 0), (1, 1), (2, 2)],
[(3, 3), (4, 4), (5, 5)],
[(6, 6), (7, 7), (8, 8)],
]
def ak_num(array, **kwargs):
"""
awkward.num() wrapper also for work in empty array
:param array: Data containing nested lists to count.
:param kwargs: keywords arguments for awkward.num().
:return: an array of integers specifying the number of elements
at a particular level. If array is empty, return empty.
"""
if len(array) == 0:
return ak.from_numpy(np.empty(0, dtype='int64'))
return ak.num(array, **kwargs)
def cluster(inter, classify_by_energy=False):
"""
Function which clusters the found clusters together.
To cluster events a weighted mean is computed for time and position.
The individual interactions are weighted by their energy.
The energy of clustered interaction is given by the total sum.
Events can be classified either by the first interaction in time in the
cluster or by the highest energy deposition.
Args:
inter (awkward.Array): Array containing at least the following
fields: x,y,z,t,ed,cluster_ids, type, parenttype, creaproc,
edproc.
Kwargs:
classify_by_energy (bool): If true events are classified
according to the properties of the highest energy deposit
within the cluster. If false cluster is classified according
to first interaction.
Returns:
awkward.Array: Clustered events with nest conform
classification.
"""
if len(inter) == 0:
result_cluster_dtype = [('x', 'float64'), ('y', 'float64'),
('z', 'float64'), ('t', 'float64'),
('ed', 'float64'), ('nestid', 'int64'),
('A', 'int64'), ('Z', 'int64')]
return ak.from_numpy(np.empty(0, dtype=result_cluster_dtype))
# Sort interactions by cluster_ids to simplify looping
inds = ak.argsort(inter['cluster_ids'])
inter = inter[inds]
# TODO: Better way to do this with awkward?
x = inter['x']
y = inter['y']
z = inter['z']
ed = inter['ed']
time = inter['t']
ci = inter['cluster_ids']
types = inter['type']
parenttype = inter['parenttype']
creaproc = inter['creaproc']
edproc = inter['edproc']
# Init result and cluster:
res = ak.ArrayBuilder()
_cluster(x, y, z, ed, time, ci, types, parenttype, creaproc, edproc,
classify_by_energy, res)
return res.snapshot()
def test_IndexedArray():
old = ak.layout.IndexedArray64(
ak.layout.Index64(np.array([1, 0], np.int64)),
ak.layout.RegularArray(
ak.from_numpy(np.arange(2 * 3 * 5).reshape(-1, 5)).layout, 3),
)
new = v1_to_v2(old)
assert v1v2_equal(old[1, 1:], new[1, 1:])
assert ak.to_list(old[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]]
assert ak.to_list(new[1, 1:]) == [[5, 6, 7, 8, 9], [10, 11, 12, 13, 14]]
with pytest.raises(IndexError):
new[1, "hello"]
with pytest.raises(IndexError):
new[1, ["hello", "there"]]
assert ak.to_list(new[0, np.newaxis]) == [[[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29]]]
assert ak.to_list(old[0, np.newaxis]) == [[[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24],
[25, 26, 27, 28, 29]]]
assert old.minmax_depth == (3, 3)
assert new.minmax_depth == (3, 3)
assert v1v2_equal(old[1, ...], new[1, ...])
assert ak.to_list(old[1, ...]) == [
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
]
assert ak.to_list(new[1, ...]) == [
[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
]
expectation = [
[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14]],
[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]],
]
assert (ak.to_list(old[[1, 0], ]) == expectation)
assert (ak.to_list(new[[1, 0], ]) == expectation)
assert ak.to_list(old[[1, 0]]) == expectation
assert ak.to_list(new[[1, 0]]) == expectation
assert ak.to_list(old[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]]
assert ak.to_list(new[1, [1, 0]]) == [[5, 6, 7, 8, 9], [0, 1, 2, 3, 4]]
assert v1v2_equal(old[1, [1, 0]], new[1, [1, 0]])
def in_sensitive_volume(events, sensitive_volumes):
"""
Function which identifies which events are inside sensitive volumes.
Further, tests if sensitive volumes overlap. Assigns volume id, and
xenon density to interactions.
Args:
events (ak.records): Awkward record of the interactions.
sensitive_volumes (dict): Dictionary of the different volumes
defined via the SensitiveVolume class.
Returns:
ak.array: Awkward array containing the event ids.
"""
if len(events) == 0:
res_det_dtype = [('xe_density', 'float64'),
('vol_id', 'int64'),
('create_S2', 'bool'),
]
return ak.from_numpy(np.empty(0, dtype=res_det_dtype))
for ind, vol in enumerate(sensitive_volumes):
res = ak.ArrayBuilder()
res = _inside_sens_vol(events['x'],
events['y'],
events['z'],
vol.roi,
vol.volume_id,
vol.xe_density,
vol.create_S2,
res)
if ind:
# Now we add the other results, but first test if
# volumes overlap. Only possible if we explicitly loop
# over everything. This reduces performance but adds layer of
# safety.
m = (result['vol_id'] > 0) & (res['vol_id'] == vol.volume_id)
if ak.any(m):
overlapping_id = result[m][0]
# Get volume name:
name = [vol.name for vol in sensitive_volumes if vol.volume_id == overlapping_id][0]
raise ValueError(f'The volume {vol.name} is overlapping with'
f' volume {name}!')
new_results = res.snapshot()
for field in result.fields:
# Workaround since we cannot sum up records-arrays anymore
result[field] = result[field] + new_results[field]
else:
# First result initiates the array
result = res.snapshot()
return result
def test_lists():
one = ak.Array([[1, 2, 3], [], [4, 5]]).layout
two = ak.Array([[1.1, 2.2], [3.3, 4.4]]).layout
three = ak.layout.EmptyArray()
four = ak.from_numpy(np.array([[10], [20]]),
regulararray=True,
highlevel=False)
assert ak.to_list(one.mergemany([two, three, four])) == [
[1.0, 2.0, 3.0],
[],
[4.0, 5.0],
[1.1, 2.2],
[3.3, 4.4],
[10.0],
[20.0],
]
assert ak.to_list(four.mergemany([three, two, one])) == [
[10.0],
[20.0],
[1.1, 2.2],
[3.3, 4.4],
[1.0, 2.0, 3.0],
[],
[4.0, 5.0],
]
one = ak.layout.ListArray64(one.starts, one.stops, one.content)
two = ak.layout.ListArray64(two.starts, two.stops, two.content)
assert ak.to_list(one.mergemany([two, three, four])) == [
[1.0, 2.0, 3.0],
[],
[4.0, 5.0],
[1.1, 2.2],
[3.3, 4.4],
[10.0],
[20.0],
]
assert ak.to_list(four.mergemany([three, two, one])) == [
[10.0],
[20.0],
[1.1, 2.2],
[3.3, 4.4],
[1.0, 2.0, 3.0],
[],
[4.0, 5.0],
]
def test_fromnumpy():
a = np.arange(2 * 3 * 5).reshape((2, 3, 5))
b = ak.from_numpy(a)
assert ak.to_list(a) == ak.to_list(b)
def test():
np_data = np.random.random(size=(4, 100 * 1024 * 1024 // 8 // 4))
array = ak.from_numpy(np_data, regulararray=False)
assert np_data.nbytes == array.nbytes
def test_toawkward0():
array = ak.from_iter([1.1, 2.2, 3.3, 4.4], highlevel=False)
assert isinstance(ak.to_awkward0(array), np.ndarray)
assert ak.to_awkward0(array).tolist() == [1.1, 2.2, 3.3, 4.4]
array = ak.from_numpy(np.arange(2 * 3 * 5).reshape(2, 3, 5),
highlevel=False).toRegularArray()
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [
[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14]],
[[15, 16, 17, 18, 19], [20, 21, 22, 23, 24], [25, 26, 27, 28, 29]],
]
array = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
array = ak.layout.ListArray64(
ak.layout.Index64(np.array([4, 999, 1], dtype=np.int64)),
ak.layout.Index64(np.array([7, 999, 3], dtype=np.int64)),
ak.layout.NumpyArray(
np.array([3.14, 4.4, 5.5, 123, 1.1, 2.2, 3.3, 321])),
)
assert isinstance(ak.to_awkward0(array), awkward0.JaggedArray)
assert ak.to_awkward0(array).tolist() == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
array = ak.from_iter(
[
{
"x": 0,
"y": []
},
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [2, 2]
},
{
"x": 3.3,
"y": [3, 3, 3]
},
],
highlevel=False,
)
assert isinstance(ak.to_awkward0(array[2]), dict)
assert ak.to_awkward0(array[2])["x"] == 2.2
assert isinstance(ak.to_awkward0(array[2])["y"], np.ndarray)
assert ak.to_awkward0(array[2])["y"].tolist() == [2, 2]
assert isinstance(ak.to_awkward0(array), awkward0.Table)
assert ak.to_awkward0(array).tolist() == [
{
"x": 0,
"y": []
},
{
"x": 1.1,
"y": [1]
},
{
"x": 2.2,
"y": [2, 2]
},
{
"x": 3.3,
"y": [3, 3, 3]
},
]
array = ak.from_iter([(0, []), (1.1, [1]), (2.2, [2, 2]),
(3.3, [3, 3, 3])],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.Table)
assert ak.to_awkward0(array).tolist() == [
(0, []),
(1.1, [1]),
(2.2, [2, 2]),
(3.3, [3, 3, 3]),
]
assert isinstance(ak.to_awkward0(array[2]), tuple)
assert ak.to_awkward0(array[2])[0] == 2.2
assert ak.to_awkward0(array[2])[1].tolist() == [2, 2]
array = ak.from_iter([0.0, [], 1.1, [1], 2.2, [2, 2], 3.3, [3, 3, 3]],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.UnionArray)
assert ak.to_awkward0(array).tolist() == [
0.0,
[],
1.1,
[1],
2.2,
[2, 2],
3.3,
[3, 3, 3],
]
array = ak.from_iter([1.1, 2.2, None, None, 3.3, None, 4.4],
highlevel=False)
assert isinstance(ak.to_awkward0(array), awkward0.IndexedMaskedArray)
assert ak.to_awkward0(array).tolist() == [
1.1, 2.2, None, None, 3.3, None, 4.4
]
content = ak.layout.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]))
index = ak.layout.Index64(np.array([3, 2, 2, 5, 0], dtype=np.int64))
array = ak.layout.IndexedArray64(index, content)
assert isinstance(ak.to_awkward0(array), awkward0.IndexedArray)
assert ak.to_awkward0(array).tolist() == [3.3, 2.2, 2.2, 5.5, 0.0]
def test_regulararray():
layout = ak.from_numpy(np.array([[1, 2, 3], [4, 5, 6]]),
regulararray=True,
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 f3(x):
return x
assert isinstance(f3(array).layout, ak.layout.VirtualArray)
assert counter[0] == 0
@numba.njit
def f1(x):
return x[1][1]
assert f1(array) == 5
assert counter[0] == 1
assert f1(array) == 5
assert counter[0] == 1
@numba.njit
def f2(x):
return x[1]
assert ak.to_list(f2(array)) == [4, 5, 6]
assert counter[0] == 1
assert ak.to_list(f2(array)) == [4, 5, 6]
assert counter[0] == 1
assert ak.to_list(f3(array)) == [[1, 2, 3], [4, 5, 6]]
def test():
array = ak.from_numpy(np.zeros((3, 0), dtype=np.int32))
buffs = ak.to_buffers(array)
new_array = ak.from_buffers(*buffs)
assert ak.to_list(new_array) == [[], [], []]
def test_awkward_from_numpy_array():
x = np.arange(12)
y = ak.from_numpy(x)
z = ak.is_none(y)
assert z.tolist() == [False] * 12
def test():
array = ak.from_numpy(np.zeros((3, 3, 5)))
flattened = ak.flatten(array, axis=None)
assert flattened.ndim == 1
def ggTauTau_inclusive_preselection(events, photons, electrons, muons, taus,
jets, dR, genPart, Category_pairsLoose,
options, debug):
"""
Performs inclusive ggTauTau preselection, requiring >=1 (leptons + tau_h).
Assumes diphoton preselection has already been applied.
Also calculates relevant event-level variables.
"""
cut_diagnostics = utils.CutDiagnostics(
events=events,
debug=debug,
cut_set="[analysis_selections.py : ggTauTau_inclusive_preselection]")
# Get number of electrons, muons, taus
selected_electrons = electrons[lepton_selections.select_electrons(
events, photons, electrons, options, debug)]
selected_muons = muons[lepton_selections.select_muons(
events, photons, muons, options, debug)]
selected_taus = taus[tau_selections.select_taus(events, photons,
selected_muons,
selected_electrons, taus,
options, debug)]
n_electrons = awkward.num(selected_electrons)
n_muons = awkward.num(selected_muons)
n_taus = awkward.num(selected_taus)
# Require >= 1 lep/tau
n_leptons_and_taus = n_electrons + n_muons + n_taus
# only events with hadronic taus (no leptonic taus!!!!!!!!!!)
atleast_one_had_tau_cut = (n_taus >= 1)
# Require OS leptons/taus for events with 2 leptons/taus
sum_charge = awkward.sum(selected_electrons.charge, axis=1) + awkward.sum(
selected_muons.charge, axis=1) + awkward.sum(selected_taus.charge,
axis=1)
charge_cut = sum_charge == 0
two_leptons = n_leptons_and_taus == 2
not_two_leptons = n_leptons_and_taus != 2
os_cut = (
two_leptons & charge_cut
) | not_two_leptons # only require 2 OS leptons if there are ==2 leptons in the event
# Select jets (don't cut on jet quantities for selection, but they will be useful for BDT training)
selected_jets = jets[jet_selections.select_jets(events, photons,
selected_electrons,
selected_muons,
selected_taus, jets,
options, debug)]
all_cuts = os_cut & atleast_one_had_tau_cut
cut_diagnostics.add_cuts([atleast_one_had_tau_cut, os_cut, all_cuts],
["N_taus >= 1", "OS dileptons", "all"])
# Keep only selected events
selected_events = events[all_cuts]
selected_photons = photons[all_cuts]
selected_electrons = selected_electrons[all_cuts]
selected_muons = selected_muons[all_cuts]
selected_taus = selected_taus[all_cuts]
selected_jets = selected_jets[all_cuts]
dR = dR[all_cuts]
# Calculate event-level variables
selected_events = lepton_selections.set_electrons(selected_events,
selected_electrons,
debug)
selected_events = lepton_selections.set_muons(selected_events,
selected_muons, debug)
selected_events = tau_selections.set_taus(selected_events, selected_taus,
debug)
selected_events = jet_selections.set_jets(selected_events, selected_jets,
options, debug)
if genPart is not None:
genPart = genPart[all_cuts]
selected_events = gen_selections.set_genZ(selected_events, genPart,
options, debug)
else:
selected_events["genZ_decayMode"] = awkward.from_numpy(
numpy.ones(len(selected_events)) * -1)
selected_events["tau_motherID"] = awkward.from_numpy(
numpy.ones(len(selected_events)) * -1)
selected_events = compound_selections.compound_selections(
selected_events, options, debug)
return selected_events
def test():
array = ak.from_numpy(np.zeros((1, 0), dtype=np.int32), regulararray=True)
packed = ak.packed(array)
assert ak.to_list(packed) == [[]]
def test_toRegularArray():
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 7, 11),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 7 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 7, 11, 0),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 7 * 11 * 0 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 7, 0, 11),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 7 * 0 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 0, 7, 11),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 0 * 7 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 0, 5, 7, 11),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 0 * 5 * 7 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 0, 3, 5, 7, 11),
np.float64)).layout.toRegularArray())))
== "2 * 0 * 3 * 5 * 7 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((0, 2, 3, 5, 7, 11),
np.float64)).layout.toRegularArray())))
== "0 * 2 * 3 * 5 * 7 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 7, 0, 11, 0),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 7 * 0 * 11 * 0 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 5, 0, 7, 11, 0),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 5 * 0 * 7 * 11 * 0 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 3, 0, 5, 7, 0, 11),
np.float64)).layout.toRegularArray())))
== "2 * 3 * 0 * 5 * 7 * 0 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((2, 0, 3, 5, 7, 0, 11),
np.float64)).layout.toRegularArray())))
== "2 * 0 * 3 * 5 * 7 * 0 * 11 * float64")
assert (str(
ak.type(
ak.Array(
ak.from_numpy(np.empty((0, 2, 3, 5, 7, 0, 11),
np.float64)).layout.toRegularArray())))
== "0 * 2 * 3 * 5 * 7 * 0 * 11 * float64")
def test_actual():
x = ak.from_numpy(np.arange(2 * 3 * 4, dtype=np.int64).reshape(2, 3, 4))
s = x[..., :0]
result = ak.zip({"q": s, "t": s})
assert str(ak.type(result)) == '2 * 3 * 0 * {"q": int64, "t": int64}'
def test():
array = np.array([1, 2, 9, 0])
nplike = ak.nplike.of(array)
ak_array = ak.from_numpy(array)
assert ak.to_list(nplike.to_rectilinear(array)) == ak.to_list(ak_array)
