大网格下4D数值积分的性能优化方案问询

问题描述 #

我正在尝试数值计算一个4维积分：

其中 d⁴x = dt d³x，函数 f(x)=f(t,x) 需要在一个大网格上求解。由于内存限制（我只有1TB RAM），我把积分在时间和空间网格上离散化，拆分成了两部分：

之后我用一个for循环计算时间求和，每个时间步都会重新计算 f(x)=f(tₙ,x)，并对其做3D FFT。

性能相关说明 #

计算 f(x)=f(t,x) 涉及一些数组操作，我用了numexpr库；FFT计算则用了pyfftw库。我没法做时间方向的向量化，因为那样FFT对应的数组会超出内存限制。

我已经尝试过用更低精度的类型（float32、complex64）来限制内存占用、减少计算时间。

每个时间步都会执行以下计算：

• 计算 f(x)=f(t,x)：需要一些numexpr数组操作，之后做逆3D FFT。
• 积分计算主体：对 f(x)=f(t,x) 执行一些数组操作，之后对结果做正3D FFT。

我用scalene和memray做了性能分析，结果显示大部分时间都消耗在FFT操作和这些特定的数组运算上。

我的问题 #

如果有人能给我一些优化这个计算性能的建议，我会非常感激。如果能推荐一些相关思路或者资料让我阅读，也会对我帮助很大。

最小工作示例 #

import numpy as np
import numexpr as ne
import pyfftw

# Minimal working example
class Function:
    def __init__(self, size=(20,20,1000), nthreads=1):
        self.a = np.random.random(size)
        
        self.tmp = np.empty(size, dtype='complex128')
        self.tmp_fftw = pyfftw.FFTW(
            self.tmp,
            self.tmp,
            axes=(0, 1, 2),
            direction="FFTW_FORWARD",
            flags=("FFTW_MEASURE",),
            threads=nthreads,
        )

    def get_function(self, t):
        ne.evaluate("exp(1j*t)*a", global_dict={"a": self.a}, out=self.tmp)
        self.tmp_fftw.execute()
        return self.tmp


class Integral:
    def __init__(self, func, size=(20,20,1000), nthreads=1):
        self.func = func
        
        self.f = np.empty(size, dtype='complex128')
        self.f_fftw = pyfftw.FFTW(
            self.f,
            self.f,
            axes=(0, 1, 2),
            direction="FFTW_FORWARD",
            flags=("FFTW_MEASURE",),
            threads=nthreads,
        )
        
        self.result = 0
    
    def get_one_time_step(self, t):
        self.f = self.func.get_function(t)
        ne.evaluate("3*f + 2", global_dict={"f": self.f}, out=self.f)
        self.f_fftw.execute()
        self.result += self.f

    def get_time_integral(self, t_grid):
        for t in t_grid:
            self.get_one_time_step(t)
        return self.result


size = (30,30,1000)
func = Function(size)
integrator = Integral(func, size)

t_grid = np.linspace(-1, 1, 21)
result = integrator.get_time_integral(t_grid)

备注：内容来源于stack exchange，提问作者maxbalrog