python - astropy.io 适合大表的高效元素访问

Question

我正在尝试使用 Python 和 astropy.io 从 FITS 文件中的二进制表中提取数据。该表包含一个事件数组，其中包含超过 200 万个事件。我想要做的是将某些事件的 TIME 值存储在一个数组中，这样我就可以对该数组进行分析。我遇到的问题是，在 fortran(使用 FITSIO)中，相同的操作在慢得多的处理器上可能需要几秒钟，而在 Python 中使用 astropy.io 的完全相同的操作需要几分钟。我想知道瓶颈到底在哪里，是否有更有效的方法来访问各个元素以确定是否将每个时间值存储在新数组中。这是我到目前为止的代码:

from astropy.io import fits

minenergy=0.3
maxenergy=0.4
xcen=20000
ycen=20000
radius=50

datafile=fits.open('datafile.fits')
events=datafile['EVENTS'].data


datafile.close()

times=[]

for i in range(len(events)):
    energy=events['PI'][i]
    if energy<maxenergy*1000:
        if energy>minenergy*1000:
            x=events['X'][i]
            y=events['Y'][i]
            radius2=(x-xcen)*(x-xcen)+(y-ycen)*(y-ycen)
            if radius2<=radius*radius:
                times.append(events['TIME'][i])

print times

如有任何帮助，我们将不胜感激。我是其他语言的 ok 程序员，但我以前真的不必担心 Python 的效率。我现在选择在 Python 中执行此操作的原因是我将 fortran 与 FITSIO 和 PGPLOT 以及 Numerical Recipes 中的一些例程一起使用，但是无法说服我在这台机器上使用的新的 fortran 编译器生成正常工作的程序(有一些 32 位与 64 位的问题，等等)。 Python 似乎具有我需要的所有功能(FITS I/O、绘图等)，但如果访问列表中的各个元素需要很长时间，我将不得不寻找其他解决方案。

非常感谢。

Answer 1

您需要使用 numpy 向量运算来执行此操作。如果没有像 numba 这样的特殊工具，在 Python 中像你所做的那样进行大循环总是很慢，因为它是一种解释型语言。您的程序应该更像:

energy = events['PI'] / 1000.
e_ok = (energy > min_energy) & (energy < max_energy)
rad2 = (events['X'][e_ok] - xcen)**2 + (events['Y'][e_ok] - ycen)**2
r_ok = rad2 < radius**2
times = events['TIMES'][e_ok][r_ok]

这应该具有与 Fortran 相当的性能。您还可以过滤整个事件表，例如:

events_filt = events[e_ok][r_ok]
times = events_filt['TIMES']