def test_where(): condition = awkward1.Array([True, False, True, False, True], check_valid=True) one = awkward1.Array([1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True) two = awkward1.Array([False, False, False, True, True], check_valid=True) three = awkward1.Array([[], [1], [2, 2], [3, 3, 3], [4, 4, 4, 4]], check_valid=True) assert awkward1.to_list(awkward1.where(condition, one, two)) == [1.1, 0.0, 3.3, 1.0, 5.5] assert awkward1.to_list(awkward1.where(condition, one, three)) == [1.1, [1], 3.3, [3, 3, 3], 5.5]
def test_issue_334(): a = awkward1.Array([1, 2, 3, 4]) b = awkward1.Array([-1]) c = awkward1.Array([True, False, True, True]) assert awkward1.where(c, a, b).tolist() == [1, -1, 3, 4] assert awkward1.where(*awkward1.broadcast_arrays(c, a, b)).tolist() == [1, -1, 3, 4] assert awkward1.where(c, a, -1).tolist() == [1, -1, 3, 4] assert awkward1.where(*awkward1.broadcast_arrays(c, a, -1)).tolist() == [1, -1, 3, 4]
def test_where(): one = awkward1.Array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True) two = awkward1.Array([ 0, 100, 200, 300, 400, 500, 600, 700, 800, 900], check_valid=True) condition = awkward1.Array([False, False, False, False, False, True, False, True, False, True], check_valid=True) assert isinstance(awkward1.where(condition)[0], awkward1.Array) assert awkward1.to_list(awkward1.where(condition)[0]) == [5, 7, 9] assert awkward1.to_list(awkward1.where(condition, one, two)) == awkward1.to_list(numpy.where(numpy.asarray(condition), numpy.asarray(one), numpy.asarray(two)))
def test(): one = awkward1.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]) two = awkward1.Array([[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]]) condition = awkward1.Array([[False, True, False], [], [True, False], [True], [False, False, True, True]]) assert awkward1.where(condition, one, two).tolist() == [[0, 1, 2.2], [], [3, 4.4], [5], [6.6, 7.7, 8, 9]]