def test_RecordArray_NumpyArray(): v1a = ak.layout.RecordArray( [ ak.layout.NumpyArray(np.array([0, 1, 2, 3, 4], dtype=np.int64)), ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])), ], ["x", "y"], ) v2a = ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array([0, 1, 2, 3, 4], dtype=np.int64) ), ak._v2.contents.numpyarray.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5]) ), ], ["x", "y"], ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist() v1b = ak.layout.RecordArray( [ ak.layout.NumpyArray(np.array([0, 1, 2, 3, 4], dtype=np.int64)), ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5])), ], None, ) v2b = ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array([0, 1, 2, 3, 4], dtype=np.int64) ), ak._v2.contents.numpyarray.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5]) ), ], None, ) assert v1v2_equal(v1b, v2b) assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b)) assert ak.to_list(v1b) == ak.to_list(v2b) assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist() v1c = ak.layout.RecordArray([], [], 10) v2c = ak._v2.contents.recordarray.RecordArray([], [], 10) assert v1v2_equal(v1c, v2c) assert v1v2_equal(v2_to_v1(v2c), v1_to_v2(v1c)) assert ak.to_list(v1c) == ak.to_list(v2c) assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist() v1d = ak.layout.RecordArray([], None, 10) v2d = ak._v2.contents.recordarray.RecordArray([], None, 10) assert v1v2_equal(v1d, v2d) assert v1v2_equal(v2_to_v1(v2d), v1_to_v2(v1d)) assert ak.to_list(v1d) == ak.to_list(v2d) assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist()
def test_ByteMaskedArray_NumpyArray(): v1a = ak.layout.ByteMaskedArray( ak.layout.Index8(np.array([1, 0, 1, 0, 1], dtype=np.int8)), ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), valid_when=True, ) v2a = ak._v2.contents.bytemaskedarray.ByteMaskedArray( ak._v2.index.Index(np.array([1, 0, 1, 0, 1], dtype=np.int8)), ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), valid_when=True, ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist() v1b = ak.layout.ByteMaskedArray( ak.layout.Index8(np.array([0, 1, 0, 1, 0], dtype=np.int8)), ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), valid_when=False, ) v2b = ak._v2.contents.bytemaskedarray.ByteMaskedArray( ak._v2.index.Index(np.array([0, 1, 0, 1, 0], dtype=np.int8)), ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), valid_when=False, ) assert v1v2_equal(v1b, v2b) assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b)) assert ak.to_list(v1b) == ak.to_list(v2b) assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
def test_UnmaskedArray(): old = ak.layout.UnmaskedArray( ak.layout.NumpyArray( np.array([[0.0, 1.1, 2.2, 3.3], [0.0, 1.1, 2.2, 3.3]]))) new = v1_to_v2(old) assert ak.to_list(old[0, 1:]) == [1.1, 2.2, 3.3] assert ak.to_list(new[0, 1:]) == [1.1, 2.2, 3.3] assert v1v2_equal(old[0, 1:], new[0, 1:]) with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[1, np.newaxis, -2]) == [2.2] assert ak.to_list(new[1, np.newaxis, np.newaxis, -2]) == [[2.2]] assert old.minmax_depth == (2, 2) assert new.minmax_depth == (2, 2) assert ak.to_list(old[1, ..., -2]) == 2.2 assert ak.to_list(new[1, ..., -2]) == 2.2 expectation = [[0.0, 1.1, 2.2, 3.3], [0.0, 1.1, 2.2, 3.3]] assert (ak.to_list(old[[1, 0], ]) == expectation) assert (ak.to_list(new[[1, 0], ]) == expectation) assert ak.to_list(new[[1, 0]]) == expectation assert v1v2_equal(old[1, ...], new[1, ...]) assert ak.to_list(old[1, [1, 0]]) == [1.1, 0.0] assert ak.to_list(new[1, [1, 0]]) == [1.1, 0.0]
def test_ListArray_RecordArray_NumpyArray(): v1a = ak.layout.ListArray64( ak.layout.Index64(np.array([4, 100, 1], dtype=np.int64)), ak.layout.Index64(np.array([7, 100, 3, 200], dtype=np.int64)), ak.layout.RecordArray( [ak.layout.NumpyArray(np.array([6.6, 4.4, 5.5, 7.7, 1.1, 2.2, 3.3, 8.8]))], ["nest"], ), ) v2a = ak._v2.contents.listarray.ListArray( ak._v2.index.Index(np.array([4, 100, 1], dtype=np.int64)), ak._v2.index.Index(np.array([7, 100, 3, 200], dtype=np.int64)), ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array([6.6, 4.4, 5.5, 7.7, 1.1, 2.2, 3.3, 8.8]) ) ], ["nest"], ), ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_EmptyArray(): v1a = ak.layout.EmptyArray() v2a = ak._v2.contents.emptyarray.EmptyArray() assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
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 test_NumpyArray(): v1a = ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3])) v2a = ak._v2.contents.numpyarray.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3])) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist() v1b = ak.layout.NumpyArray(np.arange(2 * 3 * 5, dtype=np.int64).reshape(2, 3, 5)) v2b = ak._v2.contents.numpyarray.NumpyArray( np.arange(2 * 3 * 5, dtype=np.int64).reshape(2, 3, 5) ) assert v1v2_equal(v1b, v2b) assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b)) assert ak.to_list(v1b) == ak.to_list(v2b) assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
def prepare_tuple_item(item): if ak._util.isint(item): return int(item) elif isinstance(item, slice): return item elif ak._util.isstr(item): return item elif item is np.newaxis: return item elif item is Ellipsis: return item elif isinstance(item, ak.highlevel.Array): return prepare_tuple_item(item.layout) elif isinstance(item, ak.layout.Content): return prepare_tuple_item(v1_to_v2(item)) elif isinstance(item, ak._v2.highlevel.Array): return prepare_tuple_item(item.layout) elif isinstance(item, ak._v2.contents.EmptyArray): return prepare_tuple_item(item.toNumpyArray(np.int64)) elif isinstance(item, ak._v2.contents.NumpyArray): return item.data elif isinstance(item, ak._v2.contents.Content): out = prepare_tuple_bool_to_int(prepare_tuple_nested(item)) if isinstance(out, ak._v2.contents.NumpyArray): return out.data else: return out elif isinstance(item, Iterable) and all(ak._util.isstr(x) for x in item): return list(item) elif isinstance(item, Iterable): layout = ak._v2.operations.convert.to_layout(item) as_array = layout.maybe_to_array(layout.nplike) if as_array is None: return prepare_tuple_item(layout) else: return as_array else: raise ak._v2._util.error( TypeError( "only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), " "integer/boolean arrays (possibly with variable-length nested " "lists or missing values), field name (str) or names (non-tuple " "iterable of str) are valid indices for slicing, not\n\n " + repr(item).replace("\n", "\n ") ) )
def test_UnmaskedArray_NumpyArray(): v1a = ak.layout.UnmaskedArray(ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3]))) v2a = ak._v2.contents.unmaskedarray.UnmaskedArray( ak._v2.contents.numpyarray.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3])) ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_RegularArray_RecordArray_NumpyArray(): v1a = ak.layout.RegularArray( ak.layout.RecordArray( [ak.layout.NumpyArray(np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]))], ["nest"], ), 3, ) v2a = ak._v2.contents.regulararray.RegularArray( ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6]) ) ], ["nest"], ), 3, ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist() v1b = ak.layout.RegularArray( ak.layout.RecordArray([ak.layout.EmptyArray()], ["nest"]), 0, zeros_length=10, ) v2b = ak._v2.contents.regulararray.RegularArray( ak._v2.contents.recordarray.RecordArray( [ak._v2.contents.emptyarray.EmptyArray()], ["nest"] ), 0, zeros_length=10, ) assert v1v2_equal(v1b, v2b) assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b)) assert ak.to_list(v1b) == ak.to_list(v2b) assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist()
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 test_ListOffsetArray_NumpyArray(): old = ak.layout.ListOffsetArray64( ak.layout.Index64(np.array([0, 1, 2, 3], np.int64)), ak.layout.NumpyArray( np.array([ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [11.1, 22.2, 33.3, 44.4, 55.5, 66.6], [21.1, 22.2, 23.3, 24.4, 25.5, 26.6], [31.1, 32.2, 33.3, 34.4, 35.5, 36.6], [41.1, 42.2, 43.3, 44.4, 45.5, 46.6], ])), ) new = v1_to_v2(old) assert ak.to_list(old[0, 0:]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]] assert ak.to_list(new[0, 0:]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]] assert v1v2_equal(old[0, 1:], new[0, 1:]) with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list( new[1, np.newaxis]) == [[[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]] assert ak.to_list( old[1, np.newaxis]) == [[[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]]] assert old.minmax_depth == (3, 3) assert new.minmax_depth == (3, 3) assert ak.to_list(old[1, ...]) == [[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]] assert ak.to_list(new[1, ...]) == [[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]] expectation = [ [[11.1, 22.2, 33.3, 44.4, 55.5, 66.6]], [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]], ] 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[0, [0, 0]]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert ak.to_list(new[0, [0, 0]]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ]
def test_IndexedOptionArray_NumpyArray(): v1a = ak.layout.IndexedOptionArray64( ak.layout.Index64(np.array([2, 2, -1, 1, -1, 5, 4], dtype=np.int64)), ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), ) v2a = ak._v2.contents.indexedoptionarray.IndexedOptionArray( ak._v2.index.Index(np.array([2, 2, -1, 1, -1, 5, 4], dtype=np.int64)), ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6])), ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_ListOffsetArray_NumpyArray(): v1a = ak.layout.ListOffsetArray64( ak.layout.Index64(np.array([1, 4, 4, 6], dtype=np.int64)), ak.layout.NumpyArray([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]), ) v2a = ak._v2.contents.listoffsetarray.ListOffsetArray( ak._v2.index.Index(np.array([1, 4, 4, 6], dtype=np.int64)), ak._v2.contents.numpyarray.NumpyArray([6.6, 1.1, 2.2, 3.3, 4.4, 5.5, 7.7]), ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test(): array = ak.Array([[{ "x": [[1, 2], []] }, { "x": [] }, { "x": [[3]] }]], with_name="Record") array_v2 = ak._v2.highlevel.Array(v1_to_v2(array.layout)) # Flattening inside record does not drop parameters assert "__record__" not in ak.parameters(ak.flatten(array, axis=3)) assert "__record__" not in array_v2.layout.flatten(axis=3).parameters # Flattening before record does not drop parameters assert "__record__" in ak.parameters(ak.flatten(array, axis=1)) assert "__record__" in array_v2.layout.flatten(axis=1).parameters
def test_ByteMaskedArray(): old = ak.layout.ByteMaskedArray( ak.layout.Index8(np.array([1, 1, 1], np.int8)), ak.layout.NumpyArray( np.array([ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ])), valid_when=True, ) new = v1_to_v2(old) assert ak.to_list(old[:, 5:]) == [[6.6], [6.6], [6.6]] assert ak.to_list(new[:, 5:]) == [[6.6], [6.6], [6.6]] assert v1v2_equal(old[:, 5:], new[:, 5:]) with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[1, np.newaxis]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]] assert ak.to_list(old[1, np.newaxis]) == [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6]] assert old.minmax_depth == (2, 2) assert new.minmax_depth == (2, 2) assert ak.to_list(old[0, ..., 0]) == 1.1 assert ak.to_list(new[0, ..., 0]) == 1.1 expectation = [[1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6]] assert (ak.to_list(old[[1, 0], ]) == expectation) assert (ak.to_list(new[[1, 0], ]) == expectation) assert ak.to_list(new[[1, 0]]) == expectation assert ak.to_list(old[1, [1, 0]]) == [2.2, 1.1] assert ak.to_list(new[1, [1, 0]]) == [2.2, 1.1]
def test_UnionArray_NumpyArray(): v1a = ak.layout.UnionArray8_64( ak.layout.Index8(np.array([1, 1, 0, 0, 1, 0, 1], dtype=np.int8)), ak.layout.Index64(np.array([4, 3, 0, 1, 2, 2, 4, 100], dtype=np.int64)), [ ak.layout.NumpyArray(np.array([1, 2, 3], dtype=np.int64)), ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])), ], ) v2a = ak._v2.contents.unionarray.UnionArray( ak._v2.index.Index(np.array([1, 1, 0, 0, 1, 0, 1], dtype=np.int8)), ak._v2.index.Index(np.array([4, 3, 0, 1, 2, 2, 4, 100], dtype=np.int64)), [ ak._v2.contents.numpyarray.NumpyArray(np.array([1, 2, 3], dtype=np.int64)), ak._v2.contents.numpyarray.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5])), ], ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist()
def test_RecordArray(): old = ak.layout.RecordArray( [ ak.layout.NumpyArray( np.array([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], np.int64)), ak.layout.NumpyArray( np.array([[0.0, 1.1, 2.2, 3.3, 4.4, 5.5], [0.0, 1.1, 2.2, 3.3, 4.4, 5.5]])), ], ["x", "y"], ) new = v1_to_v2(old) assert ak.to_list(old[:, 3:]) == [ { "x": [3, 4], "y": [3.3, 4.4, 5.5] }, { "x": [3, 4], "y": [3.3, 4.4, 5.5] }, ] assert ak.to_list(new[:, 3:]) == [ { "x": [3, 4], "y": [3.3, 4.4, 5.5] }, { "x": [3, 4], "y": [3.3, 4.4, 5.5] }, ] with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[1, np.newaxis]) == [{ "x": [0, 1, 2, 3, 4], "y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5] }] assert ak.to_list(old[1, np.newaxis]) == [{ "x": [0, 1, 2, 3, 4], "y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5] }] assert old.minmax_depth == (2, 2) assert new.minmax_depth == (2, 2) assert ak.to_list(old[0, ..., 0]) == {"x": 0, "y": 0.0} assert ak.to_list(new[0, ..., 0]) == {"x": 0, "y": 0.0} expectation = [ { "x": [0, 1, 2, 3, 4], "y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5] }, { "x": [0, 1, 2, 3, 4], "y": [0.0, 1.1, 2.2, 3.3, 4.4, 5.5] }, ] assert (ak.to_list(old[[1, 0], ]) == expectation) assert (ak.to_list(new[[1, 0], ]) == expectation) assert ak.to_list(new[[1, 0]]) == expectation assert ak.to_list(old[1, [1, 0]]) == [{ "x": 1, "y": 1.1 }, { "x": 0, "y": 0.0 }] assert ak.to_list(new[1, [1, 0]]) == [{ "x": 1, "y": 1.1 }, { "x": 0, "y": 0.0 }]
def test_ListArray(): old = ak.layout.ListArray64( ak.layout.Index64(np.array([0, 100, 1], np.int64)), ak.layout.Index64(np.array([3, 100, 3, 200], np.int64)), ak.layout.NumpyArray( np.array([ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ])), ) new = v1_to_v2(old) assert ak.to_list(old[0, :2]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert ak.to_list(new[0, :2]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[0, np.newaxis]) == [[ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ]] assert ak.to_list(old[0, np.newaxis]) == [[ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ]] assert old.minmax_depth == (3, 3) assert new.minmax_depth == (3, 3) assert ak.to_list(old[0, ...]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert ak.to_list(new[0, ...]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] expectation = [ [], [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ], ] 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[0, [1, 0]]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert ak.to_list(new[0, [1, 0]]) == [ [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ]
def test_BitMaskedArray(): old = ak.layout.BitMaskedArray( ak.layout.IndexU8( np.packbits( np.array( [ [ 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, ], [ 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, ], ], np.uint8, ))), ak.layout.NumpyArray( np.array([ [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ], [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ], ])), valid_when=True, length=2, lsb_order=False, ) new = v1_to_v2(old) assert v1v2_equal(old[:, :], new[:, :]) assert ak.to_list(old[:, 5:]) == [ [5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert ak.to_list(new[:, 5:]) == [ [5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] assert v1v2_equal(old[:, 5:], new[:, 5:]) with pytest.raises(IndexError): new[1, "hello"] with pytest.raises(IndexError): new[1, ["hello", "there"]] assert ak.to_list(new[1, np.newaxis]) == [[ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6 ]] assert ak.to_list(old[1, np.newaxis]) == [[ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6 ]] assert old.minmax_depth == (2, 2) assert new.minmax_depth == (2, 2) assert ak.to_list(old[0, ..., 0]) == 0.0 assert ak.to_list(new[0, ..., 0]) == 0.0 expectation = [ [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6], ] 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]]) == [1.0, 0.0] assert ak.to_list(new[1, [1, 0]]) == [1.0, 0.0] assert v1v2_equal(old[1, [1, 0]], new[1, [1, 0]])
def _getitem(self, where): if ak._util.isint(where): return self._getitem_at(where) elif isinstance(where, slice) and where.step is None: return self._getitem_range(where) elif isinstance(where, slice): return self._getitem((where,)) elif ak._util.isstr(where): return self._getitem_field(where) elif where is np.newaxis: return self._getitem((where,)) elif where is Ellipsis: return self._getitem((where,)) elif isinstance(where, tuple): if len(where) == 0: return self items = [ak._v2._slicing.prepare_tuple_item(x) for x in where] nextwhere = ak._v2._slicing.getitem_broadcast(items) next = ak._v2.contents.RegularArray( self, self.length if self._nplike.known_shape else 1, 1, None, None, self._nplike, ) out = next._getitem_next(nextwhere[0], nextwhere[1:], None) if out.length == 0: return out._getitem_nothing() else: return out._getitem_at(0) elif isinstance(where, ak.highlevel.Array): return self._getitem(where.layout) elif isinstance(where, ak.layout.Content): return self._getitem(v1_to_v2(where)) elif isinstance(where, ak._v2.highlevel.Array): return self._getitem(where.layout) elif ( isinstance(where, Content) and where._parameters is not None and (where._parameters.get("__array__") in ("string", "bytestring")) ): return self._getitem_fields(ak._v2.operations.convert.to_list(where)) elif isinstance(where, ak._v2.contents.emptyarray.EmptyArray): return where.toNumpyArray(np.int64) elif isinstance(where, ak._v2.contents.numpyarray.NumpyArray): if issubclass(where.dtype.type, np.int64): carry = ak._v2.index.Index64(where.data.reshape(-1)) allow_lazy = True elif issubclass(where.dtype.type, np.integer): carry = ak._v2.index.Index64(where.data.astype(np.int64).reshape(-1)) allow_lazy = "copied" # True, but also can be modified in-place elif issubclass(where.dtype.type, (np.bool_, bool)): if len(where.data.shape) == 1: where = self._nplike.nonzero(where.data)[0] carry = ak._v2.index.Index64(where) allow_lazy = "copied" # True, but also can be modified in-place else: wheres = self._nplike.nonzero(where.data) return self._getitem(wheres) else: raise ak._v2._util.error( TypeError( "array slice must be an array of integers or booleans, not\n\n {}".format( repr(where.data).replace("\n", "\n ") ) ) ) out = ak._v2._slicing.getitem_next_array_wrap( self._carry(carry, allow_lazy), where.shape ) if out.length == 0: return out._getitem_nothing() else: return out._getitem_at(0) elif isinstance(where, Content): return self._getitem((where,)) elif isinstance(where, Iterable) and all(ak._util.isstr(x) for x in where): return self._getitem_fields(where) elif isinstance(where, Iterable): layout = ak._v2.operations.convert.to_layout(where) as_array = layout.maybe_to_array(layout.nplike) if as_array is None: return self._getitem(layout) else: return self._getitem( ak._v2.contents.NumpyArray(as_array, None, None, layout.nplike) ) else: raise ak._v2._util.error( TypeError( "only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), " "integer/boolean arrays (possibly with variable-length nested " "lists or missing values), field name (str) or names (non-tuple " "iterable of str) are valid indices for slicing, not\n\n " + repr(where).replace("\n", "\n ") ) )
def test_BitMaskedArray_RecordArray_NumpyArray(): v1a = ak.layout.BitMaskedArray( ak.layout.IndexU8( np.packbits( np.array( [ True, True, True, True, False, False, False, False, True, False, True, False, True, ] ) ) ), ak.layout.RecordArray( [ ak.layout.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=True, length=13, lsb_order=False, ) v2a = ak._v2.contents.bitmaskedarray.BitMaskedArray( ak._v2.index.Index( np.packbits( np.array( [ True, True, True, True, False, False, False, False, True, False, True, False, True, ] ) ) ), ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=True, length=13, lsb_order=False, ) assert v1v2_equal(v1a, v2a) assert v1v2_equal(v2_to_v1(v2a), v1_to_v2(v1a)) assert ak.to_list(v1a) == ak.to_list(v2a) assert newform(json.loads(v1a.form.tojson())) == v2a.form.tolist() v1b = ak.layout.BitMaskedArray( ak.layout.IndexU8( np.packbits( np.array( [ 0, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 1, 0, ], dtype=np.uint8, ) ) ), ak.layout.RecordArray( [ ak.layout.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=False, length=13, lsb_order=False, ) v2b = ak._v2.contents.bitmaskedarray.BitMaskedArray( ak._v2.index.Index( np.packbits( np.array( [ 0, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 1, 0, ], dtype=np.uint8, ) ) ), ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=False, length=13, lsb_order=False, ) assert v1v2_equal(v1b, v2b) assert v1v2_equal(v2_to_v1(v2b), v1_to_v2(v1b)) assert ak.to_list(v1b) == ak.to_list(v2b) assert newform(json.loads(v1b.form.tojson())) == v2b.form.tolist() v1c = ak.layout.BitMaskedArray( ak.layout.IndexU8( np.packbits( np.array( [ 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, ], dtype=np.uint8, ) ) ), ak.layout.RecordArray( [ ak.layout.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=True, length=13, lsb_order=True, ) v2c = ak._v2.contents.bitmaskedarray.BitMaskedArray( ak._v2.index.Index( np.packbits( np.array( [ 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, ], dtype=np.uint8, ) ) ), ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=True, length=13, lsb_order=True, ) assert v1v2_equal(v1c, v2c) assert v1v2_equal(v2_to_v1(v2c), v1_to_v2(v1c)) assert ak.to_list(v1c) == ak.to_list(v2c) assert newform(json.loads(v1c.form.tojson())) == v2c.form.tolist() v1d = ak.layout.BitMaskedArray( ak.layout.IndexU8( np.packbits( np.array( [ 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, ], dtype=np.uint8, ) ) ), ak.layout.RecordArray( [ ak.layout.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=False, length=13, lsb_order=True, ) v2d = ak._v2.contents.bitmaskedarray.BitMaskedArray( ak._v2.index.Index( np.packbits( np.array( [ 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, ], dtype=np.uint8, ) ) ), ak._v2.contents.recordarray.RecordArray( [ ak._v2.contents.numpyarray.NumpyArray( np.array( [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, ] ) ) ], ["nest"], ), valid_when=False, length=13, lsb_order=True, ) assert v1v2_equal(v1d, v2d) assert v1v2_equal(v2_to_v1(v2d), v1_to_v2(v1d)) assert ak.to_list(v1d) == ak.to_list(v2d) assert newform(json.loads(v1d.form.tojson())) == v2d.form.tolist()
def __getitem__(self, where): try: if ak._util.isint(where): return self._getitem_at(where) elif isinstance(where, slice) and where.step is None: return self._getitem_range(where) elif isinstance(where, slice): return self.__getitem__((where,)) elif ak._util.isstr(where): return self._getitem_field(where) elif where is np.newaxis: return self.__getitem__((where,)) elif where is Ellipsis: return self.__getitem__((where,)) elif isinstance(where, tuple): if len(where) == 0: return self nextwhere = self._getitem_broadcast( [self._prepare_tuple_item(x) for x in where], self.nplike, ) next = ak._v2.contents.RegularArray(self, len(self), 1, None, None) out = next._getitem_next(nextwhere[0], nextwhere[1:], None) if len(out) == 0: raise out._getitem_nothing() else: return out._getitem_at(0) elif isinstance(where, ak.highlevel.Array): return self.__getitem__(where.layout) elif isinstance(where, ak.layout.Content): return self.__getitem__(v1_to_v2(where)) elif isinstance(where, ak._v2.contents.numpyarray.NumpyArray): if issubclass(where.dtype.type, np.int64): carry = ak._v2.index.Index64(where.data.reshape(-1)) allow_lazy = True elif issubclass(where.dtype.type, np.integer): carry = ak._v2.index.Index64( where.data.astype(np.int64).reshape(-1) ) allow_lazy = "copied" # True, but also can be modified in-place elif issubclass(where.dtype.type, (np.bool_, bool)): carry = ak._v2.index.Index64(np.nonzero(where.data.reshape(-1))[0]) allow_lazy = "copied" # True, but also can be modified in-place else: raise TypeError( "array slice must be an array of integers or booleans, not\n\n {0}".format( repr(where.data).replace("\n", "\n ") ) ) out = self._getitem_next_array_wrap( self._carry(carry, allow_lazy, NestedIndexError), where.shape ) if len(out) == 0: return out._getitem_nothing() else: return out._getitem_at(0) elif isinstance(where, Content): return self.__getitem__((where,)) elif isinstance(where, Iterable) and all(ak._util.isstr(x) for x in where): return self._getitem_fields(where) elif isinstance(where, Iterable): return self.__getitem__( v1_to_v2(ak.operations.convert.to_layout(where)) ) else: raise TypeError( "only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), " "integer/boolean arrays (possibly with variable-length nested " "lists or missing values), field name (str) or names (non-tuple " "iterable of str) are valid indices for slicing, not\n\n " + repr(where).replace("\n", "\n ") ) except NestedIndexError as err: def format_slice(x): if isinstance(x, slice): if x.step is None: return "{0}:{1}".format( "" if x.start is None else x.start, "" if x.stop is None else x.stop, ) else: return "{0}:{1}:{2}".format( "" if x.start is None else x.start, "" if x.stop is None else x.stop, x.step, ) elif isinstance(x, tuple): return "(" + ", ".join(format_slice(y) for y in x) + ")" elif isinstance(x, ak._v2.index.Index64): return str(x.data) elif isinstance(x, Content): return str(ak.Array(v2_to_v1(x))) else: return repr(x) raise IndexError( """cannot slice {0} with {1} at inner {2} of length {3}, using sub-slice {4}.{5}""".format( repr(ak.Array(v2_to_v1(self))), format_slice(where), type(err.array).__name__, len(err.array), format_slice(err.slicer), "" if err.details is None else "\n\n{0} error: {1}.".format( type(err.array).__name__, err.details ), ) )
def _prepare_tuple_item(self, item): if ak._util.isint(item): return int(item) elif isinstance(item, slice): return item elif ak._util.isstr(item): return item elif item is np.newaxis: return item elif item is Ellipsis: return item elif isinstance(item, ak.highlevel.Array): return self._prepare_tuple_item(item.layout) elif isinstance(item, ak.layout.Content): return self._prepare_tuple_item(v1_to_v2(item)) elif isinstance(item, ak._v2.contents.NumpyArray): return item.data elif isinstance( item, ( ak._v2.contents.ListOffsetArray, ak._v2.contents.ListArray, ak._v2.contents.RegularArray, ), ): return item.toListOffsetArray64(False) elif isinstance( item, ( ak._v2.contents.IndexedOptionArray, ak._v2.contents.ByteMaskedArray, ak._v2.contents.BitMaskedArray, ak._v2.contents.UnmaskedArray, ), ): return item.toIndexedOptionArray64() elif isinstance(item, Iterable) and all(ak._util.isstr(x) for x in item): return list(item) elif isinstance(item, Iterable): return self._prepare_tuple_item( v1_to_v2(ak.operations.convert.to_layout(item)) ) else: raise TypeError( "only integers, slices (`:`), ellipsis (`...`), np.newaxis (`None`), " "integer/boolean arrays (possibly with variable-length nested " "lists or missing values), field name (str) or names (non-tuple " "iterable of str) are valid indices for slicing, not\n\n " + repr(item).replace("\n", "\n ") )