在二维numpy数组上广播函数

时间:2015-05-09 14:16:54

标签: python numpy vectorization

我希望通过在 numpy 数组上计算一次函数来提高代码的速度,而不是 for循环是{{3}的函数python库。如果我有以下功能:

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

虽然pos_arcsec24000x2的二维数组,但shear_z是一个带有24000元素的一维数组。 主要问题是我想在M200=np.arange(13., 16., 0.01)conc = np.arange(3, 10, 0.01)的网格上计算此函数。我不知道如何在M200conc上广播这个二维数组的函数。运行代码需要很多。我正在寻找加速这些计算的最佳方法。

2 个答案:

答案 0 :(得分:1)

pos是形状(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)

示例:

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])

答案 1 :(得分:1)

使用numpy.linalg.norm

如果您有阵列:

import numpy as np
import numpy.linalg as la

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

然后你可以通过

确定结果向量的大小(sqrt(a ^ 2 + b ^ 2))
b = np.sqrt(la.norm(a, axis=1)

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