NavierStokes.py 58.3 KB
Newer Older
Chichi Lalescu's avatar
Chichi Lalescu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#######################################################################
#                                                                     #
#  Copyright 2015 Max Planck Institute                                #
#                 for Dynamics and Self-Organization                  #
#                                                                     #
#  This file is part of bfps.                                         #
#                                                                     #
#  bfps is free software: you can redistribute it and/or modify       #
#  it under the terms of the GNU General Public License as published  #
#  by the Free Software Foundation, either version 3 of the License,  #
#  or (at your option) any later version.                             #
#                                                                     #
#  bfps 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 General Public License  #
#  along with bfps.  If not, see <http://www.gnu.org/licenses/>       #
#                                                                     #
# Contact: Cristian.Lalescu@ds.mpg.de                                 #
#                                                                     #
#######################################################################


Cristian Lalescu's avatar
Cristian Lalescu committed
26

Cristian Lalescu's avatar
Cristian Lalescu committed
27
import sys
28
29
30
import os
import numpy as np
import h5py
31
import argparse
32

33
import bfps
34
import bfps.tools
35
from ._code import _code
36
from ._fluid_base import _fluid_particle_base
Cristian Lalescu's avatar
Cristian Lalescu committed
37

38
class NavierStokes(_fluid_particle_base):
39
40
41
42
43
    """Objects of this class can be used to generate production DNS codes.
    Any functionality that users require should be available through this class,
    in the sense that they can implement whatever they need by simply inheriting
    this class.
    """
Cristian Lalescu's avatar
Cristian Lalescu committed
44
45
    def __init__(
            self,
46
            name = 'NavierStokes-v' + bfps.__version__,
Cristian Lalescu's avatar
Cristian Lalescu committed
47
            work_dir = './',
48
            simname = 'test',
49
            fluid_precision = 'single',
Chichi Lalescu's avatar
Chichi Lalescu committed
50
            fftw_plan_rigor = 'FFTW_MEASURE',
51
            frozen_fields = False,
Cristian Lalescu's avatar
Cristian Lalescu committed
52
            use_fftw_wisdom = True,
Cristian Lalescu's avatar
Cristian Lalescu committed
53
54
            QR_stats_on = False,
            Lag_acc_stats_on = False):
Cristian Lalescu's avatar
Cristian Lalescu committed
55
        self.QR_stats_on = QR_stats_on
Cristian Lalescu's avatar
Cristian Lalescu committed
56
        self.Lag_acc_stats_on = Lag_acc_stats_on
Cristian Lalescu's avatar
Cristian Lalescu committed
57
58
        self.frozen_fields = frozen_fields
        self.fftw_plan_rigor = fftw_plan_rigor
59
60
        _fluid_particle_base.__init__(
                self,
61
                name = name + '-' + fluid_precision,
62
                work_dir = work_dir,
63
                simname = simname,
64
65
                dtype = fluid_precision,
                use_fftw_wisdom = use_fftw_wisdom)
66
67
68
        self.parameters['nu'] = 0.1
        self.parameters['fmode'] = 1
        self.parameters['famplitude'] = 0.5
69
70
        self.parameters['fk0'] = 2.0
        self.parameters['fk1'] = 4.0
71
72
73
74
        self.parameters['forcing_type'] = 'linear'
        self.parameters['histogram_bins'] = 256
        self.parameters['max_velocity_estimate'] = 1.0
        self.parameters['max_vorticity_estimate'] = 1.0
Cristian Lalescu's avatar
Cristian Lalescu committed
75
        self.parameters['max_Lag_acc_estimate'] = 1.0
Cristian Lalescu's avatar
Cristian Lalescu committed
76
        self.parameters['max_pressure_estimate'] = 1.0
77
78
79
80
        self.parameters['QR2D_histogram_bins'] = 64
        self.parameters['max_trS2_estimate'] = 1.0
        self.parameters['max_Q_estimate'] = 1.0
        self.parameters['max_R_estimate'] = 1.0
81
        self.file_datasets_grow = """
82
                //begincpp
83
84
85
86
                hid_t group;
                group = H5Gopen(stat_file, "/statistics", H5P_DEFAULT);
                H5Ovisit(group, H5_INDEX_NAME, H5_ITER_NATIVE, grow_statistics_dataset, NULL);
                H5Gclose(group);
87
88
                //endcpp
                """
89
90
        self.style = {}
        self.statistics = {}
Chichi Lalescu's avatar
Chichi Lalescu committed
91
        self.fluid_output = 'fs->write(\'v\', \'c\');\n'
92
        return None
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    def create_stat_output(
            self,
            dset_name,
            data_buffer,
            data_type = 'H5T_NATIVE_DOUBLE',
            size_setup = None,
            close_spaces = True):
        new_stat_output_txt = 'Cdset = H5Dopen(stat_file, "{0}", H5P_DEFAULT);\n'.format(dset_name)
        if not type(size_setup) == type(None):
            new_stat_output_txt += (
                    size_setup +
                    'wspace = H5Dget_space(Cdset);\n' +
                    'ndims = H5Sget_simple_extent_dims(wspace, dims, NULL);\n' +
                    'mspace = H5Screate_simple(ndims, count, NULL);\n' +
                    'H5Sselect_hyperslab(wspace, H5S_SELECT_SET, offset, NULL, count, NULL);\n')
        new_stat_output_txt += ('H5Dwrite(Cdset, {0}, mspace, wspace, H5P_DEFAULT, {1});\n' +
                                'H5Dclose(Cdset);\n').format(data_type, data_buffer)
        if close_spaces:
            new_stat_output_txt += ('H5Sclose(mspace);\n' +
                                    'H5Sclose(wspace);\n')
        return new_stat_output_txt
114
115
116
    def write_fluid_stats(self):
        self.fluid_includes += '#include <cmath>\n'
        self.fluid_includes += '#include "fftw_tools.hpp"\n'
Cristian Lalescu's avatar
Cristian Lalescu committed
117
118
        self.stat_src += """
                //begincpp
119
120
121
                hid_t stat_group;
                if (myrank == 0)
                    stat_group = H5Gopen(stat_file, "statistics", H5P_DEFAULT);
122
                fs->compute_velocity(fs->cvorticity);
123
124
125
126
                std::vector<double> max_estimate_vector;
                max_estimate_vector.resize(4);
                *tmp_vec_field = fs->cvelocity;
                switch(fs->dealias_type)
127
                {
128
                    case 0:
Cristian Lalescu's avatar
Cristian Lalescu committed
129
130
                        tmp_vec_field->compute_stats(
                            kk_two_thirds,
131
                            stat_group,
Cristian Lalescu's avatar
Cristian Lalescu committed
132
133
134
                            "velocity",
                            fs->iteration / niter_stat,
                            max_velocity_estimate/sqrt(3));
135
136
                        break;
                    case 1:
Cristian Lalescu's avatar
Cristian Lalescu committed
137
138
                        tmp_vec_field->compute_stats(
                            kk_smooth,
139
                            stat_group,
Cristian Lalescu's avatar
Cristian Lalescu committed
140
141
142
                            "velocity",
                            fs->iteration / niter_stat,
                            max_velocity_estimate/sqrt(3));
143
                        break;
144
                }
Cristian Lalescu's avatar
Cristian Lalescu committed
145
                //endcpp
146
                """
