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.

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

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.