如何优化Python多进程3D伊辛模型模拟速度？求ML样本生成方案

我用Python实现了Wolff聚类算法来模拟3D伊辛模型，生成指定温度和晶格尺寸下的磁化强度值。已经用multiprocessing做了并行化，流程是先让系统平衡，再以自相关时间为间隔采集样本。

在80×80×80晶格、临界温度4.5的条件下，代码运行耗时约5天，但查资料发现500×500×500规模的晶格通常仅需3天，希望了解更快生成磁化样本的方法。另外我刚接触C++，暂时想继续用Python实现；同时想咨询能否用机器学习生成这类样本，试过混合密度网络（Mixture Density Networks）但效果不好。

import numpy as np 
from collections import deque
import time
import matplotlib.pyplot as plt
from multiprocessing import Pool
plt.style.use('ggplot')

SAMPLES = 1000  #No of Independant samples that needs to be generated
AUTO = 20  #Autocorrelation time to get independant samples
LOOP = 10 #The same simulation is done LOOP no of times to get SAMPLES*LOOP no of samples

class Cluster3D:
    
    def __init__(self, temp):
        self.temp = temp   #Sampling temperature
        self.width = 20     #Size of the lattice
        self.state = np.ones((self.width,self.width,self.width), dtype=int)   #Creates a 3D lattice with all spins = +1
        self.length = SAMPLES*AUTO
        self.equi = 500    #The model is made to equilibriate before the samples are taken
        self.p_add = 1 - np.exp(-2/temp)     #Probability with which the clusters are formed
        self.auto = 0
        self.avg_size = 0
        
    def neighboring_sites(self, s):
        w = self.width
        return [((s[0]+1)%w, s[1], s[2]), ((s[0]-1)%w, s[1], s[2]),
                 (s[0], (s[1]+1)%w, s[2]), (s[0], (s[1]-1)%w, s[2]),
                 (s[0], s[1], (s[2]+1)%w), (s[0], s[1], (s[2]-1)%w)]
  
    def cluster_flip(self, seed):
        spin = self.state[seed]
        self.state[seed] = -spin  
        cluster_size = 1
        unvisited = deque([seed])   # use a deque to efficiently track the unvisited cluster sites
        while unvisited:   # while unvisited sites remain
            site = unvisited.pop()  # take one and remove from the unvisited list
            for nbr in self.neighboring_sites(site):
                if self.state[nbr] == spin and np.random.random() < self.p_add:
                    self.state[nbr] = -spin
                    unvisited.appendleft(nbr)
                    cluster_size += 1
        return cluster_size
    
    def wolff_cluster_move(self):
        rng = np.random.default_rng()
        seed = tuple(rng.integers(0,self.width,3))
        return self.cluster_flip(seed)
            
    def compute_magnetization(self):
        return np.sum(self.state)/self.width**3

    def compute_internal_energy(self):
        n = self.width
        l = self.state
        e = 0.0
        for i in range(0,n):
            for j in range(0,n):
                for k in range(0,n):
                    if i+1<=n-1:
                        e += -l[i,j,k]*l[i+1,j,k]
                    if j+1<=n-1:
                        e += -l[i,j,k]*l[i,j+1,k]
                    if k+1<=n-1:
                        e += -l[i,j,k]*l[i,j,k+1]
                    if i-1>=0:
                        e += -l[i,j,k]*l[i-1,j,k]
                    if j-1>=0:
                        e += -l[i,j,k]*l[i,j-1,k]
                    if k-1>=0:
                        e += -l[i,j,k]*l[i,j,k-1]
        return e

    def run_ising_wolff_mcmc(self, n):
        total = 0
        for _ in range(n):
            total += self.wolff_cluster_move()
        return total

    def sample_autocovariance(self, x):
        x_shifted = x - np.mean(x)
        return np.array([np.dot(x_shifted[:len(x)-t],x_shifted[t:])/len(x) 
                         for t in range(self.length)])

    def find_correlation_time(self, autocov):
        smaller = np.where(autocov < np.exp(-1)*autocov[0])[0]
        return smaller[0] if len(smaller) > 0 else len(autocov)

    def do_wolff(self):
        trace = np.zeros(self.length)
        energy = np.zeros(self.length)
        total_flips = 0
        self.run_ising_wolff_mcmc(self.equi)
        for i in range(self.length):    
            total_flips += self.wolff_cluster_move()
            trace[i] = self.compute_magnetization()
            energy[i] = self.compute_internal_energy()
            
        autocov = self.sample_autocovariance(np.abs(trace))
        time = self.find_correlation_time(autocov)
        self.avg_size = total_flips/self.length
        self.auto = time
        
        return trace

