python - 在二维 numpy 数组上广播函数

Question

我想通过在 numpy 数组而不是 for 循环 上计算一个函数来提高我的代码的速度是 this 的函数 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_arcsec 是一个 24000x2 的二维数组，而 shear_z 是一个包含 24000 元素的一维数组以及。主要问题 是我想在 M200=np.arange(13., 16., 0.01) 和 conc = np 的网格上计算此函数.arange(3, 10, 0.01).我不知道如何通过 M200 和 conc 广播此函数以估计此二维数组。运行代码需要很多时间。我正在寻找加速这些计算的最佳方法。

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