Cristian Lalescu's avatar
Cristian Lalescu committed
147
148
149
        if self.Lag_acc_stats_on:
            self.stat_src += """
                    //begincpp
150
                    tmp_vec_field->real_space_representation = false;
Cristian Lalescu's avatar
Cristian Lalescu committed
151
152
                    fs->compute_Lagrangian_acceleration(tmp_vec_field->get_cdata());
                    switch(fs->dealias_type)
153
                    {
Cristian Lalescu's avatar
Cristian Lalescu committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
                        case 0:
                            tmp_vec_field->compute_stats(
                                kk_two_thirds,
                                stat_group,
                                "Lagrangian_acceleration",
                                fs->iteration / niter_stat,
                                max_Lag_acc_estimate);
                            break;
                        case 1:
                            tmp_vec_field->compute_stats(
                                kk_smooth,
                                stat_group,
                                "Lagrangian_acceleration",
                                fs->iteration / niter_stat,
                                max_Lag_acc_estimate);
                            break;
170
                    }
Cristian Lalescu's avatar
Cristian Lalescu committed
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
                    tmp_scal_field->real_space_representation = false;
                    fs->compute_velocity(fs->cvorticity);
                    fs->ift_velocity();
                    fs->compute_pressure(tmp_scal_field->get_cdata());
                    switch(fs->dealias_type)
                    {
                        case 0:
                            tmp_scal_field->compute_stats(
                                kk_two_thirds,
                                stat_group,
                                "pressure",
                                fs->iteration / niter_stat,
                                max_pressure_estimate);
                            break;
                        case 1:
                            tmp_scal_field->compute_stats(
                                kk_smooth,
                                stat_group,
                                "pressure",
                                fs->iteration / niter_stat,
                                max_pressure_estimate);
                            break;
                    }
Cristian Lalescu's avatar
Cristian Lalescu committed
194
                    //endcpp
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
                    """
        self.stat_src += """
                //begincpp
                *tmp_vec_field = fs->cvorticity;
                switch(fs->dealias_type)
                {
                    case 0:
                        tmp_vec_field->compute_stats(
                            kk_two_thirds,
                            stat_group,
                            "vorticity",
                            fs->iteration / niter_stat,
                            max_vorticity_estimate/sqrt(3));
                        break;
                    case 1:
                        tmp_vec_field->compute_stats(
                            kk_smooth,
                            stat_group,
                            "vorticity",
                            fs->iteration / niter_stat,
                            max_vorticity_estimate/sqrt(3));
                        break;
                }
                //endcpp
                """
Cristian Lalescu's avatar
Cristian Lalescu committed
220
221
222
        if self.QR_stats_on:
            self.stat_src += """
                //begincpp
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
                double *trS2_Q_R_moments  = new double[10*3];
                double *gradu_moments     = new double[10*9];
                ptrdiff_t *hist_trS2_Q_R  = new ptrdiff_t[histogram_bins*3];
                ptrdiff_t *hist_gradu     = new ptrdiff_t[histogram_bins*9];
                ptrdiff_t *hist_QR2D      = new ptrdiff_t[QR2D_histogram_bins*QR2D_histogram_bins];
                double trS2QR_max_estimates[3];
                double gradu_max_estimates[9];
                trS2QR_max_estimates[0] = max_trS2_estimate;
                trS2QR_max_estimates[1] = max_Q_estimate;
                trS2QR_max_estimates[2] = max_R_estimate;
                std::fill_n(gradu_max_estimates, 9, sqrt(3*max_trS2_estimate));
                fs->compute_gradient_statistics(
                    fs->cvelocity,
                    gradu_moments,
                    trS2_Q_R_moments,
                    hist_gradu,
                    hist_trS2_Q_R,
                    hist_QR2D,
                    trS2QR_max_estimates,
                    gradu_max_estimates,
                    histogram_bins,
                    QR2D_histogram_bins);
                //endcpp
                """
Cristian Lalescu's avatar
Cristian Lalescu committed
247
248
        self.stat_src += """
                //begincpp
249
250
                if (myrank == 0)
                    H5Gclose(stat_group);
251
252
253
254
255
256
257
258
259
                if (fs->cd->myrank == 0)
                {{
                    hid_t Cdset, wspace, mspace;
                    int ndims;
                    hsize_t count[4], offset[4], dims[4];
                    offset[0] = fs->iteration/niter_stat;
                    offset[1] = 0;
                    offset[2] = 0;
                    offset[3] = 0;
Cristian Lalescu's avatar
Cristian Lalescu committed
260
                //endcpp
Cristian Lalescu's avatar
Cristian Lalescu committed
261
262
263
264
265
                """.format(self.C_dtype)
        if self.dtype == np.float32:
            field_H5T = 'H5T_NATIVE_FLOAT'
        elif self.dtype == np.float64:
            field_H5T = 'H5T_NATIVE_DOUBLE'
Cristian Lalescu's avatar
Cristian Lalescu committed
266
267
268
269
270
271
272
273
274
        if self.QR_stats_on:
            self.stat_src += self.create_stat_output(
                    '/statistics/moments/trS2_Q_R',
                    'trS2_Q_R_moments',
                    size_setup ="""
                        count[0] = 1;
                        count[1] = 10;
                        count[2] = 3;
                        """)
275
276
277
278
279
280
281
282
283
            self.stat_src += self.create_stat_output(
                    '/statistics/moments/velocity_gradient',
                    'gradu_moments',
                    size_setup ="""
                        count[0] = 1;
                        count[1] = 10;
                        count[2] = 3;
                        count[3] = 3;
                        """)
Cristian Lalescu's avatar
Cristian Lalescu committed
284
285
286
287
288
289
290
291
292
            self.stat_src += self.create_stat_output(
                    '/statistics/histograms/trS2_Q_R',
                    'hist_trS2_Q_R',
                    data_type = 'H5T_NATIVE_INT64',
                    size_setup = """
                        count[0] = 1;
                        count[1] = histogram_bins;
                        count[2] = 3;
                        """)
