_fluid_base.py 18.8 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                                 #
#                                                                     #
#######################################################################


26

27
28
from ._code import _code
from bfps import tools
29

30
import os
31
32
import numpy as np
import h5py
33

34
class _fluid_particle_base(_code):
35
36
37
38
    """This class is meant to put together all common code between the
    different C++ solvers/postprocessing tools, so that development of
    specific functionalities is not overwhelming.
    """
39
40
41
42
    def __init__(
            self,
            name = 'solver',
            work_dir = './',
43
            simname = 'test',
44
45
            dtype = np.float32,
            use_fftw_wisdom = True):
46
47
        _code.__init__(
                self,
48
49
                work_dir = work_dir,
                simname = simname)
50
        self.use_fftw_wisdom = use_fftw_wisdom
51
        self.name = name
52
        self.particle_species = 0
53
54
55
56
        if dtype in [np.float32, np.float64]:
            self.dtype = dtype
        elif dtype in ['single', 'double']:
            if dtype == 'single':
Cristian Lalescu's avatar
Cristian Lalescu committed
57
                self.dtype = np.dtype(np.float32)
58
            elif dtype == 'double':
Cristian Lalescu's avatar
Cristian Lalescu committed
59
                self.dtype = np.dtype(np.float64)
60
61
        self.rtype = self.dtype
        if self.rtype == np.float32:
Cristian Lalescu's avatar
Cristian Lalescu committed
62
            self.ctype = np.dtype(np.complex64)
63
64
            self.C_dtype = 'float'
        elif self.rtype == np.float64:
Cristian Lalescu's avatar
Cristian Lalescu committed
65
            self.ctype = np.dtype(np.complex128)
66
            self.C_dtype = 'double'
Cristian Lalescu's avatar
Cristian Lalescu committed
67
        self.parameters['dealias_type'] = 1
68
69
70
71
72
        self.parameters['dkx'] = 1.0
        self.parameters['dky'] = 1.0
        self.parameters['dkz'] = 1.0
        self.parameters['niter_todo'] = 8
        self.parameters['niter_part'] = 1
73
        self.parameters['niter_stat'] = 1
74
75
76
77
78
79
80
81
82
        self.parameters['niter_out'] = 1024
        self.parameters['nparticles'] = 0
        self.parameters['dt'] = 0.01
        self.fluid_includes = '#include "fluid_solver.hpp"\n'
        self.fluid_variables = ''
        self.fluid_definitions = ''
        self.fluid_start = ''
        self.fluid_loop = ''
        self.fluid_end  = ''
Cristian Lalescu's avatar
Cristian Lalescu committed
83
        self.fluid_output = ''
84
        self.stat_src = ''
85
        self.particle_includes = ''
86
87
88
89
90
        self.particle_variables = ''
        self.particle_definitions = ''
        self.particle_start = ''
        self.particle_loop = ''
        self.particle_end  = ''
91
        self.particle_stat_src = ''
92
        self.file_datasets_grow   = ''
93
94
95
        return None
    def finalize_code(self):
        self.includes   += self.fluid_includes
96
        self.includes   += '#include <ctime>\n'
97
98
99
100
101
102
        self.variables  += self.fluid_variables
        self.definitions+= self.fluid_definitions
        if self.particle_species > 0:
            self.includes    += self.particle_includes
            self.variables   += self.particle_variables
            self.definitions += self.particle_definitions
103
104
        self.definitions += ('int grow_single_dataset(hid_t dset, int tincrement)\n{\n' +
                             'int ndims;\n' +
105
                             'hsize_t dims[5];\n' +
106
107
108
109
110
111
112
113
114
115
116
117
118
                             'hsize_t space;\n' +
                             'space = H5Dget_space(dset);\n' +
                             'ndims = H5Sget_simple_extent_dims(space, dims, NULL);\n' +
                             'dims[0] += tincrement;\n' +
                             'H5Dset_extent(dset, dims);\n' +
                             'H5Sclose(space);\n' +
                             'return EXIT_SUCCESS;\n}\n')
        self.definitions += ('herr_t grow_statistics_dataset(hid_t o_id, const char *name, const H5O_info_t *info, void *op_data)\n{\n' +
                             'if (info->type == H5O_TYPE_DATASET)\n{\n' +
                             'hsize_t dset = H5Dopen(o_id, name, H5P_DEFAULT);\n' +
                             'grow_single_dataset(dset, niter_todo/niter_stat);\n'
                             'H5Dclose(dset);\n}\n' +
                             'return 0;\n}\n')
