Select each row of data with different column indices

[maplexgitx0302]

Dear experts,

I have a basic question, but I couldn't find the solution. I think there must be some easy way to do this, so my question is:

Suppose I have an ak.Array

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

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

If I also have another list (or array) that have the same length to len(x), let's say

y = [2,0,1] # the column index I want to choose for each row
# i.e., choose column 2 from row 0, choose column 0 from row 1, choose column 1 from row 2

Is there any way to select data so that the output would give me

result_from_x1 = ak.Array([3,4,9])
result_from_x2 = ak.Array([[5,6], [1,1], [4,5,6]])

I have tried something like x[:, y] or x[ak.zip([range(len(x), y)])], but somehow I couldn't get the desired result :( I think I might miss somewhere in the document. Thank you very much!

[agoose77]

Hi @r08222011

What you're describing here is an advanced index. NumPy has its own concept of this (that Awkward supports), e.g.

x1[[0, 1,2], [2, 0, 1]]

where each array in the subscript operation picks out a component from the corresponding dimension. In other words,

[[0, 1,2], [2, 0, 1]]

is interpreted as a one dimensional array of N dimensional tuples:

[(0, 2), (1, 0), (2, 1)] 

to build the result.

However, this means that the result is always one dimensional.

Awkward also has an additional indexing mode, if you pass in a ragged array which has the same number of dimensions as your x1 array (i.e. with length-1 lists in the final dimension), it will pick out the values you're interested in, albeit with an additional dimension

ix = ak.Array([
    [2],
    [0],
    [1]
])
x1[ix]

You can simply slice the result:

ix = ak.Array([
    [2],
    [0],
    [1]
])
x1[ix][:, 0]

[jpivarski]

Thanks, @agoose77! This is a more complete explanation of the same thing I was trying:

>>> x1[[[2], [0], [1]]][:, 0]
<Array [3, 4, 9] type='3 * int64'>

>>> x2[[[2], [0], [1]]][:, 0]
<Array [[5, 6], [1, 1], [4, 5, 6]] type='3 * var * int64'>

[maplexgitx0302]

@agoose77 and @jpivarski , thank you for your reply! These explanations are really helpful and I understand now :) Thank you very much and have a great day!

[agoose77]

Note that this differs to NumPy (our Awkward indexing "wins out"). Even if ix is regular, this works. @jpivarski I don't recall that always being the rules, but we've discussed this so many times that I've lost track of things. Are we happy that ix does not need to be ragged (apparently)?