Question

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

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

插图：

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)

Answer 1

一个解决方案（但不是最好看的编辑）是将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的速度比在stack([coords.rho, coords.phi, coords.z])上迭代该函数更快。这具有保留单位的附加优势。

请参阅this answer。

Astropy，Numpy：在坐标上应用函数非常慢

1 个答案: