Marco Spinaci
Marco Spinaci

Reputation: 1879

numpy.array select all even elements from d-dimensional array

Is there an efficient and pythonic way to select, from a d-dimensional array, all elements that have an even index, without knowing d beforehand? And all the remaining ones (i.e. all those that have at least an odd index)?

Minimal example of the first question

import numpy as np
a = np.array(range(27)).reshape((3,3,3))
a[::2,::2,::2] 
# -> array([[[ 0,  2],
#            [ 6,  8]],
#           [[18, 20],
#            [24, 26]]])

Only non-pythonic way I found for d-dimensional objects, d being variable At least for the "all even" part, the "at least one odd" still escapes me.

d = 3
a = np.array(range(3**d)).reshape([3]*d)
b = a
for i in range(d):
    b = np.take(b, np.array(range(0,b.shape[i],2)), axis=i)

The reason why I'm asking this (that maybe has already a higher-level solution) is that I want to iteratively create a large d-dimensional object of size (2**n+1, ..., 2**n+1) in n steps, at each step copying the even-indexed elements from the previous steps, something like:

for n in range(N):
    new_array = np.zeros([2**n+1]*d)
    new_array[all_even] = old_array
    new_array[at_least_one_odd] = #something else

Thanks in advance for any hint!

Upvotes: 3

Views: 7909

Answers (2)

miindlek
miindlek

Reputation: 3563

I guess you could use slice objects.

even = a[[slice(None, None, 2) for _ in range(a.ndim)]]
odd = a[[slice(1, None, 2) for _ in range(a.ndim)]]

Upvotes: 3

Divakar
Divakar

Reputation: 221614

Here's an approach using np.ix_ -

a[np.ix_(*[range(0,i,2) for i in a.shape])]

Sample runs -

In [813]: def even_idx(a):
     ...:     return a[np.ix_(*[range(0,i,2) for i in a.shape])]
     ...: 

In [814]: a = np.array(range(27)).reshape((3,3,3))

In [815]: np.allclose(a[::2,::2,::2], even_idx(a) )
Out[815]: True

In [816]: a = np.array(range(27*4)).reshape((3,3,3,4))

In [817]: np.allclose(a[::2,::2,::2,::2], even_idx(a) )
Out[817]: True

In [818]: a = np.array(range(27*4*5)).reshape((3,3,3,4,5))

In [819]: np.allclose(a[::2,::2,::2,::2,::2], even_idx(a) )
Out[819]: True

Upvotes: 2

Related Questions