dft.c

Go to the documentation of this file.

#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include <fftw3.h> 

#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);

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();
}


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);
   }
      }   
    }
  }
}


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];
    }
  }
}


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);
}


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);
      }
    }
  }
}

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);
}


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;
      }
    }
  }
}


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)

static double
ranf (void)
{
  return rand() * AM;
}