如何用Eigen::Map替代循环将FFTW数组复制到Eigen Tensor？

我编写了一段C++代码，用FFTW对Eigen三维Tensor执行3D FFT，将输出结果存回Eigen复Tensor。目前通过嵌套循环把FFTW的output_array复制到Eigen Tensor中，但输出结果存在异常：部分列被填充为0，且元素位置错位（如5e11的位置错误）。想知道是否可以用Eigen::TensorMap替代现有嵌套循环来完成数据复制？

代码示例 #

static const int nx = 4;
static const int ny = 4;
static const int nz = 4;
void r2cfft3d(Eigen::Tensor<double, 3>& rArr, Eigen::Tensor<std::complex<double>, 3>& cArr){    
       
        fftw_complex *input_array;
        input_array = (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
        memcpy(input_array, rArr.data(), nx*ny*nz * sizeof(fftw_complex));

        fftw_complex *output_array;
        output_array = (fftw_complex*) fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
        
        fftw_plan forward = fftw_plan_dft_3d(nx, ny, nz, input_array, output_array, FFTW_FORWARD, FFTW_ESTIMATE);
        
        fftw_execute(forward);
        fftw_destroy_plan(forward);
        fftw_cleanup();

        for(int i=0; i < nx; ++i){
            for(int j=0; j < ny; ++j) {
                for(int k=0; k < nz; ++k) {
                    cArr(i,j,k).real(output_array[i + nz * (j + ny * k)][REAL]);
                    cArr(i,j,k).imag(output_array[i + nz * (j + ny * k)][IMAG]);
                }
            }
        }
        
        fftw_free(input_array);
        fftw_free(output_array);
    }

当前输出 #

1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11 -9.53674e-07 -9.53674e-07 -9.53674e-07  2.86102e-06            0            0            0            0
       1e+11        1e+11        1e+11        1e+11        1e+11        5e+11        1e+11        1e+11 -4.76837e-06 -4.76837e-06 -4.76837e-06 -1.62125e-05 -1.52588e-05 -1.52588e-05 -1.52588e-05 -1.52588e-05
       1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11 -9.53674e-07 -9.53674e-07 -9.53674e-07  2.86102e-06            0            0            0            0
       1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11        1e+11 -8.58307e-06 -8.58307e-06 -8.58307e-06 -4.76837e-06 -7.62939e-06 -7.62939e-06 -7.62939e-06 -7.62939e-06

正确输出 #

1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11
1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11
1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 5e+11 1e+11 1e+11 1e+11 1e+11 1e+11
1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11 1e+11

解答 #

结论：可以用Eigen::TensorMap替代嵌套循环，但首先要解决内存布局不匹配的问题——这才是当前输出异常的核心原因。 #

问题根源

你的代码存在两个关键错误：

1. 输入数据拷贝错误：rArr是Eigen::Tensor<double,3>，每个元素占1个double；但fftw_complex是包含两个double的结构体（实部+虚部）。你用memcpy拷贝了nx*ny*nz * sizeof(fftw_complex)字节，相当于拷贝了2倍于实际需要的数据，直接导致输入乱码。
2. 索引顺序不匹配：FFTW的3D数组默认是列优先（Fortran顺序），存储顺序为k, j, i；而Eigen Tensor默认是行优先（C顺序），存储顺序为i, j, k。嵌套循环中的索引计算i + nz*(j + ny*k)完全不符合两者的布局逻辑，直接导致元素位置错位。

修正后的实现（用TensorMap替代循环）

static const int nx = 4;
static const int ny = 4;
static const int nz = 4;

void r2cfft3d(Eigen::Tensor<double, 3>& rArr, Eigen::Tensor<std::complex<double>, 3>& cArr){    
    // 1. 将实Tensor转为复Tensor（虚部补0）
    Eigen::Tensor<std::complex<double>, 3> input_complex = rArr.cast<std::complex<double>>();
    
    // 2. 分配FFTW内存
    fftw_complex *input_array = (fftw_complex*)fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
    fftw_complex *output_array = (fftw_complex*)fftw_malloc(nx*ny*nz * sizeof(fftw_complex));
    
    // 3. 用TensorMap按FFTW的列优先布局拷贝输入数据
    Eigen::TensorMap<Eigen::Tensor<std::complex<double>, 3, Eigen::ColMajor>> input_map(input_array, nx, ny, nz);
    input_map = input_complex.swap_layout(); // 转换Eigen Tensor的布局为列优先，匹配FFTW
    
    // 4. 执行FFT
    fftw_plan forward = fftw_plan_dft_3d(nx, ny, nz, input_array, output_array, FFTW_FORWARD, FFTW_ESTIMATE);
    fftw_execute(forward);
    
    // 5. 用TensorMap映射FFTW输出，转换布局为Eigen默认的行优先后赋值给cArr
    Eigen::TensorMap<Eigen::Tensor<std::complex<double>, 3, Eigen::ColMajor>> output_map(output_array, nx, ny, nz);
    cArr = output_map.swap_layout();
    
    // 6. 清理资源
    fftw_destroy_plan(forward);
    fftw_free(input_array);
    fftw_free(output_array);
}

关键说明

• Eigen::Tensor的第三个模板参数可指定布局：Eigen::RowMajor（C顺序，默认）、Eigen::ColMajor（Fortran顺序）。
• swap_layout()方法可快速转换Tensor的存储顺序，无需手动计算索引，避免出错。
• TensorMap直接映射FFTW的内存块，无需嵌套循环，效率更高且逻辑更清晰。
• 输入数据需先转为复Tensor再按FFTW布局拷贝，避免内存拷贝错误。

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