在使用FFTW的C/C++中实现MATLAB冒号运算符的等效方法

需求背景 #

将MATLAB中的零填充逻辑转写为C/C++代码，核心是实现MATLAB冒号运算符:生成连续索引的功能，并正确填充FFTW3分配的复数数组。

原MATLAB代码逻辑回顾 #

N = [Nx Ny];
M = 3 * N / 2;
    
% 创建零矩阵
fk_pad = zeros(M(1), M(2));
    
% 生成填充区域索引：两段连续区间的合并
for i = 1:2
    ind_pad{i} = [1:N(i)/2 M(i)-N(i)/2+1:M(i)];
end
    
% 将输入矩阵fk填充到目标区域
fk_pad(ind_pad{1},ind_pad{2}) = fk;

这里的冒号运算符a:b用于生成从a到b的连续整数序列，最终ind_pad存储的是两段连续索引的合并结果，用于定位零矩阵中需要填充数据的区域。

C/C++代码实现步骤 #

1. 修正FFTW数组初始化

原代码中用memset初始化fftw_complex数组是错误的——memset按字节赋值，无法正确设置双精度浮点值。正确的零初始化方式如下：

static const int nx = 8, ny = 8;
static const int Mx = 3*nx/2, My= 3*ny/2;

// 分配FFTW复数数组
fftw_complex *fk_pad = (fftw_complex*) fftw_malloc(Mx * My * sizeof(fftw_complex));

// 零初始化：逐个设置实部和虚部为0
for (int i = 0; i < Mx * My; ++i) {
    fk_pad[i][0] = 0.0; // 实部
    fk_pad[i][1] = 0.0; // 虚部
}

2. 实现冒号运算符的索引生成逻辑

MATLAB的冒号序列是1基索引，而C/C++是0基索引，需要对应转换：

• 对于x维度（Mx行）：
  • MATLAB的1:Nx/2 → C/C++的0到(nx/2)-1
  • MATLAB的Mx-Nx/2+1:Mx → C/C++的Mx - nx/2到Mx-1
• 对于y维度（My列）：
  • MATLAB的1:Ny/2 → C/C++的0到(ny/2)-1
  • MATLAB的My-Ny/2+1:My → C/C++的My - ny/2到My-1

我们可以预先生成对应的索引数组，再用于填充：

// 生成x方向目标索引（对应MATLAB的ind_pad{1}）
int x_indices[nx];
int count = 0;
// 第一段：前nx/2个索引
for (int i = 0; i < nx/2; ++i) {
    x_indices[count++] = i;
}
// 第二段：后nx/2个索引
for (int i = Mx - nx/2; i < Mx; ++i) {
    x_indices[count++] = i;
}

// 生成y方向目标索引（对应MATLAB的ind_pad{2}）
int y_indices[ny];
count = 0;
for (int i = 0; i < ny/2; ++i) {
    y_indices[count++] = i;
}
for (int i = My - ny/2; i < My; ++i) {
    y_indices[count++] = i;
}

3. 执行填充操作

假设输入的fk是nx*ny大小的fftw_complex数组（C风格行优先存储），通过双重循环将fk的元素映射到fk_pad的对应位置：

// 假设fk是已初始化的输入复数数组
fftw_complex *fk = (fftw_complex*) fftw_malloc(nx * ny * sizeof(fftw_complex));
// （这里省略fk的赋值逻辑）

// 填充fk_pad
for (int fk_row = 0; fk_row < nx; ++fk_row) {
    for (int fk_col = 0; fk_col < ny; ++fk_col) {
        // 计算fk中当前元素的索引
        int fk_idx = fk_row * ny + fk_col;
        // 获取目标填充位置的行和列
        int target_row = x_indices[fk_row];
        int target_col = y_indices[fk_col];
        // 计算fk_pad中的索引（行优先存储）
        int pad_idx = target_row * My + target_col;
        // 赋值实部和虚部
        fk_pad[pad_idx][0] = fk[fk_idx][0];
        fk_pad[pad_idx][1] = fk[fk_idx][1];
    }
}

关键说明 #

• FFTW数组采用**行优先（C风格）**存储，因此二维坐标(row, col)对应的一维索引是row * 列数 + col，和MATLAB的列优先存储逻辑不同，需注意转换。
• 不需要像MATLAB那样先存储完整的索引矩阵，直接通过循环生成索引或实时计算位置，在C/C++中效率更高。

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