Ejemplo n.º 1
0
def test_ArrayBuilder_append_2():
    # issue #415
    A = awkward1.from_numpy(numpy.array([0, 1, 2], dtype=numpy.float32))
    B = awkward1.from_numpy(numpy.array([0, 1], dtype=numpy.float32))

    builder = awkward1.ArrayBuilder()
    with builder.list():
        builder.append(A.tolist())
    with builder.list():
        builder.append(A.tolist())
    with builder.list():
        pass
    with builder.list():
        builder.append(B.tolist())

    assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [],
                                           [[0, 1]]]
    assert str(awkward1.type(builder.snapshot())) == "4 * var * var * float64"

    builder = awkward1.ArrayBuilder()
    with builder.list():
        builder.append(A)
    with builder.list():
        builder.append(A)
    with builder.list():
        pass
    with builder.list():
        builder.append(B)

    assert builder.snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [],
                                           [[0, 1]]]
    assert str(awkward1.type(builder.snapshot())) == "4 * var * var * float32"

    @numba.njit
    def f1(builder, A, B):
        builder.begin_list()
        builder.append(A)
        builder.end_list()

        builder.begin_list()
        builder.append(A)
        builder.end_list()

        builder.begin_list()
        builder.end_list()

        builder.begin_list()
        builder.append(B)
        builder.end_list()

        return builder

    assert f1(awkward1.ArrayBuilder(),
              A, B).snapshot().tolist() == [[[0, 1, 2]], [[0, 1, 2]], [],
                                            [[0, 1]]]
def test_ArrayBuilder_record():
    @numba.njit
    def f1(x):
        x.begin_record()
        x.field("x").append(1)
        x.field("y").append(1.1)
        x.end_record()
        x.begin_record()
        x.field("x").append(2)
        x.field("y").append(2.2)
        x.end_record()
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a)
    assert awkward1.to_list(a.snapshot()) == [{"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2}]
    assert awkward1.to_list(b.snapshot()) == [{"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2}]
    c = f1.py_func(a)
    assert awkward1.to_list(a.snapshot()) == [{"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2},
                                             {"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2}]
    assert awkward1.to_list(c.snapshot()) == [{"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2},
                                             {"x": 1, "y": 1.1},
                                             {"x": 2, "y": 2.2}]
def test_ArrayBuilder_refcount():
    builder = awkward1.ArrayBuilder()
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 2)

    @numba.njit
    def f1(x):
        return 3.14

    y = f1(builder)
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 2)

    @numba.njit
    def f2(x):
        return x

    y = f2(builder)
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 3)
    del y
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 2)

    @numba.njit
    def f3(x):
        return x, x

    y = f3(builder)
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 4)
    del y
    assert (sys.getrefcount(builder), sys.getrefcount(builder._layout)) == (2, 2)
def test_ArrayBuilder_append_numba5():
    @numba.njit
    def f1(builder, x):
        builder.append(x)

    @numba.njit
    def f2(builder, i):
        if i % 2 == 0:
            return 3
        else:
            return None

    @numba.njit
    def f3(builder, i):
        builder.append(f2(builder, i))

    builder = awkward1.ArrayBuilder()

    f1(builder, True)
    f1(builder, 1)
    f1(builder, 2.2)
    f3(builder, 0)
    f3(builder, 1)
    f1(builder, None)

    assert awkward1.to_list(builder.snapshot()) == [True, 1, 2.2, 3, None, None]
def test_ArrayBuilder_simple():
    @numba.njit
    def f1(x):
        x.clear()
        return 3.14

    a = awkward1.ArrayBuilder()
    f1(a)
def test_ArrayBuilder_real():
    @numba.njit
    def f1(x, z):
        x.real(1)
        x.real(2.2)
        x.real(z)
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a, numpy.array([3.5], dtype=numpy.float32)[0])
    assert awkward1.to_list(a.snapshot()) == [1, 2.2, 3.5]
    assert awkward1.to_list(b.snapshot()) == [1, 2.2, 3.5]
def test_ArrayBuilder_integer():
    @numba.njit
    def f1(x):
        x.integer(1)
        x.integer(2)
        x.integer(3)
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a)
    assert awkward1.to_list(a.snapshot()) == [1, 2, 3]
    assert awkward1.to_list(b.snapshot()) == [1, 2, 3]
def test_ArrayBuilder_boolean():
    @numba.njit
    def f1(x):
        x.boolean(True)
        x.boolean(False)
        x.boolean(False)
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a)
    assert awkward1.to_list(a.snapshot()) == [True, False, False]
    assert awkward1.to_list(b.snapshot()) == [True, False, False]
