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

任何帮助都会很感激。我在其他语言方面是一个不错的程序员，但我以前从来没有真正担心过Python的效率。我现在选择在Python中这样做的原因是，我正在使用FITSIO和PGPLOT的fortran，以及一些来自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']