diff --git a/src/dft.c b/src/dft.c new file mode 100644 index 0000000000000000000000000000000000000000..c473f78e0ac16a42c2681adc6aa161a7c5f59fce --- /dev/null +++ b/src/dft.c @@ -0,0 +1,373 @@ +/** @file dft.c + * @brief All routines related with discrete fourier transformations. + * + * We make intensive use of the fftw library.\n + * FFTW, the Fastest Fourier Transform in the West, is a collection of fast + * C routines to compute the discrete Fourier transform. + * See http://www.fftw.org/doc/ for the manual documents of FFTW version 3.3.3. + * + * Note: All calls to FFTW should be made here, so + * 1. If we decide to change our FFT library, we only have to make + * changes in this module. + * 2. We can compile a 2d version of the code that does not depend on fftw. + * by eliminating this module. + */ + +#include +#include +#include +#include + +#include + +#include "cdr.h" +#include "cstream.h" +#include "grid.h" +#include "parameters.h" +#include "proto.h" +#include "rz_array.h" +#include "species.h" + +static void renormalize (cdr_grid_t *grid, rz_array_t *var); +static double rnd_gauss (double mu, double sigma); +static double ranf (void); + +/** @brief Transform the discrete fourier transformations ???? */ +void +dft_transform (rz_array_t *in, rz_array_t *out, int sign) +{ + int i; + fftw_plan *plans; + fftw_r2r_kind kind; + + debug (2, "dft_transform (..., sign = %d)\n", sign); + + if (sign > 0) kind = FFTW_R2HC; + else kind = FFTW_HC2R; + + assert(in->dim == 3 && out->dim == 3); + assert(in->ntheta == out->ntheta && in->nr == out->nr && in->nz == out->nz); + + /* Unfortunately, the fftw planning routines are not thread-safe. */ + + /* The +4 here is because we also transform the buffer, which is + useful to easily set boundaries, etc. */ + plans = xmalloc (sizeof (fftw_plan) * (in->nz + 4)); + + for (i = 0; i < in->nz + 4; i++) { + plans[i] = fftw_plan_many_r2r (1, /* Rank */ + &in->ntheta, /* n[] */ + in->nr + 4, /* howmany */ + RZTP (in, in->r0, in->z0 + i, 0), + /* in */ + &in->ntheta, /* inembed */ + in->strides[THETA_INDX], /* istride */ + in->strides[R_INDX], /* idist */ + RZTP (out, in->r0, out->z0 + i, 0), + /* out */ + &out->ntheta, /* onembed */ + out->strides[THETA_INDX], /* ostride */ + out->strides[R_INDX], /* odist */ + &kind, /* kind */ + FFTW_ESTIMATE); /* flags */ + + } + +#pragma omp parallel + { +#pragma omp for + for (i = 0; i < in->nz + 4; i++) { + fftw_execute (plans[i]); + } + } + + + for (i = 0; i < in->nz + 4; i++) { + fftw_destroy_plan (plans[i]); + } + + free (plans); + fftw_cleanup(); +} + + +/** @brief Calculates the Fourier transform of the derivative of a function given its + Fourier transform. */ +void +dft_diff (grid_t *grid, rz_array_t *f) +{ + int ir, iz, k; + + debug (2, "dft_diff (" grid_printf_str ", ...)\n", grid_printf_args(grid)); + + /* The parallelization here is not optimal if max_ntheta is not a multiple + of 2. */ +#pragma omp parallel + { +#pragma omp for private(ir, iz) + for (k = 0; k < grid->ntheta / 2 + 1; k++) { + if (k < grid->ntheta / 2) { + iter_grid_n (grid, ir, iz, 2) { + double re, im, re_f, im_f; + + /* For k == 0, everything is real. */ + re_f = RZT (f, ir, iz, k); + im_f = (k == 0? 0: RZT (f, ir, iz, grid->ntheta - k)); + + re = re_f * wk[k] - (k == 0? 0: im_f * wk[grid->ntheta - k]); + im = (k == 0? 0: re_f * wk[grid->ntheta - k] + im_f * wk[k]); + + RZT (f, ir, iz, k) = re / r_at (ir, grid->level); + if (k != 0) + RZT (f, ir, iz, grid->ntheta - k) = im / r_at (ir, grid->level); + } + } else if ((grid->ntheta % 2) == 0) { + iter_grid_n (grid, ir, iz, 2) { + /* Also for k = ntheta / 2, everything is real. */ + RZT (f, ir, iz, k) = wk[k] * RZT (f, ir, iz, k) + / r_at (ir, grid->level); + } + } + } + } +} + + +/** @brief Calculates the @a weight for a given cdr grid and a variable ??? + * + * For a given cdr grid and a variable, calculates its "weight": + * the integral of the variable raised to power for each Fourier mode. + * weights[] must be able to contain at least cdr->ntheta values. +*/ +void +dft_weight (cdr_grid_t *cdr, rz_array_t *var, double weights[], double power) +{ + int k, ir, iz; + double pwr, tmp; + + debug (2, "dft_weigth(" grid_printf_str ", ...)\n", grid_printf_args(cdr)); + +#pragma omp parallel + { +#pragma omp for private(ir, iz, pwr, tmp) + for (k = 0; k < cdr->ntheta; k++) { + pwr = 0; + iter_grid(cdr, ir, iz) { + /* Slow, since power is usually 1 or 2, + but this is outside the main loop, so we buy generality + at the cost of speed. */ + tmp = pow(RZT(var, ir, iz, k), power); + pwr += (tmp * r_at(ir, cdr->level)); + } + weights[k] = twopi * pwr * dr[cdr->level] * dz[cdr->level]; + } + } +} + + +/** @brief Outputs to a file 'weights.tsv' the results of dft_weight. + * + * Assumes that grid->charge is calculated but in real space then it + * transforms it to Fourier space and leaves it like that: so do not + * assume that grid->charge is unchanged. + * + * grid->dens[electrons], on the other hand, is transformed forth and back. + */ +void +dft_out_weights (cdr_grid_t *grid, const char *prefix, double t) +{ + static FILE *fp_charge = NULL; + static FILE *fp_electrons = NULL; + static double *powers_charge = NULL; + static double *powers_electrons = NULL; + int i; + + if (NULL == fp_charge) { + char *fname; + asprintf (&fname, "%s/charge-weights.tsv", prefix); + fp_charge = fopen (fname, "w"); + + if (NULL == fp_charge) { + fatal ("Unable to open %s\n", fname); + return; + } + free (fname); + + asprintf (&fname, "%s/electrons-weights.tsv", prefix); + fp_electrons = fopen (fname, "w"); + + if (NULL == fp_electrons) { + fatal ("Unable to open %s\n", fname); + return; + } + free (fname); + } + + if (NULL == powers_charge) { + powers_charge = xmalloc (sizeof(double) * grid->ntheta); + powers_electrons = xmalloc (sizeof(double) * grid->ntheta); + } + + dft_transform (grid->charge, grid->charge, 1); + dft_weight (grid, grid->charge, powers_charge, 2.0); + + /* The electron density has to be transformed back, since we will still need + it. */ + dft_transform (grid->dens[electrons], grid->dens[electrons], 1); + dft_weight (grid, grid->dens[electrons], powers_electrons, 1.0); + dft_transform (grid->dens[electrons], grid->dens[electrons], -1); + renormalize(grid, grid->dens[electrons]); + + fprintf (fp_charge, "%g", t); + fprintf (fp_electrons, "%g", t); + + for (i = 0; i < grid->ntheta; i++) { + fprintf (fp_charge, "\t%g", powers_charge[i]); + fprintf (fp_electrons, "\t%g", powers_electrons[i]); + } + + fprintf (fp_charge, "\n"); + fprintf (fp_electrons, "\n"); + + fflush(fp_charge); + fflush(fp_electrons); +} + + +/** @brief Perturbs a FFT-transformed variable. + * + * Assumes that the unperturbed variable is axi-symmetrical and hence + * all modes are zero except k=0. The perturbation is then selected + * as a Gaussian random number epsilon_k times the zero-mode value of + * the variable. epsilon_k is distributed as + * epsilon_k = N(0, perturb_epsilon). + */ +void +dft_perturb (cdr_grid_t *cdr, rz_array_t *var, double *epsilon_k) +{ + int k, ir, iz; + +#pragma omp parallel + { +#pragma omp for private(ir, iz) + for (k = 1; k < cdr->ntheta; k++) { + if (k > perturb_max_k && k < (cdr->ntheta - perturb_max_k)) + continue; + + iter_grid_n (cdr, ir, iz, 2) { + RZT (var, ir, iz, k) = + /* We multiply the perturbation by r to ensure that it is + continuous at r -> 0. */ + epsilon_k[k - 1] * RZT (var, ir, iz, 0) * r_at (ir, cdr->level); + } + } + } +} + +/** @brief Takes a density in real space, transforms it to Fourier space, + * perturbs it and then transforms back to real space. */ +void +dft_dens_perturb_r (cdr_grid_t *grid, int species, double *epsilon_k) +{ + int ir, iz, k, allocated = FALSE; + cdr_grid_t *leaf; + double A; + + debug (2, "dft_dens_perturb_r(" grid_printf_str ", %s)\n", + grid_printf_args(grid), spec_index[species]->name); + + /* We normalize the perturbation using the maximum of the densities. */ + A = invpi32 * seed_N / (seed_sigma_x * seed_sigma_y * seed_sigma_z); + + if (NULL == epsilon_k) { + epsilon_k = (double*) xmalloc (sizeof(double) * (grid->ntheta - 1)); + + for (k = 1; k < grid->ntheta; k++) { + /* We divide by ntheta because the FFT is not normalized and we want + to express perturb_epsilon as a fraction of the k=0 mode. */ + epsilon_k[k - 1] = rnd_gauss (0, perturb_epsilon) / grid->ntheta / A; + } + + allocated = TRUE; + } + + iter_childs (grid, leaf) { + dft_dens_perturb_r (leaf, species, epsilon_k); + } + + dft_transform (grid->dens[species], grid->dens[species], 1); + + renormalize(grid, grid->dens[species]); + + dft_perturb (grid, grid->dens[species], epsilon_k); + dft_transform (grid->dens[species], grid->dens[species], -1); + + if (allocated) free (epsilon_k); +} + + +/** @brief Renormalize @a var for some cases. + * + * FFTW does not normalize the Fourier transform, so after transforming + * forth and back, the original values are multiplied by max_ntheta. + * We use this routine to renormalize @a var in those cases. + * + * Note that this is taken care of when calculating the charge, so + * in most cases you do not need to call this function. + */ +static void +renormalize (cdr_grid_t *grid, rz_array_t *var) +{ + int ir, iz, itheta; + +#pragma omp parallel + { +#pragma omp for private(ir, iz) + for (itheta = 0; itheta < grid->ntheta; itheta++) { + iter_grid_n (grid, ir, iz, 2) { + RZT(var, ir, iz, itheta) /= grid->ntheta; + } + } + } +} + + +/** @brief Returns a random number with a gaussian distribution centered + * around mu with width sigma. + */ +static double +rnd_gauss(double mu, double sigma) +{ + double x1, x2, w; + static double y1, y2; + static int has_more = FALSE; + + + if (has_more) { + has_more = FALSE; + return mu + y2 * sigma; + } + + do { + x1 = 2.0 * ranf() - 1.0; + x2 = 2.0 * ranf() - 1.0; + w = x1 * x1 + x2 * x2; + } while (w >= 1.0); + + w = sqrt ((-2.0 * log (w)) / w); + y1 = x1 * w; + y2 = x2 * w; + + has_more = TRUE; + + return mu + y1 * sigma; +} + +#define AM (1.0 / RAND_MAX) + +/** @brief Returns a random number uniformly distributed in [0, 1] */ +static double +ranf (void) +{ + return rand() * AM; +}