# Plot on a 2D grid a spherically symmetric function

I am trying to compute a spherically symmetric function on a 2D grid. The simple solution that I came up with is the following:

Grid_side = 10
N_pix = 100

function = lambda r : 1.0 / r**2.0

grid_x = np.linspace(-Grid_side, Grid_side, N_pix)
grid_y = np.linspace(-Grid_side, Grid_side, N_pix)

SphericallySymmetric_function_GRID = np.zeros((N_pix, N_pix))

for i in range(N_pix):

for j in range(N_pix):

SphericallySymmetric_function_GRID[i,j] = function(np.sqrt(grid_x[i]**2.0 + grid_y[j]**2.0))


I was wondering if there is a better way to do the following, taking advantage of Numpy.

No need to iterate if the function works with arrays, and you use broadcasting to populate a grid:

def fn(X,Y):
r = np.sqrt(X**2+Y**2)
return 1/r**2
fn(grid_x[:,None], grid_y[None,:])


np.meshgrid and np.mgrid (or np.ogrid) can also be used to form the X,Y grid.

fn(*np.meshgrid(grid_x, grid_y))
fn(*np.ogrid[-Grid_side:Grid_side:N_pix*1j, -Grid_side:Grid_side:N_pix*1j])


Look up their docs, and practice. Practice with simple things like grid_x[:,None]+grid_y as well.

• I think it's clearer to write np.newaxis instead of None. Dec 18 '16 at 19:28
• Do you also like it in this slicing notation: grid_x[slice(np.newaxis, -1, np.newaxis)]? :) Dec 18 '16 at 20:18
• No, that would be misleading since a new axis is not created. grid_x[:-1] would be better. Dec 18 '16 at 20:27