def test_ArrayBuilder_len():
    builder = awkward1.ArrayBuilder()
    builder.real(1.1)
    builder.real(2.2)
    builder.real(3.3)
    builder.real(4.4)
    builder.real(5.5)

    @numba.njit
    def f1(x):
        return len(x)

    assert f1(builder) == 5
def test_ArrayBuilder_append_numba4():
    @numba.njit
    def f1(array, builder):
        builder.append(array[3])
        builder.append(array[2])
        builder.append(array[2])
        builder.append(array[0])
        builder.append(array[1])
        builder.append(array[-1])

    array = 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)
    builder = awkward1.ArrayBuilder()

    f1(array, builder)

    assert awkward1.to_list(builder.snapshot()) == [[5.5], [3.3, 4.4], [3.3, 4.4], [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]]
def test_ArrayBuilder_append_numba2():
    @numba.njit
    def f1(array, builder):
        builder.append(array[3])
        builder.append(array[2])
        builder.append(array[2])
        builder.append(array[0])
        builder.append(array[1])
        builder.append(array[-1])

    array = awkward1.Array([{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}], check_valid=True)
    builder = awkward1.ArrayBuilder()

    f1(array, builder)

    assert awkward1.to_list(builder.snapshot()) == [{"x": 3.3, "y": [3, 3, 3]}, {"x": 2.2, "y": [2, 2]}, {"x": 2.2, "y": [2, 2]}, {"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 3.3, "y": [3, 3, 3]}]
def test_ArrayBuilder_tuple():
    @numba.njit
    def f1(x):
        x.begin_tuple(2)
        x.index(0).append(1)
        x.index(1).append(1.1)
        x.end_tuple()
        x.begin_tuple(2)
        x.index(0).append(2)
        x.index(1).append(2.2)
        x.end_tuple()
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a)
    assert awkward1.to_list(a.snapshot()) == [(1, 1.1), (2, 2.2)]
    assert awkward1.to_list(b.snapshot()) == [(1, 1.1), (2, 2.2)]
    c = f1.py_func(a)
    assert awkward1.to_list(a.snapshot()) == [(1, 1.1), (2, 2.2),
                                             (1, 1.1), (2, 2.2)]
    assert awkward1.to_list(c.snapshot()) == [(1, 1.1), (2, 2.2),
                                             (1, 1.1), (2, 2.2)]
def test_ArrayBuilder_list():
    @numba.njit
    def f1(x):
        x.begin_list()
        x.real(1.1)
        x.real(2.2)
        x.real(3.3)
        x.end_list()
        x.begin_list()
        x.end_list()
        x.begin_list()
        x.real(4.4)
        x.real(5.5)
        x.end_list()
        return x

    a = awkward1.ArrayBuilder()
    b = f1(a)
    assert awkward1.to_list(a.snapshot()) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]
    assert awkward1.to_list(b.snapshot()) == [[1.1, 2.2, 3.3], [], [4.4, 5.5]]

    @numba.njit
    def f2(x):
        return len(x)

    assert f2(a) == 3
    assert f2(b) == 3

    @numba.njit
    def f3(x):
        x.clear()
        return x

    c = f3(b)
    assert awkward1.to_list(a.snapshot()) == []
    assert awkward1.to_list(b.snapshot()) == []
    assert awkward1.to_list(c.snapshot()) == []
def test_ArrayBuilder():
    builder = awkward1.ArrayBuilder()
    assert sys.getrefcount(builder._layout) == 3

    @numba.njit
    def f():
        builder.append(1)
        builder.append(2)
        builder.append(3)
        return builder, builder

    @numba.njit
    def g():
        builder.append(1)
        builder.append(2)
        builder.append(3)

    b, c = f()
    assert b.snapshot().tolist() == [1, 2, 3]
    assert c.snapshot().tolist() == [1, 2, 3]
    assert builder.snapshot().tolist() == [1, 2, 3]

    assert sys.getrefcount(builder._layout) == 5

    g()
    assert b.snapshot().tolist() == [1, 2, 3, 1, 2, 3]
    assert c.snapshot().tolist() == [1, 2, 3, 1, 2, 3]
    assert builder.snapshot().tolist() == [1, 2, 3, 1, 2, 3]

    assert sys.getrefcount(builder._layout) == 5

    del b._layout
    assert sys.getrefcount(builder._layout) == 4

    del c._layout
    assert sys.getrefcount(builder._layout) == 3
