turbx 1.0.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,940 @@
1
+ import os
2
+ import re
3
+ import sys
4
+ import timeit
5
+ from concurrent.futures import ThreadPoolExecutor
6
+ from pathlib import Path, PurePosixPath
7
+
8
+ import h5py
9
+ import numpy as np
10
+ import psutil
11
+ import scipy as sp
12
+ from mpi4py import MPI
13
+ from scipy.signal import csd
14
+ from tqdm import tqdm
15
+
16
+ from .h5 import h5_print_contents
17
+ from .utils import even_print, format_time_string
18
+
19
+ # ======================================================================
20
+
21
+ def _calc_turb_cospectrum_xpln(self, **kwargs):
22
+ '''
23
+ Calculate FFT cospectrum in [z,t] at every [x,y]
24
+ - Designed for analyzing unsteady, thin planes in [x]
25
+ - Multithreaded with ThreadPoolExecutor()
26
+ - scipy.signal.csd() automatically tries to run multithreaded
27
+ - set OMP_NUM_THREADS=1 and pass 'n_threads' to as kwarg manually
28
+ '''
29
+
30
+ if (self.rank==0):
31
+ verbose = True
32
+ else:
33
+ verbose = False
34
+
35
+ if verbose: print('\n'+'rgd.calc_turb_cospectrum_xpln()'+'\n'+72*'-')
36
+ t_start_func = timeit.default_timer()
37
+
38
+ ## assert that the opened RGD has fsubtype 'unsteady' (i.e. is NOT a prime file)
39
+ if (self.fsubtype!='unsteady'):
40
+ raise ValueError
41
+ if not self.usingmpi:
42
+ raise NotImplementedError('function is not implemented for non-MPI usage')
43
+
44
+ if not h5py.h5.get_config().mpi:
45
+ if verbose: print('h5py was not compiled for parallel usage! exiting.')
46
+ sys.exit(1)
47
+
48
+ rx = kwargs.get('rx',1)
49
+ ry = kwargs.get('ry',1)
50
+ rz = kwargs.get('rz',1)
51
+ rt = kwargs.get('rt',1)
52
+
53
+ # cy = kwargs.get('cy',1) ## number of subdivisions per rank [y] range
54
+ # if not isinstance(cy,int):
55
+ # raise TypeError('cy should be an int')
56
+ # if (cy<1):
57
+ # raise TypeError('cy should be an int')
58
+
59
+ sy = kwargs.get('sy',1) ## number of [y] layers to read at a time
60
+ if not isinstance(sy,int) or (sy<1):
61
+ raise TypeError('sy should be a positive non-zero int')
62
+
63
+ n_threads = kwargs.get('n_threads',1)
64
+
65
+ ## Debug Rank:Proc Affinity
66
+ #pp = psutil.Process()
67
+ #print(f"[Rank {self.rank}] sees CPUs: {pp.cpu_affinity()} | n_threads={n_threads} | OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS')}")
68
+
69
+ #try:
70
+ # n_threads = int(os.environ.get('OMP_NUM_THREADS'))
71
+ #except TypeError: ## not set
72
+ # n_threads = os.cpu_count()
73
+
74
+ fn_h5_out = kwargs.get('fn_h5_out',None) ## filename for output HDF5 (.h5) file
75
+ overlap_fac_nom = kwargs.get('overlap_fac_nom',0.50) ## nominal windows overlap factor
76
+ n_win = kwargs.get('n_win',8) ## number of segment windows for [t] PSD calc
77
+ #window_type = kwargs.get('window_type','hann') ## 'tukey','hann'
78
+
79
+ ## only distribute data across [y]
80
+ if (rx!=1):
81
+ raise ValueError('rx!=1')
82
+ if (rz!=1):
83
+ raise ValueError('rz!=1')
84
+ if (rt!=1):
85
+ raise ValueError('rt!=1')
86
+
87
+ if not isinstance(ry,int) or (ry<1):
88
+ raise ValueError('ry should be a positive non-zero int')
89
+
90
+ ## check the choice of ranks per dimension
91
+ if (rx*ry*rz*rt != self.n_ranks):
92
+ raise ValueError('rx*ry*rz*rt != self.n_ranks')
93
+ if (rx>self.nx):
94
+ raise ValueError('rx>self.nx')
95
+ if (ry>self.ny):
96
+ raise ValueError('ry>self.ny')
97
+ if (rz>self.nz):
98
+ raise ValueError('rz>self.nz')
99
+ if (rt>self.nt):
100
+ raise ValueError('rt>self.nt')
101
+
102
+ if (self.ny%ry!=0):
103
+ raise ValueError('ny not divisible by ry')
104
+
105
+ ## distribute 4D data over ranks --> here only in [y]
106
+ ryl_ = np.array_split(np.arange(self.ny,dtype=np.int64),ry)
107
+ ryl = [[b[0],b[-1]+1] for b in ryl_ ]
108
+ ry1,ry2 = ryl[self.rank]
109
+ nyr = ry2 - ry1
110
+
111
+ ## check all [y] ranges have same size
112
+ for ryl_ in ryl:
113
+ if not (ryl_[1]-ryl_[0]==nyr):
114
+ raise ValueError('[y] chunks are not even in size')
115
+
116
+ # ## [y] sub chunk range --> cyl = list of ranges in ry1:ry2
117
+ # ## cy is the NUMBER of chunks for the rank sub-range
118
+ # cyl_ = np.array_split( np.arange(ry1,ry2) , min(cy,nyr) )
119
+ # cyl = [[b[0],b[-1]+1] for b in cyl_ ]
120
+ #
121
+ # for nyc_ in [ cyl_[1]-cyl_[0] for cyl_ in cyl ]:
122
+ # if (nyc_ < 1):
123
+ # #raise ValueError
124
+ # print(f'rank {self.rank:d}: sub-range is <1')
125
+ # self.comm.Abort(1)
126
+ #
127
+ # if 1: ## assert that [y] sub-chunk ranges are correct
128
+ #
129
+ # yi = np.arange(self.ny, dtype=np.int32)
130
+ #
131
+ # local_indices = []
132
+ # for cyl_ in cyl:
133
+ # cy1, cy2 = cyl_
134
+ # local_indices += [ yi_ for yi_ in yi[cy1:cy2] ]
135
+ #
136
+ # G = self.comm.gather([ self.rank , local_indices ], root=0)
137
+ # G = self.comm.bcast(G, root=0)
138
+ #
139
+ # all_indices = []
140
+ # for G_ in G:
141
+ # all_indices += G_[1]
142
+ # all_indices = np.array( sorted(all_indices), dtype=np.int32 )
143
+ #
144
+ # if not np.array_equal( all_indices , yi ):
145
+ # raise AssertionError
146
+
147
+ if (nyr%sy!=0):
148
+ raise ValueError('nyr not divisible by sy')
149
+
150
+ ## output filename : HDF5 (.h5)
151
+ if (fn_h5_out is None): ## automatically determine name
152
+ fname_path = os.path.dirname(self.fname)
153
+ fname_base = os.path.basename(self.fname)
154
+ fname_root, fname_ext = os.path.splitext(fname_base)
155
+ fname_root = re.findall(r'io\S+_mpi_[0-9]+', fname_root)[0]
156
+ fname_fft_h5_base = fname_root+'_fft.h5'
157
+ fn_h5_out = str(PurePosixPath(fname_path, fname_fft_h5_base))
158
+ if (Path(fn_h5_out).suffix != '.h5'):
159
+ raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' must end in .h5")
160
+ if os.path.isfile(fn_h5_out):
161
+ #if (os.path.getsize(fn_h5_out) > 8*1024**3):
162
+ # raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' exists and is >8 [GB]. exiting for your own safety.")
163
+ if (fn_h5_out == self.fname):
164
+ raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' cannot be same as input filename.")
165
+
166
+ if verbose: even_print( 'fn_h5' , self.fname )
167
+ if verbose: even_print( 'fn_h5_out' , fn_h5_out )
168
+ if verbose: print(72*'-')
169
+ self.comm.Barrier()
170
+
171
+ ## the data dictionary to be written to .h5 later
172
+ data = {}
173
+
174
+ ## infile
175
+ fsize = os.path.getsize(self.fname)/1024**3
176
+ if verbose: even_print(os.path.basename(self.fname),'%0.1f [GB]'%fsize)
177
+ if verbose: even_print('nx',f'{self.nx:d}')
178
+ if verbose: even_print('ny',f'{self.ny:d}')
179
+ if verbose: even_print('nz',f'{self.nz:d}')
180
+ if verbose: even_print('nt',f'{self.nt:d}')
181
+ if verbose: even_print('ngp',f'{self.ngp/1e6:0.1f} [M]')
182
+ #if verbose: even_print('cy',f'{cy:d}')
183
+ if verbose: even_print('sy',f'{sy:d}')
184
+ if verbose: even_print('n_ranks',f'{self.n_ranks:d}')
185
+ if verbose: even_print('n_threads',f'{n_threads:d}')
186
+ if verbose: print(72*'-')
187
+
188
+ ## 0D freestream scalars
189
+ lchar = self.lchar ; data['lchar'] = lchar
190
+ U_inf = self.U_inf ; data['U_inf'] = U_inf
191
+ rho_inf = self.rho_inf ; data['rho_inf'] = rho_inf
192
+ T_inf = self.T_inf ; data['T_inf'] = T_inf
193
+
194
+ #data['M_inf'] = self.M_inf
195
+ data['Ma'] = self.Ma
196
+ data['Pr'] = self.Pr
197
+
198
+ ## read in 1D coordinate arrays & re-dimensionalize
199
+ x = np.copy( self['dims/x'][()] * self.lchar )
200
+ y = np.copy( self['dims/y'][()] * self.lchar )
201
+ z = np.copy( self['dims/z'][()] * self.lchar )
202
+ t = np.copy( self['dims/t'][()] * self.tchar )
203
+
204
+ nx = self.nx ; data['nx'] = nx
205
+ ny = self.ny ; data['ny'] = ny
206
+ nz = self.nz ; data['nz'] = nz
207
+ nt = self.nt ; data['nt'] = nt
208
+
209
+ ## assert constant Δz
210
+ dz0 = np.diff(z)[0]
211
+ if not np.all(np.isclose(np.diff(z), dz0, rtol=1e-6)):
212
+ raise NotImplementedError('Δz not constant')
213
+ dz = np.diff(z)[0]
214
+
215
+ ## dimensional [s]
216
+ dt = self.dt * self.tchar
217
+ np.testing.assert_allclose(dt, t[1]-t[0], rtol=1e-12, atol=1e-12)
218
+
219
+ t_meas = self.duration * self.tchar
220
+ np.testing.assert_allclose(t_meas, t.max()-t.min(), rtol=1e-12, atol=1e-12)
221
+
222
+ zrange = z.max() - z.min()
223
+
224
+ data['x'] = x
225
+ data['y'] = y
226
+ data['z'] = z
227
+
228
+ data['t'] = t
229
+ data['t_meas'] = t_meas
230
+ data['dt'] = dt
231
+ data['dz'] = dz
232
+ data['zrange'] = zrange
233
+
234
+ if verbose: even_print( 'Δt/tchar' , f'{dt/self.tchar:0.8f}' )
235
+ if verbose: even_print( 'Δt' , f'{dt:0.3e} [s]' )
236
+ if verbose: even_print( 'duration/tchar' , f'{self.duration:0.1f}' )
237
+ if verbose: even_print( 'duration' , f'{self.duration*self.tchar:0.3e} [s]' )
238
+ if verbose: print(72*'-')
239
+
240
+ ## report
241
+ if verbose:
242
+ even_print('Δt' , f'{dt :0.5e} [s]' )
243
+ even_print('t_meas' , f'{t_meas:0.5e} [s]' )
244
+ even_print('Δz' , f'{dz0 :0.5e} [m]' )
245
+ even_print('zrange' , f'{zrange:0.5e} [m]' )
246
+ print(72*'-')
247
+
248
+ ## establish [t] windowing (old)
249
+ #win_len, overlap = get_overlapping_window_size(nt, n_win, overlap_fac_nom)
250
+ #overlap_fac = overlap / win_len
251
+ #tw, n_win, n_pad = get_overlapping_windows(t, win_len, overlap)
252
+ #t_meas_per_win = (win_len-1)*dt
253
+ #data['win_len'] = win_len
254
+ #data['overlap_fac'] = overlap_fac
255
+ #data['overlap'] = overlap
256
+ #data['n_win'] = n_win
257
+ #data['t_meas_per_win'] = t_meas_per_win
258
+
259
+ # ## temporal [t] frequency (f) vector (Short Time Fourier Transform)
260
+ # freq_full = sp.fft.fftfreq(n=win_len, d=dt)
261
+ # fp = np.where(freq_full>0) ## indices of positive values
262
+ # freq = np.copy(freq_full[fp])
263
+ # df = freq[1] - freq[0]
264
+ # nf = freq.size
265
+
266
+ ## establish [t] windowing & get frequency
267
+ nperseg = nt // n_win
268
+ noverlap = int(round(nperseg*overlap_fac_nom))
269
+ overlap_fac = noverlap / nperseg
270
+ fs = 1./dt ## dimensional [1/s]
271
+
272
+ ## get [freq] vector
273
+ freq,_ = csd(
274
+ np.zeros((nt,),dtype=np.float64),
275
+ np.zeros((nt,),dtype=np.float64),
276
+ fs=fs,
277
+ nperseg=nperseg,
278
+ noverlap=noverlap,
279
+ window='hann',
280
+ detrend='constant',
281
+ scaling='density',
282
+ return_onesided=True,
283
+ )
284
+
285
+ nf = freq.shape[0]
286
+ df = np.diff(freq)[0]
287
+
288
+ data['nperseg'] = nperseg
289
+ data['noverlap'] = noverlap
290
+ data['freq'] = freq
291
+ data['df'] = df
292
+ data['nf'] = nf
293
+
294
+ if verbose:
295
+ even_print('overlap_fac (nominal)' , f'{overlap_fac_nom:0.5f}' )
296
+ even_print('n_win' , f'{n_win:d}' )
297
+ even_print('nperseg' , f'{nperseg:d}' )
298
+ even_print('noverlap' , f'{noverlap:d}' )
299
+ even_print('overlap_fac' , f'{overlap_fac:0.5f}' )
300
+ print(72*'-')
301
+
302
+ if verbose:
303
+ even_print('freq min',f'{freq.min():0.1f} [Hz]')
304
+ even_print('freq max',f'{freq.max():0.1f} [Hz]')
305
+ even_print('df',f'{df:0.1f} [Hz]')
306
+ even_print('nf',f'{nf:d}')
307
+ print(72*'-')
308
+
309
+ ## spatial [x] wavenumber (kx) and wavelength (λx)
310
+ # λx = u/f
311
+ # kx = 2·π·f/u
312
+ # λx+ = λx/(ν/uτ)
313
+ # kx+ = (2·π·f/u)·(ν/uτ)
314
+
315
+ ## kx = 2·π·f/u --> [y,f]
316
+ #kx = np.copy( 2*np.pi*freq[na,:] / u_avg[:,na] )
317
+
318
+ ## λx = u/f --> [y,f]
319
+ #lx = np.copy( u_avg[:,na] / freq[na,:] )
320
+
321
+ #data['kx'] = kx
322
+ #data['lx'] = lx
323
+
324
+ ## spatial [z] wavenumber (kz) vector
325
+ kz_full = sp.fft.fftfreq(n=nz, d=dz0) * ( 2 * np.pi )
326
+ kzp = np.where(kz_full>0) ## indices of positive values
327
+ kz = np.copy(kz_full[kzp])
328
+ dkz = kz[1] - kz[0]
329
+ nkz = kz.shape[0]
330
+
331
+ ## wavenumber vector should be size nz//2-1
332
+ if (nkz!=nz//2-1):
333
+ raise ValueError
334
+
335
+ data['kz'] = kz
336
+ data['dkz'] = dkz
337
+ data['nkz'] = nkz
338
+
339
+ if verbose:
340
+ even_print('kz min',f'{kz.min():0.1f} [1/m]')
341
+ even_print('kz max',f'{kz.max():0.1f} [1/m]')
342
+ even_print('dkz',f'{dkz:0.1f} [1/m]')
343
+ even_print('nkz',f'{nkz:d}')
344
+ print(72*'-')
345
+
346
+ ## Wavelength λz = (2·π)/kz
347
+ lz = np.copy( 2 * np.pi / kz )
348
+ data['lz'] = lz
349
+
350
+ # ===
351
+
352
+ ## cospectrum pairs
353
+ ## [ (str:var1, bool:ρ_weighting) , (str:var2, bool:ρ_weighting) ]
354
+ fft_combis = [
355
+
356
+ [ ('u',True) , ('v',True) ], ## [ ρ·u″ , ρ·v″ ]
357
+ [ ('u',True) , ('u',True) ], ## [ ρ·u″ , ρ·u″ ]
358
+ [ ('v',True) , ('v',True) ], ## [ ρ·v″ , ρ·v″ ]
359
+ [ ('w',True) , ('w',True) ], ## [ ρ·w″ , ρ·w″ ]
360
+
361
+ [ ('u',False) , ('v',False) ], ## [ u′ , v′ ]
362
+ [ ('u',False) , ('u',False) ], ## [ u′ , u′ ]
363
+ [ ('v',False) , ('v',False) ], ## [ v′ , v′ ]
364
+ [ ('w',False) , ('w',False) ], ## [ w′ , w′ ]
365
+
366
+ [ ('p',False) , ('u',False) ], ## [ p′ , u′ ]
367
+ [ ('p',False) , ('v',False) ], ## [ p′ , v′ ]
368
+ [ ('p',False) , ('w',False) ], ## [ p′ , w′ ]
369
+
370
+ [ ('p',False) , ('u',True) ], ## [ p′ , ρ·u″ ]
371
+ [ ('p',False) , ('v',True) ], ## [ p′ , ρ·v″ ]
372
+ [ ('p',False) , ('w',True) ], ## [ p′ , ρ·w″ ]
373
+
374
+ [ ('p',False) , ('p',False) ], ## [ p′ , p′ ]
375
+ [ ('T',False) , ('T',False) ], ## [ T′ , T′ ]
376
+ [ ('rho',False) , ('rho',False) ], ## [ ρ′ , ρ′ ]
377
+
378
+ ]
379
+
380
+ ## generate FFT cospectrum scalar names
381
+ scalars = []
382
+ for fft_combi in fft_combis:
383
+ if not isinstance(fft_combi,list):
384
+ raise RuntimeError
385
+ if not len(fft_combi)==2:
386
+ raise RuntimeError
387
+
388
+ sL = fft_combi[0][0]
389
+ do_density_weighting_L = fft_combi[0][1]
390
+ if do_density_weighting_L:
391
+ sLs = f'r{sL}II'
392
+ else:
393
+ sLs = f'{sL}I'
394
+
395
+ sR = fft_combi[1][0]
396
+ do_density_weighting_R = fft_combi[1][1]
397
+ if do_density_weighting_R:
398
+ sRs = f'r{sR}II'
399
+ else:
400
+ sRs = f'{sR}I'
401
+
402
+ scalars.append(f'{sLs}_{sRs}')
403
+
404
+ ## generate avg scalar names
405
+ scalars_Re_avg = []
406
+ scalars_Fv_avg = []
407
+ for fft_combi in fft_combis:
408
+ sL = fft_combi[0][0]
409
+ do_density_weighting_L = fft_combi[0][1]
410
+ sR = fft_combi[1][0]
411
+ do_density_weighting_R = fft_combi[1][1]
412
+
413
+ if do_density_weighting_L or do_density_weighting_R:
414
+ if 'rho' not in scalars_Re_avg:
415
+ scalars_Re_avg.append('rho')
416
+
417
+ if do_density_weighting_L:
418
+ if sL not in scalars_Fv_avg:
419
+ scalars_Fv_avg.append(sL)
420
+ else:
421
+ if sL not in scalars_Re_avg:
422
+ scalars_Re_avg.append(sL)
423
+
424
+ if do_density_weighting_R:
425
+ if sR not in scalars_Fv_avg:
426
+ scalars_Fv_avg.append(sR)
427
+ else:
428
+ if sR not in scalars_Re_avg:
429
+ scalars_Re_avg.append(sR)
430
+
431
+ ## numpy formatted arrays: buffers for PSD & other data (rank-local)
432
+ Ekz = np.zeros(shape=(nyr,nkz ) , dtype={'names':scalars , 'formats':[ np.dtype(np.complex128) for s in scalars ] })
433
+ Ef = np.zeros(shape=(nyr,nf ) , dtype={'names':scalars , 'formats':[ np.dtype(np.complex128) for s in scalars ] })
434
+ covariance = np.zeros(shape=(nyr, ) , dtype={'names':scalars , 'formats':[ np.dtype(np.float64) for s in scalars ] })
435
+ avg_Re = np.zeros(shape=(nyr, ) , dtype={'names':scalars_Re_avg , 'formats':[ np.dtype(np.float64) for s in scalars_Re_avg ] })
436
+ avg_Fv = np.zeros(shape=(nyr, ) , dtype={'names':scalars_Fv_avg , 'formats':[ np.dtype(np.float64) for s in scalars_Fv_avg ] })
437
+
438
+ if verbose:
439
+ even_print('n turb spectrum scalar combinations' , '%i'%(len(fft_combis),))
440
+ print(72*'-')
441
+
442
+ ## window for [z] -- rectangular because [z] is assumed periodic already
443
+ window_z = np.ones(nz,dtype=np.float64)
444
+ sum_sqrt_win_z = np.sum(np.sqrt(window_z))
445
+ mean_sq_win_z = np.mean(window_z**2)
446
+ if verbose:
447
+ #even_print('sum(sqrt(window_z))' , '%0.5f'%(sum_sqrt_win_z,))
448
+ even_print('sum(sqrt(window_z)) / nz' , '%0.5f'%(sum_sqrt_win_z/nz,))
449
+ even_print('mean(window_z**2)' , '%0.5f'%(mean_sq_win_z,))
450
+
451
+ # ## window function for [t]
452
+ # if (window_type=='tukey'):
453
+ # window_t = sp.signal.windows.tukey(win_len,alpha=0.5,sym=False) ## α=0:rectangular, α=1:Hann
454
+ # elif (window_type=='hann'):
455
+ # window_t = sp.signal.windows.hann(win_len,sym=False)
456
+ # elif (window_type is None):
457
+ # window_t = np.ones(win_len, dtype=np.float64)
458
+ # else:
459
+ # raise ValueError
460
+ #
461
+ # if verbose:
462
+ # even_print('window type [t]', '\'%s\''%str(window_type))
463
+ #
464
+ # ## sum of sqrt of window: needed for normalization
465
+ # sum_sqrt_win_t = np.sum(np.sqrt(window_t))
466
+ # if verbose:
467
+ # #even_print('sum(sqrt(window_t))' , '%0.5f'%(sum_sqrt_win_t,))
468
+ # even_print('sum(sqrt(window_t)) / win_len', '%0.5f'%(sum_sqrt_win_t/win_len,))
469
+
470
+ #if verbose: print(72*'-')
471
+
472
+ # ==============================================================
473
+ # check memory
474
+ # ==============================================================
475
+
476
+ hostname = MPI.Get_processor_name()
477
+ mem_free_gb = psutil.virtual_memory().free / 1024**3
478
+ G = self.comm.gather([ self.rank , hostname , mem_free_gb ], root=0)
479
+ G = self.comm.bcast(G, root=0)
480
+
481
+ host_mem = {}
482
+ for rank, host, mem in G:
483
+ if host not in host_mem or mem < host_mem[host]:
484
+ host_mem[host] = mem
485
+ total_free = sum(host_mem.values())
486
+
487
+ if verbose:
488
+ print(72*'-')
489
+ for key,value in host_mem.items():
490
+ even_print(f'RAM free {key}', f'{int(np.floor(value)):d} [GB]')
491
+ even_print('RAM free (local,min)', f'{int(np.floor(min(host_mem.values()))):d} [GB]')
492
+ even_print('RAM free (global)', f'{int(np.floor(total_free)):d} [GB]')
493
+
494
+ shape_read = (nx,sy,nz,nt) ## local
495
+ if verbose: even_print('read shape (local)', f'[{nx:d},{sy:d},{nz:d},{nt:d}]')
496
+ data_gb = np.dtype(np.float64).itemsize * np.prod(shape_read) / 1024**3
497
+ if verbose: even_print('read size (global)', f'{int(np.ceil(data_gb*ry)):d} [GB]')
498
+
499
+ if verbose: even_print('read size (global) ×8', f'{int(np.ceil(data_gb*ry*8)):d} [GB]')
500
+ ram_usage_est = data_gb*ry*8/total_free
501
+ if verbose: even_print('RAM usage estimate', f'{100*ram_usage_est:0.1f} [%]')
502
+
503
+ self.comm.Barrier()
504
+ if (ram_usage_est>0.80):
505
+ print('RAM consumption might be too high. exiting.')
506
+ self.comm.Abort(1)
507
+
508
+ # ==============================================================
509
+ # main loop
510
+ # ==============================================================
511
+
512
+ if verbose:
513
+ progress_bar = tqdm(
514
+ #total=len(fft_combis)*cy,
515
+ total=len(fft_combis)*(nyr//sy),
516
+ ncols=100,
517
+ desc='fft',
518
+ leave=True,
519
+ file=sys.stdout,
520
+ mininterval=0.1,
521
+ smoothing=0.,
522
+ #bar_format="\033[B{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}\033[A\n\b",
523
+ bar_format="{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}",
524
+ ascii="░█",
525
+ colour='#FF6600',
526
+ )
527
+
528
+ for cci,cc in enumerate(fft_combis): ## fft pairs
529
+
530
+ if verbose: tqdm.write(72*'-')
531
+
532
+ scalar_L = cc[0][0]
533
+ do_density_weighting_L = cc[0][1]
534
+ scalar_R = cc[1][0]
535
+ do_density_weighting_R = cc[1][1]
536
+
537
+ if do_density_weighting_L:
538
+ sLs = f'r{scalar_L}II'
539
+ sLsF = f'ρ·{scalar_L}″'
540
+ else:
541
+ sLs = f'{scalar_L}I'
542
+ sLsF = f'{scalar_L}′'
543
+
544
+ if do_density_weighting_R:
545
+ sRs = f'r{scalar_R}II'
546
+ sRsF = f'ρ·{scalar_R}″'
547
+ else:
548
+ sRs = f'{scalar_R}I'
549
+ sRsF = f'{scalar_R}′'
550
+
551
+ msg = f'[{sLsF},{sRsF}]'
552
+
553
+ dset_L = self[f'data/{scalar_L}']
554
+ dset_R = self[f'data/{scalar_R}']
555
+ dset_rho = self['data/rho']
556
+
557
+ scalar = scalars[cci]
558
+
559
+ ## assert scalar name
560
+ if (f'{sLs}_{sRs}' != scalar):
561
+ raise RuntimeError(f"'{sLs}_{sRs}' != '{scalar}'")
562
+
563
+ ## [y] loop outer (grid subchunks within rank)
564
+ # for cyl_ in cyl:
565
+ # cy1, cy2 = cyl_
566
+ # nyc = cy2 - cy1
567
+
568
+ for ci in range(nyr//sy):
569
+
570
+ ## buffers
571
+ data_L = np.zeros((nx,sy,nz,nt), dtype=np.float64)
572
+ data_R = np.zeros((nx,sy,nz,nt), dtype=np.float64)
573
+ if do_density_weighting_L or do_density_weighting_R:
574
+ rho = np.zeros((nx,sy,nz,nt), dtype=np.float64)
575
+ else:
576
+ rho = None
577
+
578
+ cy1 = ry1 + ci*sy
579
+ cy2 = cy1 + sy
580
+ nyc = cy2 - cy1
581
+
582
+ self.comm.Barrier()
583
+ t_start = timeit.default_timer()
584
+
585
+ ## read data L
586
+ n_scalars_read = 1 ## initialize
587
+ scalar_str = scalar_L
588
+ with dset_L.collective:
589
+ data_L[:,:,:,:] = dset_L[:,:,cy1:cy2,:].T
590
+
591
+ ## read data R (if != data L)
592
+ if (scalar_L==scalar_R):
593
+ data_R[:,:,:,:] = data_L[:,:,:,:]
594
+ else:
595
+ n_scalars_read += 1
596
+ scalar_str += f',{scalar_R}'
597
+ with dset_R.collective:
598
+ data_R[:,:,:,:] = dset_R[:,:,cy1:cy2,:].T
599
+
600
+ ## read ρ
601
+ if do_density_weighting_L or do_density_weighting_R:
602
+ n_scalars_read += 1
603
+ scalar_str += ',ρ'
604
+ with dset_rho.collective:
605
+ rho[:,:,:,:] = dset_rho[:,:,cy1:cy2,:].T
606
+ else:
607
+ rho = None
608
+
609
+ self.comm.Barrier()
610
+ t_delta = timeit.default_timer() - t_start
611
+ data_gb = n_scalars_read * ( nx * ry * (cy2-cy1) * nz * nt * dset_L.dtype.itemsize ) / 1024**3
612
+ if verbose:
613
+ tqdm.write(even_print(f'read: {scalar_str}', '%0.3f [GB] %0.3f [s] %0.3f [GB/s]'%(data_gb,t_delta,(data_gb/t_delta)), s=True))
614
+
615
+ ## data_L and data_R should be [nx,nyc,nz,nt] where nyc is the chunk [y] range
616
+ if ( data_L.shape != (nx,nyc,nz,nt) ) or ( data_R.shape != (nx,nyc,nz,nt) ):
617
+ print(f'rank {self.rank:d}: shape violation')
618
+ self.comm.Abort(1)
619
+ if (rho is not None) and ( rho.shape != (nx,nyc,nz,nt) ):
620
+ print(f'rank {self.rank:d}: shape violation')
621
+ self.comm.Abort(1)
622
+
623
+ # === redimensionalize
624
+
625
+ if scalar_L in ['u','v','w',]:
626
+ data_L *= U_inf
627
+ elif scalar_L in ['p',]:
628
+ data_L *= rho_inf * U_inf**2
629
+ elif scalar_L in ['T',]:
630
+ data_L *= T_inf
631
+ elif scalar_L in ['rho',]:
632
+ data_L *= rho_inf
633
+ else:
634
+ raise RuntimeError
635
+
636
+ if scalar_R in ['u','v','w',]:
637
+ data_R *= U_inf
638
+ elif scalar_R in ['p',]:
639
+ data_R *= rho_inf * U_inf**2
640
+ elif scalar_R in ['T',]:
641
+ data_R *= T_inf
642
+ elif scalar_R in ['rho',]:
643
+ data_R *= rho_inf
644
+ else:
645
+ raise RuntimeError
646
+
647
+ if (rho is not None):
648
+ rho *= rho_inf
649
+
650
+ # === compute mean-removed data
651
+
652
+ ## Reynolds avg(□) or Favre avg(ρ·□)/avg(ρ) in [t]
653
+ if do_density_weighting_L:
654
+ rho_avg = np.mean( rho , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
655
+ data_L_avg = np.mean( rho*data_L , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
656
+ data_L_avg /= rho_avg
657
+ else:
658
+ data_L_avg = np.mean( data_L , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
659
+
660
+ ## Reynolds avg(□) or Favre avg(ρ·□)/avg(ρ) in [t]
661
+ if do_density_weighting_R:
662
+ rho_avg = np.mean( rho , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
663
+ data_R_avg = np.mean( rho*data_R , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
664
+ data_R_avg /= rho_avg
665
+ else:
666
+ data_R_avg = np.mean( data_R , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
667
+
668
+ ## Reynolds prime □′ or Favre prime □″
669
+ data_L -= data_L_avg
670
+ data_R -= data_R_avg
671
+
672
+ ## LEFT
673
+ ## Assert stationarity + definition Re/Fv averaging
674
+ ## avg(□′)==0 or avg(ρ·□″)==0
675
+ if do_density_weighting_L:
676
+ a_ = np.mean(rho*data_L, axis=3, dtype=np.float64, keepdims=True)
677
+ else:
678
+ a_ = np.mean(data_L, axis=3, dtype=np.float64, keepdims=True)
679
+ if not np.allclose( a_, np.zeros_like(a_), atol=1e-6 ):
680
+ print(f'rank {self.rank:d}: avg(□′)!=0 or avg(ρ·□″)!=0')
681
+ self.comm.Abort(1)
682
+
683
+ ## RIGHT
684
+ ## Assert stationarity + definition Re/Fv averaging
685
+ ## avg(□′)==0 or avg(ρ·□″)==0
686
+ if do_density_weighting_R:
687
+ a_ = np.mean(rho*data_R, axis=3, dtype=np.float64, keepdims=True)
688
+ else:
689
+ a_ = np.mean(data_R, axis=3, dtype=np.float64, keepdims=True)
690
+ if not np.allclose( a_, np.zeros_like(a_), atol=1e-6 ):
691
+ print(f'rank {self.rank:d}: avg(□′)!=0 or avg(ρ·□″)!=0')
692
+ self.comm.Abort(1)
693
+
694
+ ## covariance: <□′·□′> OR <ρ□″·ρ□″> etc. --> note that this is NOT the typical Favre <ρ·□″□″> except in special cases
695
+ if do_density_weighting_L and do_density_weighting_R:
696
+ covariance_ = np.mean( rho*data_L * rho*data_R , axis=3 , dtype=np.float64, keepdims=True)
697
+ elif do_density_weighting_L and not do_density_weighting_R:
698
+ covariance_ = np.mean( rho*data_L * data_R , axis=3 , dtype=np.float64, keepdims=True)
699
+ elif not do_density_weighting_L and do_density_weighting_R:
700
+ covariance_ = np.mean( data_L * rho*data_R , axis=3 , dtype=np.float64, keepdims=True)
701
+ elif not do_density_weighting_L and not do_density_weighting_R:
702
+ covariance_ = np.mean( data_L * data_R , axis=3 , dtype=np.float64, keepdims=True)
703
+ else:
704
+ raise RuntimeError
705
+
706
+ ## write this chunk/scalar's covariance to covariance buffer
707
+ ## avg over [x,z] : [x,y,z,1] --> [y]
708
+ yiA = cy1 - ry1
709
+ yiB = cy2 - ry1
710
+ covariance[scalar][yiA:yiB] = np.squeeze( np.mean( covariance_ , axis=(0,2,3) , dtype=np.float64) )
711
+
712
+ ## write (rank-local) 1D [y] averages
713
+ if do_density_weighting_L or do_density_weighting_R:
714
+ avg_Re['rho'][yiA:yiB] = np.squeeze( np.mean( rho_avg , axis=(0,2,3) , dtype=np.float64) )
715
+
716
+ if do_density_weighting_L:
717
+ avg_Fv[scalar_L][yiA:yiB] = np.squeeze( np.mean( data_L_avg , axis=(0,2,3) , dtype=np.float64) )
718
+ else:
719
+ avg_Re[scalar_L][yiA:yiB] = np.squeeze( np.mean( data_L_avg , axis=(0,2,3) , dtype=np.float64) )
720
+
721
+ if do_density_weighting_R:
722
+ avg_Fv[scalar_R][yiA:yiB] = np.squeeze( np.mean( data_R_avg , axis=(0,2,3) , dtype=np.float64) )
723
+ else:
724
+ avg_Re[scalar_R][yiA:yiB] = np.squeeze( np.mean( data_R_avg , axis=(0,2,3) , dtype=np.float64) )
725
+
726
+ # ===============================================================================
727
+ # At this point you have 4D [x,y,z,t] □′ or □″ data
728
+ # ===============================================================================
729
+
730
+ def __fft_z_thread_kernel(xi,ti,yii,do_density_weighting_L,do_density_weighting_R):
731
+
732
+ ## 1D [z] □′ or ρ·□″ vectors
733
+ if do_density_weighting_L:
734
+ uL = rho[xi,yii,:,ti] * data_L[xi,yii,:,ti]
735
+ else:
736
+ uL = data_L[xi,yii,:,ti]
737
+
738
+ if do_density_weighting_R:
739
+ uR = rho[xi,yii,:,ti] * data_R[xi,yii,:,ti]
740
+ else:
741
+ uR = data_R[xi,yii,:,ti]
742
+
743
+ ## One-sided amplitude spectra
744
+ A1 = sp.fft.fft( uL * window_z )[kzp] / nz
745
+ A2 = sp.fft.fft( uR * window_z )[kzp] / nz
746
+
747
+ #P = 2. * np.real(A1*np.conj(A2)) / ( dkz * mean_sq_win_z )
748
+
749
+ ## One-sided complex cross-spectral density in [kz]
750
+ P = 2. * A1 * np.conj(A2) / ( dkz * mean_sq_win_z )
751
+
752
+ return xi,ti,P
753
+
754
+ def __fft_t_thread_kernel(xi,zi,yii,do_density_weighting_L,do_density_weighting_R):
755
+
756
+ ## 1D [t] □′ or ρ·□″ vectors
757
+ if do_density_weighting_L:
758
+ uL = rho[xi,yii,zi,:] * data_L[xi,yii,zi,:]
759
+ else:
760
+ uL = data_L[xi,yii,zi,:]
761
+
762
+ if do_density_weighting_R:
763
+ uR = rho[xi,yii,zi,:] * data_R[xi,yii,zi,:]
764
+ else:
765
+ uR = data_R[xi,yii,zi,:]
766
+
767
+ ## ## OLD with manual windowing
768
+ ## uL_, nw, n_pad = get_overlapping_windows(uL, win_len, overlap)
769
+ ## uR_, nw, n_pad = get_overlapping_windows(uR, win_len, overlap)
770
+ ##
771
+ ## ## STFT buffer
772
+ ## E_ijk_buf = np.zeros((nw,nf), dtype=np.float64)
773
+ ##
774
+ ## ## compute fft for each overlapped window segment
775
+ ## for wi in range(nw):
776
+ ## A1 = sp.fft.fft( uL_[wi,:] * window_t )[fp] / sum_sqrt_win_t
777
+ ## A2 = sp.fft.fft( uR_[wi,:] * window_t )[fp] / sum_sqrt_win_t
778
+ ## E_ijk_buf[wi,:] = 2. * np.real(A1*np.conj(A2)) / df
779
+ ##
780
+ ## ## mean across short time FFT (STFT) segments
781
+ ## E_ijk = np.mean(E_ijk_buf, axis=0, dtype=np.float64)
782
+
783
+ ## One-sided complex cross-spectral density in [f] (using Welch's method)
784
+ _,P = csd(
785
+ uL,uR,
786
+ fs=fs,
787
+ nperseg=nperseg,
788
+ noverlap=noverlap,
789
+ window='hann',
790
+ detrend='constant',
791
+ scaling='density',
792
+ return_onesided=True,
793
+ )
794
+
795
+ ### Real part = co-spectral density (in-phase contribution to covariance)
796
+ #P = np.real(P)
797
+ #return xi,zi,P
798
+
799
+ return xi,zi,P
800
+
801
+ # ===============================================================================
802
+
803
+ ## [y] loop inner ([y] indices within subdivision within rank)
804
+ for yi in range(cy1,cy2):
805
+
806
+ yii = yi - cy1 ## chunk local
807
+ yiii = yi - ry1 ## rank local
808
+
809
+ ## PSD buffers for [y] loop inner
810
+ E_xt = np.zeros((nx,nt,nkz) , dtype=np.complex128) ## [x,t] range for FFT(z)
811
+ E_xz = np.zeros((nx,nz,nf ) , dtype=np.complex128) ## [x,z] range for FFT(t)
812
+
813
+ # ===========================================================================
814
+ # FFT(z) : loop over [x,t]
815
+ # ===========================================================================
816
+
817
+ ## concurrent/threaded execution for fft(z)
818
+ tasks = [(xi,ti,yii,do_density_weighting_L,do_density_weighting_R) for xi in range(nx) for ti in range(nt)]
819
+ with ThreadPoolExecutor(max_workers=n_threads) as executor:
820
+ results = executor.map(lambda task: __fft_z_thread_kernel(*task,), tasks)
821
+ for xi,ti,P in results:
822
+ E_xt[xi,ti,:] = P
823
+
824
+ ## for xi in range(nx):
825
+ ## for ti in range(nt):
826
+ ## ...
827
+
828
+ ## avg in [x,t] & write in rank context
829
+ Ekz[scalar][yiii,:] = np.mean(E_xt, axis=(0,1))
830
+
831
+ # ===========================================================================
832
+ # FFT(t) : loop over [x,z], use windows
833
+ # ===========================================================================
834
+
835
+ ## concurrent/threaded execution for fft(t)
836
+ tasks = [(xi,zi,yii,do_density_weighting_L,do_density_weighting_R) for xi in range(nx) for zi in range(nz)]
837
+ with ThreadPoolExecutor(max_workers=n_threads) as executor:
838
+ results = executor.map(lambda task: __fft_t_thread_kernel(*task,), tasks)
839
+ for xi,zi,P in results:
840
+ E_xz[xi,zi,:] = P
841
+
842
+ ## for xi in range(nx):
843
+ ## for zi in range(nz):
844
+ ## ...
845
+
846
+ ## avg in [x,z] & write in rank context
847
+ Ef[scalar][yiii,:] = np.mean(E_xz, axis=(0,1))
848
+
849
+ self.comm.Barrier()
850
+ t_delta = timeit.default_timer() - t_start
851
+ if verbose: tqdm.write(even_print(msg, format_time_string(t_delta), s=True))
852
+ if verbose: progress_bar.update() ## (scalar, [y] chunk) progress
853
+
854
+ #break ## debug --> only do one round in [y] sub-chunk loop
855
+
856
+ if verbose: progress_bar.close()
857
+ self.comm.Barrier()
858
+ if verbose: print(72*'-')
859
+
860
+ # ==============================================================
861
+ # write HDF5 (.h5) file
862
+ # ==============================================================
863
+
864
+ ## overwrite outfile!
865
+ ## open on rank 0 and write attributes, dimensions, etc.
866
+ if (self.rank==0):
867
+ with h5py.File(fn_h5_out, 'w') as hfw:
868
+
869
+ ## write floats,ints as top-level attributes
870
+ for key,val in data.items():
871
+ if isinstance(data[key], (int,np.int32,np.int64)):
872
+ hfw.attrs[key] = val
873
+ elif isinstance(data[key], (float,np.float32,np.float64)):
874
+ hfw.attrs[key] = val
875
+ elif isinstance(data[key], np.ndarray):
876
+ pass
877
+ else:
878
+ print(f'key {key} is type {str(type(data[key]))}')
879
+ self.comm.Abort(1)
880
+
881
+ ## write numpy arrays
882
+ hfw.create_dataset( 'dims/x' , data=x ) ## [m]
883
+ hfw.create_dataset( 'dims/y' , data=y ) ## [m]
884
+ hfw.create_dataset( 'dims/z' , data=z ) ## [m]
885
+ hfw.create_dataset( 'dims/t' , data=t ) ## [s]
886
+ hfw.create_dataset( 'dims/freq' , data=freq ) ## [1/s] | [Hz]
887
+ hfw.create_dataset( 'dims/kz' , data=kz ) ## [1/m]
888
+ hfw.create_dataset( 'dims/lz' , data=lz ) ## [m]
889
+
890
+ ## initialize datasets
891
+ for scalar in scalars:
892
+ hfw.create_dataset( f'covariance/{scalar}' , shape=(ny,) , dtype=np.float64 , chunks=None , data=np.full((ny,),0.,np.float64) )
893
+ hfw.create_dataset( f'Ekz/{scalar}' , shape=(ny,nkz) , dtype=np.complex128 , chunks=(1,nkz) , data=np.full((ny,nkz),0.,np.complex128) )
894
+ hfw.create_dataset( f'Ef/{scalar}' , shape=(ny,nf) , dtype=np.complex128 , chunks=(1,nf) , data=np.full((ny,nf),0.,np.complex128) )
895
+
896
+ ## initialize datasets 1D [y] mean
897
+ for scalar in avg_Re.dtype.names:
898
+ hfw.create_dataset( f'avg/Re/{scalar}', shape=(ny,), dtype=np.float64, chunks=None, data=np.full((ny,),0.,np.float64) )
899
+ for scalar in avg_Fv.dtype.names:
900
+ hfw.create_dataset( f'avg/Fv/{scalar}', shape=(ny,), dtype=np.float64, chunks=None, data=np.full((ny,),0.,np.float64) )
901
+
902
+ self.comm.Barrier()
903
+
904
+ with h5py.File(fn_h5_out, 'a', driver='mpio', comm=self.comm) as hfw:
905
+
906
+ ## collectively write covariance,Ekz,Ef
907
+ for scalar in scalars:
908
+ dset = hfw[f'covariance/{scalar}']
909
+ with dset.collective:
910
+ dset[ry1:ry2] = covariance[scalar][:]
911
+ dset = hfw[f'Ekz/{scalar}']
912
+ with dset.collective:
913
+ dset[ry1:ry2,:] = Ekz[scalar][:,:]
914
+ dset = hfw[f'Ef/{scalar}']
915
+ with dset.collective:
916
+ dset[ry1:ry2,:] = Ef[scalar][:,:]
917
+
918
+ ## collectively write 1D [y] avgs (Reynolds,Favre)
919
+ for scalar in avg_Re.dtype.names:
920
+ dset = hfw[f'avg/Re/{scalar}']
921
+ with dset.collective:
922
+ dset[ry1:ry2] = avg_Re[scalar][:]
923
+ for scalar in avg_Fv.dtype.names:
924
+ dset = hfw[f'avg/Fv/{scalar}']
925
+ with dset.collective:
926
+ dset[ry1:ry2] = avg_Fv[scalar][:]
927
+
928
+ ## report file contents
929
+ self.comm.Barrier()
930
+ if (self.rank==0):
931
+ even_print( os.path.basename(fn_h5_out) , f'{(os.path.getsize(fn_h5_out)/1024**2):0.1f} [MB]' )
932
+ print(72*'-')
933
+ with h5py.File(fn_h5_out,'r') as hfr:
934
+ h5_print_contents(hfr)
935
+ self.comm.Barrier()
936
+
937
+ if verbose: print(72*'-')
938
+ if verbose: print('total time : rgd.calc_turb_cospectrum_xpln() : %s'%format_time_string((timeit.default_timer() - t_start_func)))
939
+ if verbose: print(72*'-')
940
+ return