Add `xt::ix_`

[mirmohammad]

xtensor is missing cross product indices generation (numpy.ix_) https://numpy.org/doc/stable/reference/generated/numpy.ix_.html

I am posting this code as a suggestion. It works for my use cases but definitely requires more testing. The idea is to return a vector of indices (that are cross product of all the inputs) used with index_view on an xarray.

in xt namespace:

template <typename... INDICES>
std::vector<std::array<std::size_t, sizeof...(INDICES)>> ix_(INDICES &&...indices)
{
    constexpr std::size_t N = sizeof...(INDICES);
    std::vector<std::array<std::size_t, N>> ix_indices;
    auto mesh = xt::meshgrid(std::forward<INDICES>(indices)...);
    auto stack = std::apply([=](auto &&...m)
                            { return xt::stack(xt::xtuple(m...), N).reshape({-1, (int)N}); },
                            mesh);
    auto start = xt::axis_begin(stack, 0);
    auto end = xt::axis_end(stack, 0);
    while (start != end)
    {
        ix_indices.push_back({});
        std::copy(start->begin(), start->end(), ix_indices.back().begin());
        start++;
    }
    return ix_indices;
}

Example Usage 1:

xt::xarray<int> some_indices{ 0, 1, 2, 3 };
auto indices = xt::ix_(some_indices, xt::arange(2));
auto selected_scores = xt::index_view(all_scores, indices);

Example Usage 2:

xt::xarray<int> some_indices{ 0, 1 };
auto indices = xt::ix_(some_indices, some_indices, some_indices);
auto selected_scores = xt::index_view(all_scores, indices);

[mirmohammad]

@AntoinePrv Would you mind checking my solution out?

[AntoinePrv]

Hi @mirmohammad,

What are you trying to achieve? From NumPy doc, it seems np.ix_ is simpler than what you are doing. It looks like one should go through the inputs and apply a reshape to add a number of 1s.

Perhaps what you are trying to do an advanced indexing where you apply a 1D index selection on each axis. As you noticed, there is no easy/performant way of doing in xtensor so far.

I ended up doing something similar when I needed: build the expanded list of indices on way or another, then pass it to xt::index_view.