¿Cómo convertir el segmento ND a índices en NumPy?

Question

Dada cualquier N-tupla de sectores (también conocido como sector ND) en NumPy, ¿cómo convertirlo en los índices correspondientes de la matriz ND representada como una tupla de matrices 1D (índices a lo largo de cada eje)? Por ejemplo, si tenemos el np.nd_slice_to_indexessiguiente código:

import numpy as np
print(np.nd_slice_to_indexes(np.s_[1 : 3]))
print(np.nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))

debería imprimir

(array([1, 2]),)
(array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))

Es común que NumPy represente índices de la matriz ND como N-tupla de matrices 1-D de la misma longitud (cada elemento de la k-thmatriz en la tupla representa el siguiente índice a lo largo de la k-ésima dimensión). Por ejemplo, np.nonzerodevuelve tal N-tupla en código

print(np.nonzero([[0, 1, 1], [1, 1, 0]])) # Non-zero elements in 2D array.
# (array([0, 0, 1, 1], dtype=int64), array([1, 2, 0, 1], dtype=int64))

Se debe lograr el mismo comportamiento que en la función Pythonic a continuación, pero de una manera más eficiente (de rendimiento):

¡Pruébelo en línea!

import numpy as np

def nd_slice_to_indexes(nd_slice):
    assert type(nd_slice) in [tuple, slice], type(nd_slice)
    if type(nd_slice) is not tuple:
        nd_slice = (nd_slice,)
    def iter_slices(slices):
        if len(slices) == 0:
            yield ()
        else:
            for i in range(slices[0].start, slices[0].stop, slices[0].step or 1):
                for r in iter_slices(slices[1:]):
                    yield (i,) + r
    *res, = np.vstack(list(iter_slices(nd_slice))).T
    return tuple(res)

print(nd_slice_to_indexes(np.s_[1 : 3]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2, 8 : 14 : 3]))
# (array([1, 2]),)
# (array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))
# (array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]), array([5, 5, 7, 7, 9, 9, 5, 5, 7, 7, 9, 9]), array([ 8, 11,  8, 11,  8, 11,  8, 11,  8, 11,  8, 11]))

Answer 1

Gracias a la sugerencia de @hpaulj resolvió la tarea de manera eficiente usando np.mgrid .

¡Pruébelo en línea!

import numpy as np

def nd_slice_to_indexes(nd_slice):
    grid = np.mgrid[{tuple: nd_slice, slice: (nd_slice,)}[type(nd_slice)]]
    return tuple(grid[i].ravel() for i in range(grid.shape[0]))
    
print(nd_slice_to_indexes(np.s_[1 : 3]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2]))
print(nd_slice_to_indexes(np.s_[1 : 3, 5 : 11 : 2, 8 : 14 : 3]))
# (array([1, 2]),)
# (array([1, 1, 1, 2, 2, 2]), array([5, 7, 9, 5, 7, 9]))
# (array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2]), array([5, 5, 7, 7, 9, 9, 5, 5, 7, 7, 9, 9]), array([ 8, 11,  8, 11,  8, 11,  8, 11,  8, 11,  8, 11]))