293
294
295
296
297
298
299
300
301
302
            self.stat_src += self.create_stat_output(
                    '/statistics/histograms/velocity_gradient',
                    'hist_gradu',
                    data_type = 'H5T_NATIVE_INT64',
                    size_setup = """
                        count[0] = 1;
                        count[1] = histogram_bins;
                        count[2] = 3;
                        count[3] = 3;
                        """)
Cristian Lalescu's avatar
Cristian Lalescu committed
303
304
305
306
307
308
309
310
311
            self.stat_src += self.create_stat_output(
                    '/statistics/histograms/QR2D',
                    'hist_QR2D',
                    data_type = 'H5T_NATIVE_INT64',
                    size_setup = """
                        count[0] = 1;
                        count[1] = QR2D_histogram_bins;
                        count[2] = QR2D_histogram_bins;
                        """)
312
        self.stat_src += '}\n'
Cristian Lalescu's avatar
Cristian Lalescu committed
313
314
315
        if self.QR_stats_on:
            self.stat_src += """
                //begincpp
316
317
318
319
320
                delete[] trS2_Q_R_moments;
                delete[] gradu_moments;
                delete[] hist_trS2_Q_R;
                delete[] hist_gradu;
                delete[] hist_QR2D;
Cristian Lalescu's avatar
Cristian Lalescu committed
321
322
                //endcpp
                """
323
324
325
        return None
    def fill_up_fluid_code(self):
        self.fluid_includes += '#include <cstring>\n'
326
327
328
        self.fluid_variables += (
                'fluid_solver<{0}> *fs;\n'.format(self.C_dtype) +
                'field<{0}, FFTW, THREE> *tmp_vec_field;\n'.format(self.C_dtype) +
Cristian Lalescu's avatar
Cristian Lalescu committed
329
                'field<{0}, FFTW, ONE> *tmp_scal_field;\n'.format(self.C_dtype) +
330
331
                'kspace<FFTW, SMOOTH> *kk_smooth;\n' +
                'kspace<FFTW, TWO_THIRDS> *kk_two_thirds;\n')
Chichi Lalescu's avatar
Chichi Lalescu committed
332
333
334
335
336
337
        self.fluid_definitions += """
                    typedef struct {{
                        {0} re;
                        {0} im;
                    }} tmp_complex_type;
                    """.format(self.C_dtype)
338
        self.write_fluid_stats()
Cristian Lalescu's avatar
Cristian Lalescu committed
339
340
341
342
        if self.dtype == np.float32:
            field_H5T = 'H5T_NATIVE_FLOAT'
        elif self.dtype == np.float64:
            field_H5T = 'H5T_NATIVE_DOUBLE'
343
        self.fluid_start += """
Cristian Lalescu's avatar
Cristian Lalescu committed
344
345
                //begincpp
                char fname[512];
346
                fs = new fluid_solver<{0}>(
Cristian Lalescu's avatar
Cristian Lalescu committed
347
348
                        simname,
                        nx, ny, nz,
349
                        dkx, dky, dkz,
350
351
                        dealias_type,
                        {1});
352
353
354
355
                tmp_vec_field = new field<{0}, FFTW, THREE>(
                        nx, ny, nz,
                        MPI_COMM_WORLD,
                        {1});
Cristian Lalescu's avatar
Cristian Lalescu committed
356
357
358
359
                tmp_scal_field = new field<{0}, FFTW, ONE>(
                        nx, ny, nz,
                        MPI_COMM_WORLD,
                        {1});
360
361
362
363
364
365
                kk_smooth = new kspace<FFTW, SMOOTH>(
                        tmp_vec_field->clayout,
                        fs->dkx, fs->dky, fs->dkz);
                kk_two_thirds = new kspace<FFTW, TWO_THIRDS>(
                        tmp_vec_field->clayout,
                        fs->dkx, fs->dky, fs->dkz);
Cristian Lalescu's avatar
Cristian Lalescu committed
366
367
368
                fs->nu = nu;
                fs->fmode = fmode;
                fs->famplitude = famplitude;
Cristian Lalescu's avatar
Cristian Lalescu committed
369
370
371
                fs->fk0 = fk0;
                fs->fk1 = fk1;
                strncpy(fs->forcing_type, forcing_type, 128);
Chichi Lalescu's avatar
Chichi Lalescu committed
372
                fs->iteration = iteration;
Cristian Lalescu's avatar
Cristian Lalescu committed
373
                fs->read('v', 'c');
374
375
                //endcpp
                """.format(self.C_dtype, self.fftw_plan_rigor, field_H5T)
376
        self.fluid_start += self.store_kspace
Chichi Lalescu's avatar
Chichi Lalescu committed
377
378
379
380
381
382
        if not self.frozen_fields:
            self.fluid_loop = 'fs->step(dt);\n'
        else:
            self.fluid_loop = ''
        self.fluid_loop += ('if (fs->iteration % niter_out == 0)\n{\n' +
                            self.fluid_output + '\n}\n')
Cristian Lalescu's avatar
Cristian Lalescu committed
383
384
        self.fluid_end = ('if (fs->iteration % niter_out != 0)\n{\n' +
                          self.fluid_output + '\n}\n' +
385
386
                          'delete fs;\n' +
                          'delete tmp_vec_field;\n' +
Cristian Lalescu's avatar
Cristian Lalescu committed
387
                          'delete tmp_scal_field;\n' +
388
389
                          'delete kk_smooth;\n' +
                          'delete kk_two_thirds;\n')
390
        return None
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
    def add_3D_rFFTW_field(
            self,
            name = 'rFFTW_acc'):
        if self.dtype == np.float32:
            FFTW = 'fftwf'
        elif self.dtype == np.float64:
            FFTW = 'fftw'
        self.fluid_variables += '{0} *{1};\n'.format(self.C_dtype, name)
        self.fluid_start += '{0} = {1}_alloc_real(2*fs->cd->local_size);\n'.format(name, FFTW)
        self.fluid_end   += '{0}_free({1});\n'.format(FFTW, name)
        return None
    def add_interpolator(
            self,
            interp_type = 'spline',
            neighbours = 1,
            smoothness = 1,
            name = 'field_interpolator',
408
409
410
411
412
            field_name = 'fs->rvelocity',
            class_name = 'rFFTW_interpolator'):
        self.fluid_includes += '#include "{0}.hpp"\n'.format(class_name)
        self.fluid_variables += '{0} <{1}, {2}> *{3};\n'.format(
                class_name, self.C_dtype, neighbours, name)
413
414
415
416
417
418
419
        self.parameters[name + '_type'] = interp_type
        self.parameters[name + '_neighbours'] = neighbours
        if interp_type == 'spline':
            self.parameters[name + '_smoothness'] = smoothness
            beta_name = 'beta_n{0}_m{1}'.format(neighbours, smoothness)
        elif interp_type == 'Lagrange':
            beta_name = 'beta_Lagrange_n{0}'.format(neighbours)
