Writing leafs with fixed number of objects

[kjmeagher]

I am trying to use uproot to replicate the structure of data that is currently being written in root. Each leaf has 3 integers.
The when opened in uproot the branch has a typename uint16_t[][3] and interpretation AsJagged(AsDtype("('>u2', (3,))")) and I can read the data as expected. However, I cannot figure how to write such a tree with uproot. I couldn't find anything in the uproot documentation that addresses this. This is what I tried so far.

import uproot
import awkward as ak
outfile = uproot.recreate('outfile.root')

desc = {'vector' : 'var * 3 * uint16'}
tree = outfile.mktree("Tree", desc, title="Title")

d = {
     'vector': ak.Array([ 
                    [ [1,1,1], [2,2,2] ,[3,3,3] ],
                    [ [4,4,4], [5,5,5] ] 
                    ])
     }
print(d)
tree.extend(d)

[jpivarski]

You need to make the ak.Array regular (or just use a NumPy array, which are all regular). The mktree descriptor is good, but data passed to extend is not cast as the descriptor (unless it's pure Python lists of int and float); it's checked against the descriptor and will complain if there's a mismatch.

Fixed-length lists are a different type from variable-length lists. In your descriptor, the "var" refers to a variable-length dimension and the "3" refers to a fixed-length dimension.

When you construct an ak.Array from Python data, it is all variable-length because Python lists don't have fixed lengths.

>>> varlen = ak.Array([[[1, 1, 1], [2, 2, 2]], [[4, 4, 4], [5, 5, 5]]])
>>> varlen
<Array [[[1, 1, 1], [2, 2, ... 4], [5, 5, 5]]] type='2 * var * var * int64'>

Note the "var * var" in the type.

When you construct an ak.Array from a NumPy array, it's all regular-length because NumPy dimensions have fixed lengths.

>>> reglen = ak.Array(np.array([[[1, 1, 1], [2, 2, 2]], [[4, 4, 4], [5, 5, 5]]]))
>>> reglen
<Array [[[1, 1, 1], [2, 2, ... 4], [5, 5, 5]]] type='2 * 2 * 3 * int64'>

Note the "2 * 3" in the type.

You can also turn variable-length dimensions into fixed-length dimensions with ak.to_regular (there's also an ak.from_regular).

>>> ak.to_regular(varlen, axis=1)
<Array [[[1, 1, 1], [2, 2, ... 4], [5, 5, 5]]] type='2 * 2 * var * int64'>
>>> ak.to_regular(ak.to_regular(varlen, axis=1), axis=2)
<Array [[[1, 1, 1], [2, 2, ... 4], [5, 5, 5]]] type='2 * 2 * 3 * int64'>

Another mismatch is the numeric type: "int64" versus "uint16". If you construct the ak.Array from a NumPy array with the desired dtype, the ak.Array will have that dtype, too. If you have an ak.Array made from Python data, it will default to 64-bit numbers. You can change that with ak.values_astype.

It looks like you want mixed variable-length and regular-length: the only way to do that is with an Awkward Array. (Python lists will be interpreted as all-variable and NumPy arrays are all-regular.)

[kjmeagher]

Thanks for the quick reply. This makes more sense now, but I still seem to have problems, I am not sure if I am still doing anything wrong or if I encountered a bug.

If I run this:

import uproot
import awkward as ak
outfile = uproot.recreate('outfile.root')

desc = {'vector' : 'var * 3 * uint16'}
tree = outfile.mktree("Tree", desc, title="Title")

x = ak.Array([ [ [1,1,1], [2,2,2] ,[3,3,3] ],
               [ [4,4,4], [5,5,5]],
               [ [6,6,6]],
               [ [7,7,7], [8,8,8]]
               ])

x = ak.to_regular(x, axis=2)
x = ak.values_astype(x,'uint16')
tree.extend({'vector':x})
outfile.close()

It runs without error but I can't read the output.

with uproot.open('outfile.root') as f:
    f["Tree"].show()
    f["Tree"]['vector'].array()

gives this output

ame                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
nvector              | int32_t                  | AsDtype('>i4')
vector               | uint16_t[][3]            | AsJagged(AsDtype("('>u2', (...
Traceback (most recent call last):
  File "/usr/lib/python3.10/site-packages/uproot/interpretation/numerical.py", line 340, in basket_array
    output = data.view(dtype).reshape((-1,) + shape)
ValueError: cannot reshape array of size 8 into shape (3)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/kjm/icecube/icetray/uprootwriter/printfile.py", line 5, in <module>
    f["Tree"]['vector'].array()
  File "/usr/lib/python3.10/site-packages/uproot/behaviors/TBranch.py", line 2078, in array
    _ranges_or_baskets_to_arrays(
  File "/usr/lib/python3.10/site-packages/uproot/behaviors/TBranch.py", line 3502, in _ranges_or_baskets_to_arrays
    uproot.source.futures.delayed_raise(*obj)
  File "/usr/lib/python3.10/site-packages/uproot/source/futures.py", line 36, in delayed_raise
    raise exception_value.with_traceback(traceback)
  File "/usr/lib/python3.10/site-packages/uproot/behaviors/TBranch.py", line 3446, in basket_to_array
    basket_arrays[basket.basket_num] = interpretation.basket_array(
  File "/usr/lib/python3.10/site-packages/uproot/interpretation/jagged.py", line 174, in basket_array
    content = self._content.basket_array(
  File "/usr/lib/python3.10/site-packages/uproot/interpretation/numerical.py", line 342, in basket_array
    raise ValueError(
ValueError: basket 0 in tree/branch /Tree;1:vector has the wrong number of bytes (16) for interpretation AsDtype("('>u2', (3,))")
in file outfile.root

If i try to read this file in ROOT i get the following output.

root [1] Tree->Scan()
***********************************************
*    Row   * Instance * nvector.n * vector.ve *
***********************************************
Warning in <TBasket::ReadBasketBuffers>: basket:vector has fNevBuf=4 but fEntryOffset=0, pos=310, len=145, fNbytes=130, fObjlen=72, trying to repair
*        0 *        0 *         3 *         1 *
*        0 *        1 *         3 *         1 *
*        0 *        2 *         3 *         1 *
*        0 *        3 *         3 *         2 *
*        0 *        4 *         3 *         2 *
*        0 *        5 *         3 *         2 *
*        0 *        6 *         3 *         3 *
*        0 *        7 *         3 *         3 *
*        0 *        8 *         3 *         3 *
*        1 *        0 *         2 *     60416 *
*        1 *        1 *         2 *         0 *
*        1 *        2 *         2 *     18541 *
*        1 *        3 *         2 *     19645 *
*        1 *        4 *         2 *      9472 *
*        1 *        5 *         2 *     18688 *
*        2 *        0 *         1 *         0 *
*        2 *        1 *         1 *       130 *
*        2 *        2 *         1 *      1004 *
*        3 *        0 *         2 *         0 *
*        3 *        1 *         2 *       130 *
*        3 *        2 *         2 *      1004 *
*        3 *        3 *         2 *         0 *
*        3 *        4 *         2 *        72 *
*        3 *        5 *         2 *     27980 *
***********************************************

It looks like after the first row the contents is random memory.