Fast iteration over vectors in numpy multidimensional array
I am writing python + numpy + cython code and am trying to find the most elegant and efficient way to do the following iteration over an array:
Let's say I have a function f (x, y) that takes a vector x of form (3) and a vector y of form (10) and returns a vector of form (10,). Now I have two arrays X and Y of the form sx + (3,) and sy + (10,), where sx and sy are two forms that can be translated together (i.e. either sx == sy or when the axis is different, one of the two has a length of 1, in which case it will be repeated). I want to create an array Z that is of the form zs + (10,) where zs is the sx cast form with sy. Each 10-dimensional vector in Z is equal to f (x, y) vectors x and y at their corresponding locations in X and Y.
I looked at np.nditer , and while it works fine with cython (see at the bottom of the page for links), it doesn't seem to allow iterating over vectors from a multidimensional array instead of elements. I also looked at index grids , but the problem is that cython indexing is only fast when the number of indices is equal to the dimension of the array, and are stored as cython integers instead of python tuples.
Any help is greatly appreciated!
source to share
You describe what Numpy calls a generic generic FUNCTION, or gufunc. As the name suggests, this is a ufuncs extension. You probably want to start by reading these two pages:
The second example uses Cython and has some gufunc stuff. To go all the way down the gufunc road, you will need to read the relevant section in the numpy C API documentation:
I don't know of any example of gufuncs coding in Cython, although it shouldn't be too hard to do with the examples above. If you want to look at gufuncs coded in C, you can look at the source code np.linalg
here , although it can be a daunting experience. A while ago, I missed my local Python user group to death talking about extending numpy with C, which was mostly about writing gufuncs in C, the slides of that conversation, and a sample Python module providing the new gufunc can be found here .
source to share
If you want to stick with nditer
, approximate sizes can be used here. Here's pure Python, but t21 doesn't have to be hard (although it still has a tuple iterator). I am borrowing ideas from ndindex
as described in shallow iteration with nditer
The idea is to find a general broadcast form sz
and build an iterator on it multi_index
.
I use as_strided
to extend to X
and Y
to the views used and pass the appropriate vectors (actually (1,n)
arrays) to the function f(x,y)
.
import numpy as np
from numpy.lib.stride_tricks import as_strided
def f(x,y):
# sample that takes (10,) and (3,) arrays, and returns (10,) array
assert x.shape==(1,10), x.shape
assert y.shape==(1,3), y.shape
z = x*10 + y.mean()
return z
def brdcast(X, X1):
# broadcast X to shape of X1 (keep last dim of X)
# modeled on np.broadcast_arrays
shape = X1.shape + (X.shape[-1],)
strides = X1.strides + (X.strides[-1],)
X1 = as_strided(X, shape=shape, strides=strides)
return X1
def F(X, Y):
X1, Y1 = np.broadcast_arrays(X[...,0], Y[...,0])
Z = np.zeros(X1.shape + (10,))
it = np.nditer(X1, flags=['multi_index'])
X1 = brdcast(X, X1)
Y1 = brdcast(Y, Y1)
while not it.finished:
I = it.multi_index + (None,)
Z[I] = f(X1[I], Y1[I])
it.iternext()
return Z
sx = (2,3) # works with (2,1)
sy = (1,3)
# X, Y = np.ones(sx+(10,)), np.ones(sy+(3,))
X = np.repeat(np.arange(np.prod(sx)).reshape(sx)[...,None], 10, axis=-1)
Y = np.repeat(np.arange(np.prod(sy)).reshape(sy)[...,None], 3, axis=-1)
Z = F(X,Y)
print Z.shape
print Z[...,0]
source to share