runner_doiact_fft.c 10.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2016 Matthieu Schaller (matthieu.schaller@durham.ac.uk)
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 ******************************************************************************/

/* Config parameters. */
#include "../config.h"

Matthieu Schaller's avatar
Matthieu Schaller committed
23
#ifdef HAVE_FFTW
24
#include <fftw3.h>
Matthieu Schaller's avatar
Matthieu Schaller committed
25
26
#endif

27
28
29
30
/* This object's header. */
#include "runner_doiact_fft.h"

/* Local includes. */
31
32
#include "engine.h"
#include "error.h"
33
#include "kernel_long_gravity.h"
34
#include "runner.h"
35
#include "space.h"
Matthieu Schaller's avatar
Matthieu Schaller committed
36
#include "timers.h"
37

38
39
#ifdef HAVE_FFTW

40
41
42
43
44
45
46
47
48
49
/**
 * @brief Returns 1D index of a 3D NxNxN array using row-major style.
 *
 * @param i Index along x.
 * @param j Index along y.
 * @param k Index along z.
 * @param N Size of the array along one axis.
 */
__attribute__((always_inline)) INLINE static int row_major_id(int i, int j,
                                                              int k, int N) {
50
  return ((i % N) * N * N + (j % N) * N + (k % N));
51
52
53
54
55
56
57
58
59
}

/**
 * @brief Assigns a given multipole to a density mesh using the CIC method.
 *
 * @param m The #multipole.
 * @param rho The density mesh.
 * @param N the size of the mesh along one axis.
 * @param fac The width of a mesh cell.
Matthieu Schaller's avatar
Matthieu Schaller committed
60
 * @param dim The dimensions of the simulation box.
61
 */
62
__attribute__((always_inline)) INLINE static void multipole_to_mesh_CIC(
63
64
    const struct gravity_tensors* m, double* rho, int N, double fac,
    const double dim[3]) {
65

66
67
68
69
70
71
  /* Box wrap the multipole's position */
  const double CoM_x = box_wrap(m->CoM[0], 0., dim[0]);
  const double CoM_y = box_wrap(m->CoM[1], 0., dim[1]);
  const double CoM_z = box_wrap(m->CoM[2], 0., dim[2]);

  int i = (int)(fac * CoM_x);
72
  if (i >= N) i = N - 1;
73
  const double dx = fac * CoM_x - i;
74
75
  const double tx = 1. - dx;

76
  int j = (int)(fac * CoM_y);
77
  if (j >= N) j = N - 1;
78
  const double dy = fac * CoM_y - j;
79
80
  const double ty = 1. - dy;

81
  int k = (int)(fac * CoM_z);
82
  if (k >= N) k = N - 1;
83
  const double dz = fac * CoM_z - k;
84
85
  const double tz = 1. - dz;

86
87
88
89
90
91
#ifdef SWIFT_DEBUG_CHECKS
  if (i < 0 || i >= N) error("Invalid multipole position in x");
  if (j < 0 || j >= N) error("Invalid multipole position in y");
  if (k < 0 || k >= N) error("Invalid multipole position in z");
#endif

92
  /* CIC ! */
93
94
95
96
97
98
99
100
  rho[row_major_id(i + 0, j + 0, k + 0, N)] += m->m_pole.M_000 * tx * ty * tz;
  rho[row_major_id(i + 0, j + 0, k + 1, N)] += m->m_pole.M_000 * tx * ty * dz;
  rho[row_major_id(i + 0, j + 1, k + 0, N)] += m->m_pole.M_000 * tx * dy * tz;
  rho[row_major_id(i + 0, j + 1, k + 1, N)] += m->m_pole.M_000 * tx * dy * dz;
  rho[row_major_id(i + 1, j + 0, k + 0, N)] += m->m_pole.M_000 * dx * ty * tz;
  rho[row_major_id(i + 1, j + 0, k + 1, N)] += m->m_pole.M_000 * dx * ty * dz;
  rho[row_major_id(i + 1, j + 1, k + 0, N)] += m->m_pole.M_000 * dx * dy * tz;
  rho[row_major_id(i + 1, j + 1, k + 1, N)] += m->m_pole.M_000 * dx * dy * dz;
101
102
103
104
105
106
107
108
109
110
}

/**
 * @brief Computes the potential on a multipole from a given mesh using the CIC
 * method.
 *
 * @param m The #multipole.
 * @param pot The potential mesh.
 * @param N the size of the mesh along one axis.
 * @param fac width of a mesh cell.
Matthieu Schaller's avatar
Matthieu Schaller committed
111
 * @param dim The dimensions of the simulation box.
112
 */
