[Python] Efficient and readable way to get N-dimensional index array in C-order using NumPy

When I need to generate an N-dimensional index array in C-order, I’ve tried a few different NumPy approaches.

The fastest for larger arrays but less readable:

np.stack(np.meshgrid(*[np.arange(i, dtype=dtype) for i in sizes], indexing="ij"), axis=-1).reshape(-1, len(sizes))


More readable with good performance:

np.ascontiguousarray(np.indices(sizes, dtype=dtype).reshape(len(sizes), -1).T)


Here I'm not sure if the ascontiguousarray copy is actually necessary, or if there's a better way to make sure the result is in C-order without forcing a copy.

Most readable, but by far the slowest:

np.vstack([*np.ndindex(sizes)], dtype=dtype)


The iterator conversion is quite slow for larger arrays.

Is there a built-in more straightforward and readable way to achieve this that matches the performance of np.meshgrid or np.indices using NumPy? If not, can either the meshgrid or indices approaches be optimized to avoid unnecessary memory copies (like ascontiguousarray) while still making sure the array is C-contiguous?

Example:

sizes = (3, 1, 2)
idx = np.ascontiguousarray(np.indices(sizes).reshape(len(sizes), -1).T)
print(idx)
print(f"C_CONTIGUOUS: {idx.flags['C_CONTIGUOUS']}")
# [[0 0 0]
# [0 0 1]
# [1 0 0]
# [1 0 1]
# [2 0 0]
# [2 0 1]]
# C_CONTIGUOUS: True

Continue reading...