使用切片的numpy数组索引numpy数组

Question

(Edit: I wrote a solution basing on hpaulj's answer, see code at the bottom of this post)

I wrote a function that subdivides an n-dimensional array into smaller ones such that each of the subdivisions has max_chunk_size elements in total.

Since I need to subdivide many arrays of same shapes and then perform operations on the corresponding chunks, it doesn't actually operate on the data rather than creates an array of "indexers", ie an array of (slice(x1, x2), slice(y1, y2), ...) objects (see the code below). With these indexers I can retrieve subdivisions by calling the_array[indexer[i]] (see examples below).

Also, the array of these indexers has same number of dimensions as input and divisions are aligned along corresponding axes, ie blocks the_array[indexer[i,j,k]] and the_array[indexer[i+1,j,k]] are adjusent along the 0-axis, etc.

I was expecting that I should also be able to concatenate these blocks by calling the_array[indexer[i:i+2,j,k]] and that the_array[indexer] would return just the_array , however such calls result in an error:

IndexError: arrays used as indices must be of integer (or boolean) type

Is there a simple way around this error?

Here's the code:

import numpy as np
import itertools

def subdivide(shape, max_chunk_size=500000):
    shape = np.array(shape).astype(float)
    total_size = shape.prod()

    # calculate maximum slice shape:
    slice_shape = np.floor(shape * min(max_chunk_size / total_size, 1.0)**(1./len(shape))).astype(int)

    # create a list of slices for each dimension:
    slices = [[slice(left, min(right, n)) \
      for left, right in zip(range(0, n, step_size), range(step_size, n + step_size, step_size))] \
        for n, step_size in zip(shape.astype(int), slice_shape)]

    result = np.empty(reduce(lambda a,b:a*len(b), slices, 1), dtype=np.object)
    for i, el in enumerate(itertools.product(*slices)): result[i] = el
    result.shape = np.ceil(shape / slice_shape).astype(int)
    return result

Here's an example usage:

>>> ar = np.arange(90).reshape(6,15)
>>> ar
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44],
       [45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74],
       [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89]])

>>> slices = subdivide(ar.shape, 16)
>>> slices
array([[(slice(0, 2, None), slice(0, 6, None)),
        (slice(0, 2, None), slice(6, 12, None)),
        (slice(0, 2, None), slice(12, 15, None))],
       [(slice(2, 4, None), slice(0, 6, None)),
        (slice(2, 4, None), slice(6, 12, None)),
        (slice(2, 4, None), slice(12, 15, None))],
       [(slice(4, 6, None), slice(0, 6, None)),
        (slice(4, 6, None), slice(6, 12, None)),
        (slice(4, 6, None), slice(12, 15, None))]], dtype=object)

>>> ar[slices[1,0]]
array([[30, 31, 32, 33, 34, 35],
       [45, 46, 47, 48, 49, 50]])
>>> ar[slices[0,2]]
array([[12, 13, 14],
       [27, 28, 29]])
>>> ar[slices[2,1]]
array([[66, 67, 68, 69, 70, 71],
       [81, 82, 83, 84, 85, 86]])

>>> ar[slices[:2,1:3]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: arrays used as indices must be of integer (or boolean) type

---

Here's a solution based on hpaulj's answer:

import numpy as np
import itertools

class Subdivision():
    def __init__(self, shape, max_chunk_size=500000):
        shape = np.array(shape).astype(float)
        total_size = shape.prod()

        # calculate maximum slice shape:
        slice_shape = np.floor(shape * min(max_chunk_size / total_size, 1.0)**(1./len(shape))).astype(int)

        # create a list of slices for each dimension:
        slices = [[slice(left, min(right, n)) \
          for left, right in zip(range(0, n, step_size), range(step_size, n + step_size, step_size))] \
            for n, step_size in zip(shape.astype(int), slice_shape)]

        self.slices = \
            np.array(list(itertools.product(*slices)), \
                     dtype=np.object).reshape(tuple(np.ceil(shape / slice_shape).astype(int)) + (len(shape),))

    def __getitem__(self, args):
        if type(args) != tuple: args = (args,)

        # turn integer index into equivalent slice
        args = tuple(slice(arg, arg + 1 if arg != -1 else None) if type(arg) == int else arg for arg in args)

        # select the slices
        # always select all elements from the last axis (which contains slices for each data dimension)
        slices = self.slices[args + ((slice(None),) if Ellipsis in args else (Ellipsis, slice(None)))]

        return np.ix_(*tuple(np.r_[tuple(slices[tuple([0] * i + [slice(None)] + \
                                                      [0] * (len(slices.shape) - 2 - i) + [i])])] \
                                for i in range(len(slices.shape) - 1)))

Example usage:

>>> ar = np.arange(90).reshape(6,15)
>>> ar
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44],
       [45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74],
       [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89]])

