Repository navigation
ENH: create a method to index N-dim tensors using 1D index #23992
Description
Proposed new feature or change:
I work with geospatial data that requires tensors with many dimensions. One challenge I used to have was when I had to implement a __getitem__ to access those tensors using a 1D index. One use case is machine learning, when one needs to feed models sequentially with sub-tensors of the original one. A simple example is a matrix of shape 2x2 that has the positions (0, 0), (0, 1), (1, 0), (1, 1), and I want to access it via indices [0, 1, 2, 3]. So, when I say __getitem__(2), I want the position (1, 0).
The internals of numpy probably has such information because arrays are usually stored contiguously in memory, but I couldn't find a way to access this information.
I came up with this function that calculates tensor indices given the 1D index and the shape of the tensor:
from typing import Tuple
import numpy as np
def f(i: int, s: Tuple[int]) -> Tuple[int]:
return (i,) if len(s) == 1 else f(i // s[-1], s[:-1]) + (i % s[-1],)How to use it:
tensor = np.arange(12).reshape((3, 4))
row_id, col_id = f(i=6, s=(3, 4))Would numpy like to include this method in its code? If yes, I can check how to submit a PR. If not, I would appreciate an indication of a more suitable library (scipy, maybe xarray).
Daniel
Activity
rkern commentedon Jun 20, 2023
Actually accessing the data like that is already implemented with the
ndarray.flatproperty.Given the availability of that property, I don't think we have many use cases for computing the ndim indices one by one in Python, so I would not add a function to numpy's API to do that.
dancivitarese commentedon Jun 20, 2023
Hi @rkern, I agree entirely with you. I would like to know about this function earlier. For some reason, it's not very popular because I looked into this and found nothing about it.
dancivitarese commentedon Jun 20, 2023
Hi @rkern, I was thinking here, and maybe this
flatfunction does not solve the problem entirely. I need to get the coordinates of my index, not the value. Is there a way to get that in numpy?rkern commentedon Jun 20, 2023
See
np.unravel_index().dancivitarese commentedon Jun 21, 2023
Thank you very much!