展开NumPy einsum操作：获取多维向量的索引映射关系

索引映射关系拆解 #

给定的einsum操作核心逻辑可拆解为以下规则：

• 输入：两个同形状的5维数组，分别是原向量vector和其共轭np.conj(vector)
• 索引标记仅用于维度对应关系，与数字大小无关：
  • 原向量的索引[0,1,2,3,4]对应其5个维度：vector[a,b,c,d,e]（a/b/c/d/e取值0或1）
  • 共轭向量的索引[0,1,2,8,9]对应其5个维度：np.conj(vector)[a,b,c,f,g]（8/9是区分标记，实际对应共轭向量的第3/4维）
  • 输出索引[3,4,8,9]对应输出数组的4个维度：out[d,e,f,g]

求和消去的索引 #

索引0,1,2（即a,b,c）同时出现在两个输入的索引列表中，但未出现在输出索引里——这意味着要对这三个维度的所有可能取值做累加求和，最终得到输出数组。

循环展开实现 #

因为每个维度长度都是2（取值0和1），可将einsum操作展开为嵌套循环，等价代码如下：

import numpy as np

# 初始化输入向量（示例）
vector = np.random.rand(2,2,2,2,2) + 1j*np.random.rand(2,2,2,2,2)
# 初始化输出数组
out = np.zeros((2,2,2,2), dtype=np.complex128)

# 嵌套循环实现einsum逻辑
for d in range(2):
    for e in range(2):
        for f in range(2):
            for g in range(2):
                acc = 0.0 + 0.0j
                # 对求和维度a,b,c遍历累加
                for a in range(2):
                    for b in range(2):
                        for c in range(2):
                            acc += vector[a,b,c,d,e] * np.conj(vector[a,b,c,f,g])
                out[d,e,f,g] = acc

结果索引的影响因素 #

输出数组的每个位置out[d,e,f,g]，由以下两类向量元素共同决定：

• 原向量中**前3维任意、后2维为(d,e)**的所有元素：vector[a,b,c,d,e]（a/b/c取0或1）
• 共轭向量中**前3维与原向量完全相同、后2维为(f,g)**的所有元素：np.conj(vector[a,b,c,f,g])
• 对所有a/b/c的组合，将对应元素相乘后累加，就得到该位置的输出值。

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