[FIXED] Numpy array slicing to return sliced array and corresponding array indices

Issue

I'm trying to generate two numpy arrays from one. One which is a slice slice of an original array, and another which represents the indexes which can be used to look up the values produced. The best way I can explain this is by example:

import numpy as np

original = np.array([
    [5, 3, 7, 3, 2],
    [8, 4, 22, 6, 4],
])

sliced_array = original[:,::3]
indices_of_slice = None # here, what command to use?

for val, idx in zip(np.nditer(sliced_array), np.nditer(indices_of_slice)):
    # I want indices_of_slice to behave the following way:
    assert val == original[idx], "Error. This implementation is not correct. "

Ultimately what I'm aiming for is an array which I can iterate through with nditer, and a corresponding array indices_of_slices, which returns the original lookup indices (i,j,...). Then, the value of the sliced array should be equal to value of the original array in index (i,j,...).

Main question is: Can I return both the new sliced array as well as the indices of the values when slicing a numpy array? Hopefully all is clear!

Edit: Here are the expected printouts of the two arrays:

# print(sliced_array)
# >>> [[5 3]
# >>>  [8 6]]

# expected result of 
# print(indices_of_slice)
# >>> [[(0 0) (0 3)]
# >>>  [(1 0) (1 3)]]

Solution

You can use numpy's slice np.s_[] with a tiny bit of gymnastics to get the indices you are looking for:

slc = np.s_[:, ::3]

shape = original.shape
ix = np.unravel_index(np.arange(np.prod(shape)).reshape(shape)[slc], shape)

>>> ix
(array([[0, 0],
        [1, 1]]),
 array([[0, 3],
        [0, 3]]))

>>> original[ix]
array([[5, 3],
       [8, 6]])

>>> original[slc]
array([[5, 3],
       [8, 6]])

Note that this works with slices that have some reverse direction:

slc = np.s_[:, ::-2]

# ... (as above)

>>> ix
(array([[0, 0, 0],
        [1, 1, 1]]),
 array([[4, 2, 0],
        [4, 2, 0]]))

>>> np.array_equal(original[ix], original[slc])
True

Answered By - Pierre D