119
120
121
122
123
        self.definitions += ('herr_t grow_particle_datasets(hid_t g_id, const char *name, const H5L_info_t *info, void *op_data)\n{\n' +
                             'std::string full_name;\n' +
                             'hsize_t dset;\n')
        for key in ['state', 'velocity', 'acceleration']:
            self.definitions += ('full_name = (std::string(name) + std::string("/{0}"));\n'.format(key) +
124
                                 'if (H5Lexists(g_id, full_name.c_str(), H5P_DEFAULT))\n{\n' +
125
126
                                 'dset = H5Dopen(g_id, full_name.c_str(), H5P_DEFAULT);\n' +
                                 'grow_single_dataset(dset, niter_todo/niter_part);\n' +
127
                                 'H5Dclose(dset);\n}\n')
128
        self.definitions += ('full_name = (std::string(name) + std::string("/rhs"));\n' +
Cristian Lalescu's avatar
Cristian Lalescu committed
129
                             'if (H5Lexists(g_id, full_name.c_str(), H5P_DEFAULT))\n{\n' +
130
131
                             'dset = H5Dopen(g_id, full_name.c_str(), H5P_DEFAULT);\n' +
                             'grow_single_dataset(dset, 1);\n' +
Cristian Lalescu's avatar
Cristian Lalescu committed
132
                             'H5Dclose(dset);\n}\n' +
133
                             'return 0;\n}\n')
134
135
        self.definitions += ('int grow_file_datasets()\n{\n' +
                             'int file_problems = 0;\n' +
136
                             self.file_datasets_grow +
137
                             'return file_problems;\n'
138
                             '}\n')
Chichi Lalescu's avatar
Chichi Lalescu committed
139
        self.definitions += 'void do_stats()\n{\n' + self.stat_src + '}\n'
140
        self.definitions += 'void do_particle_stats()\n{\n' + self.particle_stat_src + '}\n'
141
        # take care of wisdom
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
        if self.use_fftw_wisdom:
            if self.dtype == np.float32:
                fftw_prefix = 'fftwf_'
            elif self.dtype == np.float64:
                fftw_prefix = 'fftw_'
            self.main_start += """
                        //begincpp
                        if (myrank == 0)
                        {{
                            char fname[256];
                            sprintf(fname, "%s_fftw_wisdom.txt", simname);
                            {0}import_wisdom_from_filename(fname);
                        }}
                        {0}mpi_broadcast_wisdom(MPI_COMM_WORLD);
                        //endcpp
                        """.format(fftw_prefix)
            self.main_end = """
                        //begincpp
                        {0}mpi_gather_wisdom(MPI_COMM_WORLD);
                        MPI_Barrier(MPI_COMM_WORLD);
                        if (myrank == 0)
                        {{
                            char fname[256];
                            sprintf(fname, "%s_fftw_wisdom.txt", simname);
                            {0}export_wisdom_to_filename(fname);
                        }}
                        //endcpp
                        """.format(fftw_prefix) + self.main_end
170
171
172
        self.main        = self.fluid_start
        if self.particle_species > 0:
            self.main   += self.particle_start
173
174
175
        self.main       += """
                           //begincpp
                           int data_file_problem;
176
                           clock_t time0, time1;
177
                           double time_difference, local_time_difference;
178
                           time0 = clock();
179
180
181
182
183
184
185
186
187
188
                           if (myrank == 0) data_file_problem = grow_file_datasets();
                           MPI_Bcast(&data_file_problem, 1, MPI_INT, 0, MPI_COMM_WORLD);
                           if (data_file_problem > 0)
                           {
                               std::cerr << data_file_problem << " problems growing file datasets.\\ntrying to exit now." << std::endl;
                               MPI_Finalize();
                               return EXIT_SUCCESS;
                           }
                           //endcpp
                           """
189
190
191
192
        output_time_difference = ('time1 = clock();\n' +
                                  'local_time_difference = ((unsigned int)(time1 - time0))/((double)CLOCKS_PER_SEC);\n' +
                                  'time_difference = 0.0;\n' +
                                  'MPI_Allreduce(&local_time_difference, &time_difference, ' +
193
                                      '1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);\n' +
194
195
196
197
                                  'if (myrank == 0) std::cout << "iteration " ' +
                                      '<< iteration << " took " ' +
                                      '<< time_difference/nprocs << " seconds" << std::endl;\n' +
                                  'time0 = time1;\n')
198
199
200
201
202
        self.main       += 'for (int max_iter = iteration+niter_todo; iteration < max_iter; iteration++)\n'
        self.main       += '{\n'
        self.main       += 'if (iteration % niter_stat == 0) do_stats();\n'
        if self.particle_species > 0:
            self.main       += 'if (iteration % niter_part == 0) do_particle_stats();\n'
203
204
            self.main   += self.particle_loop
        self.main       += self.fluid_loop
205
        self.main       += output_time_difference
