Plot on a 2D grid a spherically symmetric function

Question

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.

Answer 1

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.

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