def objUnion(objs, sortBy=None):
objs = tuple(objs)
assert objs
assert all(isinstance(obj, awkward.JaggedArray) for obj in objs)
assert all(obj.shape == objs[0].shape for obj in objs)
tags = awkward.JaggedArray.concatenate(
[
awkward.JaggedArray(obj.starts, obj.stops,
np.full(obj.shape, i, dtype=np.intp))
for i, obj in enumerate(objs)
],
axis=1,
)
index = awkward.JaggedArray.concatenate(
[
awkward.JaggedArray(obj.starts, obj.stops,
np.arange(len(obj.content), dtype=np.intp))
for obj in objs
],
axis=1,
)
if sortBy:
if not callable(sortBy):
sortBy = attrgetter(sortBy)
shuf = awkward.JaggedArray.concatenate(list(map(sortBy, objs)),
axis=1).argsort()
tags = tags[shuf]
index = index[shuf]
# index = tags.copy(content=index.content) # reduces memory footprint
return tags.copy(content=awkward.UnionArray(tags.content, index.content,
[obj.content for obj in objs]))
def mergeArray(a1, a2):
'''
Merge two arrays into one, e.g. electrons and muons
'''
import awkward
a1_tags = awkward.JaggedArray(a1.starts, a1.stops, np.full(len(a1.content), 0, dtype=np.int64))
a1_index = awkward.JaggedArray(a1.starts, a1.stops, np.arange(len(a1.content), dtype=np.int64))
a2_tags = awkward.JaggedArray(a2.starts, a2.stops, np.full(len(a2.content), 1, dtype=np.int64))
a2_index = awkward.JaggedArray(a2.starts, a2.stops, np.arange(len(a2.content), dtype=np.int64))
tags = awkward.JaggedArray.concatenate([a1_tags, a2_tags], axis=1)
index = awkward.JaggedArray.concatenate([a1_index, a2_index], axis=1)
return awkward.JaggedArray(tags.starts, tags.stops, awkward.UnionArray(tags.content, index.content, [a1.content, a2.content]))
def recurse(layout):
if isinstance(layout, awkward1.partition.PartitionedArray):
return awkward0.ChunkedArray(
[recurse(x) for x in layout.partitions])
elif isinstance(layout, awkward1.layout.NumpyArray):
return numpy.asarray(layout)
elif isinstance(layout, awkward1.layout.EmptyArray):
return numpy.array([])
elif isinstance(layout, awkward1.layout.RegularArray):
# content, size
if keeplayout:
raise ValueError(
"awkward0 has no equivalent of RegularArray; "
"try keeplayout=False")
offsets = numpy.arange(0,
(len(layout) + 1)*layout.size,
layout.size)
return awkward0.JaggedArray.fromoffsets(
offsets, recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListArray32):
# starts, stops, content
return awkward0.JaggedArray(
numpy.asarray(layout.starts),
numpy.asarray(layout.stops),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListArrayU32):
# starts, stops, content
return awkward0.JaggedArray(
numpy.asarray(layout.starts),
numpy.asarray(layout.stops),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListArray64):
# starts, stops, content
return awkward0.JaggedArray(
numpy.asarray(layout.starts),
numpy.asarray(layout.stops),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListOffsetArray32):
# offsets, content
return awkward0.JaggedArray.fromoffsets(
numpy.asarray(layout.offsets),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListOffsetArrayU32):
# offsets, content
return awkward0.JaggedArray.fromoffsets(
numpy.asarray(layout.offsets),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ListOffsetArray64):
# offsets, content
return awkward0.JaggedArray.fromoffsets(
numpy.asarray(layout.offsets),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.Record):
# istuple, numfields, field(i)
out = []
for i in range(layout.numfields):
content = layout.field(i)
if isinstance(content, (awkward1.layout.Content,
awkward1.layout.Record)):
out.append(recurse(content))
else:
out.append(content)
if layout.istuple:
return tuple(out)
else:
return dict(zip(layout.keys(), out))
elif isinstance(layout, awkward1.layout.RecordArray):
# istuple, numfields, field(i)
if layout.numfields == 0 and len(layout) != 0:
raise ValueError(
"cannot convert zero-field, nonzero-length RecordArray "
"to awkward0.Table (limitation in awkward0)")
keys = layout.keys()
values = [recurse(x) for x in layout.contents]
pairs = collections.OrderedDict(zip(keys, values))
out = awkward0.Table(pairs)
if layout.istuple:
out._rowname = "tuple"
return out
elif isinstance(layout, awkward1.layout.UnionArray8_32):
# tags, index, numcontents, content(i)
return awkward0.UnionArray(numpy.asarray(layout.tags),
numpy.asarray(layout.index),
[recurse(x) for x in layout.contents])
elif isinstance(layout, awkward1.layout.UnionArray8_U32):
# tags, index, numcontents, content(i)
return awkward0.UnionArray(numpy.asarray(layout.tags),
numpy.asarray(layout.index),
[recurse(x) for x in layout.contents])
elif isinstance(layout, awkward1.layout.UnionArray8_64):
# tags, index, numcontents, content(i)
return awkward0.UnionArray(numpy.asarray(layout.tags),
numpy.asarray(layout.index),
[recurse(x) for x in layout.contents])
elif isinstance(layout, awkward1.layout.IndexedOptionArray32):
# index, content
index = numpy.asarray(layout.index)
toosmall = (index < -1)
if toosmall.any():
index = index.copy()
index[toosmall] = -1
return awkward0.IndexedMaskedArray(index, recurse(layout.content))
elif isinstance(layout, awkward1.layout.IndexedOptionArray64):
# index, content
index = numpy.asarray(layout.index)
toosmall = (index < -1)
if toosmall.any():
index = index.copy()
index[toosmall] = -1
return awkward0.IndexedMaskedArray(index, recurse(layout.content))
elif isinstance(layout, awkward1.layout.IndexedArray32):
# index, content
return awkward0.IndexedArray(numpy.asarray(layout.index),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.IndexedArrayU32):
# index, content
return awkward0.IndexedArray(numpy.asarray(layout.index),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.IndexedArray64):
# index, content
return awkward0.IndexedArray(numpy.asarray(layout.index),
recurse(layout.content))
elif isinstance(layout, awkward1.layout.ByteMaskedArray):
# mask, content, valid_when
return awkward0.MaskedArray(numpy.asarray(layout.mask),
recurse(layout.content),
maskedwhen=(not layout.valid_when))
elif isinstance(layout, awkward1.layout.BitMaskedArray):
# mask, content, valid_when, length, lsb_order
return awkward0.BitMaskedArray(numpy.asarray(layout.mask),
recurse(layout.content),
maskedwhen=(not layout.valid_when),
lsborder=layout.lsb_order)
elif isinstance(layout, awkward1.layout.UnmaskedArray):
# content
return recurse(layout.content) # no equivalent in awkward0
else:
raise AssertionError(
"missing converter for {0}".format(type(layout).__name__))