420
        self.fluid_start += '{0} = new {1}<{2}, {3}>(fs, {4}, {5});\n'.format(
421
                name,
422
                class_name,
423
424
                self.C_dtype,
                neighbours,
Cristian Lalescu's avatar
Cristian Lalescu committed
425
426
                beta_name,
                field_name)
427
        self.fluid_end += 'delete {0};\n'.format(name)
428
        return None
429
430
    def add_particles(
            self,
431
            integration_steps = 2,
Cristian Lalescu's avatar
Cristian Lalescu committed
432
            kcut = None,
433
            interpolator = 'field_interpolator',
434
            frozen_particles = False,
435
            acc_name = None,
436
            class_name = 'particles'):
437
438
439
440
441
442
443
444
445
        """Adds code for tracking a series of particle species, each
        consisting of `nparticles` particles.

        :type integration_steps: int, list of int
        :type kcut: None (default), str, list of str
        :type interpolator: str, list of str
        :type frozen_particles: bool
        :type acc_name: str

446
447
448
        .. warning:: if not None, kcut must be a list of decreasing
                     wavenumbers, since filtering is done sequentially
                     on the same complex FFTW field.
449
        """
Cristian Lalescu's avatar
Cristian Lalescu committed
450
451
452
453
454
        if self.dtype == np.float32:
            FFTW = 'fftwf'
        elif self.dtype == np.float64:
            FFTW = 'fftw'
        s0 = self.particle_species
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
        if type(integration_steps) == int:
            integration_steps = [integration_steps]
        if type(kcut) == str:
            kcut = [kcut]
        if type(interpolator) == str:
            interpolator = [interpolator]
        nspecies = max(len(integration_steps), len(interpolator))
        if type(kcut) == list:
            nspecies = max(nspecies, len(kcut))
        if len(integration_steps) == 1:
            integration_steps = [integration_steps[0] for s in range(nspecies)]
        if len(interpolator) == 1:
            interpolator = [interpolator[0] for s in range(nspecies)]
        if type(kcut) == list:
            if len(kcut) == 1:
                kcut = [kcut[0] for s in range(nspecies)]
        assert(len(integration_steps) == nspecies)
        assert(len(interpolator) == nspecies)
        if type(kcut) == list:
            assert(len(kcut) == nspecies)
        for s in range(nspecies):
Cristian Lalescu's avatar
Cristian Lalescu committed
476
            neighbours = self.parameters[interpolator[s] + '_neighbours']
477
478
            if type(kcut) == list:
                self.parameters['tracers{0}_kcut'.format(s0 + s)] = kcut[s]
Cristian Lalescu's avatar
Cristian Lalescu committed
479
            self.parameters['tracers{0}_interpolator'.format(s0 + s)] = interpolator[s]
480
            self.parameters['tracers{0}_acc_on'.format(s0 + s)] = int(not type(acc_name) == type(None))
Cristian Lalescu's avatar
Cristian Lalescu committed
481
482
            self.parameters['tracers{0}_integration_steps'.format(s0 + s)] = integration_steps[s]
            self.file_datasets_grow += """
483
                        //begincpp
484
                        group = H5Gopen(particle_file, "/tracers{0}", H5P_DEFAULT);
485
                        grow_particle_datasets(group, "", NULL, NULL);
486
                        H5Gclose(group);
487
                        //endcpp
Cristian Lalescu's avatar
Cristian Lalescu committed
488
489
490
491
492
493
494
                        """.format(s0 + s)

        #### code that outputs statistics
        output_vel_acc = '{\n'
        # array for putting sampled velocity in
        # must compute velocity, just in case it was messed up by some
        # other particle species before the stats
495
        output_vel_acc += 'fs->compute_velocity(fs->cvorticity);\n'
496
497
        if not type(kcut) == list:
            output_vel_acc += 'fs->ift_velocity();\n'
Cristian Lalescu's avatar
Cristian Lalescu committed
498
499
500
501
        if not type(acc_name) == type(None):
            # array for putting sampled acceleration in
            # must compute acceleration
            output_vel_acc += 'fs->compute_Lagrangian_acceleration({0});\n'.format(acc_name)
502
        for s in range(nspecies):
503
504
505
            if type(kcut) == list:
                output_vel_acc += 'fs->low_pass_Fourier(fs->cvelocity, 3, {0});\n'.format(kcut[s])
                output_vel_acc += 'fs->ift_velocity();\n'
Cristian Lalescu's avatar
Cristian Lalescu committed
506
            output_vel_acc += """
507
                {0}->read_rFFTW(fs->rvelocity);
508
                ps{1}->sample({0}, "velocity");
Cristian Lalescu's avatar
Cristian Lalescu committed
509
510
511
                """.format(interpolator[s], s0 + s)
            if not type(acc_name) == type(None):
                output_vel_acc += """
512
                    {0}->read_rFFTW({1});
513
                    ps{2}->sample({0}, "acceleration");
Cristian Lalescu's avatar
Cristian Lalescu committed
514
515
516
517
                    """.format(interpolator[s], acc_name, s0 + s)
        output_vel_acc += '}\n'

        #### initialize, stepping and finalize code
518
519
520
521
522
523
524
        if not type(kcut) == list:
            update_fields = ('fs->compute_velocity(fs->cvorticity);\n' +
                             'fs->ift_velocity();\n')
            self.particle_start += update_fields
            self.particle_loop  += update_fields
        else:
            self.particle_loop += 'fs->compute_velocity(fs->cvorticity);\n'
525
        self.particle_includes += '#include "{0}.hpp"\n'.format(class_name)
Cristian Lalescu's avatar
Cristian Lalescu committed
526
527
528
        self.particle_stat_src += (
                'if (ps0->iteration % niter_part == 0)\n' +
                '{\n')
529
        for s in range(nspecies):
530
            neighbours = self.parameters[interpolator[s] + '_neighbours']
Cristian Lalescu's avatar
Cristian Lalescu committed
531
            self.particle_start += 'sprintf(fname, "tracers{0}");\n'.format(s0 + s)
532
            self.particle_end += ('ps{0}->write();\n' +
Cristian Lalescu's avatar
Cristian Lalescu committed
533
                                  'delete ps{0};\n').format(s0 + s)
534
535
            self.particle_variables += '{0}<VELOCITY_TRACER, {1}, {2}> *ps{3};\n'.format(
                    class_name,
536
537
                    self.C_dtype,
                    neighbours,
Cristian Lalescu's avatar
Cristian Lalescu committed
538
                    s0 + s)
