How To Use Numpy.frompyfunc To Return An Array Of Elements Instead Of Array Of Arrays?

I am using the PLegendre function from the SHTOOLS package. It returns an array of Legendre polynomials for a particular argument. PLegendre(lmax,x) returns an array of Legendre po

Solution 1:

I think it can not be done directly as already pointed out here in the case of np.vectorize which does almost the same thing. Note that your code is not faster then an ordinary for loop by using np.frompyfunc ... the code only looks nicer.

However, what you can do is using np.vstack instead of the list comprehension

a = legendre(3,np.linspace(0,1,4))
np.vstack(a)

Solution 2:

np.frompyfunc is compiled, so I'd have to dig into the source to see exactly what it is doing. But it appears to assume that the func output is an (inscrutable) Python object.

foo1 = np.frompyfunc(np.arange,1,1)
foo2 = np.vectorize(np.arange,otypes='O')

These 2 functions produce the same outputs, though foo1 is faster.

foo1(np.arange(4))

produces arrays with different sizes

array([array([], dtype=int32), array([0]), array([0, 1]), array([0, 1, 2])], dtype=object)

where as foo1(np.ones((4,)) are all the same, and in theory could be stacked.

There's no attempt, during the loop or after, to test whether the objects are arrays (or lists) and whether they can be combined into a single higher dimensional array.

plonser's use of vstack is a good idea. In fact frompyfunc plus vstack is faster than the more common list comprehension plus vstack.

In [54]: timeit np.vstack([np.arange(i) for i in 10*np.ones((10,))])
10000 loops, best of 3: 169 µs per loop

In [55]: timeit np.vstack(foo1(10*np.ones((10,))))
10000 loops, best of 3: 127 µs per loop