Poisson equation solver in 3-D Sine transform of r2r type using FFTW in Fortran

2
votes

I have written following code in Fortran to solve Poisson's equation using r2r (real to real) type FFTW sin transform. In this code, first I have done normal FFTW i.e. r2c (real to complex) type of a mathematical function 27*sin(3x)*sin(3y)*sin(3z) and then divide it by 27 (3*3+3*3+3*3) to calculate second derivative of input function. A 3-D plot of amplitude of input function shows it's correct amplitude.amplitude along z-axis against x-,y- co-ordinate.
Then inverse FFTW of c2r (complex to real) type regenerates the input function but the amplitude now reduced to 1 as shown in 3-D plot 2. This says my code for Poisson equation solver works fine in case of normal FFTW.

Program Derivative

! To run this code: gcc dst_3d.c -c -std=c99 && gfortran derivative.f95 dst_3d.o -lfftw3 && ./a.out

implicit none

include "fftw3.f"

integer ( kind = 4 ), parameter :: Nx = 64
integer ( kind = 4 ), parameter :: Ny = 64
integer ( kind = 4 ), parameter :: Nz = 64

real ( kind = 8 ), parameter :: pi=3.14159265358979323846d0

integer ( kind = 4 ) i,j,k
real ( kind = 8 ) Lx,Ly,Lz,dx,dy,dz,kx,ky,kz
real ( kind = 8 ) x(Nx),y(Ny),z(Nz)

real ( kind = 8 ) in_dst(Nx,Ny,Nz),out_dst(Nx,Ny,Nz) ! DST
real ( kind = 8 ) in_k_dst(Nx,Ny,Nz),out_k_dst(Nx,Ny,Nz) ! DST

real ( kind = 8 ) in_dft(Nx,Ny,Nz),out_dft(Nx,Ny,Nz) ! DFT
complex ( kind = 8 ) in_k_dft(Nx/2+1,Ny,Nz),out_k_dft(Nx/2+1,Ny,Nz) ! DFT
integer ( kind = 8 ) plan_forward,plan_backward ! DFT

! System Size.
Lx = 2.0d0*pi; Ly = 2.0d0*pi; Lz = 2.0d0*pi 

! Grid Resolution.
dx = Lx/dfloat(Nx); dy = Ly/dfloat(Ny); dz = Lz/dfloat(Nz)

! =================================== INPUT ===========================================

! Initial Profile Details.
do i = 1, Nx
  x(i) = dfloat(i-1)*dx
  do j = 1, Ny
    y(j) = dfloat(j-1)*dy
    do k = 1, Nz
      z(k) = dfloat(k-1)*dz
      in_dst(i,j,k) = 27.0d0*sin(3.0d0*x(i))*sin(3.0d0*y(j))*sin(3.0d0*z(k))
      in_dft(i,j,k) = in_dst(i,j,k)
      write(10,*) x(i), y(j), z(k), in_dft(i,j,k)
    enddo
  enddo
enddo

! =================================== 3D DFT ===========================================

  call dfftw_plan_dft_r2c_3d_ (plan_forward, Nx, Ny, Nz, in_dft, in_k_dft, FFTW_ESTIMATE)
  call dfftw_execute_ (plan_forward)
  call dfftw_destroy_plan_ (plan_forward)