>>> subdiv = Subdivision(ar.shape, 16)
>>> ar[subdiv[...]]
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44],
       [45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74],
       [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89]])

>>> ar[subdiv[0]]
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]])

>>> ar[subdiv[:2,1]]
array([[ 6,  7,  8,  9, 10, 11],
       [21, 22, 23, 24, 25, 26],
       [36, 37, 38, 39, 40, 41],
       [51, 52, 53, 54, 55, 56]])

>>> ar[subdiv[2,:3]]
array([[60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74],
       [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89]])

>>> ar[subdiv[...,:2]]
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
       [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41],
       [45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71],
       [75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86]])

Answer 1

Your slices produce 2x6 and 2x3 arrays.

In [36]: subslice=slices[:2,1:3]
In [37]: subslice[0,0]
Out[37]: array([slice(0, 2, None), slice(6, 12, None)], dtype=object)

In [38]: ar[tuple(subslice[0,0])]
Out[38]: 
array([[ 6,  7,  8,  9, 10, 11],
       [21, 22, 23, 24, 25, 26]])

My numpy version expects me to turn the subslice into a tuple. This is the same as

ar[slice(0,2), slice(6,12)]
ar[:2, 6:12]

That's just the basic syntax of indexing and slicing. ar is 2d, so ar[(i,j)] requires a 2 element tuple - of slices, lists, arrays, or integers. It won't work with an array of slice objects.

How ever it is possible to concatenate the results into a larger array. That can be done after indexing or the slices can be converted into indexing lists.

np.bmat for example concatenates together a 2d arangement of arrays:

In [42]: np.bmat([[ar[tuple(subslice[0,0])], ar[tuple(subslice[0,1])]], 
                  [ar[tuple(subslice[1,0])],ar[tuple(subslice[1,1])]]])
Out[42]: 
matrix([[ 6,  7,  8,  9, 10, 11, 12, 13, 14],
        [21, 22, 23, 24, 25, 26, 27, 28, 29],
        [36, 37, 38, 39, 40, 41, 42, 43, 44],
        [51, 52, 53, 54, 55, 56, 57, 58, 59]])

You could generalize this. It just uses hstack and vstack on the nested lists. The result is np.matrix but can be converted back to array .

The other approach is to use tools like np.arange , np.r_ , np.xi_ to create index arrays. It'll take some playing around to generate an example.

To combine the [0,0] and [0,1] subslices:

In [64]: j = np.r_[subslice[0,0,1],subslice[0,1,1]]
In [65]: i = np.r_[subslice[0,0,0]]

In [66]: i,j
Out[66]: (array([0, 1]), array([ 6,  7,  8,  9, 10, 11, 12, 13, 14]))
In [68]: ix = np.ix_(i,j)
In [69]: ix
Out[69]: 
(array([[0],
        [1]]), array([[ 6,  7,  8,  9, 10, 11, 12, 13, 14]]))

In [70]: ar[ix]
Out[70]: 
array([[ 6,  7,  8,  9, 10, 11, 12, 13, 14],
       [21, 22, 23, 24, 25, 26, 27, 28, 29]])

Or with i = np.r_[subslice[0,0,0], subslice[1,0,0]] , ar[np.ix_(i,j)] produces the 4x9 array.