fftw

#include <fftw3-mpi.h>
#include <mpi.h>
#include <iostream>
#include <cmath>
#include <functional>

using namespace std;

static double main_function(const double x1, const double x2, const double x3) {
    return exp(sin(x1 - x2 + x3));
}
static double derivative_x1(const double x1, const double x2, const double x3) {
    return exp(sin(x1 - x2 + x3)) * cos(x1 - x2 + x3);
}
static double derivative_x2(const double x1, const double x2, const double x3) {
    return -exp(sin(x1 - x2 + x3)) * cos(x1 - x2 + x3);
}
static double derivative_x3(const double x1, const double x2, const double x3) {
    return exp(sin(x1 - x2 + x3)) * cos(x1 - x2 + x3);
}

struct FFTW_ptrs {
    ptrdiff_t alloc_local, local_n0, local_0_start;
    FFTW_ptrs(ptrdiff_t _alloc_local, ptrdiff_t _local_n0, ptrdiff_t _local_0_start)
    : alloc_local(_alloc_local),
      local_n0(_local_n0),
      local_0_start(_local_0_start)
    {}
};

void find_fourie_derivative(function<double(const double, const double, const double)> *derivatives,
                            const ptrdiff_t N,
                            FFTW_ptrs ptrs,
                            const double range_from,
                            const double range_to)
{
    fftw_plan forward, backward;
    double *real          = fftw_alloc_real(ptrs.alloc_local * 2);
    fftw_complex *complex = fftw_alloc_complex(ptrs.alloc_local);
    forward  = fftw_mpi_plan_dft_r2c_3d(N, N, N, real, complex, MPI_COMM_WORLD, FFTW_MEASURE);
    backward = fftw_mpi_plan_dft_c2r_3d(N, N, N, complex, real, MPI_COMM_WORLD, FFTW_MEASURE);
    double coefficient = 0;
    auto indices = new double[N];
    for (ptrdiff_t i = 0; i <= N / 2; i++) {
        indices[i] = i;
    }
    for (ptrdiff_t i = N / 2 + 1; i < N; i++) {
        indices[i] = i - N;
    }
    for (int parameter = 0; parameter < 3; parameter++) {
        //filling values of main function
        for (ptrdiff_t i = 0; i < ptrs.local_n0; i++) {
            for (ptrdiff_t j = 0; j < N; j++) {
                for (ptrdiff_t k = 0; k < N; k++) {
                    double current_x1 = range_to * (ptrs.local_0_start + i) / N;
                    double current_x2 = range_to * j / N;
                    double current_x3 = range_to * k / N;
                    real[(i * N + j) * (2 * (N / 2 + 1)) + k] = main_function(current_x1, current_x2, current_x3);
                }
            }
        }
        fftw_execute(forward);

        for (ptrdiff_t i = 0; i < ptrs.local_n0; i++) {
            for (ptrdiff_t j = 0; j < N; j++) {
                for (ptrdiff_t k = 0; k < N/2 + 1; k++) {
                    switch (parameter) {
                        case 0:
                            coefficient = indices[ptrs.local_0_start + i];
                            break;
                        case 1:
                            coefficient = indices[j];
                            break;
                        case 2:
                            coefficient = k;
                        default:
                            break;
                    }
                    complex[(i * N + j) * (2 * (N / 2 + 1)) + k][0] *=  coefficient;
                    complex[(i * N + j) * (2 * (N / 2 + 1)) + k][1] *= -coefficient;
                    swap(complex[(i * N + j) * (2 * (N / 2 + 1)) + k][0],
                            complex[(i * N + j) * (2 * (N / 2 + 1)) + k][1]);
                }
            }
        }

        fftw_execute(backward);

        for (ptrdiff_t i = 0; i < ptrs.local_n0; i++)
            for (ptrdiff_t j = 0; j < N; j++)
                for (ptrdiff_t k = 0; k < N; k++)
                    real[(i * N + j) * (2 * (N / 2 + 1)) + k] /= N*N*N;

        double max_diff = 0.0;

        for (ptrdiff_t i = 0; i < ptrs.local_n0; i++)
            for (ptrdiff_t j = 0; j < N; j++)
                for (ptrdiff_t k = 0; k < N; k++) {
                    double current_x1 = range_to * (ptrs.local_0_start + i) / N;
                    double current_x2 = range_to * j / N;
                    double current_x3 = range_to * k / N;
                    max_diff = max(abs(real[(i * N + j) * (2 * (N / 2 + 1)) + k] -
                                       derivatives[parameter](current_x1, current_x2, current_x3)), max_diff);
                }
        cout << "Difference in derivative for x1: " << max_diff << endl;
    }
    fftw_free(real);
    fftw_free(complex);
    fftw_destroy_plan(forward);
    fftw_destroy_plan(backward);
}

int main(int argc, char **argv)
{
    if (argc != 2) {
        cout << "Wrong number of arguments, program will be terminated\n";
        return -1;
    }
    const ptrdiff_t N = strtol(argv[1], nullptr, 10);
    ptrdiff_t alloc_local, local_n0, local_0_start;
    alloc_local = fftw_mpi_local_size_3d(N, N, N / 2 + 1,
                                         MPI_COMM_WORLD, &local_n0, &local_0_start);

    MPI_Init(&argc, &argv);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    fftw_mpi_init();
    FFTW_ptrs ptrs(alloc_local, local_n0, local_0_start);
    const double rangeFrom = 0.0, rangeTo = 2 * M_PI;
    function<double(const double, const double, const double)> derivatives[3];
    derivatives[0] = derivative_x1;
    derivatives[1] = derivative_x2;
    derivatives[2] = derivative_x3;
    find_fourie_derivative(derivatives, N, ptrs, rangeFrom, rangeTo);
    MPI_Finalize();
    return 0;
}