python – Astropy,Numpy：在坐标上应用函数非常慢

2019年7月28日 308次阅读

我有一个包含在单个天文坐标对象中的大量坐标.我想将函数并行应用于每个坐标并生成相同形状的输出数组 – 但这很慢.

(在我的例子中,该函数是一个采用半中心坐标并输出与空间中该点相关的“亮度”的模型.)

插图：

In [339]: type(data)
Out[339]: astropy.coordinates.builtin_frames.galactocentric.Galactocentric

In [340]: data.shape, data.size              # Not that big, really
Out[340]: ((21, 21, 31), 13671)

In [341]: data[0,0,0]                        # An example of a single coordinate
Out[341]: 
<Galactocentric Coordinate (galcen_distance=8.3 kpc, galcen_ra=266d24m18.36s, galcen_dec=-28d56m10.23s, z_sun=27.0 pc, roll=0.0 deg): (rho, phi, z) in (kpc, deg, kpc)
    ( 8.29995608,  180.,  0.027)>

In [342]: func = vectorize(lambda coord: 0)  # Dummy function

In [343]: %time func(data).shape
CPU times: user 33.2 s, sys: 88.1 ms, total: 33.3 s
Wall time: 33.4 s
Out[343]: (21, 21, 31)

我怀疑这很慢,因为在每次迭代时,新的坐标对象在被传递到矢量化函数(discussion)之前被初始化.

解决方案可能是在应用函数之前将坐标对象转换为普通的numpy数组,丢弃单元信息和元数据(因为单元是同质的).

但是,我找不到办法做到这一点.

我该怎么做呢？如果转换为vanilla numpy数据类型是最佳解决方案,那么如何实现？

谢谢！

最小的工作示例：

from numpy import *
from astropy import units as u
from astropy.coordinates import Galactocentric

# Generate lots of coordinates
x = linspace(0, 1, 1e3)*u.pc
data = Galactocentric(x=x, y=0*u.pc, z=0*u.pc)

@vectorize
def func(coord):
    '''ultimately in terms of coord.x, coord.y, coord.z...'''
    return 0

# timeit
func(data)

最佳答案 一个解决方案(但不是最好看的编辑)是将astropy坐标转换为numpy数组,然后像numpy一样正常进行.可以通过分别提取每个坐标组件来完成此转换：

coords_np = stack([coords.rho, coords.phi, coords.z]).value

(由于生成的数组将具有混合单位,因此我们通过获取.value来丢弃单位.)

现在,坐标三元组(rho,phi,z)沿着新轴,

>>> coords_np[:,0,0,0]
array([  <rho>,  <phi>,    <z>])

并且您可以应用您的函数(rho,phi,z) – > x到coords_np如下：

scalar_field = apply_along_axis(func, 0, coords_np)

此结果相当于执行func(coords)(直接在星座坐标上),但速度更快.

编辑：如果可能,通过矢量化函数完全避免apply_along_axis,而不是将其应用于每个坐标.例如,如果函数类似于lambda rho,phi,z：rho ** 2 z ** 2,那么简单地计算coords就会快得多.rho ** 2 coords.z ** 2比迭代它更快堆栈功能([coords.rho,coords.phi,coords.z])如上所述.这具有保留单元的附加优点.

见this answer.

点赞