Transpose a part of an Array

[alpetukhov]

Hello everyone,

I'm trying to get the particle track data from the .root files with the uproot, but I guess the question would be more appropriate here.

The Array I'm getting from the .root file with the iterate() function is structured like:

[Event [ Hit variable [ Hit number ] ] ]

and I want to get a structure that is

[ Event [ Hit number [ hit variable ] ] ]

I suppose it would be more clear with an example. I've got

[
  [ 'hit_x' : [ 0, 1, 0],
    'hit_y' : [12, 0.3, 7],
    'hit_z': [3, 8, 6]],
  
  ['hit_x' : [ 1, 1],
    'hit_y' : [3, 7],
    'hit_z': [0,4]]
]

and I want to get

[
  [[0, 12, 3], [1, 0.3, 8], [0, 7, 6]],
  [[1, 3, 0], [1, 7, 4]]
]

So not only I want to transpose, I want to

1. Drop the names of the hit variables in the Array
2. Deal with the fact that number of track is not fixed so the innermost dimension in the initial Array is changing

---

I've tried the most obvious thing - the numpy transpose function:

np.transpose(hit_array)

But got the following error:

TypeError: no implementation found for 'numpy.transpose' on types that implement __array_function__: [<class 'awkward.highlevel.Array'>]

So what would be the most appropriate way to do this transformation? I'm doing it so I can write the resulting array as a 'hit_variables' column of the pandas dataset to do some more transformations and write it to the .hdf file. This is only the part of the whole transformation process which is quite memory consuming and I'm to to this python way to handle big datases so I want to make sure I would create as little intermediate data as possible.

Thanks in advance,
Aleksandr

[agoose77]

Given that you're using uproot, I believe you're obtaining arrays of the form

>>> arr = ak.Array([
...   { 'hit_x' : [ 0, 1, 0],
...     'hit_y' : [12, 0.3, 7],
...     'hit_z': [3, 8, 6]},
...
...   {'hit_x' : [ 1, 1],
...     'hit_y' : [3, 7],
...     'hit_z': [0,4]}
... ])

If you just wanted the records to have scalar fields, then one way to achieve what you want is to re-zip the array, i.e.

>>> fields = ak.unzip(arr)
>>> names = ak.fields(arr)
>>> rezipped = ak.zip(dict(zip(names, fields)))
>>> rezipped.tolist()
[[{'hit_x': 0, 'hit_y': 12.0, 'hit_z': 3},
  {'hit_x': 1, 'hit_y': 0.3, 'hit_z': 8},
  {'hit_x': 0, 'hit_y': 7.0, 'hit_z': 6}],
 [{'hit_x': 1, 'hit_y': 3.0, 'hit_z': 0},
  {'hit_x': 1, 'hit_y': 7.0, 'hit_z': 4}]]

Then you can unzip the records

>>> ak.unzip(rezipped)
(<Array [[0, 1, 0], [1, 1]] type='2 * var * int64'>,
 <Array [[12, 0.3, 7], [3, 7]] type='2 * var * float64'>,
 <Array [[3, 8, 6], [0, 4]] type='2 * var * int64'>)

However, you are asking about how to concatenate the fields as columns. To do this, we do not need to re-zip the array. First, we need to unzip the original array:

>>> x, y, z = ak.unzip(arr)

What we want to do is call np.stack, but currently Awkward does not implement support for it. Instead, we can add a new innermost dimension to our array, and then call ak.concatenate:

>>> new_dim = (..., np.newaxis)
>>> transposed = ak.concatenate((x[new_dim], y[new_dim], z[new_dim]), axis=-1)
>>> transposed.tolist()
[[[0.0, 12.0, 3.0], [1.0, 0.3, 8.0], [0.0, 7.0, 6.0]],
 [[1.0, 3.0, 0.0], [1.0, 7.0, 4.0]]]

Note that this will create intermediate data because concatenate creates a copy. In general this is pretty unavoidable; Awkward does have constructs to reference existing arrays, but when your array data size is similar to the size of the index, it's probably not worth doing (and is likely slower, because it adds another layer of indirection for computation).

There might be a faster way to do this (if that matters), depending upon how the data are serialised & deserialised; if hit_x, hit_y, and hit_z are different branches, then I believe that they are deserialised as separate arrays. If this weren't the case, then the fields of the record array would just be strided views on one big array. In that case, to obtain the transpose one would just take a RegularArray over the underlying data.

[jpivarski]

@agoose77's answer is exactly what I would have said: this is a good way of solving the problem. The only thing I would add for @alpetukhov is that this problem is not really "transposing," since a record of three fields is a very different thing from a list of length three. ak.unzip breaks down the record structure and ak.concatenate builds up a list structure, and that's why this works. :)

[alpetukhov]

Thanks for such a swift reply! That was exactly what I needed.
One small question though. I don't want to hardcode the variables, so is it memory safe to do something like this?

>>> new_dim = (..., np.newaxis)
>>> fields_to_transpose = [x[new_dim] for x in ak.unzip(arr)]
>>> transposed = `ak.concatenate(fields_to_transpose,` axis=-1)

[jpivarski]

I'm not sure what you're asking by "memory safe," but the above safely avoids Python for loops over large datasets. We assume that the number of fields in a record is small (a few thousand or less), which is not expensive to iterate over (for x in ak.unzip(arr)), but the number of items in an array is large (~billions). The above does not iterate over the array.

The ak.unzip and x[..., np.newaxis] doesn't even copy all the data (though it occurs to me that this particular slice could be optimized more than it is), but the ak.concatenate does need to copy data, since it's collating it from many sources into a single output array. (Concatenation of more than two sources is better optimized than it used to be—it used to do it two arrays at a time, but now it copies everything into the output array in a single pass.)

If by "memory safe," you mean memory leaks and such, there aren't any known ways to cause memory leaks, though a few bugs involving memory leaks have been fixed in the past. Memory leaks are bugs that we should fix, not an indicator that you're doing anything wrong.