如何用awkward数组处理四缪子的Z玻色子衰变组合筛选？

你遇到的类型不匹配问题，本质是用ak.combinations处理四缪子内部配对时，生成的嵌套结构和原有四缪子组合的类型不兼容。四缪子分成两个Z候选对只有三种固定分法，直接构造这三种配对结构即可高效解决问题，具体步骤如下：

1. 构造三种Z对组合 #

假设你已经通过ak.combinations生成了四缪子组合数组four_muons（字段为mu1, mu2, mu3, mu4），直接基于这个数组构造三种符合H→Z→μμ衰变逻辑的配对：

# 配对方式1：(mu1, mu2) 作为第一个Z候选，(mu3, mu4)作为第二个
pairing1 = ak.zip({
    "z1": ak.zip({"mu": [four_muons.mu1, four_muons.mu2]}),
    "z2": ak.zip({"mu": [four_muons.mu3, four_muons.mu4]})
})

# 配对方式2：(mu1, mu3) 和 (mu2, mu4)
pairing2 = ak.zip({
    "z1": ak.zip({"mu": [four_muons.mu1, four_muons.mu3]}),
    "z2": ak.zip({"mu": [four_muons.mu2, four_muons.mu4]})
})

# 配对方式3：(mu1, mu4) 和 (mu2, mu3)
pairing3 = ak.zip({
    "z1": ak.zip({"mu": [four_muons.mu1, four_muons.mu4]}),
    "z2": ak.zip({"mu": [four_muons.mu2, four_muons.mu3]})
})

# 把三种配对合并到一个数组，方便统一做筛选
all_pairings = ak.concatenate([pairing1, pairing2, pairing3], axis=1)

2. 实现电荷与质量筛选逻辑 #

中性Z玻色子衰变的两个缪子电荷总和必须为0，且不变质量要落在Z的质量窗口（通常80-100GeV）：

# 计算每个Z候选的总电荷，筛选出两个Z候选电荷都为0的情况
z1_charge_ok = (ak.sum(all_pairings.z1.mu.charge, axis=2) == 0)
z2_charge_ok = (ak.sum(all_pairings.z2.mu.charge, axis=2) == 0)
charge_cut = z1_charge_ok & z2_charge_ok

# 定义不变质量计算函数
def calc_invariant_mass(mu_pair):
    total_px = ak.sum(mu_pair.mu.px, axis=2)
    total_py = ak.sum(mu_pair.mu.py, axis=2)
    total_pz = ak.sum(mu_pair.mu.pz, axis=2)
    total_e = ak.sum(mu_pair.mu.energy, axis=2)
    return np.sqrt(total_e**2 - total_px**2 - total_py**2 - total_pz**2)

# 计算两个Z候选的质量并做窗口筛选
z1_mass = calc_invariant_mass(all_pairings.z1)
z2_mass = calc_invariant_mass(all_pairings.z2)
mass_cut = (z1_mass > 80) & (z1_mass < 100) & (z2_mass > 80) & (z2_mass < 100)

3. 提取有效组合 #

合并筛选条件，提取符合要求的配对和对应的四缪子组合：

# 合并电荷与质量筛选条件
valid_mask = charge_cut & mass_cut

# 得到符合条件的Z对组合
valid_z_pairs = all_pairings[valid_mask]

# 若需要还原回四缪子组合，可从有效配对中提取
valid_four_muons = ak.zip({
    "mu1": valid_z_pairs.z1.mu[:, 0],
    "mu2": valid_z_pairs.z1.mu[:, 1],
    "mu3": valid_z_pairs.z2.mu[:, 0],
    "mu4": valid_z_pairs.z2.mu[:, 1]
})

关键注意点 #

• 不要用ak.combinations生成四缪子内部的两两配对：这会产生6种组合，但我们只需要符合Z衰变逻辑的3种电荷共轭配对，直接构造固定配对不仅高效，还能避免嵌套结构导致的类型不匹配问题。
• 利用Awkward的ak.zip和数组广播特性，批量处理所有组合，避免循环操作，保证处理效率。

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