206
        self.main       += '}\n'
207
208
        self.main       += 'do_stats();\n'
        self.main       += 'do_particle_stats();\n'
209
210
211
212
        if self.particle_species > 0:
            self.main   += self.particle_end
        self.main       += self.fluid_end
        return None
213
214
215
216
217
218
219
220
221
222
223
    def read_rfield(
            self,
            field = 'velocity',
            iteration = 0,
            filename = None):
        if type(filename) == type(None):
            filename = os.path.join(
                    self.work_dir,
                    self.simname + '_r' + field + '_i{0:0>5x}'.format(iteration))
        return np.memmap(
                filename,
224
                dtype = self.dtype,
225
226
227
                shape = (self.parameters['nz'],
                         self.parameters['ny'],
                         self.parameters['nx'], 3))
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
    def transpose_frame(
            self,
            field = 'velocity',
            iteration = 0,
            filename = None,
            ofile = None):
        Rdata = self.read_rfield(
                field = field,
                iteration = iteration,
                filename = filename)
        new_data = np.zeros(
                (3,
                 self.parameters['nz'],
                 self.parameters['ny'],
                 self.parameters['nx']),
243
                dtype = self.dtype)
244
245
246
247
248
249
        for i in range(3):
            new_data[i] = Rdata[..., i]
        if type(ofile) == type(None):
            ofile = os.path.join(
                    self.work_dir,
                    self.simname + '_r' + field + '_i{0:0>5x}_3xNZxNYxNX'.format(iteration))
Cristian Lalescu's avatar
Cristian Lalescu committed
250
251
        else:
            new_data.tofile(ofile)
252
        return new_data
253
254
255
256
257
258
259
260
    def plot_vel_cut(
            self,
            axis,
            field = 'velocity',
            iteration = 0,
            yval = 13,
            filename = None):
        axis.set_axis_off()
261
        Rdata0 = self.read_rfield(field = field, iteration = iteration, filename = filename)
262
263
264
265
266
267
268
269
270
271
        energy = np.sum(Rdata0[:, yval, :, :]**2, axis = 2)*.5
        axis.imshow(energy, interpolation='none')
        axis.set_title('{0}'.format(np.average(Rdata0[..., 0]**2 +
                                               Rdata0[..., 1]**2 +
                                               Rdata0[..., 2]**2)*.5))
        return Rdata0
    def generate_vector_field(
            self,
            rseed = 7547,
            spectra_slope = 1.,
272
            amplitude = 1.,
273
274
275
276
            iteration = 0,
            field_name = 'vorticity',
            write_to_file = False):
        np.random.seed(rseed)
277
        Kdata00 = tools.generate_data_3D(
Cristian Lalescu's avatar
Cristian Lalescu committed
278
279
280
                self.parameters['nz']//2,
                self.parameters['ny']//2,
                self.parameters['nx']//2,
281
282
                p = spectra_slope,
                amplitude = amplitude).astype(self.ctype)
283
        Kdata01 = tools.generate_data_3D(
Cristian Lalescu's avatar
Cristian Lalescu committed
284
285
286
                self.parameters['nz']//2,
                self.parameters['ny']//2,
                self.parameters['nx']//2,
287
288
                p = spectra_slope,
                amplitude = amplitude).astype(self.ctype)
289
        Kdata02 = tools.generate_data_3D(
Cristian Lalescu's avatar
Cristian Lalescu committed
290
291
292
                self.parameters['nz']//2,
                self.parameters['ny']//2,
                self.parameters['nx']//2,
293
294
                p = spectra_slope,
                amplitude = amplitude).astype(self.ctype)
295
296
297
298
299
300
        Kdata0 = np.zeros(
                Kdata00.shape + (3,),
                Kdata00.dtype)
        Kdata0[..., 0] = Kdata00
        Kdata0[..., 1] = Kdata01
        Kdata0[..., 2] = Kdata02
301
        Kdata1 = tools.padd_with_zeros(
302
                Kdata0,
Chichi Lalescu's avatar
Chichi Lalescu committed
303
                self.parameters['nz'],
304
                self.parameters['ny'],
305
306
307
308
309
310
311
312
                self.parameters['nx'])
        if write_to_file:
            Kdata1.tofile(
                    os.path.join(self.work_dir,
                                 self.simname + "_c{0}_i{1:0>5x}".format(field_name, iteration)))
        return Kdata1
    def generate_tracer_state(
            self,
313
            rseed = None,
314
315
316
            iteration = 0,
            species = 0,
            write_to_file = False,
317
            ncomponents = 3,
318
319
320
321
            testing = False,
            data = None):
        if (type(data) == type(None)):
            if not type(rseed) == type(None):
