如何在numpy中创建数组元组的3D数组，其中3D数组的每个元素就像np.where的输出？

Question

I have a 3D array where each axis is the same length (~100) and corresponds to a location in 3D space. Each element of this array has a value (float64). I am trying to create an array identifying neighbors where each element of this new 3D array consists of a tuple of arrays, similar to the output of np.where so that I can loop through each element of the 3D array and perform an operation on all the neighbors.

I know np.where for a 3D array would output a 3-tuple of arrays, where the length of each array is the same and is the number of elements satisfying the given condition. I'm trying to do something similar, but in my case I know that the length of these arrays is 6, since each element of the 3D array has 6 immediate neighbors. I want to set each element of the 3D neighbors array to be a tuple of 3 arrays of length 6.

What I would like to do is something like:

>>> values = np.zeros((100,100,100))
>>> neighbors = np.zeros((100,100,100))
>>> for i in range(neighbors.shape[0]):
>>>     for j in range(neighbors.shape[1]):
>>>         for k in range(neighbors.shape[2]):
>>>             neighbors[i,j,k] = (np.array([i-1,i+1,i,i,i,i]),np.array([j,j,j-1,j+1,j,j]),np.array([k,k,k,k,k-1,k+1]))
>>> neighbors[5,6,9]
(array([4,6,5,5,5,5]), array([6,6,5,7,6,6]), array([9,9,9,9,8,10]))

whereas what I actually get is the following error:

>>> values = np.zeros((100,100,100))
>>> neighbors = np.zeros((100,100,100))
>>> for i in range(neighbors.shape[0]):
>>>     for j in range(neighbors.shape[1]):
>>>         for k in range(neighbors.shape[2]):
>>>             neighbors[i,j,k] = (np.array([i-1,i+1,i,i,i,i]),np.array([j,j,j-1,j+1,j,j]),np.array([k,k,k,k,k-1,k+1]))
ValueError: setting an array element with a sequence.

Any way to set elements of an array as a tuple of arrays?

Answer 1

To store arbitrary Python objects in the cells of a NumPy array, change the dtype to object : Here, you would need to change

neighbors = np.zeros((100,100,100))

to

neighbors = np.zeros((100,100,100), dtype='O')

By default np.zeros creates an array of floating point dtype, which means every cell of the array must hold a floating point value, hence the error message:

ValueError: setting an array element with a sequence.

---

Beware that using arrays of object dtype inhibits NumPy from using fast numerical methods -- the code will be slower than equivalent NumPy code that uses arrays of native (eg floating point) dtype.

For example, you might be better off creating an array of floating point dtype and shape (3,6,100,100,100) where the first axis (of length 3) allows you to select amongst the three subarrays,

(array([4,6,5,5,5,5]), array([6,6,5,7,6,6]), array([9,9,9,9,8,10]))

and the second axis (of length 6) allows you to select amongst the 6 values in the subarray:

import numpy as np
d, h, w = 100, 100, 100
I, J, K = np.mgrid[:d, :h, :w]
neighbors = np.array([(I-1,I+1,I,I,I,I), (J,J,J-1,J+1,J,J), (K,K,K,K,K-1,K+1)])

Then neighbors is an array with a native NumPy dtype:

In [32]: neighbors.dtype
Out[32]: dtype('int64')

In [33]: neighbors.shape
Out[33]: (3, 6, 100, 100, 100)

Here are the three subarrays, corresponding to your neighbors[5,6,9] :

In [36]: neighbors[:,:,5,6,9]
Out[36]: 
array([[ 4,  6,  5,  5,  5,  5],
       [ 6,  6,  5,  7,  6,  6],
       [ 9,  9,  9,  9,  8, 10]])

Here is the second of these subarrays, corresponding to your neighbors[5,6,9][1] :

In [37]: neighbors[1,:,5,6,9]
Out[37]: array([6, 6, 5, 7, 6, 6])

And here is the fourth value in the second subarray, corresponding to your neighbors[5,6,9][1][3] :

In [38]: neighbors[1,3,5,6,9]
Out[38]: 7