mwe.f08

! gfortran -O3 -fopenmp -o mwe mwe.f08
program mwe
    use, intrinsic :: iso_c_binding
    use omp_lib
    implicit none

    integer, parameter :: N = 729, threadNum = 1, iterations = 1000
    integer :: i, j, k
    double precision :: start, finish
    complex(C_DOUBLE_COMPLEX), parameter :: im = (0.d0, 1.d0)
    double precision, parameter :: dx = 1.d0/dble(N), dy = dx, dt = dx
    double precision :: getFreq

    complex(C_DOUBLE_COMPLEX), allocatable, dimension(:,:) :: A, T

    allocate(A(N,N))
    allocate(T(N,N))

    A = 1.d0

    call omp_set_num_threads(threadNum)

    ! T(x,y) = e^(-i/2 (px^2 + py^2) dt)
    !$OMP PARALLEL DO
    do j = 1, N
        do i = 1, N
            T(i,j) = exp(-im/2.d0 * (getFreq(i, dx)**2.d0 + getFreq(j, dy)**2.d0) * dt)
        end do
    end do
    !$OMP END PARALLEL DO

    start = omp_get_wtime()
    do k=1,iterations
        !$OMP PARALLEL DO
        do j = 1, N
            do i = 1, N
                A(i,j) = T(i,j) * A(i,j)
            end do
        end do
        !$OMP END PARALLEL DO
        !A = T * A
    end do
    finish = omp_get_wtime()
    print '("Time = ",f12.3," seconds.")',finish-start
end program

! returns appropriate frequency for given index:
! i/N, if i <= N/2
! (i-N)/N = i/N - 1, if N/2 < i <= N
double precision function getFreq(i, d)
    integer, intent(in) :: i
    double precision, intent(in) :: d
    if (i.le.N/2) then
        getFreq = 2.d0 * pi * dble(i-1)/(d * dble(N))
    else
        getFreq = 2.d0 * pi * dble(i-N-1)/(d * dble(N))
    endif
end function