322
                np.random.seed(rseed)
323
324
325
326
327
            #point with problems: 5.37632864e+00,   6.10414710e+00,   6.25256493e+00]
            data = np.zeros(self.parameters['nparticles']*ncomponents).reshape(-1, ncomponents)
            data[:, :3] = np.random.random((self.parameters['nparticles'], 3))*2*np.pi
        else:
            assert(data.shape == (self.parameters['nparticles'], ncomponents))
328
        if testing:
329
330
            #data[0] = np.array([3.26434, 4.24418, 3.12157])
            data[0] = np.array([ 0.72086101,  2.59043666,  6.27501953])
331
        with h5py.File(os.path.join(self.work_dir, self.simname + '.h5'), 'r+') as data_file:
332
333
            time_chunk = 2**20 // (8*ncomponents*
                                   self.parameters['nparticles'])
334
            time_chunk = max(time_chunk, 1)
335
336
            dset = data_file.create_dataset(
                    '/particles/tracers{0}/state'.format(species),
337
                    (1,
338
339
                     self.parameters['nparticles'],
                     ncomponents),
340
                    chunks = (time_chunk, self.parameters['nparticles'], ncomponents),
341
342
                    maxshape = (None, self.parameters['nparticles'], ncomponents),
                    dtype = np.float64)
343
            dset[0] = data
344
345
346
347
        if write_to_file:
            data.tofile(
                    os.path.join(
                        self.work_dir,
348
                        "tracers{0}_state_i{1:0>5x}".format(species, iteration)))
349
350
351
352
353
354
        return data
    def generate_initial_condition(self):
        self.generate_vector_field(write_to_file = True)
        for species in range(self.particle_species):
            self.generate_tracer_state(
                    species = species,
355
                    write_to_file = False)
356
        return None
357
358
359
    def get_kspace(self):
        kspace = {}
        if self.parameters['dealias_type'] == 1:
360
361
362
            kMx = self.parameters['dkx']*(self.parameters['nx']//2 - 1)
            kMy = self.parameters['dky']*(self.parameters['ny']//2 - 1)
            kMz = self.parameters['dkz']*(self.parameters['nz']//2 - 1)
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        else:
            kMx = self.parameters['dkx']*(self.parameters['nx']//3 - 1)
            kMy = self.parameters['dky']*(self.parameters['ny']//3 - 1)
            kMz = self.parameters['dkz']*(self.parameters['nz']//3 - 1)
        kspace['kM'] = max(kMx, kMy, kMz)
        kspace['dk'] = min(self.parameters['dkx'],
                           self.parameters['dky'],
                           self.parameters['dkz'])
        nshells = int(kspace['kM'] / kspace['dk']) + 2
        kspace['nshell'] = np.zeros(nshells, dtype = np.int64)
        kspace['kshell'] = np.zeros(nshells, dtype = np.float64)
        kspace['kx'] = np.arange( 0,
                                  self.parameters['nx']//2 + 1).astype(np.float64)*self.parameters['dkx']
        kspace['ky'] = np.arange(-self.parameters['ny']//2 + 1,
377
                                  self.parameters['ny']//2 + 1).astype(np.float64)*self.parameters['dky']
378
379
        kspace['ky'] = np.roll(kspace['ky'], self.parameters['ny']//2+1)
        kspace['kz'] = np.arange(-self.parameters['nz']//2 + 1,
380
                                  self.parameters['nz']//2 + 1).astype(np.float64)*self.parameters['dkz']
381
382
        kspace['kz'] = np.roll(kspace['kz'], self.parameters['nz']//2+1)
        return kspace
383
    def write_par(self, iter0 = 0):
384
385
386
387
388
        assert (self.parameters['niter_todo'] % self.parameters['niter_stat'] == 0)
        assert (self.parameters['niter_todo'] % self.parameters['niter_out']  == 0)
        assert (self.parameters['niter_todo'] % self.parameters['niter_part'] == 0)
        assert (self.parameters['niter_out']  % self.parameters['niter_stat'] == 0)
        assert (self.parameters['niter_out']  % self.parameters['niter_part'] == 0)
389
        _code.write_par(self, iter0 = iter0)
390
391
392
393
394
395
        with h5py.File(os.path.join(self.work_dir, self.simname + '.h5'), 'r+') as ofile:
            ofile['field_dtype'] = np.dtype(self.dtype).str
            kspace = self.get_kspace()
            for k in kspace.keys():
                ofile['kspace/' + k] = kspace[k]
            nshells = kspace['nshell'].shape[0]
396
            ofile.close()
397
        return None
398
399
400
401
402
    def specific_parser_arguments(
            self,
            parser):
        _code.specific_parser_arguments(self, parser)
        return None
403