def test_ArrayBuilder_append():
    array = 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)

    builder = awkward1.ArrayBuilder()
    builder.append(array, 3)
    builder.append(array, 2)
    builder.append(array, 2)
    builder.append(array, 0)
    builder.append(array, 1)
    builder.append(array, -1)
    assert awkward1.to_list(builder.snapshot()) == [[5.5], [3.3, 4.4], [3.3, 4.4], [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9]]

    builder.extend(array)
    assert awkward1.to_list(builder.snapshot()) == [[5.5], [3.3, 4.4], [3.3, 4.4], [0.0, 1.1, 2.2], [], [6.6, 7.7, 8.8, 9.9], [0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]]

    builder = awkward1.ArrayBuilder()
    builder.null()
    builder.null()
    builder.null()
    builder.append(array, 3)
    builder.append(array, 2)
    builder.append(array, 2)
    builder.append(array, -1)

    assert awkward1.to_list(builder.snapshot()) == [None, None, None, [5.5], [3.3, 4.4], [3.3, 4.4], [6.6, 7.7, 8.8, 9.9]]

    builder.null()
    assert awkward1.to_list(builder.snapshot()) == [None, None, None, [5.5], [3.3, 4.4], [3.3, 4.4], [6.6, 7.7, 8.8, 9.9], None]

    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([[3.3, 2.2, 1.1, 0.0], [5.5, 4.4], [], [6.6]], check_valid=True)

    builder = awkward1.ArrayBuilder()
    builder.append(one, 2)
    builder.append(two, 1)
    builder.append(one, 0)
    builder.append(two, -1)
    builder.append(one, -1)

    assert awkward1.to_list(builder.snapshot()) == [[3.3, 4.4], [5.5, 4.4], [0.0, 1.1, 2.2], [6.6], [6.6, 7.7, 8.8, 9.9]]

    builder = awkward1.ArrayBuilder()
    builder.null()
    builder.append(one, 2)
    builder.null()
    builder.append(two, 1)
    builder.null()
    assert awkward1.to_list(builder.snapshot()) == [None, [3.3, 4.4], None, [5.5, 4.4], None]

    builder = awkward1.ArrayBuilder()
    builder.string("hello")
    builder.append(one, 2)
    builder.string("there")
    builder.append(one, 0)
    assert awkward1.to_list(builder.snapshot()) == ["hello", [3.3, 4.4], "there", [0.0, 1.1, 2.2]]

    builder = awkward1.ArrayBuilder()
    builder.null()
    builder.string("hello")
    builder.null()
    builder.append(one, 2)
    builder.null()
    builder.string("there")
    builder.append(one, 0)
    assert awkward1.to_list(builder.snapshot()) == [None, "hello", None, [3.3, 4.4], None, "there", [0.0, 1.1, 2.2]]

    builder = awkward1.ArrayBuilder()
    builder.append(one, 2)
    builder.string("there")
    builder.append(one, 0)
    assert awkward1.to_list(builder.snapshot()) == [[3.3, 4.4], "there", [0.0, 1.1, 2.2]]

    builder = awkward1.ArrayBuilder()
    builder.null()
    builder.append(one, 2)
    builder.null()
    builder.string("there")
    builder.null()
    builder.append(one, 0)
    assert awkward1.to_list(builder.snapshot()) == [None, [3.3, 4.4], None, "there", None, [0.0, 1.1, 2.2]]

    array = awkward1.Array(["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"], check_valid=True)
    builder = awkward1.ArrayBuilder()
    builder.begin_list()
    builder.append(array, 1)
    builder.append(array, 2)
    builder.append(array, 3)
    builder.end_list()
    builder.begin_list()
    builder.end_list()
    builder.begin_list()
    builder.append(array, 4)
    builder.append(array, 5)
    builder.end_list()
    assert awkward1.to_list(builder.snapshot()) == [["one", "two", "three"], [], ["four", "five"]]

    builder.append(array, -1)
    assert awkward1.to_list(builder.snapshot()) == [["one", "two", "three"], [], ["four", "five"], "nine"]

    array = awkward1.Array([{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}], check_valid=True)
    builder = awkward1.ArrayBuilder()
    builder.append(array[2])
    builder.append(array[2])
    builder.append(array[1])
    builder.append(array[-1])
    tmp = builder.snapshot()
    assert awkward1.to_list(tmp) == [{"x": 2.2, "y": [2, 2]}, {"x": 2.2, "y": [2, 2]}, {"x": 1.1, "y": [1]}, {"x": 3.3, "y": [3, 3, 3]}]
    assert isinstance(tmp.layout, awkward1.layout.IndexedArray64)
    assert isinstance(tmp.layout.content, awkward1.layout.RecordArray)

    builder.extend(array)
    tmp = builder.snapshot()
    assert awkward1.to_list(tmp) == [{"x": 2.2, "y": [2, 2]}, {"x": 2.2, "y": [2, 2]}, {"x": 1.1, "y": [1]}, {"x": 3.3, "y": [3, 3, 3]}, {"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}]
    assert isinstance(tmp.layout, awkward1.layout.IndexedArray64)
    assert isinstance(tmp.layout.content, awkward1.layout.RecordArray)

    builder.append(999)
    tmp = builder.snapshot()
    assert awkward1.to_list(tmp) == [{"x": 2.2, "y": [2, 2]}, {"x": 2.2, "y": [2, 2]}, {"x": 1.1, "y": [1]}, {"x": 3.3, "y": [3, 3, 3]}, {"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}, 999]
    assert isinstance(tmp.layout, awkward1.layout.UnionArray8_64)

    builder.append([1, 2, 3, 4, 5])
    tmp = builder.snapshot()
    assert awkward1.to_list(tmp) == [{"x": 2.2, "y": [2, 2]}, {"x": 2.2, "y": [2, 2]}, {"x": 1.1, "y": [1]}, {"x": 3.3, "y": [3, 3, 3]}, {"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [2, 2]}, {"x": 3.3, "y": [3, 3, 3]}, 999, [1, 2, 3, 4, 5]]
    assert isinstance(tmp.layout, awkward1.layout.UnionArray8_64)

    array1 = awkward1.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]], check_valid=True)

    builder = awkward1.ArrayBuilder()
    builder.append(array1, 2)
    builder.append(array1, 1)
    builder.append(array1, 0)
    array2 = builder.snapshot()
    assert isinstance(array2.layout.content, awkward1.layout.ListOffsetArray64)

    builder = awkward1.ArrayBuilder()
    builder.append(array2, 2)
    builder.append(array2, 1)
    builder.append(array2, 0)
    array3 = builder.snapshot()
    assert isinstance(array3.layout.content, awkward1.layout.ListOffsetArray64)

    builder = awkward1.ArrayBuilder()
    builder.append(array3, 2)
    builder.append(array3, 1)
    builder.append(array3, 0)
    array4 = builder.snapshot()
    assert isinstance(array4.layout.content, awkward1.layout.ListOffsetArray64)
Ejemplo n.º 16
0
def fill_record(inpLCIO, collection, nevents):
    """this function reads all events in LCIO file and put them into awkward array"""

    ## open LCIO file
    reader = IOIMPL.LCFactory.getInstance().createLCReader()
    reader.open(inpLCIO)

    ## create awkward array
    b = ak.ArrayBuilder()

    ## start looping
    nEvt = 0
    for evt in reader:
        nEvt += 1
        if nEvt > nevents:
            break

        ## First thing: MC particle collection
        b.begin_list()
        mcparticle = evt.getCollection("MCParticle")

        ## fill energy for each MCParticle (in this case just an incoming photon)
        for enr in mcparticle:
            b.begin_record()
            b.field("E")
            b.real(enr.getEnergy())
            ## calculate polar angle theta and fill
            b.field("theta")
            pVec = enr.getMomentum()
            theta = math.pi / 2.00 - math.atan(pVec[2] / pVec[1])
            b.real(theta)
            b.end_record()

        ## ECAL barrel collection
        ecalBarrel = evt.getCollection(collection)
        cellIDString = ecalBarrel.getParameters().getStringVal(
            "CellIDEncoding")
        decoder = CellIDDecoder(cellIDString)
        ##

        for hit in ecalBarrel:

            l = decoder.layer(
                hit.getCellID0())  ## get the layer information from CellID0
            e = hit.getEnergy()
            pos = hit.getPosition()

            ## start filling a record with all relevant information
            b.begin_record()
            b.field("x")
            b.real(pos[0])
            b.field("y")
            b.real(pos[1])
            b.field("z")
            b.real(pos[2])
            b.field("e")
            b.real(e * 1000)
            b.field("layer")
            b.integer(l)
            b.end_record()

        b.end_list()

    ### Example:
    # Get the incident energy of the first event --> b[0].E
    # Get the x positions of the the first event --> b[0].x

    return b
def test_builder_string():
    builder = awkward1.ArrayBuilder()

    builder.bytestring(b"one")
    builder.bytestring(b"two")
    builder.bytestring(b"three")

    a = builder.snapshot()
    if py27:
        assert str(a) == "['one', 'two', 'three']"
    else:
        assert str(a) == "[b'one', b'two', b'three']"
    assert awkward1.to_list(a) == [b'one', b'two', b'three']
    assert awkward1.to_json(a) == '["one","two","three"]'
    if py27:
        assert repr(a) == "<Array ['one', 'two', 'three'] type='3 * bytes'>"
    else:
        assert repr(a) == "<Array [b'one', b'two', b'three'] type='3 * bytes'>"
    assert repr(a.type) == "3 * bytes"

    builder = awkward1.ArrayBuilder()

    builder.string("one")
    builder.string("two")
    builder.string("three")

    a = builder.snapshot()
    if py27:
        assert str(a) == "[u'one', u'two', u'three']"
    else:
        assert str(a) == "['one', 'two', 'three']"
    assert awkward1.to_list(a) == ['one', 'two', 'three']
    assert awkward1.to_json(a) == '["one","two","three"]'
    if py27:
        assert repr(
            a) == "<Array [u'one', u'two', u'three'] type='3 * string'>"
    else:
        assert repr(a) == "<Array ['one', 'two', 'three'] type='3 * string'>"
    assert repr(a.type) == "3 * string"

    builder = awkward1.ArrayBuilder()

    builder.begin_list()
    builder.string("one")
    builder.string("two")
    builder.string("three")
    builder.end_list()

    builder.begin_list()
    builder.end_list()

    builder.begin_list()
    builder.string("four")
    builder.string("five")
    builder.end_list()

    a = builder.snapshot()
    if py27:
        assert str(a) == "[[u'one', u'two', u'three'], [], [u'four', u'five']]"
    else:
        assert str(a) == "[['one', 'two', 'three'], [], ['four', 'five']]"
    assert awkward1.to_list(a) == [['one', 'two', 'three'], [],
                                   ['four', 'five']]
    assert awkward1.to_json(a) == '[["one","two","three"],[],["four","five"]]'
    assert repr(a.type) == "3 * var * string"
def test():
    out = awkward1.ArrayBuilder()

    out.begin_record()
    if True:
        out.field("x")
        out.integer(3)
        out.field("extreme")
        out.begin_record()
        if True:
            out.field("pt")
            out.real(3.3)
            out.field("charge")
            out.integer(-1)
            out.field("iso")
            out.integer(100)
        out.end_record()
    out.end_record()

    out.begin_record()
    if True:
        out.field("x")
        out.integer(3)
    out.end_record()

    ss = out.snapshot()
    assert awkward1.to_list(ss) == [{
        "x": 3,
        "extreme": {
            "pt": 3.3,
            "charge": -1,
            "iso": 100
        }
    }, {
        "x": 3,
        "extreme": None
    }]

    assert awkward1.to_list(
        awkward1.Array([{
            "x": 3,
            "extreme": {
                "pt": 3.3,
                "charge": -1,
                "iso": 100
            }
        }, {
            "x": 3
        }],
                       check_valid=True)) == [{
                           "x": 3,
                           "extreme": {
                               "pt": 3.3,
                               "charge": -1,
                               "iso": 100
                           }
                       }, {
                           "x": 3,
                           "extreme": None
                       }]

    assert awkward1.to_list(
        awkward1.Array([{
            "x": 3,
            "extreme": {
                "pt": 3.3,
                "charge": -1,
                "iso": 100
            }
        }, {
            "x": 3,
            "what": 3
        }],
                       check_valid=True)) == [{
                           "x": 3,
                           "extreme": {
                               "pt": 3.3,
                               "charge": -1,
                               "iso": 100
                           },
                           "what": None
                       }, {
                           "x": 3,
                           "extreme": None,
                           "what": 3
                       }]