539
            self.particle_start += ('ps{0} = new {1}<VELOCITY_TRACER, {2}, {3}>(\n' +
Cristian Lalescu's avatar
Cristian Lalescu committed
540
                                    'fname, particle_file, {4},\n' +
541
                                    'niter_part, tracers{0}_integration_steps);\n').format(
Cristian Lalescu's avatar
Cristian Lalescu committed
542
                                            s0 + s,
543
                                            class_name,
544
545
                                            self.C_dtype,
                                            neighbours,
Cristian Lalescu's avatar
Cristian Lalescu committed
546
547
548
                                            interpolator[s])
            self.particle_start += ('ps{0}->dt = dt;\n' +
                                    'ps{0}->iteration = iteration;\n' +
549
                                    'ps{0}->read();\n').format(s0 + s)
Cristian Lalescu's avatar
Cristian Lalescu committed
550
            if not frozen_particles:
551
552
553
554
                if type(kcut) == list:
                    update_field = ('fs->low_pass_Fourier(fs->cvelocity, 3, {0});\n'.format(kcut[s]) +
                                    'fs->ift_velocity();\n')
                    self.particle_loop += update_field
555
                self.particle_loop += '{0}->read_rFFTW(fs->rvelocity);\n'.format(interpolator[s])
556
                self.particle_loop += 'ps{0}->step();\n'.format(s0 + s)
557
            self.particle_stat_src += 'ps{0}->write(false);\n'.format(s0 + s)
Cristian Lalescu's avatar
Cristian Lalescu committed
558
559
        self.particle_stat_src += output_vel_acc
        self.particle_stat_src += '}\n'
560
        self.particle_species += nspecies
561
        return None
562
563
564
565
    def get_cache_file_name(self):
        return os.path.join(self.work_dir, self.simname + '_cache.h5')
    def get_cache_file(self):
        return h5py.File(self.get_postprocess_file_name(), 'r')
566
    def get_postprocess_file_name(self):
567
        return self.get_cache_file_name()
Cristian Lalescu's avatar
Cristian Lalescu committed
568
569
    def get_postprocess_file(self):
        return h5py.File(self.get_postprocess_file_name(), 'r')
Chichi Lalescu's avatar
Chichi Lalescu committed
570
    def compute_statistics(self, iter0 = 0, iter1 = None):
571
572
573
574
575
576
577
578
579
580
581
582
        """Run basic postprocessing on raw data.
        The energy spectrum :math:`E(t, k)` and the enstrophy spectrum
        :math:`\\frac{1}{2}\omega^2(t, k)` are computed from the

        .. math::

            \sum_{k \\leq \\|\\mathbf{k}\\| \\leq k+dk}\\hat{u_i} \\hat{u_j}^*, \\hskip .5cm
            \sum_{k \\leq \\|\\mathbf{k}\\| \\leq k+dk}\\hat{\omega_i} \\hat{\\omega_j}^*

        tensors, and the enstrophy spectrum is also used to
        compute the dissipation :math:`\\varepsilon(t)`.
        These basic quantities are stored in a newly created HDF5 file,
583
        ``simname_cache.h5``.
584
        """
585
586
        if len(list(self.statistics.keys())) > 0:
            return None
Cristian Lalescu's avatar
Cristian Lalescu committed
587
        self.read_parameters()
588
        with self.get_data_file() as data_file:
589
590
            if 'moments' not in data_file['statistics'].keys():
                return None
591
592
            iter0 = min((data_file['statistics/moments/velocity'].shape[0] *
                         self.parameters['niter_stat']-1),
593
                        iter0)
Chichi Lalescu's avatar
Chichi Lalescu committed
594
595
596
597
            if type(iter1) == type(None):
                iter1 = data_file['iteration'].value
            else:
                iter1 = min(data_file['iteration'].value, iter1)
598
599
            ii0 = iter0 // self.parameters['niter_stat']
            ii1 = iter1 // self.parameters['niter_stat']
600
            self.statistics['kshell'] = data_file['kspace/kshell'].value
Cristian Lalescu's avatar
Cristian Lalescu committed
601
            self.statistics['nshell'] = data_file['kspace/nshell'].value
Chichi Lalescu's avatar
Chichi Lalescu committed
602
603
604
            for kk in [-1, -2]:
                if (self.statistics['kshell'][kk] == 0):
                    self.statistics['kshell'][kk] = np.nan
605
            self.statistics['kM'] = data_file['kspace/kM'].value
Chichi Lalescu's avatar
Chichi Lalescu committed
606
            self.statistics['dk'] = data_file['kspace/dk'].value
607
            computation_needed = True
608
609
610
611
            pp_file = h5py.File(self.get_postprocess_file_name(), 'a')
            if 'ii0' in pp_file.keys():
                computation_needed =  not (ii0 == pp_file['ii0'].value and
                                           ii1 == pp_file['ii1'].value)
612
                if computation_needed:
613
                    for k in ['t', 'vel_max(t)', 'renergy(t)',
Cristian Lalescu's avatar
Cristian Lalescu committed
614
615
                              'energy(t)', 'enstrophy(t)',
                              'energy(k)', 'enstrophy(k)',
616
617
                              'energy(t, k)',
                              'enstrophy(t, k)',
Cristian Lalescu's avatar
Cristian Lalescu committed
618
                              'R_ij(t)',
619
                              'ii0', 'ii1', 'iter0', 'iter1']:
620
621
                        if k in pp_file.keys():
                            del pp_file[k]
622
            if computation_needed:
623
624
625
626
627
628
629
                pp_file['iter0'] = iter0
                pp_file['iter1'] = iter1
                pp_file['ii0'] = ii0
                pp_file['ii1'] = ii1
                pp_file['t'] = (self.parameters['dt']*
                                self.parameters['niter_stat']*
                                (np.arange(ii0, ii1+1).astype(np.float)))
Cristian Lalescu's avatar
Cristian Lalescu committed
630
                phi_ij = data_file['statistics/spectra/velocity_velocity'][ii0:ii1+1]
Cristian Lalescu's avatar
Cristian Lalescu committed
631
632
                pp_file['R_ij(t)'] = np.sum(phi_ij, axis = 1)
                energy_tk = (
Cristian Lalescu's avatar
Cristian Lalescu committed
633
634
635
                    phi_ij[:, :, 0, 0] +
                    phi_ij[:, :, 1, 1] +
                    phi_ij[:, :, 2, 2])/2
Cristian Lalescu's avatar
Cristian Lalescu committed
636
                pp_file['energy(t)'] = np.sum(energy_tk, axis = 1)
Cristian Lalescu's avatar
Cristian Lalescu committed
637
                pp_file['energy(k)'] = np.mean(energy_tk, axis = 0)*(4*np.pi*self.statistics['kshell']**2) / (self.statistics['dk']*self.statistics['nshell'])
