Applying Numpy broadcasting on function involving linear algebra

Question

I would like to use numpy broadcasting feature on mathematical function which involves linear algebra (bivariate gaussian distribution without the denominator part) . The Minimal, Reproducible Example of my code is this:

I have the following function

import numpy as np
def  gaussian(x):
    mu = np.array([[2],
                   [2]])
    sigma = np.array([[10, 0],
                      [0, 10]])
    xm = x - mu
    result = np.exp((-1/2) * xm.T @ np.linalg.inv(sigma) @ xm)
    return result

The function assumes that x is a 2x1 array. My aim is to use the function to generate a 2D array where the individual elements are products of the function. I apply this function as follows:

x, y = np.arange(5), np.arange(5)
xLen, yLen = len(x), len(y)
z = np.zeros((yLen, xLen))

for y_index in range(yLen):
        for x_index in range(xLen):
            element = np.array([[x[x_index]],
                                [y[y_index]]])
            result = gaussian(element)
            z[y_index][x_index] = result

This works but as you can see, I use two for loops for indexing. I am aware that this is bad practise and when working with bigger arrays it is terribly slow. I would like to solve this with numpy broadcasting feature. I attempted the following code:

X, Y = np.meshgrid(x, y, indexing= 'xy')
element = np.array([[X],
                    [Y]])
Z = gaussian(element)

But I am getting this error: ValueError: operands could not be broadcast together with shapes (2,1,5,5) (2,1) for line xm = x - mu of the function. I understand this error to a certain extent.

In addition, even if I solved this I would be getting another error: ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 5 is different from 2) for the result = np.exp((-1/2) * xm.T @ np.linalg.inv(sigma) @ xm) line of the fuction. Again, I understand why. xm would no longer be 2x1 array and multiplying it with sigma , which is 2x2, would not work.

Does anyone have a suggestion on how to modify my function so the broadcasting implementation works?

Answer 1

The following may work. Two things to note:

• I'm using np.einsum for the vector-matrix-vector multiplication. There may be faster ways, but this can nicely handle the other dimensions that are to be broadcasted.
• From my experience, for larger arrays, using broadcasting with 3+ dimensions, things may not be faster than a simple nested loop. I haven't dug into this: perhaps the calculations were done on the wrong dimension (the column- versus row-wise issue), which would slow things down. So perhaps by tweaking or playing around with the dimension order, things could be sped up

Setup code

nx, ny = 5, 5
x, y = np.arange(nx), np.arange(ny)
X, Y = np.meshgrid(x, y, indexing= 'xy')
element = np.array([[X],
                    [Y]])
# Stack X and Y into a nx x ny x 2 array
XY = np.dstack([X, Y])

New function

def  gaussian(x):
    # Note that I have removed the extra dimension: 
    # mu is a simple array of shape (2,)
    # This is no problem, since we're using einsum
    # for the matrix multiplication
    mu = np.array([2, 2])
    sigma = np.array([[10, 0],
                      [0, 10]])
    # Broadcast xm to x's shape: (nx, ny, 2)
    xm = x - mu[..., :]
    invsigma = np.linalg.inv(sigma)
    # Compute the (double) matrix multiplication
    # Leave the first two dimension (ab) alone
    # The other dimensions will sum up to a single scalar
    # and thus only the ab dimensions are there in the output
    alpha = np.einsum('abi,abj,ji->ab', xm, xm, invsigma)
    result = np.exp((-1/2) * alpha)
    # The shape of result is (nx, ny)
    return result

And then call:

gaussian(XY)

Obviously, please double check. I did one brief check, which seems to be correct, but transcription errors may eg have swapped dimensions.

Answer 2

So a (2,1) input returns a (1,1) result:

In [83]: gaussian(np.ones((2,1)))
Out[83]: array([[0.90483742]])

Adding some leading dimensions:

In [84]: gaussian(np.ones((3,4,2,1)))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [84], in <cell line: 1>()
----> 1 gaussian(np.ones((3,4,2,1)))

Input In [80], in gaussian(x)
      4 sigma = np.array([[10, 0],
      5                   [0, 10]])
      6 xm = x - mu
----> 7 result = np.exp((-1/2) * xm.T @ np.linalg.inv(sigma) @ xm)
      8 return result

ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 3)

x-mu works because (3,4,2,1) broadcasts with (2,1)

The error occurs in (-1/2) * xm.T @ np.linalg.inv(sigma)

np.linalg.inv(sigma) is (2,2)

xm is (3,4,2,1) , so its transpose is (1,2,4,3).

If instead the arrays are (3,4,1,2) @ (2,2) @ (3,4,2,1) the result should be (3,4,1,1).

So let's refine the transpose:

def  gaussian(x):
    mu = np.array([[2],
                   [2]])
    sigma = np.array([[10, 0],
                      [0, 10]])
    xm = x - mu
    xmt =xm.swapaxes(-2,-1)
    result = np.exp((-1/2) * xmt @ np.linalg.inv(sigma) @ xm)
    return result

Now it works for both the original (2,1), and any other (n,m,2,1) shape:

In [87]: gaussian(np.ones((3,4,2,1))).shape
Out[87]: (3, 4, 1, 1)

In [88]: gaussian(np.ones((2,1))).shape
Out[88]: (1, 1)