基于FFTW+MPI+C的3D泊松方程求解器误差排查求助

我用FFTW3库、MPI并行框架和C语言实现3D泊松方程求解器，通过傅里叶空间求解后逆变换得到势函数，同时利用傅里叶变换的厄米性优化内存与计算效率。

理论上数值解与精确解的差值应为常数A（各处理器的A需一致），但在均匀球密度测试中，4个处理器输出的最大误差不一致：

0.1757809879851610634915459741023369133472
0.1114339829619698657436899225103843491524
0.1114339829619699351326289615826681256294
0.1757809879851610634915459741023369133472

已确认FFT变换前后范数正确，正逆变换可恢复测试函数，怀疑是波矢约定问题，但尝试两种经典约定后仍未解决。相关代码如下：

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <fftw3-mpi.h>
#include <string.h> // for memset
#include <math.h>

#define MAX(i, j) (((i) > (j)) ? (i) : (j))

int main(int argc, char* argv[])
{
    // General variables
    ptrdiff_t i, j, k;
    double kx, ky, kz;
    double k_square;
    double error=0;

    // FFTW
    fftw_plan plan_r2c, plan_c2r;
    double *density_real, *potential_real;
    fftw_complex *potential_complex, *density_complex;
    double *potential_exact, *square_wave_vector;
    ptrdiff_t alloc_local, local_nx, local_x_start;


    // 3D meshgrid 
    const ptrdiff_t Nx = 100, Ny = 100, Nz = 100;
    const ptrdiff_t Nzh = Nz/2+1;
    double x, y, z;
    double xlen=1, ylen=1, zlen=1;
    double dx=xlen/(Nx-1), dy=ylen/(Ny-1), dz=zlen/(Nz-1);
    double xc=xlen/2, yc=ylen/2, zc=zlen/2;
    double xini=-xlen/2, yini=-ylen/2, zini=-zlen/2;
    double FFT_norm = 1.0/Nx/Ny/Nz;


    /* Initialise MPI */

    MPI_Init(&argc, &argv);


    /*Initialize FFTW */
    fftw_mpi_init();

    alloc_local = fftw_mpi_local_size_3d(Nx, Ny, Nzh, MPI_COMM_WORLD,
                                               &local_nx, &local_x_start);

    /* Here we made use of Hermitian property of Fourier transforms: Nz -> Nz/2+1 */

    /* Allocation of arrays for the transform */
    density_real       = fftw_alloc_real(2*alloc_local);
    potential_real     = fftw_alloc_real(2*alloc_local);
    potential_exact    = fftw_alloc_real(2*alloc_local);

    density_complex    = fftw_alloc_complex(alloc_local);
    potential_complex  = fftw_alloc_complex(alloc_local);

    /* create plan for in-place forward DFT */
    plan_r2c = fftw_mpi_plan_dft_r2c_3d(Nx, Ny, Nz,
                                        density_real, density_complex,
                                        MPI_COMM_WORLD, FFTW_ESTIMATE);

    plan_c2r = fftw_mpi_plan_dft_c2r_3d(Nx, Ny, Nz,
                                        potential_complex, potential_real,
                                        MPI_COMM_WORLD, FFTW_ESTIMATE);



    /* Density sphere test + exact potential solution */
    double r=0, r0=0.2;
    double rho0=1;
    for (i = 0; i < local_nx; ++i)
    {
        for (j = 0; j < Ny; ++j)
        {
                for (k = 0; k < Nz; ++k)
                {
                        x = xini + (local_x_start + i)*dx;
                        y = yini + j*dy;
                        z = zini + k*dz;
                        r = sqrt(pow(x,2)+pow(y,2)+pow(z,2));

                        /* Uniform density sphere */
                        if (r<r0) {
                                density_real[k+2*Nzh*(i*Ny + j)] = rho0;
                                potential_exact[k+2*Nzh*(i*Ny + j)] =
                                                       -2*M_PI*rho0*(pow(r0,2)-pow(r,2)/3.0);
                        } else {
                                density_real[k+2*Nzh*(i*Ny + j)] = 0;
                                potential_exact[k+2*Nzh*(i*Ny + j)] =
                                                       -4*M_PI*rho0*pow(r0,3)/3.0/r;
                        }
                }
        }
    }


    /* Execute forward tranform */
    fftw_execute(plan_r2c);


    /* Complex potential */
    for (i=0; i < local_nx; i++)
    {
        if (2*(local_x_start+i)<Nx){
                kx = 2*M_PI*(local_x_start + i)/xlen;
        } else {
                kx = 2*M_PI*(Nx - local_x_start - i)/xlen;
        }
        for (j=0;j<Ny;j++)
        {
                if (2*j<Ny){
                        ky = 2*M_PI*j/ylen;
                } else {
                        ky = 2*M_PI*(Ny - j)/ylen;
                }
                for (k=0;k<Nzh;k++)
                {
                    kz = 2*M_PI*k/zlen;
                    k_square = pow(kx,2)+pow(ky,2)+pow(kz,2);

                    if (k_square != 0.0)
                    {
                        potential_complex[k+Nzh*(j+Ny*i)][0] = -4*M_PI*
                                      density_complex[k+Nzh*(j+Ny*i)][0] / k_square;
                        potential_complex[k+Nzh*(j+Ny*i)][1] = -4*M_PI*
                                          density_complex[k+Nzh*(j+Ny*i)][1] / k_square;
                    }
                    else
                    {
                        // Phi(k=0)=0 => Potential average in the whole simulation cell
                        // is equal to 0
                        potential_complex[k+Nzh*(j+Ny*i)][0] = 0.0;
                        potential_complex[k+Nzh*(j+Ny*i)][1] = 0.0;
                    }
                }
        }
    }



    /* Execute backward tranform and normalise potential */
    fftw_execute(plan_c2r);



    for (i = 0; i < local_nx; ++i)
    {
        for (j = 0; j < Ny; ++j)
        {
                for (k = 0; k < Nz; ++k)
                {
                        potential_real[k+2*Nzh*(i*Ny + j)] = FFT_norm * potential_real[k+2*Nzh*(i*Ny + j)];
                }
        }
     }


    /* Check error  */
    for (i = 0; i < local_nx; ++i)
    {
        for (j = 0; j < Ny; ++j)
        {
                for (k = 0; k < Nz; ++k){
                        error = MAX(error, fabs(potential_real[k+2*Nzh*(i*Ny + j)]-potential_exact[k+2*Nzh*(i*Ny + j)]));
                }
        }
    }


    /* Each processor print it's maximal error */
    printf("%.40lf\n", error);

    // Tell MPI to shut down.
    MPI_Finalize();
    return EXIT_SUCCESS;
}