Cristian Lalescu's avatar
Cristian Lalescu committed
638
                enstrophy_tk = (
639
640
641
                    data_file['statistics/spectra/vorticity_vorticity'][ii0:ii1+1, :, 0, 0] +
                    data_file['statistics/spectra/vorticity_vorticity'][ii0:ii1+1, :, 1, 1] +
                    data_file['statistics/spectra/vorticity_vorticity'][ii0:ii1+1, :, 2, 2])/2
Cristian Lalescu's avatar
Cristian Lalescu committed
642
                pp_file['enstrophy(t)'] = np.sum(enstrophy_tk, axis = 1)
Cristian Lalescu's avatar
Cristian Lalescu committed
643
                pp_file['enstrophy(k)'] = np.mean(enstrophy_tk, axis = 0)*(4*np.pi*self.statistics['kshell']**2) / (self.statistics['dk']*self.statistics['nshell'])
Cristian Lalescu's avatar
Cristian Lalescu committed
644
                pp_file['vel_max(t)'] = data_file['statistics/moments/velocity'][ii0:ii1+1, 9, 3]
645
                pp_file['renergy(t)'] = data_file['statistics/moments/velocity'][ii0:ii1+1, 2, 3]/2
646
                if 'trS2_Q_R' in data_file['statistics/moments'].keys():
647
                    pp_file['mean_trS2(t)'] = data_file['statistics/moments/trS2_Q_R'][:, 1, 0]
648
            for k in ['t',
Cristian Lalescu's avatar
Cristian Lalescu committed
649
650
651
652
653
                      'energy(t)',
                      'energy(k)',
                      'enstrophy(t)',
                      'enstrophy(k)',
                      'R_ij(t)',
654
                      'vel_max(t)',
655
656
                      'renergy(t)',
                      'mean_trS2(t)']:
657
658
                if k in pp_file.keys():
                    self.statistics[k] = pp_file[k].value
659
660
            self.compute_time_averages()
        return None
661
662
663
664
665
666
667
668
669
670
    def compute_Reynolds_stress_invariants(
            self):
        Rij = self.statistics['R_ij(t)']
        Rij /= (2*self.statistics['energy(t)'][:, None, None])
        Rij[:, 0, 0] -= 1./3
        Rij[:, 1, 1] -= 1./3
        Rij[:, 2, 2] -= 1./3
        self.statistics['I2(t)'] = np.sqrt(np.einsum('...ij,...ij', Rij, Rij, optimize = True) / 6)
        self.statistics['I3(t)'] = np.cbrt(np.einsum('...ij,...jk,...ki', Rij, Rij, Rij, optimize = True) / 6)
        return None
671
    def compute_time_averages(self):
672
673
674
        """Compute easy stats.

        Further computation of statistics based on the contents of
675
        ``simname_cache.h5``.
676
677
678
679
680
681
        Standard quantities are as follows
        (consistent with [Ishihara]_):

        .. math::

            U_{\\textrm{int}}(t) = \\sqrt{\\frac{2E(t)}{3}}, \\hskip .5cm
Cristian Lalescu's avatar
Cristian Lalescu committed
682
683
684
            L_{\\textrm{int}} = \\frac{\pi}{2U_{int}^2} \\int \\frac{dk}{k} E(k), \\hskip .5cm
            T_{\\textrm{int}} =
            \\frac{L_{\\textrm{int}}}{U_{\\textrm{int}}}
685
686
687
688
689
690
691
692

            \\eta_K = \\left(\\frac{\\nu^3}{\\varepsilon}\\right)^{1/4}, \\hskip .5cm
            \\tau_K = \\left(\\frac{\\nu}{\\varepsilon}\\right)^{1/2}, \\hskip .5cm
            \\lambda = \\sqrt{\\frac{15 \\nu U_{\\textrm{int}}^2}{\\varepsilon}}

            Re = \\frac{U_{\\textrm{int}} L_{\\textrm{int}}}{\\nu}, \\hskip
            .5cm
            R_{\\lambda} = \\frac{U_{\\textrm{int}} \\lambda}{\\nu}
Cristian Lalescu's avatar
Cristian Lalescu committed
693
694
695
696
697
698
699
700

        .. [Ishihara] T. Ishihara et al,
                      *Small-scale statistics in high-resolution direct numerical
                      simulation of turbulence: Reynolds number dependence of
                      one-point velocity gradient statistics*.
                      J. Fluid Mech.,
                      **592**, 335-366, 2007
        """
Chichi Lalescu's avatar
Chichi Lalescu committed
701
702
703
704
        self.statistics['Uint(t)'] = np.sqrt(2*self.statistics['energy(t)'] / 3)
        for key in ['energy',
                    'enstrophy',
                    'mean_trS2',
Cristian Lalescu's avatar
Cristian Lalescu committed
705
                    'Uint']:
706
707
            if key + '(t)' in self.statistics.keys():
                self.statistics[key] = np.average(self.statistics[key + '(t)'], axis = 0)
Cristian Lalescu's avatar
Cristian Lalescu committed
708
        self.statistics['vel_max'] = np.max(self.statistics['vel_max(t)'])
709
710
711
712
713
714
715
        for suffix in ['', '(t)']:
            self.statistics['diss'    + suffix] = (self.parameters['nu'] *
                                                   self.statistics['enstrophy' + suffix]*2)
            self.statistics['etaK'    + suffix] = (self.parameters['nu']**3 /
                                                   self.statistics['diss' + suffix])**.25
            self.statistics['tauK'    + suffix] =  (self.parameters['nu'] /
                                                    self.statistics['diss' + suffix])**.5
Chichi Lalescu's avatar
Chichi Lalescu committed
716
717
718
719
720
721
            self.statistics['lambda' + suffix] = (15 * self.parameters['nu'] *
                                                  self.statistics['Uint' + suffix]**2 /
                                                  self.statistics['diss' + suffix])**.5
            self.statistics['Rlambda' + suffix] = (self.statistics['Uint' + suffix] *
                                                   self.statistics['lambda' + suffix] /
                                                   self.parameters['nu'])
722
723
            self.statistics['kMeta' + suffix] = (self.statistics['kM'] *
                                                 self.statistics['etaK' + suffix])
Cristian Lalescu's avatar
Cristian Lalescu committed
724
725
            if self.parameters['dealias_type'] == 1:
                self.statistics['kMeta' + suffix] *= 0.8
Cristian Lalescu's avatar
Cristian Lalescu committed
726
        self.statistics['Lint'] = ((np.pi /
Cristian Lalescu's avatar
Cristian Lalescu committed
727
728
729
730
731
732
                                    (2*self.statistics['Uint']**2)) *
                                   np.nansum(self.statistics['energy(k)'] /
                                                self.statistics['kshell']))
        self.statistics['Re'] = (self.statistics['Uint'] *
                                 self.statistics['Lint'] /
                                 self.parameters['nu'])
