天文网格单元掩码高效生成：替代循环解决内存占用问题

拥有含1200万个星系的天区坐标目录，包含ra（赤经）、dec（赤纬，垂直视线方向）与红移（沿视线方向）数据：

• 垂直视线方向用healpy工具像素化，得到索引数组res，ra[res[j]]可获取第j个垂直单元内所有星系的赤经
• 沿视线方向的距离chi通过以下代码分箱：

  bins = np.linspace(np.min(chi),np.max(chi),nzbin)
  hist, edges = np.histogram(chi, bins=bins)
  

当前通过双层循环创建布尔掩码数组mask_grid，标记每个网格单元包含的天体：

mask_list = []
for i in range(nzbin-1):
    for j in range(len(res)):
        mask = (np.min(ra[res[j]]) <= ra ) & ( ra <= np.max(ra[res[j]])) & (np.min(dec[res[j]]) <= dec) & (dec <= np.max(dec[res[j]])) & (chi >= edges[i]) & (chi < edges[i+1]) 
        mask_list += [mask]
mask_grid = np.vstack(mask_list)  

后续遍历mask_grid计算单元属性，但当nzbin增大至5000时出现严重内存不足问题，需替换循环实现，优化内存与效率。

核心思路是不存储庞大的二维布尔掩码数组，而是给每个星系分配所属的网格单元标签，后续通过标签直接分组处理，彻底解决内存问题。

步骤1：给每个星系标记垂直视线方向的像素ID #

利用res数组直接给每个星系打上所属的healpy像素索引标签，避免重复计算ra/dec范围：

import numpy as np

# 初始化像素ID数组，长度等于星系总数
pixel_id = np.zeros(len(ra), dtype=int)
# 为每个healpy像素内的星系分配对应ID
for j, galaxy_indices in enumerate(res):
    pixel_id[galaxy_indices] = j

步骤2：给每个星系标记视线方向的分箱ID #

用np.digitize直接获取每个星系对应的视线分箱索引，与原掩码的chi范围逻辑完全一致：

# 获取每个星系的视线分箱ID（范围0到nzbin-2）
zbin_id = np.digitize(chi, bins=edges) - 1
# 修正chi等于最大值时的溢出问题
zbin_id[zbin_id == nzbin-1] = nzbin-2

步骤3：按需计算单元属性 #

通过pixel_id和zbin_id的组合定位目标单元，直接提取对应星系数据，无需遍历所有单元：

# 获取所有存在天体的唯一网格单元
unique_cells = np.unique(np.stack([pixel_id, zbin_id], axis=1), axis=0)

for cell in unique_cells:
    j, i = cell
    # 提取该单元的所有星系索引
    galaxy_idx = np.where((pixel_id == j) & (zbin_id == i))[0]
    # 获取该单元的星系数据
    ra_masked = ra[galaxy_idx]
    dec_masked = dec[galaxy_idx]
    chi_masked = chi[galaxy_idx]
    # 后续计算物理属性...

如果允许使用pandas，分组效率会更高：

import pandas as pd

df = pd.DataFrame({
    'ra': ra, 'dec': dec, 'chi': chi,
    'pixel_id': pixel_id, 'zbin_id': zbin_id
})

# 按网格单元分组计算
for (pixel_j, zbin_i), group in df.groupby(['pixel_id', 'zbin_id']):
    ra_masked = group['ra'].values
    # 计算属性...

• 内存占用：原方案中mask_grid是(N_cells, 1.2e7)的布尔数组，仅5000个nzbin+1024个healpy像素就需约60TB内存；优化后仅需存储两个长度为1.2e7的整数数组，总内存约96MB
• 准确性：直接用healpy的像素索引标记，避免了原方案中用ra/dec极值判断可能出现的跨像素错误
• 效率：省去了双层循环中重复计算ra/dec极值的冗余操作，整体计算速度大幅提升

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