Changing shape based on rows with equal values

[CaueSousa]

Hello, dear community!

I'm trying to figure out if is possible to do something with the help of Awkward.

I have an ak array with '21104832 * 23 * float64' dimensions, where some lines share the same values on the first two columns:

Column 1	Column 2	...	Column 23
x1	y1	...	arbitrary
x2	y2	...	arbitrary
x2	y2	...	arbitrary
x2	y2	...	arbitrary
x3	y3	...	arbitrary
x4	y4	...	arbitrary
x4	y4	...	arbitrary
x4	y4	...	arbitrary
x4	y4	...	arbitrary

The lines that share the same values on the first two columns are in a sequence.

I would like to know if there's a way to change this ak array in order to keep these lines with same values on the first two columns separated from the other "groups".

[agoose77]

Hi @CaueSousa! I'm interpreting your question as:

You want to split an array of length 21104832 into sub-lists, where each sub-list contains the rows with identical leading two columns. Is this correct?

One thing that would be helpful to know; is it ever possible that the first column and second column of two rows are not both equal or non-equal? i.e., is this

10 1000 ...
11 1000 ...

or this

11 1001 ...
11 1000 ...

possible?

Additionally, does the order of this array matter? Can we re-order the rows, as long as the groups are preserved?

[jpivarski]

We sent at the same time, and it looks like you had a different take on that this problem is.

[agoose77]

@jpivarski your answer is where my mind was going! I was unsure as to the requirement to split on both columns, but I accidentally found the solution :)

@CaueSousa I wasn't initially sure how easy this would be. It turns out that there's a neat trick (or at least, I think it's a trick!) that let's do this using run_lengths without any complex slicing.

First, let's define an example array

>>> array = ak.Array(
    [
        # A
        [0, 0, 1],
        [0, 0, 2],
        # B
        [1, 1, 3],
        # C
        [2, 2, 4],
        # D
        [2, 3, 5],
        [2, 3, 6],
    ]
)

Let's find the lengths of the "runs" of same-value items in the first column:

>>> runs_in_col_0 = ak.run_lengths(array[:, 0])
>>> runs_in_col_0
[2,
 1,
 3]

We can use this to introduce a new dimension that groups our rows by these runs in column 0:

>>> array_by_col_0 = ak.unflatten(array, runs_in_col_0, axis=0)
>>> array_by_col_0
[[[0, 0, 1], [0, 0, 2]],
 [[1, 1, 3]],
 [[2, 2, 4], [2, 3, 5], [2, 3, 6]]]

Then we can compute the runs in the second column of this unflattened array

>>> runs_in_col_1_of_grouped_by_0 = ak.run_lengths(array_by_col_0[..., 1])
>>> runs_in_col_1_of_grouped_by_0
[[2],
 [1],
 [1, 2]]

You can see that this array has the same structure as our unflattened array_by_col_0 array. This is because run_lengths always operates at axis=-1, which means that ak.run_lengths(array_by_col_0[..., 1]) will always stop counting a run if the sublist in array_by_col_0[..., 1] ends. Let's remind ourselves what array_by_col_0[..., 1] looks like:

>>> array_by_col_0[..., 1]
[[0, 0],
 [1],
 [2, 3, 3]]

Visually, our runs should be [[2], [1], [1, 2]], which is what we computed for runs_in_col_1_of_grouped_by_0.

Notice that the sums of runs_in_col_1_of_grouped_by_0 and runs_in_col_0 are both 6. If we flatten runs_in_col_1_of_grouped_by_0

>>> runs_in_col_1_of_grouped_by_0_flat = ak.flatten(runs_in_col_1_of_grouped_by_0, axis=None)
>>> runs_in_col_1_of_grouped_by_0_flat
[2,
 1,
 1,
 2]

we can use the resulting array of counts to unflatten the original array:

>>> ak.unflatten(array, runs_in_col_1_of_grouped_by_0_flat, axis=0)
[[[0, 0, 1], [0, 0, 2]],
 [[1, 1, 3]],
 [[2, 2, 4]],
 [[2, 3, 5], [2, 3, 6]]]

That's the trick — using the counts computed over the temporary array_by_col_0 array to unflatten the original array. We can do this only because the total number of rows is unchanged!

There are other ways to do this, but this approach seems the most performant to me.

[CaueSousa]

Hi, @agoose77

Your assumptions were correct! In my context, each row carries information about an object called track(the trajectory of a charged particle) and the jet('spray' of particles) in which this track is. The first two columns are the information of the jets, and so, lines with the same values on the first two columns are inside the same jet.

This piece does the trick precisely!

@agoose77 , @jpivarski you're wizards, guys. Thank you so much, seriously!

[jpivarski]

Oddly enough, there is a function for this. You can use ak.run_lengths to determine how many values are equal (contiguously), and then ak.unflatten to make groups of those sizes. You might need to flatten and select an individual column before applying this operation, though.

The links to documentation point to similar examples.