Chichi Lalescu's avatar
Chichi Lalescu committed
733
734
        self.statistics['Tint'] = self.statistics['Lint'] / self.statistics['Uint']
        self.statistics['Taylor_microscale'] = self.statistics['lambda']
735
        return None
736
737
738
739
740
    def set_plt_style(
            self,
            style = {'dashes' : (None, None)}):
        self.style.update(style)
        return None
Chichi Lalescu's avatar
Chichi Lalescu committed
741
742
743
744
    def read_cfield(
            self,
            field_name = 'vorticity',
            iteration = 0):
Cristian Lalescu's avatar
Cristian Lalescu committed
745
746
747
748
749
750
        """read the Fourier representation of a vector field.

        Read the binary file containing iteration ``iteration`` of the
        field ``field_name``, and return it as a properly shaped
        ``numpy.memmap`` object.
        """
Chichi Lalescu's avatar
Chichi Lalescu committed
751
752
753
754
755
756
757
758
759
        return np.memmap(
                os.path.join(self.work_dir,
                             self.simname + '_{0}_i{1:0>5x}'.format('c' + field_name, iteration)),
                dtype = self.ctype,
                mode = 'r',
                shape = (self.parameters['ny'],
                         self.parameters['nz'],
                         self.parameters['nx']//2+1,
                         3))
760
761
762
763
    def write_par(
            self,
            iter0 = 0,
            particle_ic = None):
764
        _fluid_particle_base.write_par(self, iter0 = iter0)
Cristian Lalescu's avatar
Cristian Lalescu committed
765
        with h5py.File(self.get_data_file_name(), 'r+') as ofile:
766
767
            kspace = self.get_kspace()
            nshells = kspace['nshell'].shape[0]
Cristian Lalescu's avatar
Cristian Lalescu committed
768
            vec_stat_datasets = ['velocity', 'vorticity']
Cristian Lalescu's avatar
Cristian Lalescu committed
769
            scal_stat_datasets = []
770
            for k in vec_stat_datasets:
771
772
773
                time_chunk = 2**20 // (
                        self.dtype.itemsize*3*
                        self.parameters['nx']*self.parameters['ny'])
774
                time_chunk = max(time_chunk, 1)
775
776
777
778
                ofile.create_dataset('statistics/0slices/' + k + '/real',
                                     (1, self.parameters['ny'], self.parameters['nx'], 3),
                                     chunks = (time_chunk, self.parameters['ny'], self.parameters['nx'], 3),
                                     maxshape = (None, self.parameters['ny'], self.parameters['nx'], 3),
779
                                     dtype = self.dtype)
780
781
            if self.Lag_acc_stats_on:
                vec_stat_datasets += ['Lagrangian_acceleration']
Cristian Lalescu's avatar
Cristian Lalescu committed
782
                scal_stat_datasets += ['pressure']
783
            for k in vec_stat_datasets:
784
785
786
787
788
789
                time_chunk = 2**20//(8*3*3*nshells)
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/spectra/' + k + '_' + k,
                                     (1, nshells, 3, 3),
                                     chunks = (time_chunk, nshells, 3, 3),
                                     maxshape = (None, nshells, 3, 3),
790
                                     dtype = np.float64)
791
792
793
794
795
796
                time_chunk = 2**20//(8*4*10)
                time_chunk = max(time_chunk, 1)
                a = ofile.create_dataset('statistics/moments/' + k,
                                     (1, 10, 4),
                                     chunks = (time_chunk, 10, 4),
                                     maxshape = (None, 10, 4),
797
                                     dtype = np.float64)
798
799
800
801
802
803
804
805
806
807
808
809
                time_chunk = 2**20//(8*4*self.parameters['histogram_bins'])
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/histograms/' + k,
                                     (1,
                                      self.parameters['histogram_bins'],
                                      4),
                                     chunks = (time_chunk,
                                               self.parameters['histogram_bins'],
                                               4),
                                     maxshape = (None,
                                                 self.parameters['histogram_bins'],
                                                 4),
810
                                     dtype = np.int64)
Cristian Lalescu's avatar
Cristian Lalescu committed
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
            for k in scal_stat_datasets:
                time_chunk = 2**20//(8*nshells)
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/spectra/' + k + '_' + k,
                                     (1, nshells),
                                     chunks = (time_chunk, nshells),
                                     maxshape = (None, nshells),
                                     dtype = np.float64)
                time_chunk = 2**20//(8*10)
                time_chunk = max(time_chunk, 1)
                a = ofile.create_dataset('statistics/moments/' + k,
                                     (1, 10),
                                     chunks = (time_chunk, 10),
                                     maxshape = (None, 10),
                                     dtype = np.float64)
                time_chunk = 2**20//(8*self.parameters['histogram_bins'])
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/histograms/' + k,
                                     (1,
                                      self.parameters['histogram_bins']),
                                     chunks = (time_chunk,
                                               self.parameters['histogram_bins']),
                                     maxshape = (None,
                                                 self.parameters['histogram_bins']),
                                     dtype = np.int64)
836
837
838
839
840
841
842
843
844
845
846
847
848
            if self.QR_stats_on:
                time_chunk = 2**20//(8*3*self.parameters['histogram_bins'])
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/histograms/trS2_Q_R',
                                     (1,
                                      self.parameters['histogram_bins'],
                                      3),
                                     chunks = (time_chunk,
                                               self.parameters['histogram_bins'],
                                               3),
                                     maxshape = (None,
                                                 self.parameters['histogram_bins'],
                                                 3),
849
                                     dtype = np.int64)
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
                time_chunk = 2**20//(8*9*self.parameters['histogram_bins'])
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/histograms/velocity_gradient',
                                     (1,
                                      self.parameters['histogram_bins'],
                                      3,
                                      3),
                                     chunks = (time_chunk,
                                               self.parameters['histogram_bins'],
                                               3,
                                               3),
                                     maxshape = (None,
                                                 self.parameters['histogram_bins'],
                                                 3,
                                                 3),
865
                                     dtype = np.int64)
866
867
868
869
870
871
                time_chunk = 2**20//(8*3*10)
                time_chunk = max(time_chunk, 1)
                a = ofile.create_dataset('statistics/moments/trS2_Q_R',
                                     (1, 10, 3),
                                     chunks = (time_chunk, 10, 3),
                                     maxshape = (None, 10, 3),
872
                                     dtype = np.float64)
873
874
875
876
877
878
                time_chunk = 2**20//(8*9*10)
                time_chunk = max(time_chunk, 1)
                a = ofile.create_dataset('statistics/moments/velocity_gradient',
                                     (1, 10, 3, 3),
                                     chunks = (time_chunk, 10, 3, 3),
                                     maxshape = (None, 10, 3, 3),
879
                                     dtype = np.float64)
