Using FFTW library to take a 2D FFT of an Eigen tensor in C++

40 Views Asked by At

So I am trying to take an FFT along a specific direction of an Eigen tensor, similarly to my code in MATLAB.

In MATLAB:

Nx = 8;  
Ny = 6; 
Nz= 8;
Lx =6;
Ly = 6;
xi_x =  (2*pi)/Lx;
yi_y =  (2*pi)/Ly;
xi  = ((0:Nx-1)/Nx)*(2*pi);
yi  = ((0:Ny-1)/Ny)*(2*pi);
x =  xi/xi_x;
y =  yi/yi_y;

zlow = 0;%a
zupp =6; %b
Lz = (zupp-zlow);
zgl = (1/2)*(((zupp-zlow)*zgl) + (zupp+zlow));
[X,Z,Y] = meshgrid(x,zgl,y);
%ICs
A = 2*pi / Lx;
B = 2*pi / Ly;
u = (Z-zlow) .* (Z-zupp) .* sin(A*X).* sin(B*Y);
%Take FFT along only X and Y direction
uh =fft(fft(u,[],2),[],3);

The meshgrid function in MATLAB works in such a way that this line [code][X,Z,Y] = meshgrid(x,zgl,y); [/code] returns a 3d grid with a size of z-by-x-by-y. Then when taking my FFT along X which is the second element:

uhx =fft(u,[],2);

and along Y which is 3rd element:

uhx =fft(u,[],3);

Now, I am trying to do something similar in C++. I can already perform the following in C++:

fft(fft(fft(u,[],1),[],2),[],3);

which is by using fftw_plan_dft_3d and this works perfectly on an Eigen tensor. However, when attempting to take a 1D FFT along say the X direction only the code returns just incorrect results.

For example, I define my function using Eigen tensor as the following:

static const int nx = 8; 
static const int ny = 6;  
static const int nz = 8; 
using namespace std;
using namespace Eigen;

double SpatialMesh3DTrans(double dx,double dy, double dz, double zlow, double zupp,Eigen::Tensor<double, 3>& xarr, Eigen::Tensor<double, 3>& yarr, Eigen::Tensor<double, 3>& zarr);
int main(){

double Lx = 6.;
double zlow = 0.;
double zupp = 6.;   
double Ly = 6.;
double Lz = (zupp-zlow);
double dx = Lx / nx;
double dy = Ly / ny;
double dz = Lz / nz;

Eigen::Tensor<double, 3> eXX((nz+1),nx,ny); 
eXX.setZero(); 

Eigen::Tensor<double, 3> eYY((nz+1),nx,ny);   //tensor(row,col,matrix)
eYY.setZero(); 
Eigen::Tensor<double, 3> eZZ((nz+1),nx,ny);   
eZZ.setZero(); 

double kmax = SpatialMesh3DTrans(dx,dy,dz,zlow,zupp,eXX,eYY,eZZ);

Eigen::Tensor<double, 3> uFun((nz+1),nx,ny);   //tensor(row,col,matrix)
uFun.setZero();
 
    double A = 2*EIGEN_PI/Lx;
    double A1 = 2*EIGEN_PI/Ly;
   
        for(int x = 0; x< nx; x++){//cols
            for(int y = 0; y< ny; y++){
                for(int z = 0; z< nz+1; z++){
                uFun(z,x,y) = ( (eZZ(z,x,y)-zlow) * (eZZ(z,x,y)-zupp) * sin(A*eXX(z,x,y) ) * sin(A1 *eYY(z,x,y) ) );
                
            }
        }
    }
}

double SpatialMesh3DTrans(double dx,double dy, double dz, double zlow, double zupp,Eigen::Tensor<double, 3>& xarr, Eigen::Tensor<double, 3>& yarr, Eigen::Tensor<double, 3>& zarr){
    for(int x = 0; x< nx; x++){
        for(int y = 0; y< ny; y++){
            for(int z = 0; z< nz+1; z++){
                xarr(z,x,y)  = x*dx;
                yarr(z,x,y)  = y*dy;
                zarr(z,x,y)  = 1. * cos(((z) * EIGEN_PI )/nz); 
                zarr(z,x,y) = (1./2)*(((zupp-zlow)*zarr(z,x,y) ) + (zupp+zlow));
            }
        }       
    }
    
        if (dx < dy){
            return 1/(2*dx);
        }else {
            return 1/(2*dy);
        }
}

How can I take a 1D FFT of Eigen tensor? Also, I am able to do this using Eigen matrix and taking a 1D FFT along X of 2D Eigen matrix but not an Eigen tensor for some reason!

EDIT: This is my first attempt in taking that is not really working:

void r2cfft2d(Eigen::Tensor<double, 3>& rArr, Eigen::Tensor<std::complex<double>, 3>& cArr){    
double* input_array = fftw_alloc_real(nx*ny*(nz+1));
        
memcpy(input_array, rArr.data(), (nx*ny*(nz+1))* sizeof(double));
fftw_complex *output_array;
output_array = (fftw_complex*) fftw_malloc((nx*ny*(nz+1)) * sizeof(fftw_complex));
int constexpr n[]{ny,nx};

fftw_plan forward = fftw_plan_many_dft_r2c(2,n, (nz+1), input_array,n,1,(ny*nx), output_array,n,1,(ny*nx), FFTW_MEASURE);
 fftw_execute(forward);

memcpy(cArr.data(),output_array, (nx*ny*(nz+1)) * sizeof(fftw_complex));
fftw_free(input_array);
fftw_free(output_array);
}
0

There are 0 best solutions below