Python 3.7: Modelling a 2D Gaussian equation using a Numpy meshgrid and arrays without iterating through each point

Question

I am currently trying to write my own 2D Gaussian function as a coding exercise, and have been able to create the following script:

import numpy as np
import matplotlib.pyplot as plt

def Gaussian2D_v1(coords=None,  # x and y coordinates for each image.
                  amplitude=1,  # Highest intensity in image.
                  xo=0,  # x-coordinate of peak centre.
                  yo=0,  # y-coordinate of peak centre.
                  sigma_x=1,  # Standard deviation in x.
                  sigma_y=1,  # Standard deviation in y.
                  rho=0,  # Correlation coefficient.
                  offset=0):  # Offset from zero (background radiation).
    x, y = coords

    xo = float(xo)
    yo = float(yo)

    # Create covariance matrix
    mat_cov = [[sigma_x**2, rho * sigma_x * sigma_y],
               [rho * sigma_x * sigma_y, sigma_y**2]]
    mat_cov = np.asarray(mat_cov)
    # Find its inverse
    mat_cov_inv = np.linalg.inv(mat_cov)

    G_array = []
    # Calculate pixel by pixel
    # Iterate through row last
    for i in range(0, np.shape(y)[0]):
        # Iterate through column first
        for j in range(0, np.shape(x)[1]):
            mat_coords = np.asarray([[x[i, j]-xo],
                                     [y[i, j]-xo]])
            G = (amplitude * np.exp(-0.5*np.matmul(np.matmul(mat_coords.T,
                                                             mat_cov_inv),
                                                   mat_coords)) + offset)
            G_array.append(G)

    G_array = np.asarray(G_array)
    G_array = G_array.reshape(64, 64)
    return G_array.ravel()


coords = np.meshgrid(np.arange(0, 64), np.arange(0, 64))
model_1 = Gaussian2D_v1(coords,
                        amplitude=20,
                        xo=32,
                        yo=32,
                        sigma_x=6,
                        sigma_y=3,
                        rho=0.8,
                        offset=20).reshape(64, 64)

plt.figure(figsize=(5, 5)).add_axes([0,
                                     0,
                                     1,
                                     1])
plt.contourf(model_1)

The code as it is works, but as you can see, I am currently iterating through the mesh grid one point at a time, and appending each point to a list, which is then converted to an array and re-shaped to give the 2D Gaussian distribution.

How can I modify the script to forgo using a nested "for" loop and have the program consider the whole meshgrid for matrix calculations? Is such a method possible?

Thanks!

Answer 1

Of course there is a solution, numpy is all about array operations and vectorization of the code! np.matmul can take args with more than 2 dimensions and apply the matrix multiplication on the last two axes only (and this calculation in parallel over the others axes). However, making sure of the right axes order can get tricky.

Here is your edited code:

import numpy as np
import matplotlib.pyplot as plt

def Gaussian2D_v1(coords,  # x and y coordinates for each image.
                  amplitude=1,  # Highest intensity in image.
                  xo=0,  # x-coordinate of peak centre.
                  yo=0,  # y-coordinate of peak centre.
                  sigma_x=1,  # Standard deviation in x.
                  sigma_y=1,  # Standard deviation in y.
                  rho=0,  # Correlation coefficient.
                  offset=0):  # Offset from zero (background radiation).
    x, y = coords

    xo = float(xo)
    yo = float(yo)

    # Create covariance matrix
    mat_cov = [[sigma_x**2, rho * sigma_x * sigma_y],
               [rho * sigma_x * sigma_y, sigma_y**2]]
    mat_cov = np.asarray(mat_cov)
    # Find its inverse
    mat_cov_inv = np.linalg.inv(mat_cov)

    # PB We stack the coordinates along the last axis
    mat_coords = np.stack((x - xo, y - yo), axis=-1)

    G = amplitude * np.exp(-0.5*np.matmul(np.matmul(mat_coords[:, :, np.newaxis, :],
                                                    mat_cov_inv),
                                          mat_coords[..., np.newaxis])) + offset
    return G.squeeze()



coords = np.meshgrid(np.arange(0, 64), np.arange(0, 64))
model_1 = Gaussian2D_v1(coords,
                        amplitude=20,
                        xo=32,
                        yo=32,
                        sigma_x=6,
                        sigma_y=3,
                        rho=0.8,
                        offset=20)
plt.figure(figsize=(5, 5)).add_axes([0, 0, 1, 1])
plt.contourf(model_1)

So, the equation is exp(-0.5 * (X - µ)' Cinv (X - µ) ), where X is our coordinate matrix, µ the mean (x0, y0) and Cinv the inverse covariance matrix (and ' is a transpose). In the code, I stack both meshgrids to a new matrix so that: mat_coords has a shape of (Ny, Nx, 2). In the first np.matmul call, I add a new axis so that the shapes go like :(Ny, Nx, 1, 2) * (2, 2) = (Ny, Nx, 1, 2). As you see, the matrix multiplication is done on the two last axes, in parallel on the other. Then, I add a new axis so that: (Ny, Nx, 1, 2) * (Ny, Nx, 2, 1) = (Ny, Nx, 1, 1). The np.squeeze() call returns a version without the two last singleton axes.