Issue concatenating jagged arrays with axis=1

[carlgwilliam]

Hi,

As discussed on stack overflow

https://stackoverflow.com/questions/58440970/how-to-combine-two-uproot-jagged-arrays-into-one

in some cases I find that I concatenate inner jagged arrays with axis=1.

Here is a MWE for concatenating elecs and muons to make leps

import uproot as up
import awkward as awk

t = up.open("http://scikit-hep.org/uproot/examples/HZZ-objects.root")["events"]

azip = awk.array.jagged.JaggedArray.zip
acat = awk.concatenate

adict = t.arrays(["muonp4", "muoniso", "electronp4", "electroniso"], namedecode="utf-8")

mdict = {k.replace("muon", ""):v for k,v in adict.items() if k.startswith("muon")}
edict = {k.replace("electron", ""):v for k,v in adict.items() if k.startswith("electron")}

muons = azip(mdict)
elecs = azip(edict)

leps = acat([elecs, muons], axis=1)

gives

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-34-6ffe1cebff62> in <module>
----> 1 leps = acat([elecs, muons], axis=1)

/opt/anaconda3/lib/python3.7/site-packages/awkward/__init__.py in concatenate(arrays, axis)
     20 
     21 def concatenate(arrays, axis=0):
---> 22     return AwkwardArray.concatenate(arrays, axis=axis)
     23 
     24 from awkward.generate import fromiter

/opt/anaconda3/lib/python3.7/site-packages/awkward/util.py in <lambda>(*args, **kwargs)
     46     def __get__(self, ins, typ):
     47         if ins is None:
---> 48             return lambda *args, **kwargs: self.fcn(True, typ, *args, **kwargs)
     49         else:
     50             return lambda *args, **kwargs: self.fcn(False, ins, *args, **kwargs)

/opt/anaconda3/lib/python3.7/site-packages/awkward/array/base.py in concatenate(isclassmethod, cls_or_self, arrays, axis)
    589             return type(arrays[0])._concatenate_axis0(arrays)
    590         elif axis == 1:
--> 591             return type(arrays[0])._concatenate_axis1(arrays)
    592         else:
    593             raise NotImplementedError("axis > 1")

/opt/anaconda3/lib/python3.7/site-packages/awkward/array/jagged.py in _concatenate_axis1(cls, arrays)
   1705 
   1706         else:
-> 1707             raise NotImplementedError("concatenate with axis=1 is not implemented for " + type(arrays[0]).__name__)
   1708 
   1709         return arrays[0].__class__(starts[::n_arrays], stops[n_arrays-1::n_arrays], content)

NotImplementedError: concatenate with axis=1 is not implemented for JaggedArray

[jpivarski]

See PR #206.

This wasn't a JaggedArray of numbers as presented on SO, but probably because our error message wrongly said you couldn't do that (it was constructed without enough precision).

This was a JaggedArray of a Table of ObjectArrays, which wasn't handled for some trivial reasons, so I fixed them in the associated PR. This sort of thing needs to be handled in more generality (I'm thinking of the Awkward 1.0 project) so that we don't have to fix special cases like this.

Since you're concatenating Tables, you can keep track of the type of lepton by assigning a field before concatenating. That can be useful...

Let me know if the PR works for you!

[carlgwilliam]

Hi Jim,

Thanks and sorry for not being clear. This indeed seems to fix my issue.

Cheers,

Carl

[jpivarski]

Great!