我在单个天体坐标对象中包含大量坐标。我想对每个坐标并行应用一个函数并生成一个相同形状的输出数组——但这很慢。
(在我的例子中，该函数是一个采用银河同心坐标并输出与空间中该点相关的“亮度”的模型。)
插图:
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 而不是在 stack([coords.rho, coords.phi, coords.z]) 上迭代该函数作为多于。这具有保留单位的额外优势。
参见 this answer .
关于python - Astropy，Numpy : Applying function over coordinates is very slow，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/47468165/