A question about "record" combinations

[draghuram]

Please see the following code where I am using "zip" to build a record and then create combinations.

import awkward as ak

A1 = ak.Array([
        [1, 2],
        [3, 4, 5],
        [6]
    ])

A2 = ak.Array([
        [11, 21],
        [31, 41, 51],
        [61]
    ])

R = ak.zip({"a1": A1, "a2": A2})
print("R")
R.show()

C = ak.combinations(R, 2)
print("C")
C.show()

print(C[0][0][0])

I see the following response when I run this code:

[[{a1: 1, a2: 11}, {a1: 2, a2: 21}],
 [{a1: 3, a2: 31}, {a1: 4, a2: 41}, {a1: 5, a2: 51}],
 [{a1: 6, a2: 61}]]
[[({a1: 1, a2: 11}, {a1: 2, a2: 21})],
 [({a1: 3, a2: 31}, {a1: 4, a2: 41}), ({...}, ...), ({a1: 4, ...}, {...})],
 []]
Traceback (most recent call last):
  File "test_ak.py", line 23, in <module>
    print(C[0][0][0])
          ~~~~~~~^^^
  File ".../site-packages/awkward/highlevel.py", line 1754, in __getitem__
    out = self._layout[where]
          ~~~~~~~~~~~~^^^^^^^
  File ".../site-packages/awkward/record.py", line 138, in __getitem__
    return self._getitem(where)
           ^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/awkward/record.py", line 142, in _getitem
    raise IndexError("scalar Record cannot be sliced by an integer")
IndexError: scalar Record cannot be sliced by an integer

This error occurred while attempting to slice

    <Record ({a1: 1, ...}, {...}) type='({a1: int64, a2: int64}, {a1: int64...'>

with

    0

It is possible that I am missing something here but I thought that "combinations()" should add an extra dimension to the array as it generates list of tuples. But from the error message above, it seems that a tuple of records is also treated as a record. Is that correct?

I also see a "tolist()" method on the record and perhaps that will convert the combination tuple to a list but would it be possible to convert all record tuples to an extra dimension (axis)? I am hoping that this would allow me to create a mask such as "> 0.5" or something like that against entire axis.

I apologize if I haven't explained the issue properly. Any help is appreciated.

Thanks,
Raghu

[agoose77]

The Array.show() function accepts a type boolean that makes it easier to see the structure of the array:

>>> C.show(type=True)
type: 3 * var * (
    {
        a1: int64,
        a2: int64
    },
    {
        a1: int64,
        a2: int64
    }
)
[[({a1: 1, a2: 11}, {a1: 2, a2: 21})],
 [({a1: 3, a2: 31}, {a1: 4, a2: 41}), ({...}, ...), ({a1: 4, ...}, {...})],
 []]

The ak.combinations operation does not introduce new dimensions; the axis from which the combinations are taken is re-used as the result dimension (of a new length!). Note that tuples (and records) are distinct from array dimensions. So, your C array has the same number of dimensions as the R array (2).

If you want to pull out the first slot of the resulting tuple, you can use a string representing the index, e.g.

>>> C[0, 0, "0"]
{a1: 1, a2: 11}

If you would like to name these fields, and obtain a record instead of a tuple, you can pass a fields argument to ak.combinations:

>>> C = ak.combinations(R, 2, fields=["first", "second"])
>>> C.show(type=True)
type: 3 * var * {
    first: {
        a1: int64,
        a2: int64
    },
    second: {
        a1: int64,
        a2: int64
    }
}
[[{first: {a1: 1, a2: 11}, second: {a1: 2, ...}}],
 [{first: {a1: 3, a2: 31}, second: {a1: 4, ...}}, ..., {first: {...}, ...}],
 []]

[draghuram]

Thanks. It is clear to me now that ak.combinations doesn't introduce new dimension. Thanks for the clarification. Is there any way to run a filter on each of the combinations? For example, here is one combination:

({a1: 3, a2: 31}, {a1: 4, a2: 41})

Is there a way to find all combinations where at least one "a1" value is > 2? I can write code that does this check for each combination but I am wondering if it is possible to do it as a vectorized operation, making it more efficient. Thanks.

[agoose77]

Yes! Something like this?

C = ak.combinations(R, 2)
C_is_good = (C["0", "a1"] > 2) | (C["1", "a1"] > 2)
C_good = C[C_is_good]

[draghuram]

Great, this does the trick. I am curious how C["0", "a1"] works. It seems to directly access records in axis 1?

[agoose77]

In Awkward, field selection is orthogonal (independent of) dimension selection. You can place the field names anywhere in the slice with respect to the dimension index, as long as you don't re-order the strings. The API reference describes this here!

C["0", "a1"]

is equivalent to

C["0"]["a1"]

and gives you the array of "a1" values for the "left" combinations.

[draghuram]

Got it. Thanks a lot for very quick responses.