113
__attribute__((always_inline)) INLINE static void mesh_to_multipole_CIC(
114
115
116
117
118
119
120
    struct gravity_tensors* m, const double* pot, int N, double fac,
    const double dim[3]) {

  /* Box wrap the multipole's position */
  const double CoM_x = box_wrap(m->CoM[0], 0., dim[0]);
  const double CoM_y = box_wrap(m->CoM[1], 0., dim[1]);
  const double CoM_z = box_wrap(m->CoM[2], 0., dim[2]);
121

122
  int i = (int)(fac * CoM_x);
123
  if (i >= N) i = N - 1;
124
  const double dx = fac * CoM_x - i;
125
126
  const double tx = 1. - dx;

127
  int j = (int)(fac * CoM_y);
128
  if (j >= N) j = N - 1;
129
  const double dy = fac * CoM_y - j;
130
131
  const double ty = 1. - dy;

132
  int k = (int)(fac * CoM_z);
133
  if (k >= N) k = N - 1;
134
  const double dz = fac * CoM_z - k;
135
136
  const double tz = 1. - dz;

137
138
139
140
141
142
#ifdef SWIFT_DEBUG_CHECKS
  if (i < 0 || i >= N) error("Invalid multipole position in x");
  if (j < 0 || j >= N) error("Invalid multipole position in y");
  if (k < 0 || k >= N) error("Invalid multipole position in z");
#endif

143
  /* CIC ! */
144
145
146
147
148
149
150
151
  m->pot.F_000 += pot[row_major_id(i + 0, j + 0, k + 0, N)] * tx * ty * tz;
  m->pot.F_000 += pot[row_major_id(i + 0, j + 0, k + 1, N)] * tx * ty * dz;
  m->pot.F_000 += pot[row_major_id(i + 0, j + 1, k + 0, N)] * tx * dy * tz;
  m->pot.F_000 += pot[row_major_id(i + 0, j + 1, k + 1, N)] * tx * dy * dz;
  m->pot.F_000 += pot[row_major_id(i + 1, j + 0, k + 0, N)] * dx * ty * tz;
  m->pot.F_000 += pot[row_major_id(i + 1, j + 0, k + 1, N)] * dx * ty * dz;
  m->pot.F_000 += pot[row_major_id(i + 1, j + 1, k + 0, N)] * dx * dy * tz;
  m->pot.F_000 += pot[row_major_id(i + 1, j + 1, k + 1, N)] * dx * dy * dz;
152
153
}

154
155
#endif

Matthieu Schaller's avatar
Matthieu Schaller committed
156
157
158
159
160
161
162
/**
 * @brief Computes the potential on the top multipoles using a Fourier transform
 *
 * @param r The #runner task
 * @param timer Are we timing this ?
 */