do i = 1, Nx/2+1
  do j = 1, Ny/2
    do k = 1, Nz/2
      kx = 2.0d0*pi*dfloat(i-1)/Lx
      ky = 2.0d0*pi*dfloat(j-1)/Ly
      kz = 2.0d0*pi*dfloat(k-1)/Lz
      out_k_dft(i,j,k) = in_k_dft(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo 
    do k = Nz/2+1, Nz
      kx = 2.0d0*pi*dfloat(i-1)/Lx
      ky = 2.0d0*pi*dfloat(j-1)/Ly
      kz = 2.0d0*pi*dfloat((k-1)-Nz)/Lz
      out_k_dft(i,j,k) = in_k_dft(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
  enddo
  do j = Ny/2+1, Ny
    do k = 1, Nz/2
      kx = 2.0d0*pi*dfloat(i-1)/Lx
      ky = 2.0d0*pi*dfloat((j-1)-Ny)/Ly
      kz = 2.0d0*pi*dfloat(k-1)/Lz
      out_k_dft(i,j,k) = in_k_dft(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
    do k = Nz/2+1, Nz
      kx = 2.0d0*pi*dfloat(i-1)/Lx
      ky = 2.0d0*pi*dfloat((j-1)-Ny)/Ly
      kz = 2.0d0*pi*dfloat((k-1)-Nz)/Lz
      out_k_dft(i,j,k) = in_k_dft(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
  enddo   
enddo

out_k_dft(1,1,1) = in_k_dft(1,1,1)

  call dfftw_plan_dft_c2r_3d_ (plan_backward, Nx, Ny, Nz, out_k_dft, out_dft, FFTW_ESTIMATE)
  call dfftw_execute_ (plan_backward)
  call dfftw_destroy_plan_ (plan_backward)

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    out_dft(i,j,k) = out_dft(i,j,k)/dfloat(Nx*Ny*Nz)
    write(20,*) x(i), y(j), z(k), out_dft(i,j,k)    
    enddo
  enddo   
enddo

! =================================== 3D DST ===========================================

  call Forward_FFT(Nx, Ny, Nz, in_dst, in_k_dst)

do k = 1, Nz
  do j = 1, Ny/2
    do i = 1, Nx/2
      kz = 2.0d0*pi*dfloat((k-1))/Lz
      ky = 2.0d0*pi*dfloat((j-1))/Ly
      kx = 2.0d0*pi*dfloat((i-1))/Lx
      out_k_dst(i,j,k) = in_k_dst(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo 
    do i = Nx/2+1, Nx
      kz = 2.0d0*pi*dfloat((k-1))/Lz
      ky = 2.0d0*pi*dfloat((j-1))/Ly
      kx = 2.0d0*pi*dfloat(Nx-(i-1))/Lx
      out_k_dst(i,j,k) = in_k_dst(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
  enddo  
  do j = Ny/2+1, Ny
    do i = 1, Nx/2
      kz = 2.0d0*pi*dfloat((k-1))/Lz
      ky = 2.0d0*pi*dfloat(Ny-(j-1))/Ly
      kx = 2.0d0*pi*dfloat((i-1))/Lx
      out_k_dst(i,j,k) = in_k_dst(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
    do i = Nx/2+1, Nx
      kz = 2.0d0*pi*dfloat((k-1))/Lz
      ky = 2.0d0*pi*dfloat(Ny-(j-1))/Ly
      kx = 2.0d0*pi*dfloat(Nx-(i-1))/Lx
      out_k_dst(i,j,k) = in_k_dst(i,j,k)/(kx*kx+ky*ky+kz*kz)
    enddo
  enddo   
enddo

out_k_dst(1,1,1) = in_k_dst(1,1,1)

  call Backward_FFT(Nx, Ny, Nz, out_k_dst, out_dst)  

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    out_dst(i,j,k) = out_dst(i,j,k)/dfloat(8*Nx*Ny*Nz)
    write(30,*) x(i), y(j), z(k), out_dst(i,j,k)
    enddo
  enddo   
enddo

end program Derivative

! ================================== FFTW SUBROUTINES 
====================================================

! ================================================================= !
! Wrapper Subroutine to call forward fftw c functions for 3d arrays !
! ================================================================= !

subroutine Forward_FFT(Nx, Ny, Nz, in, out)
implicit none
integer ( kind = 4 ) Nx,Ny,Nz,i,j,k
real ( kind = 8 ) in(Nx, Ny, Nz),out(Nx, Ny, Nz),dum(Nx*Ny*Nz)

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    dum(((i-1)*Ny+(j-1))*Nz+k) = in(i,j,k)
    enddo
  enddo
enddo

  call dst3f(Nx, Ny, Nz, dum)

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    out(i,j,k) = dum(((i-1)*Ny+(j-1))*Nz+k)
    enddo
  enddo
enddo

end subroutine

! ================================================================== !
! Wrapper Subroutine to call backward fftw c functions for 3d arrays !
! ================================================================== !

subroutine Backward_FFT(Nx, Ny, Nz, in, out)
implicit none
integer ( kind = 4 ) Nx,Ny,Nz,i,j,k
real ( kind = 8 ) in(Nx, Ny, Nz),out(Nx, Ny, Nz),dum(Nx*Ny*Nz)

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    dum(((i-1)*Ny+(j-1))*Nz+k) = in(i,j,k)
    enddo
  enddo
enddo

  call dst3b(Nx, Ny, Nz, dum)

do i = 1, Nx
  do j = 1, Ny
    do k = 1, Nz
    out(i,j,k) = dum(((i-1)*Ny+(j-1))*Nz+k)
    enddo
  enddo
enddo

end subroutine
! ==================================================================

This code uses below C-wrapper to calculate forward 3D FFTW sine transform and backward 3D FFTW sine transform,

#include <fftw3.h>

int dst3f_(int *n0, int *n1, int *n2, double *in3cs)
{
    double *out3cs;
    out3cs = (double*) fftw_malloc(sizeof(double) * (*n0) * (*n1) * (*n2));
    fftw_plan p3cs;
    p3cs = fftw_plan_r2r_3d(*n0, *n1, *n2, in3cs, out3cs, FFTW_RODFT10, FFTW_RODFT10, FFTW_RODFT10, FFTW_ESTIMATE);

    fftw_execute(p3cs);
    fftw_destroy_plan(p3cs);

    for(int i0=0;i0<*n0;i0++) {
        for(int i1=0;i1<*n1;i1++) {
            for(int i2=0;i2<*n2;i2++) {
                in3cs[(i0*(*n1)+i1)*(*n2)+i2] = out3cs[(i0*(*n1)+i1)*(*n2)+i2];
            }
        }
    }

    return 0;
}

int dst3b_(int *n0, int *n1, int *n2, double *in3cs)
{
    double *out3cs;
    out3cs = (double*) fftw_malloc(sizeof(double) * (*n0) * (*n1) * (*n2));
    fftw_plan p3cs;
    p3cs = fftw_plan_r2r_3d(*n0, *n1, *n2, in3cs, out3cs, FFTW_RODFT01, FFTW_RODFT01, FFTW_RODFT01, FFTW_ESTIMATE);

    fftw_execute(p3cs);
    fftw_destroy_plan(p3cs);

    for(int i0=0;i0<*n0;i0++) {
        for(int i1=0;i1<*n1;i1++) {
            for(int i2=0;i2<*n2;i2++) {
                in3cs[(i0*(*n1)+i1)*(*n2)+i2] = out3cs[(i0*(*n1)+i1)*(*n2)+i2];
            }
        }
    }

    return 0;
}

Then I have tried to solve the same Poisson equation now using sine transform FFTW i.e. r2r (real to real) type. When I do a 3-D plot of output as shown here 3, the amplitude is now reduced to less than 1. I couldn't find out where is my mistake in the code because of which this amplitude is reduced in case of sine transform.

1
cfortranfftfftwdft
Welcome, please always use tag fortran. Note that magic numbers 4 and 8 in kind=4 and kind=8 are ugly and not portable. I did loads of similar stuff and mostly my eigenvalues were wrong. Double check them, try how different wavelenght inputs are affected. Make sure which wavelenght is stored in which array index in the transformed array. – Vladimir F
Furthermore, your test for the c2r - r2c version is not complete, it is very incomplete in fact. You must test a right hand side with multiple wavenumbers present. Not just a simple sine wave. One sine wave demonstrates you have one eigenvalue correct. You must test several of them. – Vladimir F

1 Answers

1
votes

The use of real-to-real transforms to solve the Poisson's equation is very attractive, as it allows using various boundary conditions. Nevertheless, the sensing points do not exacly correspond to those considered for periodic boundary conditions.

For periodic boundary conditions, sensing points are on a regular grid, say 0,1..,n-1. If the size of the unit cell is Lx, then the spacing between the points is Lx/n. End of the story.

Now, let's consider boundary conditions such that the DST III is to be used for forward transform, i.e. flag RODFT01. As shown there, it is signaled in the documentation of FFTW that :

FFTW_RODFT01 (DST-III): odd around j=-1 and even around j=n-1.

The sensing points are still 0,1..,n-1. If the length is Lx, the spacing is still dx=Lx/(n). But the input functions and base functions of the DST III are odd around j=-1 and even around j=n-1. This explains a discrepency of magnitude:

in_dst(i,j,k) = 27.0d0*sin(3.0d0*x(i))*sin(3.0d0*y(j))*sin(3.0d0*z(k))

Indeed, this input is not odd around i=-1 and even around i=n-1. It is odd around both i=0 and i=n. As a consequence, the following actions could be helpful:

  • Making sure that the considered forward transform is the right one. For odd-odd boundary conditions, it could be FFTW_RODFT00 (DST-I): odd around j=-1 and odd around j=n..
  • Evaluate the input at the correct sensing points. For DST-I, using n points, dx=Lx/(n+1) and x_i=dx*(i+1) for i=0,1,..,n-1 if the function is odd at 0 and Lx.
  • Compute the frequencies in the transformed space. The way the inverse transform writes helps finding these frequencies. Since the inverse of DST I is DST I, the first frequency (k=0) corresponds to a period of 2Lx; the second (k=1) corresponds to Lx, k-st frequencies corresponds to a period of 2Lx/(k+1).

Finally, a scalling of the output may be needed, as the FFTW transforms are not normalized. For DST-I, FFTW_RODFT00, it's N=2(n+1). The advice of VladimirF undoubtly is a good one. Indeed, while testing a single frequency is ideal to understand and implement the algorithm, the final test must cover multiple frequencies to ensure that the program is reliable.