Astropy，Numpy：在坐标上应用函数是非常s的

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）之前被初始化。在

一种解决方案可能是在应用函数之前将坐标对象转换为普通的numy数组，丢弃单元信息和元数据（因为单元是同质的）。在

但是，我找不到办法。在

---

我该怎么做？如果转换为普通的numpy数据类型是最好的解决方案，那么如何实现呢？在

谢谢！在

---

最小工作示例：

^{pr2}$

Answer 1

一个解决方案（但不是最好的see编辑）是将astropy坐标转换为numpy数组，然后按正常方式处理numpy。这种转换可以通过分别提取每个坐标分量来完成：

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

（由于生成的数组将具有混合单元，因此我们通过使用.value来丢弃单元。）

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

^{pr2}$

您可以将函数(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要比上面迭代{}快得多。这有一个额外的优势，即保存单位。在

见this answer。在