broadcasting a function on a 2-dimensional numpy array

问题

I would like to improve the speed of my code by computing a function once on a numpy array instead of a for loop is over a function of this python library. If I have a function as following:

import numpy as np
import galsim
from math import *
M200=1e14
conc=6.9
def func(M200, conc):
    halo_z=0.2
    halo_pos =[1200., 3769.7]
    halo_pos = galsim.PositionD(x=halo_pos_arcsec[0],y=halo_pos_arcsec[1])
    nfw      = galsim.NFWHalo(mass=M200, conc=conc, redshift=halo_z,halo_pos=halo_pos, omega_m = 0.3, omega_lam =0.7)
    for i in range(len(shear_z)):
       shear_pos=galsim.PositionD(x=pos_arcsec[i,0],y=pos_arcsec[i,1])
       model_g1, model_g2  = nfw.getShear(pos=self.shear_pos, z_s=shear_z[i])
    l=np.sum(model_g1-model_g2)/sqrt(np.pi)
    return l

While pos_arcsec is a two-dimensional array of 24000x2 and shear_z is a 1D array with 24000 elements as well. The main problem is that I want to calculate this function on a grid where M200=np.arange(13., 16., 0.01) and conc = np.arange(3, 10, 0.01). I don't know how to broadcast this function to be estimated for this two dimensional array over M200 and conc. It takes a lot to run the code. I am looking for the best approaches to speed up these calculations.

回答1:

This here should work when pos is an array of shape (n,2)

import numpy as np

def f(pos, z):
    r=np.sqrt(pos[...,0]**2+pos[...,1]**2)
    return np.log(r)*(z+1)

Example:

z = np.arange(10)
pos = np.arange(20).reshape(10,2)

f(pos,z)
# array([  0.        ,   2.56494936,   5.5703581 ,   8.88530251,
#         12.44183436,  16.1944881 ,  20.11171117,  24.17053133,
#         28.35353608,  32.64709419])

回答2:

Use numpy.linalg.norm

If you have an array:

import numpy as np
import numpy.linalg as la

a = np.array([[3, 4], [5, 12], [7, 24]])

then you can determine the magnitude of the resulting vector (sqrt(a^2 + b^2)) by

b = np.sqrt(la.norm(a, axis=1)

>>> print b
array([  5.,  15.  25.])

来源：https://stackoverflow.com/questions/30140966/broadcasting-a-function-on-a-2-dimensional-numpy-array