如何加速匹配相似亮度与距离的星系数据的Python代码？

问题描述 #

我需要对比模拟星系（S）与观测星系（O），从观测星系列表中为每个模拟星系找出亮度（Br）和距离（D）都在特定范围内的相似星系集合。目前已实现可运行的代码，但运行耗时约9-10小时，希望找到更高效的优化方法。

原实现代码 #

list1= flux_hb
list2= flux_inp_ha

flux_MC=[[]]*len(list1) # map flux from MAMBO to COSMOS (mambo--> (list of cosmos galaxies))
noise=[[]]*len(list1)
z_MC=[[]]*len(list1)

index=np.zeros(len(list1))
flux_ha_temp=np.zeros(len(list1))

for i in tqdm(range(len(list1))):
    
    temp_flux_diff = abs(list1[i] - list2)/list1[i] #Matrix subtraction
    temp_z_diff = abs(z_geo[i] - zinp)/z_geo[i]
    
    for j in range(len(temp_flux_diff)):
        
        if (temp_flux_diff[j]<0.01 and temp_z_diff[j]<0.01):
            flux_MC[i].append(list2[j])
            noise[i].append(snr_ha[j])
            z_MC[i].append(zinp[j])

原实现思路 #

• 创建空列表存储筛选后的星系信息
• 遍历每个模拟星系，计算其与所有观测星系的亮度、距离差值比例
• 遍历所有观测星系，筛选出差值比例小于阈值的星系并存入对应列表

优化更新（更新1） #

根据建议将第二层循环替换为numpy的向量化操作，修改后代码如下：

index = np.where(np.logical_and(temp_flux_diff<0.01, temp_z_diff<0.01))[0]

flux_MC[i].append(flux_inp_ha[index])
noise[i].append(snr_ha[index])
z_MC[i].append(zinp[index])

修改后代码运行耗时缩短至约1小时，效率大幅提升。

内容的提问来源于stack exchange，提问作者Aryan Bansal