使用Numpy实现PIMC QAVRP时遇维度索引错误求助

问题背景 #

在实现基于路径积分蒙特卡洛（Path-Integral Monte Carlo）的量子退火算法PIMC QAVRP时，遇到Numpy维度索引错误，暂时不关注复数转实数的警告，需要修复该问题以获取VRP的最优解S[0]。

原实现代码 #

import numpy as np

#Calculo da energia potencial do sistema
def pot_energy(S, d):
    return np.sum(d * np.einsum('...i,...j,ij->...ij', S, S, np.ones_like(d)), axis=(-1, -2))

#Calculo da energia potencial do sistema
def kin_energy(S):
    Hkin = np.zeros_like(S)
    Hkin[1:] += S[:-1] * S[1:]
    Hkin[-1] += S[-1] * S[0]
    return np.sum(Hkin)

#Calculo da transformação de Suzuki-Trotter para uma replica
def suzuki_trotter_replica(S, J, dt):
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    np.fft.fftn(S, axes=[-2, -1])
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    np.fft.ifftn(S, axes=[-2, -1])
    S = S * np.exp(-1j * dt).astype(float)
    S[1:] *= S[:-1]
    S[-1] *= S[0]
    S = S * np.exp(-1j * dt).astype(float)
    return S

#Calculo da transformação de Suzuki-Trotter para todas as replicas
def suzuki_trotter(S, J, dt):
    for i in range(S.shape[0]):
        S[i] = suzuki_trotter_replica(S[i], J, dt)
    return S


#ALGORITIMO DE EXECUÇÃO DO QAVRP
def pimc_qa_vrp(d, num_steps, num_replicas, beta, dt):
    # Inicializa a matriz de spins aleatóriamente
    S = np.random.choice([-1, 1], size=(num_replicas, *d.shape))
    # Inicializa os arrays de energia
    H_pot = np.zeros(num_replicas)
    H_kin = np.zeros(num_replicas)
    # Loop principal
    for step in range(num_steps):
        # Aplica a transformação de Suzuki-Trotter
        S = suzuki_trotter(S, beta, dt)
        # Calcula a energia potencial do sistema
        for i in range(num_replicas):
            H_pot[i] = np.sum(pot_energy(S[i],d))
            H_kin[i] = kin_energy(S[i])
        # Calcula a energia Total
        H = H_pot / num_replicas - np.sum(H_kin) * beta / num_replicas
        # Aplica o QAVRP
        S = np.sign(np.tanh(-beta * (H - H.min()) / 2) + np.random.rand(*S[:, 0:1, :, :].shape + (1,)) - 0.5)
        '''
        Essa expressão implementa a transição quântica aleatória, onde o estado atual é atualizado para um novo estado aleatório 
        com uma probabilidade que depende da diferença de energia potencial entre os dois estados. Se a energia potencial do novo
        estado é menor do que a do estado atual, o novo estado é aceito. Caso contrário, o novo estado é aceito com uma probabilidade 
        que depende da temperatura.
        '''
    return S[0]


d = np.array([[0, 1, np.sqrt(2)],
              [1, 0, 1],
              [np.sqrt(2), 1, 0]])

num_steps = 1000
num_replicas = 100
beta = 2.0
dt = 0.1

solution = pimc_qa_vrp(d, num_steps, num_replicas, beta, dt)
print(solution)

报错信息 #

C:\Users\vitor\OneDrive\Documentos\qiskit\test.py:16: ComplexWarning: Casting complex values to real discards the imaginary part
  S = (S * np.exp(-1j * J * dt / 2)).astype(float)
C:\Users\vitor\OneDrive\Documentos\qiskit\test.py:18: ComplexWarning: Casting complex values to real discards the imaginary part
  S = (S * np.exp(-1j * J * dt / 2)).astype(float)
C:\Users\vitor\OneDrive\Documentos\qiskit\test.py:20: ComplexWarning: Casting complex values to real discards the imaginary part
  S = S * np.exp(-1j * dt).astype(float)
C:\Users\vitor\OneDrive\Documentos\qiskit\test.py:23: ComplexWarning: Casting complex values to real discards the imaginary part
  S = S * np.exp(-1j * dt).astype(float)

Traceback (most recent call last):
  File "C:\Users\vitor\OneDrive\Documentos\qiskit\test.py", line 70, in <module>
    solution = pimc_qa_vrp(d, num_steps, num_replicas, beta, dt)
  File "C:\Users\vitor\OneDrive\Documentos\qiskit\test.py", line 51, in pimc_qa_vrp
    S = np.sign(np.tanh(-beta * (H - H.min()) / 2) + np.random.rand(*S[:, 0:1, :, :].shape + (1,)) - 0.5)
