Introduction

numpy.ndarray has a mechanism called fancy index. It's very convenient, but the nested index structure is contrary to my intuition and is difficult to use. As an example, let's quote the following code from the previous article.

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
a[[0, 1], [2, 2]] #=> [3, 6]

This, the index for extracting the [0, 2] element and the [1, 2] element is ([0, 1], [2, 2]), but I am ʻa [[ I want to write 0, 2], [1, 2]] `(well, maybe this is nested because I want to use it with broadcasts and slices, but I want to write it intuitively). Zip functions can be used for such nesting:

a[tuple(zip([0, 2], [1, 2]))] #=> array([3, 6])
#The zip object cannot be indexed directly, so it is tupled.

However, numpy.ndarray may be a higher-order tensor, and a simple zip may not be enough, so I want a function that recursively replaces.

Implementation

Since it is difficult to verbalize the requirements, let's show the result of trial and error first.

import collections

def deepzip(*x):
  if isinstance(x[0][0], collections.abc.Iterable):
    return zip(*map(lambda y: deepzip(*y), x))
  else:
    return zip(*x)
def invert_index(idx):
  return tuple(deepzip(*idx))

a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
a[invert_index(([[0, 2], [1, 2]], [[1, 0], [1, 1]]))] #=> array([[3, 6], [4, 5]])

Since it is judged by whether it is ʻIterable, it is OK if the index is not a list but a tuple or ndarray. x [0] [0]` is not good, but I omitted the illegal processing such as the depth is not uniform. Both input and output (if they are consistent) can be any number of floors ...

Please let me know if there is a pattern that does not work or if you should do this.

[PYTHON] Recursive zip

Introduction

Implementation