Operations on 'N' dimensional numpy arrays

Question

I am attempting to generalize some Python code to operate on arrays of arbitrary dimension. The operations are applied to each vector in the array. So for a 1D array, there is simply one operation, for a 2-D array it would be both row and column-wise (linearly, so order does not matter). For example, a 1D array (a) is simple:

b = operation(a)

where 'operation' is expecting a 1D array. For a 2D array, the operation might proceed as

for ii in range(0,a.shape[0]):
    b[ii,:] = operation(a[ii,:])
for jj in range(0,b.shape[1]):
    c[:,ii] = operation(b[:,ii])

I would like to make this general where I do not need to know the dimension of the array beforehand, and not have a large set of if/elif statements for each possible dimension. Solutions that are general for 1 or 2 dimensions are ok, though a completely general solution would be preferred. In reality, I do not imagine needing this for any dimension higher than 2, but if I can see a general example I will learn something!

Extra information: I have a matlab code that uses cells to do something similar, but I do not fully understand how it works. In this example, each vector is rearranged (basically the same function as fftshift in numpy.fft). Not sure if this helps, but it operates on an array of arbitrary dimension.

function aout=foldfft(ain)
nd = ndims(ain);
for k = 1:nd
    nx = size(ain,k);
    kx = floor(nx/2);
    idx{k} = [kx:nx 1:kx-1];
end
aout = ain(idx{:});

Answer 1

If you are looking for a programmatic way to index the k-th dimension an n -dimensional array, then numpy.take might help you.

An implementation of foldfft is given below as an example:

In[1]:
import numpy as np

def foldfft(ain):
    result = ain
    nd = len(ain.shape)
    for k in range(nd):
        nx = ain.shape[k]
        kx = (nx+1)//2
        shifted_index = list(range(kx,nx)) + list(range(kx))
        result = np.take(result, shifted_index, k)
    return result

a = np.indices([3,3])
print("Shape of a = ", a.shape)
print("\nStarting array:\n\n", a)
print("\nFolded array:\n\n", foldfft(a))

Out[1]:
Shape of a =  (2, 3, 3)

Starting array:

 [[[0 0 0]
  [1 1 1]
  [2 2 2]]

 [[0 1 2]
  [0 1 2]
  [0 1 2]]]

Folded array:

 [[[2 0 1]
  [2 0 1]
  [2 0 1]]

 [[2 2 2]
  [0 0 0]
  [1 1 1]]]

Answer 2

You could use numpy.ndarray.flat , which allows you to linearly iterate over an dimensional numpy array. Your code should then look something like this:

b = np.asarray(x)
for i in range(len(x.flat)):
    b.flat[i] = operation(x.flat[i])

Answer 3

In Octave, your MATLAB code does:

octave:19> size(ain)
ans =
   2   3   4
octave:20> idx
idx = 
{
  [1,1] =
     1   2
  [1,2] =
     1   2   3
  [1,3] =
     2   3   4   1
}

and then it uses the idx cell array to index ain . With these dimensions it 'rolls' the size 4 dimension.

For 5 and 6 the index lists would be:

 2   3   4   5   1
 3   4   5   6   1   2

The equivalent in numpy is:

In [161]: ain=np.arange(2*3*4).reshape(2,3,4)
In [162]: idx=np.ix_([0,1],[0,1,2],[1,2,3,0])
In [163]: idx
Out[163]: 
(array([[[0]],

        [[1]]]), array([[[0],
         [1],
         [2]]]), array([[[1, 2, 3, 0]]]))
In [164]: ain[idx]
Out[164]: 
array([[[ 1,  2,  3,  0],
        [ 5,  6,  7,  4],
        [ 9, 10, 11,  8]],

       [[13, 14, 15, 12],
        [17, 18, 19, 16],
        [21, 22, 23, 20]]])

Besides the 0 based indexing, I used np.ix_ to reshape the indexes. MATLAB and numpy use different syntax to index blocks of values.

The next step is to construct [0,1],[0,1,2],[1,2,3,0] with code, a straight forward translation.

I can use np.r_ as a short cut for turning 2 slices into an index array:

In [201]: idx=[]
In [202]: for nx in ain.shape:
    kx = int(np.floor(nx/2.))
    kx = kx-1;
    idx.append(np.r_[kx:nx, 0:kx])
   .....:     
In [203]: idx
Out[203]: [array([0, 1]), array([0, 1, 2]), array([1, 2, 3, 0])]

and pass this through np.ix_ to make the appropriate index tuple:

In [204]: ain[np.ix_(*idx)]
Out[204]: 
array([[[ 1,  2,  3,  0],
        [ 5,  6,  7,  4],
        [ 9, 10, 11,  8]],

       [[13, 14, 15, 12],
        [17, 18, 19, 16],
        [21, 22, 23, 20]]])

In this case, where 2 dimensions don't roll anything, slice(None) could replace those:

In [210]: idx=(slice(None),slice(None),[1,2,3,0])
In [211]: ain[idx]

======================

np.roll does:

indexes = concatenate((arange(n - shift, n), arange(n - shift)))
res = a.take(indexes, axis)

np.apply_along_axis is another function that constructs an index array (and turns it into a tuple for indexing).

Answer 4

The folks above provided multiple appropriate solutions. For completeness, here is my final solution. In this toy example for the case of 3 dimensions, the function 'ops' replaces the first and last element of a vector with 1.

import numpy as np

def ops(s):
    s[0]=1
    s[-1]=1
    return s

a = np.random.rand(4,4,3)
print '------'
print 'Array a'
print a
print '------'
for ii in np.arange(a.ndim):
    a = np.apply_along_axis(ops,ii,a)
    print '------'
    print ' Axis',str(ii)
    print a
    print '------'
    print ' '

The resulting 3D array has a 1 in every element on the 'border' with the numbers in the middle of the array unchanged. This is of course a toy example; however ops could be any arbitrary function that operates on a 1D vector.

Flattening the vector will also work; I chose not to pursue that simply because the book-keeping is more difficult and apply_along_axis is the simplest approach.

apply_along_axis reference page