How to apply function which returns vector to each numpy array element (and get array with higher dimension)

Question

Let's write it directly in code

Note: I edited mapper (original example use x -> (x, 2 * x, 3 * x) just for example), to generic blackbox function, which cause the troubles.

import numpy as np

def blackbox_fn(x): #I can't be changed!
    assert np.array(x).shape == (), "I'm a fussy little function!"
    return np.array([x, 2*x, 3*x])

# let's have 2d array
arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)

# each element should be mapped to vector
def mapper(x, blackbox_fn):
    # there is some 3rdparty non-trivial function, returning np.array
    # in examples returns np.array((x, 2 * x, 3 * x))
    # but still this 3rdparty function operates only on scalar values
    return vectorized_blackbox_fn(x)

So for input 2d array

array([[0, 1],
       [2, 3]], dtype=uint8)

I would like to get 3d array

array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]], dtype=uint8)

I can write naive algorithm using for loop

# result should be 3d array, last dimension is same as mapper result size
arr3d = np.empty(arr2d.shape + (3,), dtype=np.uint8)
for y in range(arr2d.shape[1]):
    for x in xrange(arr2d.shape[0]):
        arr3d[x, y] = mapper(arr2d[x, y])

But is seems quite slow for large arrays. I know there is np.vectorize, but using

np.vectorize(mapper)(arr2d)

not work, because of

ValueError: setting an array element with a sequence.

(seems that vectorize can't change dimension) Is there some better (numpy idiomatic and faster) solution?

Actually, vectorize will not help you as it is essentially also looping. See the notes in the documentation: "The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.". Without knowing the "general non-trivial function" it will be probably be hard to help you... — jotasi
– jotasi, Commented Sep 1, 2017 at 12:04
does blackbox_fn always output an array of the same size? — Daniel F
– Daniel F, Commented Sep 1, 2017 at 13:10
You'll probably need some sort of np.nditer solution, but I'm not an expert there. — Daniel F
– Daniel F, Commented Sep 1, 2017 at 13:36

hpaulj · Accepted Answer · 2017-09-01 16:33:57Z

np.vectorize with the new signature option can handle this. It doesn't improve the speed, but makes the dimensional bookkeeping easier.

In [159]: def blackbox_fn(x): #I can't be changed!
     ...:     assert np.array(x).shape == (), "I'm a fussy little function!"
     ...:     return np.array([x, 2*x, 3*x])
     ...:

The documentation for signature is a bit cryptic. I've worked with it before, so made a good first guess:

In [161]: f = np.vectorize(blackbox_fn, signature='()->(n)')
In [162]: f(np.ones((2,2)))
Out[162]: 
array([[[ 1.,  2.,  3.],
        [ 1.,  2.,  3.]],

       [[ 1.,  2.,  3.],
        [ 1.,  2.,  3.]]])

With your array:

In [163]: arr2d = np.array(list(range(4)), dtype=np.uint8).reshape(2, 2)
In [164]: f(arr2d)
Out[164]: 
array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]])
In [165]: _.dtype
Out[165]: dtype('int32')

The dtype is not preserved, because your blackbox_fn doesn't preserve it. As a default vectorize makes a test calculation with the first element, and uses its dtype to determine the result's dtype. It is possible to specify return dtype with the otypes parameter.

It can handle arrays other than 2d:

In [166]: f(np.arange(3))
Out[166]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6]])
In [167]: f(3)
Out[167]: array([3, 6, 9])

With a signature vectorize is using a Python level iteration. Without a signature it uses np.frompyfunc, with a bit better performance. But as long as blackbox_fn has to be called for element of the input, we can't improve the speed by much (at most 2x).

np.frompyfunc returns a object dtype array:

In [168]: fpy = np.frompyfunc(blackbox_fn, 1,1)
In [169]: fpy(1)
Out[169]: array([1, 2, 3])
In [170]: fpy(np.arange(3))
Out[170]: array([array([0, 0, 0]), array([1, 2, 3]), array([2, 4, 6])], dtype=object)
In [171]: np.stack(_)
Out[171]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6]])
In [172]: fpy(arr2d)
Out[172]: 
array([[array([0, 0, 0]), array([1, 2, 3])],
       [array([2, 4, 6]), array([3, 6, 9])]], dtype=object)

stack can't remove the array nesting in this 2d case:

In [173]: np.stack(_)
Out[173]: 
array([[array([0, 0, 0]), array([1, 2, 3])],
       [array([2, 4, 6]), array([3, 6, 9])]], dtype=object)

but we can ravel it, and stack. It needs a reshape:

In [174]: np.stack(__.ravel())
Out[174]: 
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6],
       [3, 6, 9]])

Speed tests:

In [175]: timeit f(np.arange(1000))
14.7 ms ± 322 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [176]: timeit fpy(np.arange(1000))
4.57 ms ± 161 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [177]: timeit np.stack(fpy(np.arange(1000).ravel()))
6.71 ms ± 207 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [178]: timeit np.array([blackbox_fn(i) for i in np.arange(1000)])
6.44 ms ± 235 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Having your function return a list instead of any array might make reassembling the result easier, and maybe even faster

def foo(x):
    return [x, 2*x, 3*x]

or playing about with the frompyfunc parameters;

def foo(x):
    return x, 2*x, 3*x   # return a tuple
In [204]: np.stack(np.frompyfunc(foo, 1,3)(arr2d),2)
Out[204]: 
array([[[0, 0, 0],
        [1, 2, 3]],

       [[2, 4, 6],
        [3, 6, 9]]], dtype=object)

10x speed up - I'm surprised:

In [212]: foo1 = np.frompyfunc(foo, 1,3)
In [213]: timeit np.stack(foo1(np.arange(1000)),1)
428 µs ± 17.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Nils Werner · Accepted Answer · 2017-09-01 12:25:26Z

You can use basic NumPy broadcasting for these kind of "outer products"

np.arange(3)[:, None] * np.arange(2)
# array([[0, 0],
#        [0, 1],
#        [0, 2]])

In your case it would be

def mapper(x):
    return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))

note the .transpose() is only needed if you specifically need the new axis to be at the end.

And it is almost 3x as fast as stacking 3 separate multiplications:

def mapper(x):
    return (np.arange(3)[:, None, None] * x).transpose((1, 2, 0))


def mapper2(x):
    return np.stack((x, 2 * x, 3 * x), axis = -1)

a = np.arange(30000).reshape(-1, 30)

%timeit mapper(a)   # 48.2 µs ± 417 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit mapper2(a)  # 137 µs ± 3.57 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

zipa · Accepted Answer · 2017-09-01 12:53:42Z

0

I might be getting this wrong, but comprehension does the job:

a = np.array([[0, 1],
     [2, 3]])


np.array([[[j, j*2, j*3] for j in i] for i in a ])
#[[[0 0 0]
#  [1 2 3]]
#
# [[2 4 6]
#  [3 6 9]]]

answered Sep 1, 2017 at 12:53

zipa

28k6 gold badges45 silver badges62 bronze badges

1 Comment

farincz Over a year ago

I just updated question, because it wasn't clear what is the point of question.

Collectives™ on Stack Overflow

How to apply function which returns vector to each numpy array element (and get array with higher dimension)

3 Answers 3

Comments

Comments

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related