def generator(temp):
    state = Cluster3D(temp)
    raw_spins = state.do_wolff()
    nets = [raw_spins[-AUTO*i] for i in range(SAMPLES)]
    return abs(np.array(nets))   

if __name__ == '__main__':
    
    u1 = time.time()
    temp = 5.0      #Temperature at which the samples are generated
    
    temps = [temp for i in range(LOOP)]
    raw_spin = []
    pool = Pool()
    data = pool.map(generator, temps)     #Runs the algorithm for the same temperature for LOOP no of times across multiple CPU cores
    raw_spin.append(data)
    spins = np.array(data).flatten()    #Size of spins is SAMPLES*LOOP
   
    u2 = time.time()

    print('Time Taken to Simulate - ', abs(u2-u1))

一、Python实现的性能优化方向 #

1. 核心算法的数值计算优化 #

• 预计算邻域索引：neighboring_sites每次都重新计算模运算，可提前生成整个晶格的邻域索引矩阵，后续直接索引访问，避免重复计算。
• JIT编译核心函数：用numba对cluster_flip、wolff_cluster_move这类循环密集型函数做即时编译，能大幅降低Python解释器的开销。
• 向量化能量计算：替换compute_internal_energy的三重循环，用NumPy的np.roll实现相邻自旋的乘积求和，示例：

  def compute_internal_energy(self):
      l = self.state
      energy = -np.sum(l * np.roll(l, 1, axis=0))
      energy += -np.sum(l * np.roll(l, 1, axis=1))
      energy += -np.sum(l * np.roll(l, 1, axis=2))
      return energy
  

• 优化自旋存储：用np.int8替代np.int存储自旋状态，减少内存占用，提升缓存命中率。

2. 并行与采样策略优化 #

• 减少进程初始化开销：当前Pool.map每次启动新实例并重复平衡过程，可改为单个进程内完成多次采样，或用imap_unordered让结果返回后立即处理，提高资源利用率。
• 动态调整自相关时间：临界温度附近自相关时间会显著增大，动态计算后再确定采样间隔，避免不必要的采样步骤；同时可通过监测磁化强度波动判断平衡状态，避免过度平衡。
• 跳过冗余计算：如果仅需磁化强度样本，可在do_wolff中跳过能量计算，节省时间。

3. 随机数生成优化 #

在__init__中初始化一次np.random.default_rng()，避免每次wolff_cluster_move都新建生成器的开销。

二、机器学习生成样本的可行思路 #

混合密度网络效果不佳，主要因为临界温度附近的伊辛模型样本具有强空间相关性和非高斯分布特性，可尝试以下方向：

1. 归一化流（Normalizing Flows） #

通过一系列可逆变换将简单分布（如高斯分布）映射到伊辛模型的平衡分布，能精确建模复杂概率分布，训练稳定且可估计样本对数似然，便于评估生成质量。

2. CNN-based生成式对抗网络（GAN） #

临界温度下的自旋构型具有分形特征，CNN可有效捕捉这类空间相关性，直接生成自旋构型后提取磁化强度。已有研究证明这种方法在伊辛模型样本生成上的有效性。

3. 基于物理的ML方法 #

• 神经Metropolis采样：用强化学习优化MCMC的采样策略，加速平衡态收敛；
• 小晶格预训练迁移：先在小晶格上训练模型，再迁移到大晶格，利用伊辛模型的尺度不变性减少训练数据需求。

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