CODEWITHSUNDEEP
pythonarraysnumpy
lewisxy
lewisxy

Reputation: 323

numpy broadcast of vectorized function

I have a function f that compute the inner-product-like value of 2 vectors (the function is symmetrical, so f(x, y) = f(y, x)). It take 2 1d array of same size and output a single value. Now I have a data matrix X of shape (n, d) (each row represents an input vector), and I would like to compute a matrix K with shape (n, n) such that K[i][j] = f(X[i], X[j]). One way to do this is to use the following code.

import numpy as np
# ... some other code
K = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        K[i][j] = f(X[i], X[j])

This works, but it's neither elegant nor performant. Another way to do this is using vectorize from numpy

import numpy as np
# ... some other code
vf = np.vectorize(f, signature="(n),(n)->()")

# this does not work (output shape (n,))
#K = vf(X, X)

# this works (output shape (n, n))
K = vf(X, X[: np.newaxis])

# this also works (output shape (n, n), result same as above)
#K = vf(X[: np.newaxis], X)

If we put in an array of (n, d) and another array of (n, d), it will output an 1 dimensional array of (n,), essentially taking each of the n rows of both matrix as input and thus produce n output. However, if we transform it to X[: np.newaxis], which has shape (n, 1, d), it works, but why it works?

I already checked this and this, but still can't figure out.

Upvotes: 0

Views: 369

Answers (2)

Akshay Sehgal
Akshay Sehgal

Reputation: 19322

Another example in addition to @hpaulj's answer

Broadcasting can be used to vectorize your code efficiently. Let's take an example. I have an array (10,2) and I want to take the Euclidean distance between each of the points.

X = np.random.uniform(-5,5,(10,2))

One way to do this is to broadcast my complete formula for Euclidean distance using NumPy function which supports broadcasting.

np.sqrt(np.sum(np.square(np.subtract(X[:,None,:], X[None,:,:])), axis=-1))

#This returns a (10,10) matrix, with euclidean distance of each of the 10 points to the other

Similarly, another thing you could do is vectorize your function to behave similar to a numpy function and support broadcasting.

def euclidean(point1, point2):
    return ((point1[0]-point2[0])**2 + (point1[1]-point2[1])**2)**(1/2)

euclidean_v = np.vectorize(euclidean, signature="(n),(n)->()")

euclidean_v(X[:,None,:], X[None,:,:])

#Again returns a (10,10)

In both cases -

         |--- (10,1,2)--|
(10,2)---|              |---(10,10,2)--->(10,10)
         |--- (1,10,2)--|

Upvotes: 1

hpaulj
hpaulj

Reputation: 231385

We can illustrate this broadcasting with add:

In [111]: x = np.arange(12).reshape(3,4)
In [112]: x
Out[112]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [113]: x + x[:,None]
Out[113]: 
array([[[ 0,  2,  4,  6],
        [ 4,  6,  8, 10],
        [ 8, 10, 12, 14]],

       [[ 4,  6,  8, 10],
        [ 8, 10, 12, 14],
        [12, 14, 16, 18]],

       [[ 8, 10, 12, 14],
        [12, 14, 16, 18],
        [16, 18, 20, 22]]])
In [114]: _.shape
Out[114]: (3, 3, 4)
In [115]: (x + x[:,None]).sum(axis=-1)
Out[115]: 
array([[12, 28, 44],
       [28, 44, 60],
       [44, 60, 76]])

The broadcasting goes: (3,4) and (3,1,4) => (1,3,4) and (3,1,4) => (3,3,4). Then we sum on the last axis to get (3,3)

Doing the same thing as you did:

In [116]: vf = np.vectorize(lambda a,b:(a+b).sum(), signature="(n),(n)->()")
In [117]: vf(x, x[:,None])
Out[117]: 
array([[12, 28, 44],
       [28, 44, 60],
       [44, 60, 76]])

But as noted, np.vectorize is not a speed tool, especially when using signature:

In [118]: timeit (x + x[:,None]).sum(axis=-1)
13.8 µs ± 363 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [119]: timeit vf(x, x[:,None])
278 µs ± 7.68 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Without adding the dimension to the 2nd x:

In [120]: vf(x, x)
Out[120]: array([12, 44, 76])
In [121]: (x + x).sum(axis=-1)
Out[121]: array([12, 44, 76])

Upvotes: 1

Related Questions