IndexError: too many indices for array: array is 3-dimensional, but 4 were indexed

错误原因分析 #

• 初始的S是3维数组，形状为(num_replicas, d.shape[0], d.shape[1])（示例中为(100,3,3)）
• 代码中S[:, 0:1, :, :]尝试对4维数组进行索引，但实际S只有3维，直接触发IndexError
• 同时，H是一维数组（长度为num_replicas），无法直接和3维的随机数组进行广播运算

解决方案 #

1. 修正QAVRP更新步骤的维度问题 #

将一维的H扩展维度以匹配S的3维形状，同时生成和S形状一致的随机数组，确保广播运算正常：

# 扩展H的维度，使其与S的形状兼容（从(100,)变为(100,1,1)）
H_expanded = H[:, np.newaxis, np.newaxis]
# 生成和S形状完全相同的随机数组
rand_vals = np.random.rand(*S.shape)
# 计算新的自旋状态S
S = np.sign(np.tanh(-beta * (H_expanded - H.min()) / 2) + rand_vals - 0.5)

2. 修复Suzuki-Trotter变换中的无效操作 #

原代码中np.fft.fftn和np.ifftn的结果没有赋值回S，导致这两步傅里叶变换完全无效，需要补充赋值：

def suzuki_trotter_replica(S, J, dt):
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    S = np.fft.fftn(S, axes=[-2, -1])  # 将傅里叶变换结果赋值回S
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    S = np.fft.ifftn(S, axes=[-2, -1])  # 将逆傅里叶变换结果赋值回S
    S = S * np.exp(-1j * dt).astype(float)
    S[1:] *= S[:-1]
    S[-1] *= S[0]
    S = S * np.exp(-1j * dt).astype(float)
    return S

修复后的完整代码 #

import numpy as np

#Calculo da energia potencial do sistema
def pot_energy(S, d):
    return np.sum(d * np.einsum('...i,...j,ij->...ij', S, S, np.ones_like(d)), axis=(-1, -2))

#Calculo da energia cinética do sistema
def kin_energy(S):
    Hkin = np.zeros_like(S)
    Hkin[1:] += S[:-1] * S[1:]
    Hkin[-1] += S[-1] * S[0]
    return np.sum(Hkin)

#Calculo da transformação de Suzuki-Trotter para uma replica
def suzuki_trotter_replica(S, J, dt):
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    S = np.fft.fftn(S, axes=[-2, -1])
    S = (S * np.exp(-1j * J * dt / 2)).astype(float)
    S = np.fft.ifftn(S, axes=[-2, -1])
    S = S * np.exp(-1j * dt).astype(float)
    S[1:] *= S[:-1]
    S[-1] *= S[0]
    S = S * np.exp(-1j * dt).astype(float)
    return S

#Calculo da transformação de Suzuki-Trotter para todas as replicas
def suzuki_trotter(S, J, dt):
    for i in range(S.shape[0]):
        S[i] = suzuki_trotter_replica(S[i], J, dt)
    return S


#ALGORITIMO DE EXECUÇÃO DO QAVRP
def pimc_qa_vrp(d, num_steps, num_replicas, beta, dt):
    # Inicializa a matriz de spins aleatóriamente
    S = np.random.choice([-1, 1], size=(num_replicas, *d.shape))
    # Inicializa os arrays de energia
    H_pot = np.zeros(num_replicas)
    H_kin = np.zeros(num_replicas)
    # Loop principal
    for step in range(num_steps):
        # Aplica a transformação de Suzuki-Trotter
        S = suzuki_trotter(S, beta, dt)
        # Calcula a energia potencial do sistema
        for i in range(num_replicas):
            H_pot[i] = np.sum(pot_energy(S[i],d))
            H_kin[i] = kin_energy(S[i])
        # Calcula a energia Total
        H = H_pot / num_replicas - np.sum(H_kin) * beta / num_replicas
        # Aplica o QAVRP - 修正维度问题
        H_expanded = H[:, np.newaxis, np.newaxis]
        rand_vals = np.random.rand(*S.shape)
        S = np.sign(np.tanh(-beta * (H_expanded - H.min()) / 2) + rand_vals - 0.5)
    return S[0]


d = np.array([[0, 1, np.sqrt(2)],
              [1, 0, 1],
              [np.sqrt(2), 1, 0]])

num_steps = 1000
num_replicas = 100
beta = 2.0
dt = 0.1

solution = pimc_qa_vrp(d, num_steps, num_replicas, beta, dt)
print(solution)

内容的提问来源于stack exchange，提问作者Vitor Bernstorff Clemes