def getitem_enter(array, where): if len(where) == 0: return array arraylen = 0 for x in where: if isinstance(x, numpy.ndarray) and len(x.shape) == 1: if issubclass(x.dtype.type, (numpy.bool_, numpy.bool)): arraylen = max(arraylen, numpy.count_nonzero(x)) else: arraylen = max(arraylen, len(x)) newwhere = [] for x in where: if isinstance(x, numpy.ndarray) and len(x.shape) == 1 and issubclass( x.dtype.type, (numpy.bool, numpy.bool_)): newwhere.append(numpy.nonzero(x)[0]) elif isinstance(x, int) and arraylen != 0: newwhere.append(numpy.full(arraylen, x, int)) elif isinstance(x, numpy.ndarray) and x.shape == (1, ): newwhere.append(numpy.full(arraylen, x, int)) else: newwhere.append(x) fake = getitem_next(awkward0.JaggedArray([0], [len(array)], array), newwhere, None) if isinstance(fake, numpy.ndarray): return fake[0] else: return fake.content[fake.starts[0]:fake.stops[-1]]
def getitem_intarray_none(array, head, tail, advanced): offsets = numpy.full(len(array.starts) + 1, 999, int) offsets[0] = 0 index = numpy.full(len(head) * len(array.starts), 999, int) nextadvanced = numpy.full(len(index), 999, int) k = 0 for i in range(len(array.starts)): length = array.stops[i] - array.starts[i] for j in range(len(head)): norm = head[j] if norm < 0: norm += length if norm < 0 or norm >= length: raise IndexError( "advanced index is out of bounds in JaggedArray") index[k] = array.starts[i] + norm nextadvanced[k] = j k += 1 offsets[i + 1] = k starts = offsets[:-1] stops = offsets[1:] next = getitem_next(array.content[index], tail, nextadvanced) return awkward0.JaggedArray(starts, stops, next)
def test_jagged_zip(self): a = awkward0.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]]) b = awkward0.JaggedArray([1, 5, 5], [4, 5, 7], [999, 10, 20, 30, 999, 40, 50, 999]) c = numpy.array([100, 200, 300]) d = 1000 assert awkward0.JaggedArray.zip(one=a, two=b).tolist() == [[{"one": 1.1, "two": 10}, {"one": 2.2, "two": 20}, {"one": 3.3, "two": 30}], [], [{"one": 4.4, "two": 40}, {"one": 5.5, "two": 50}]] assert awkward0.JaggedArray.zip(one=b, two=a).tolist() == [[{"one": 10, "two": 1.1}, {"one": 20, "two": 2.2}, {"one": 30, "two": 3.3}], [], [{"one": 40, "two": 4.4}, {"one": 50, "two": 5.5}]] assert awkward0.JaggedArray.zip(one=b, two=c).tolist() == [[{"one": 10, "two": 100}, {"one": 20, "two": 100}, {"one": 30, "two": 100}], [], [{"one": 40, "two": 300}, {"one": 50, "two": 300}]] assert awkward0.JaggedArray.zip(one=b, two=d).tolist() == [[{"one": 10, "two": 1000}, {"one": 20, "two": 1000}, {"one": 30, "two": 1000}], [], [{"one": 40, "two": 1000}, {"one": 50, "two": 1000}]] assert a.zip(b).tolist() == [[(1.1, 10), (2.2, 20), (3.3, 30)], [], [(4.4, 40), (5.5, 50)]] assert b.zip(a).tolist() == [[(10, 1.1), (20, 2.2), (30, 3.3)], [], [(40, 4.4), (50, 5.5)]] assert b.zip(c).tolist() == [[(10, 100), (20, 100), (30, 100)], [], [(40, 300), (50, 300)]] assert b.zip(d).tolist() == [[(10, 1000), (20, 1000), (30, 1000)], [], [(40, 1000), (50, 1000)]]
def getitem_slice2(array, head, tail, advanced): if (head.start is None or head.start == 0) and head.stop is None: next = getitem_next(array.content, tail, advanced) return awkward0.JaggedArray(array.starts, array.stops, next) starts = numpy.full(len(array.starts), 999, int) stops = numpy.full(len(array.starts), 999, int) for i in range(len(array.starts)): length = array.stops[i] - array.starts[i] a, b = head.start, head.stop if a is None: a = 0 elif a < 0: a += length if b is None: b = length elif b < 0: b += length if b <= a: a, b = 0, 0 if a < 0: a = 0 elif a > length: a = length if b < 0: b = 0 elif b > length: b = length starts[i] = array.starts[i] + a stops[i] = array.starts[i] + b next = getitem_next(array.content, tail, advanced) return awkward0.JaggedArray(starts, stops, next)
def test_jagged_structure1d(self): a = awkward0.JaggedArray([[0, 3, 3], [5, 8, 10]], [[3, 3, 5], [8, 10, 10]], [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]) assert a.tolist() == a.structure1d().tolist() a = awkward0.JaggedArray([[[0], [3], [3]], [[5], [8], [10]]], [[[3], [3], [5]], [[8], [10], [10]]], [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]) assert a.tolist() == a.structure1d().tolist()
def getitem_slice3(array, head, tail, advanced): if head.step == 0: raise ValueError offsets = numpy.full(len(array.starts) + 1, 999, int) offsets[0] = 0 index = numpy.full(len(array.content), 999, int) # too big, but okay k = 0 for i in range(len(array.starts)): length = array.stops[i] - array.starts[i] a, b, c = head.start, head.stop, head.step if c is None: c = 1 if a is None and c > 0: a = 0 elif a is None: a = length - 1 elif a < 0: a += length if b is None and c > 0: b = length elif b is None: b = -1 elif b < 0: b += length if c > 0: if b <= a: a, b = 0, 0 if a < 0: a = 0 elif a > length: a = length if b < 0: b = 0 elif b > length: b = length else: if a <= b: a, b = 0, 0 if a < -1: a = -1 elif a >= length: a = length - 1 if b < -1: b = -1 elif b >= length: b = length - 1 for j in range(a, b, c): index[k] = array.starts[i] + j k += 1 offsets[i + 1] = k starts = offsets[:-1] stops = offsets[1:] next = getitem_next(array.content[index[:k]], tail, spread_advanced(starts, stops, advanced)) return awkward0.JaggedArray(starts, stops, next)