Multidimensional Matrix Operations with Numpy Broadcasting

Illustration of stacked tensors — Tensor manipulation…

Numpy's broadcasting feature can be somewhat confusing for new users of this library, but it allows for very clean, elegant and FUN coding, so it is definitely worth the effort of getting used to. In this short article, I wanted to show a nice application of broadcasting to save some for loops and even computation time.

Let's start with a simple case.

The basic concept is that when adding or multiplying two vectors of sizes (m,1) and (1,m), numpy will broadcast (duplicate the vector) so that it allows the calculation. For example, multiplying a vector [1,2,3,4,…10] with a transposed version of itself will yield the multiplication table. For example:

A (3,1) vector multiplied by its (1,3) transpose produces a 3 by 3 multiplication table — Broadcasting a vector into a matrix. A miniature multiplication table.

In this example, we multiply a one-dimensional vector (V) of size (3,1) and the transposed version of it, which is of size (1,3), and get back a (3,3) matrix, which is the outer product of V.

If you still find this confusing, the next illustration breaks down the process into 2 steps, making it clearer:

Both vectors are first duplicated into 3 by 3 matrices, then multiplied element-wise — Broadcasting in slow motion.

You can think of broadcasting as simply duplicating both our vectors into a (3,3) matrix, and then performing element-wise multiplication.

We have just broadcast a 1-dimensional array into a 2-dimensional matrix. However, we could use this to broadcast a 2-dimensional array (or matrix) into a 3-dimensional one (a tensor). Just imagine that you needed to perform V*V.T for multiple vectors of the same length. You could simply stack them together, and get all the results in one pass.

Case study

In this case, I am looking to generate a Euclidean distance matrix for the iris data set. The Euclidean equation is:

d(p, q) = \sqrt{\sum_{i}^{n} (q_i - p_i)^2}

Obtaining the table could obviously be performed using two nested for loops:

result = []
for i,_ in enumerate(data):
    for j,_ in enumerate(data):
        result[i,j] = np.sqrt(np.sum((data[i,:]-data[j,:])**2))

However, it can also be performed using matrix operations (which are both about 100 times faster, and much cooler).

We take the data set (represented as 'A' in the illustration),

A tall rectangle labeled A, with m rows and n columns — A rectangle representing the iris data set, in which case m=150 and n=4.

create a 3D copy of it, and rotate it to the orientation shown as 'B' in the illustration.

The data set as a 3D slab labeled B, rotated — We can use numpy's rot90 function to rotate a matrix. See code below.

We then create another copy and rotate it as represented by 'C'.

A second copy of the data set rotated flat, labeled C — Copy and rotate again.

B − C will generate (via broadcasting!) a 3D cube ('D'), sized (m,m,n), which represents the calculation \((q_i - p_i)\) for each instance (m) and feature (n) of the original data set vectors (n=4 in the iris data set).

A cube labeled D of size m by m by n, summed along the feature axis — Notice that we are summing up the features for all instances at once!

Then, we simply raise the entire cube to the power of 2, sum the result along the axis representing the features (see illustration above), and take the square root.

The result is the distance matrix ('E')!

A square m by m matrix labeled E — This is the Euclidean distance matrix. It might seem like it only contains the letter 'E', but in fact it holds the distance between all instance pairs. Twice. The diagonal is the distance between every instance and itself, and if it's not equal to zero, then you should double-check your code…

If you got lost in all these matrix somersaults, go back and look at the equation for calculating the Euclidean matrix. All we did was solve this equation for several vectors simultaneously.

Code

import numpy as np
from sklearn.datasets import load_iris

# Load data:
data, _ = load_iris(return_X_y=True)

def euc_matrix(A):
    # generate distance matrix:
    B = np.rot90(A[:,:,None],1,(1,2)) #[:,:,None] is needed to add a dimension
    C = np.rot90(B,1,(0,1))
    return np.sqrt(np.sum((B-C)**2, axis=2))

mat = euc_matrix(data)

Summary

Broadcasting is a very clean and efficient method whenever you perform the same calculation on several instances. As an exercise, I encourage you to try using broadcasting on two 3D arrays (let's say you have n vectors, belonging to k classes). This means you will need to broadcast into the 4th dimension. It is impossible to imagine, of course, but implementing it follows the same logic.

Michael Chein, data scientist. Home · More posts