880
881
882
883
884
885
886
887
888
889
890
891
                time_chunk = 2**20//(8*self.parameters['QR2D_histogram_bins']**2)
                time_chunk = max(time_chunk, 1)
                ofile.create_dataset('statistics/histograms/QR2D',
                                     (1,
                                      self.parameters['QR2D_histogram_bins'],
                                      self.parameters['QR2D_histogram_bins']),
                                     chunks = (time_chunk,
                                               self.parameters['QR2D_histogram_bins'],
                                               self.parameters['QR2D_histogram_bins']),
                                     maxshape = (None,
                                                 self.parameters['QR2D_histogram_bins'],
                                                 self.parameters['QR2D_histogram_bins']),
892
                                     dtype = np.int64)
Cristian Lalescu's avatar
Cristian Lalescu committed
893
894
895
        if self.particle_species == 0:
            return None

896
897
898
899
900
901
        if type(particle_ic) == type(None):
            pbase_shape = (self.parameters['nparticles'],)
            number_of_particles = self.parameters['nparticles']
        else:
            pbase_shape = particle_ic.shape[:-1]
            assert(particle_ic.shape[-1] == 3)
902
903
904
905
906
907
            if len(pbase_shape) == 1:
                number_of_particles = pbase_shape[0]
            else:
                number_of_particles = 1
                for val in pbase_shape[1:]:
                    number_of_particles *= val
908

Cristian Lalescu's avatar
Cristian Lalescu committed
909
910
911
        with h5py.File(self.get_particle_file_name(), 'a') as ofile:
            for s in range(self.particle_species):
                ofile.create_group('tracers{0}'.format(s))
912
                time_chunk = 2**20 // (8*3*number_of_particles)
Cristian Lalescu's avatar
Cristian Lalescu committed
913
                time_chunk = max(time_chunk, 1)
914
915
916
917
                dims = ((1,
                         self.parameters['tracers{0}_integration_steps'.format(s)]) +
                        pbase_shape + (3,))
                maxshape = (h5py.h5s.UNLIMITED,) + dims[1:]
918
919
920
921
                if len(pbase_shape) > 1:
                    chunks = (time_chunk, 1, 1) + dims[3:]
                else:
                    chunks = (time_chunk, 1) + dims[2:]
922
                bfps.tools.create_alloc_early_dataset(
Cristian Lalescu's avatar
Cristian Lalescu committed
923
924
925
                        ofile,
                        '/tracers{0}/rhs'.format(s),
                        dims, maxshape, chunks)
926
927
928
929
                if len(pbase_shape) > 1:
                    chunks = (time_chunk, 1) + pbase_shape[1:] + (3,)
                else:
                    chunks = (time_chunk, pbase_shape[0], 3)
930
                bfps.tools.create_alloc_early_dataset(
Cristian Lalescu's avatar
Cristian Lalescu committed
931
932
                        ofile,
                        '/tracers{0}/state'.format(s),
933
934
                        (1,) + pbase_shape + (3,),
                        (h5py.h5s.UNLIMITED,) + pbase_shape + (3,),
935
                        chunks)
936
937
938
                # "velocity" is sampled, single precision is enough
                # for the results we are interested in.
                bfps.tools.create_alloc_early_dataset(
Cristian Lalescu's avatar
Cristian Lalescu committed
939
940
                        ofile,
                        '/tracers{0}/velocity'.format(s),
941
942
                        (1,) + pbase_shape + (3,),
                        (h5py.h5s.UNLIMITED,) + pbase_shape + (3,),
943
                        chunks,
944
                        dset_dtype = h5py.h5t.IEEE_F32LE)
Cristian Lalescu's avatar
Cristian Lalescu committed
945
                if self.parameters['tracers{0}_acc_on'.format(s)]:
946
                    bfps.tools.create_alloc_early_dataset(
Cristian Lalescu's avatar
Cristian Lalescu committed
947
948
                            ofile,
                            '/tracers{0}/acceleration'.format(s),
949
950
                            (1,) + pbase_shape + (3,),
                            (h5py.h5s.UNLIMITED,) + pbase_shape + (3,),
951
                            chunks,
952
                            dset_dtype = h5py.h5t.IEEE_F32LE)
953
        return None
Cristian Lalescu's avatar
Cristian Lalescu committed
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
    def add_particle_fields(
            self,
            interp_type = 'spline',
            kcut = None,
            neighbours = 1,
            smoothness = 1,
            name = 'particle_field',
            field_class = 'rFFTW_interpolator',
            acc_field_name = 'rFFTW_acc'):
        self.fluid_includes += '#include "{0}.hpp"\n'.format(field_class)
        self.fluid_variables += field_class + '<{0}, {1}> *vel_{2}, *acc_{2};\n'.format(
                self.C_dtype, neighbours, name)
        self.parameters[name + '_type'] = interp_type
        self.parameters[name + '_neighbours'] = neighbours
        if interp_type == 'spline':
            self.parameters[name + '_smoothness'] = smoothness
            beta_name = 'beta_n{0}_m{1}'.format(neighbours, smoothness)
        elif interp_type == 'Lagrange':
            beta_name = 'beta_Lagrange_n{0}'.format(neighbours)
        if field_class == 'rFFTW_interpolator':
            self.fluid_start += ('vel_{0} = new {1}<{2}, {3}>(fs, {4}, fs->rvelocity);\n' +
                                 'acc_{0} = new {1}<{2}, {3}>(fs, {4}, {5});\n').format(name,
                                                                                   field_class,
                                                                                   self.C_dtype,
                                                                                   neighbours,
                                                                                   beta_name,
                                                                                   acc_field_name)
        elif field_class == 'interpolator':
            self.fluid_start += ('vel_{0} = new {1}<{2}, {3}>(fs, {4});\n' +
                                 'acc_{0} = new {1}<{2}, {3}>(fs, {4});\n').format(name,
                                                                                   field_class,
                                                                                   self.C_dtype,
                                                                                   neighbours,
                                                                                   beta_name,
                                                                                   acc_field_name)
        self.fluid_end += ('delete vel_{0};\n' +
                           'delete acc_{0};\n').format(name)
        update_fields = 'fs->compute_velocity(fs->cvorticity);\n'
        if not type(kcut) == type(None):
            update_fields += 'fs->low_pass_Fourier(fs->cvelocity, 3, {0});\n'.format(kcut)
        update_fields += ('fs->ift_velocity();\n' +
                          'fs->compute_Lagrangian_acceleration(acc_{0}->field);\n').format(name)
        self.fluid_start += update_fields
        self.fluid_loop += update_fields
        return None
999
    def specific_parser_arguments(
1000
            self,
For faster browsing, not all history is shown. View entire blame