Nested cartesian product for deep arrays

[nsmith-]

Consider

import awkward as ak

a = ak.Array(
    [
        [[1, 2], [3]],
        [[4], [5, 6]],
    ]
)
b = ak.Array(
    [
        [7],
        [8, 9, 10],
    ]
)
c = ak.Array(
    [
        [
            [[(1, 7)], [(2, 7)]],
            [[(3, 7)]],
        ],
        [
            [[(4, 8), (4, 9), (4, 10)]],
            [[(5, 8), (5, 9), (5, 10)], [(6, 8), (6, 9), (6, 10)]],
        ],
    ]
)

I would expect some incantation of ak.cartesian([a, b], ...) to produce c but so far have not been successful. I think such a desire could be keyed by nested=True, axis=-1, i.e. differing depths can in principle be supported when nesting. But perhaps the full generality of ak.cartesian can only be specified with a tuple of axes?

[jpivarski]

I wouldn't want to complicate the interface to ak.cartesian any more than it already is. I don't think we're looking at a case of wanting to do a Cartesian product at two different axis values. (Which one would be used in the output? The first? The deepest?) I think we're looking at a case of wanting to make b enough like a that ak.carteisan applies without modification. Whatever that transformation is, it's likely to be useful for other things besides Cartesian products, so deeping it factored out of the ak.cartesian interface would be a good thing.

To start with, we have to figure out what we need to do to b to get it to fit with a. In this case, we need to slice it like the following:

>>> b[[[0, 0], [1, 1]]]
<Array [[[7], [7]], ... 8, 9, 10], [8, 9, 10]]] type='2 * 2 * var * int64'>

That is, the [7] in index 0 should be replicated for each of a's two lists in its first element and the [8, 9, 10] in index 1 should be replicated for each of a's two lists in its second element.

And then we can use it:

>>> c = ak.cartesian((a, b[[[0, 0], [1, 1]]]), axis=2, nested=True)
>>> c.type
2 * var * var * var * (int64, int64)
>>> c.tolist()
[
    [
        [[(1, 7)], [(2, 7)]],
        [[(3, 7)]]
    ],
    [
        [[(4, 8), (4, 9), (4, 10)]],
        [[(5, 8), (5, 9), (5, 10)], [(6, 8), (6, 9), (6, 10)]]
    ]
]

Then the question becomes, how do we get [[0, 0], [1, 1]] from a?

Maybe you can think of what, in general, that function should be named or described as—either for making [[0, 0], [1, 1]] or for making b[[[0, 0], [1, 1]]] directly (depending on what concept it's addressing). It's something like broadcasting, but more general than broadcasting (which is already being applied).

For getting this indexer in a not-too-ridiculous way, there's

>>> ak.ones_like(ak.num(a, axis=2)) * np.arange(len(a))
<Array [[0, 0], [1, 1]] type='2 * var * int64'>

The ak.num function is being abused here: we're not interested in the number of elements in each list, but at axis=2, it produces the right number of items (and calculating ak.num is probably less expensive than, say, ak.local_index at axis=1). Then the ak.ones_like function makes all the values 1 while maintaining structure. Then multiplying by np.arange(len(a)) gets the right item to duplicate in each slot.