Convert a numpy array to an array of numpy arrays

Question

How can I convert numpy array a to numpy array b in a (num)pythonic way. Solution should ideally work for arbitrary dimensions and array lengths.

Answer 1

Based on hpaulj I provide a litte more generic solution. a is an array of dimension N which shall be converted to an array b of dimension N1 with dtype object holding arrays of dimension (N-N1). In the example N equals 5 and N1 equals 3.

import numpy as np

N=5
N1=3

#create array a with dimension N
a=np.random.random(np.random.randint(2,20,size=N))

a_shape=a.shape
b_shape=a_shape[:N1] # shape of array b
b_arr_shape=a_shape[N1:] # shape of arrays in b

#Solution 1 with list() method (faster)
b=np.empty(np.prod(b_shape),dtype=object) #init b
b[:]=list(a.reshape((-1,)+b_arr_shape))
b=b.reshape(b_shape)

print "Dimension of b: {}".format(len(b.shape)) # dim of b
print "Dimension of array in b: {}".format(len(b[0,0,0].shape)) # dim of arrays in b

#Solution 2 with ndindex loop (slower)
b=np.empty(b_shape,dtype=object)
for index in np.ndindex(b_shape):
  b[index]=a[index]

print "Dimension of b: {}".format(len(b.shape)) # dim of b
print "Dimension of array in b: {}".format(len(b[0,0,0].shape)) # dim of arrays in b

Answer 2

For a start:

In [638]: a=np.arange(12).reshape(2,3,2)
In [639]: b=np.empty((2,3),dtype=object)
In [640]: for index in np.ndindex(b.shape):
    b[index]=a[index]
   .....:     
In [641]: b
Out[641]: 
array([[array([0, 1]), array([2, 3]), array([4, 5])],
       [array([6, 7]), array([8, 9]), array([10, 11])]], dtype=object)

It's not ideal since it uses iteration. But I wonder whether it is even possible to access the elements of b in any other way. By using dtype=object you break the basic vectorization that numpy is known for. b is essentially a list with numpy multiarray shape overlay. dtype=object puts an impenetrable wall around those size 2 arrays.

For example, a[:,:,0] gives me all the even numbers, in a (2,3) array. I can't get those numbers from b with just indexing. I have to use iteration:

[b[index][0] for index in np.ndindex(b.shape)]
# [0, 2, 4, 6, 8, 10]

---

np.array tries to make the highest dimension array that it can, given the regularity of the data. To fool it into making an array of objects, we have to give an irregular list of lists or objects. For example we could:

mylist = list(a.reshape(-1,2)) # list of arrays
mylist.append([]) # make the list irregular
b = np.array(mylist)  # array of objects
b = b[:-1].reshape(2,3)  # cleanup

---

The last solution suggests that my first one can be cleaned up a bit:

b = np.empty((6,),dtype=object)
b[:] = list(a.reshape(-1,2))
b = b.reshape(2,3)

I suspect that under the covers, the list() call does an iteration like

[x for x in a.reshape(-1,2)]

So time wise it might not be much different from the ndindex time.

---

One thing that I wasn't expecting about b is that I can do math on it, with nearly the same generality as on a:

b-10
b += 10
b *= 2

---

An alternative to an object dtype would be a structured dtype, e.g.

In [785]: b1=np.zeros((2,3),dtype=[('f0',int,(2,))])

In [786]: b1['f0'][:]=a

In [787]: b1
Out[787]: 
array([[([0, 1],), ([2, 3],), ([4, 5],)],
       [([6, 7],), ([8, 9],), ([10, 11],)]], 
      dtype=[('f0', '<i4', (2,))])

In [788]: b1['f0']
Out[788]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]]])

In [789]: b1[1,1]['f0']
Out[789]: array([8, 9])

And b and b1 can be added: b+b1 (producing an object dtype). Curiouser and curiouser!