void runner_do_grav_fft(struct runner* r, int timer) {
163

Matthieu Schaller's avatar
Matthieu Schaller committed
164
#ifdef HAVE_FFTW
165

166
167
  const struct engine* e = r->e;
  const struct space* s = e->s;
Matthieu Schaller's avatar
Matthieu Schaller committed
168
  const double a_smooth = e->gravity_properties->a_smooth;
169
170
  const double box_size = s->dim[0];
  const int cdim[3] = {s->cdim[0], s->cdim[1], s->cdim[2]};
171
  const double dim[3] = {s->dim[0], s->dim[1], s->dim[2]};
172

Matthieu Schaller's avatar
Matthieu Schaller committed
173
  TIMER_TIC;
174
175
176
177
178
179

  if (cdim[0] != cdim[1] || cdim[0] != cdim[2]) error("Non-square mesh");

  /* Some useful constants */
  const int N = cdim[0];
  const int N_half = N / 2;
180
  const double cell_fac = N / box_size;
181
182
183

  /* Recover the list of top-level multipoles */
  const int nr_cells = s->nr_cells;
184
  struct gravity_tensors* restrict multipoles = s->multipoles_top;
185

186
#ifdef SWIFT_DEBUG_CHECKS
187
188
189
  const struct cell* cells = s->cells_top;
  const integertime_t ti_current = e->ti_current;

190
  /* Make sure everything has been drifted to the current point */
191
  for (int i = 0; i < nr_cells; ++i)
192
    if (cells[i].ti_old_multipole != ti_current)
193
194
      error("Top-level multipole %d not drifted", i);
#endif
195
196

  /* Allocates some memory for the density mesh */
197
  double* restrict rho = fftw_malloc(sizeof(double) * N * N * N);
198
199
200
  if (rho == NULL) error("Error allocating memory for density mesh");

  /* Allocates some memory for the mesh in Fourier space */
201
202
  fftw_complex* restrict frho =
      fftw_malloc(sizeof(fftw_complex) * N * N * (N_half + 1));
203
204
205
206
  if (frho == NULL)
    error("Error allocating memory for transform of density mesh");

  /* Prepare the FFT library */
207
208
209
210
  fftw_plan forward_plan = fftw_plan_dft_r2c_3d(
      N, N, N, rho, frho, FFTW_ESTIMATE | FFTW_DESTROY_INPUT);
  fftw_plan inverse_plan = fftw_plan_dft_c2r_3d(
      N, N, N, frho, rho, FFTW_ESTIMATE | FFTW_DESTROY_INPUT);
211
212

  /* Do a CIC mesh assignment of the multipoles */
213
  bzero(rho, N * N * N * sizeof(double));
214
  for (int i = 0; i < nr_cells; ++i)
215
    multipole_to_mesh_CIC(&multipoles[i], rho, N, cell_fac, dim);
216
217
218
219
220
221

  /* Fourier transform to go to magic-land */
  fftw_execute(forward_plan);

  /* frho now contains the Fourier transform of the density field */
  /* frho contains NxNx(N/2+1) complex numbers */
222

223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
  /* Some common factors */
  const double green_fac = -1. / (M_PI * box_size);
  const double a_smooth2 = 4. * M_PI * a_smooth * a_smooth / ((double)(N * N));
  const double k_fac = M_PI / (double)N;

  /* Now de-convolve the CIC kernel and apply the Green function */
  for (int i = 0; i < N; ++i) {

    /* kx component of vector in Fourier space and 1/sinc(kx) */
    const int kx = (i > N_half ? i - N : i);
    const double kx_d = (double)kx;
    const double fx = k_fac * kx_d;
    const double sinc_kx_inv = (kx != 0) ? fx / sin(fx) : 1.;

    for (int j = 0; j < N; ++j) {

      /* ky component of vector in Fourier space and 1/sinc(ky) */
      const int ky = (j > N_half ? j - N : j);
      const double ky_d = (double)ky;
      const double fy = k_fac * ky_d;
      const double sinc_ky_inv = (ky != 0) ? fy / sin(fy) : 1.;

      for (int k = 0; k < N_half + 1; ++k) {

        /* kz component of vector in Fourier space and 1/sinc(kz) */
        const int kz = (k > N_half ? k - N : k);
        const double kz_d = (double)kz;
250
        const double fz = k_fac * kz_d;
251
        const double sinc_kz_inv = (kz != 0) ? fz / (sin(fz) + FLT_MIN) : 1.;
252
253
254
255
256
257
258
259

        /* Norm of vector in Fourier space */
        const double k2 = (kx_d * kx_d + ky_d * ky_d + kz_d * kz_d);

        /* Avoid FPEs... */
        if (k2 == 0.) continue;

        /* Green function */
260
261
        double W;
        fourier_kernel_long_grav_eval(k2 * a_smooth2, &W);
262
        const double green_cor = green_fac * W / (k2 + FLT_MIN);
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291

        /* Deconvolution of CIC */
        const double CIC_cor = sinc_kx_inv * sinc_ky_inv * sinc_kz_inv;
        const double CIC_cor2 = CIC_cor * CIC_cor;
        const double CIC_cor4 = CIC_cor2 * CIC_cor2;

        /* Combined correction */
        const double total_cor = green_cor * CIC_cor4;

        /* Apply to the mesh */
        const int index = N * (N_half + 1) * i + (N_half + 1) * j + k;
        frho[index][0] *= total_cor;
        frho[index][1] *= total_cor;
      }
    }
  }

  /* Correct singularity at (0,0,0) */
  frho[0][0] = 0.;
  frho[0][1] = 0.;

  /* Fourier transform to come back from magic-land */
  fftw_execute(inverse_plan);

  /* rho now contains the potential */
  /* This array is now again NxNxN real numbers */

  /* Get the potential from the mesh using CIC */
  for (int i = 0; i < nr_cells; ++i)
292
    mesh_to_multipole_CIC(&multipoles[i], rho, N, cell_fac, dim);
293
294
295
296
297
298
299
300

  /* Clean-up the mess */
  fftw_destroy_plan(forward_plan);
  fftw_destroy_plan(inverse_plan);
  fftw_free(rho);
  fftw_free(frho);

  /* Time the whole thing */
Matthieu Schaller's avatar
Matthieu Schaller committed
301
302
303
304
305
  if (timer) TIMER_TOC(timer_dograv_top_level);

#else
  error("No FFTW library found. Cannot compute periodic long-range forces.");
#endif
306
}
307

308
#ifdef HAVE_FFTW
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
void print_array(double* array, int N) {

  for (int k = N - 1; k >= 0; --k) {
    printf("--- z = %d ---------\n", k);
    for (int j = N - 1; j >= 0; --j) {
      for (int i = 0; i < N; ++i) {
        printf("%f ", array[i * N * N + j * N + k]);
      }
      printf("\n");
    }
  }
}

void print_carray(fftw_complex* array, int N) {

  for (int k = N - 1; k >= 0; --k) {
    printf("--- z = %d ---------\n", k);
    for (int j = N - 1; j >= 0; --j) {
      for (int i = 0; i < N; ++i) {
        printf("(%f %f) ", array[i * N * N + j * N + k][0],
               array[i * N * N + j * N + k][1]);
      }
      printf("\n");
    }
  }
}
335
#endif /* HAVE_FFTW */