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