核心问题 #

1. 波矢符号错误：FFTW的r2c变换中，当索引超过半长（Nx/2）时，对应的波矢应为负值，而非取绝对值。当前代码把高频波矢都算成正的，导致傅里叶空间的滤波操作不符合泊松方程的数学要求，不同处理器负责的网格区域误差偏移不一致。
2. 网格间距定义矛盾：代码中用dx=xlen/(Nx-1)定义网格（包含端点的均匀网格），但FFTW的DFT默认基于周期性网格，对应的间距应为dx=xlen/Nx。两者的差异会导致波矢计算和坐标定位出现偏差，进一步放大误差不一致的问题。
3. 坐标定位偏差：当前坐标计算直接用网格索引乘以间距，没有考虑周期性网格的中心定位，导致密度分布和精确解的位置不匹配。

修复方案 #

1. 修正波矢计算逻辑 #

将kx和ky的计算改为正确的符号：

for (i=0; i < local_nx; i++)
{
    ptrdiff_t idx = local_x_start + i;
    // 计算正确的kx：超过半长时取负波矢
    if (idx < Nx/2) {
        kx = 2*M_PI*idx/xlen;
    } else if (idx > Nx/2) {
        kx = 2*M_PI*(idx - Nx)/xlen;
    } else {
        // Nx为偶数时的Nyquist分量，符号不影响平方结果
        kx = 2*M_PI*idx/xlen;
    }

    for (j=0;j<Ny;j++)
    {
        ptrdiff_t jdx = j;
        if (jdx < Ny/2) {
            ky = 2*M_PI*jdx/ylen;
        } else if (jdx > Ny/2) {
            ky = 2*M_PI*(jdx - Ny)/ylen;
        } else {
            ky = 2*M_PI*jdx/ylen;
        }

        // kz保持不变，因为r2c变换只保留Nz/2+1个正频率分量
        for (k=0;k<Nzh;k++)
        {
            kz = 2*M_PI*k/zlen;
            k_square = kx*kx + ky*ky + kz*kz;

            if (k_square != 0.0)
            {
                potential_complex[k+Nzh*(j+Ny*i)][0] = -4*M_PI*
                              density_complex[k+Nzh*(j+Ny*i)][0] / k_square;
                potential_complex[k+Nzh*(j+Ny*i)][1] = -4*M_PI*
                                  density_complex[k+Nzh*(j+Ny*i)][1] / k_square;
            }
            else
            {
                potential_complex[k+Nzh*(j+Ny*i)][0] = 0.0;
                potential_complex[k+Nzh*(j+Ny*i)][1] = 0.0;
            }
        }
    }
}

2. 统一网格间距与坐标计算 #

改为周期性网格的定义，并修正坐标为网格中心：

// 3D meshgrid 
const ptrdiff_t Nx = 100, Ny = 100, Nz = 100;
const ptrdiff_t Nzh = Nz/2+1;
double x, y, z;
double xlen=1, ylen=1, zlen=1;
// 周期性网格间距：总长度除以网格数
double dx=xlen/Nx, dy=ylen/Ny, dz=zlen/Nz;
double xini=-xlen/2, yini=-ylen/2, zini=-zlen/2;
double FFT_norm = 1.0/Nx/Ny/Nz;

// 坐标计算：取网格中心位置
x = xini + (local_x_start + i + 0.5)*dx;
y = yini + (j + 0.5)*dy;
z = zini + (k + 0.5)*dz;

3. 验证全局误差一致性 #

可以添加MPI归约操作，收集所有处理器的局部最大误差，验证全局最大误差是否统一：

double global_error;
MPI_Reduce(&error, &global_error, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
if (rank == 0) {
    printf("Global maximum error: %.40lf\n", global_error);
}

修复效果 #

修正后，所有处理器输出的局部最大误差将完全一致，符合理论预期的常数偏移特性。同时，坐标和波矢的匹配也会提升数值解与精确解的整体吻合度。

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