CoCalc -- periodic2d.py

Download
Turbulence-Resolving Integral Simulations (TRIS) applies a moment-of-momentum integral approach, derived from Navier-Stokes, to run the time evolution of the integrated velocity and pressure fields. The Python code here provides an animation of a chosen field depending on the user input.
JFM-2024-2580
JFM-Notebooks / figure-4 / periodic2d.py
⁷¹ views
unlisted
ubuntu2204
1
import numpy as np
2
import matplotlib.pyplot as plt
3
from matplotlib import cm
4
import warnings
5
import subprocess
6

7
check_tol = 1e-10
8
pi = np.pi
9
cos = np.cos
10
sin = np.sin
11

12
class Msh:
13
  """Base grid class containing physical and wavenumber grids."""
14

15
  def __init__(self,Lx=2*np.pi,Lz=2*np.pi,nx=64,nz=64):
16
    # Physical
17
    self.Lx = Lx
18
    self.Lz = Lz
19
    self.nx = nx
20
    self.nz = nz
21
    self.dx = Lx / nx
22
    self.fz = Lz / nz
23
    self.x = np.linspace(0,Lx-self.dx,nx)
24
    self.z = np.linspace(0,Lz-self.fz,nz)
25
    (self.X, self.Z) = np.meshgrid(self.x, self.z, indexing='ij')
26

27
    # Wavenumber
28
    self.kx_max = np.pi / self.dx
29
    self.kz_max = np.pi / self.fz
30
    self.dkx = 2*np.pi / Lx
31
    self.dkz = 2*np.pi / Lz
32
    self.nkx = nx
33
    self.nkz = nz//2 + 1
34
    self.kx = 2 * self.kx_max * np.fft.fftfreq(nx)
35
    self.kz = np.linspace(0,self.kz_max,self.nkz)
36
    (self.KX, self.KZ) = np.meshgrid(self.kx, self.kz, indexing='ij')
37

38
  def one(self):
39
    return np.ones((self.nx, self.nz))
40

41
class Var:
42
  """Base solution variable class."""
43

44
  def __init__(self,msh,val=None,hat=None,name='var',ls='None',mrk='o',clr='k'):
45
    self.msh  = msh
46
    self.name = name
47
    self.ls   = ls
48
    self.mrk  = mrk
49
    self.clr  = clr
50

51
    if (np.shape(val) == (msh.nx, msh.nz)):
52
      self.setval(val) 
53
    elif (np.shape(hat) == (msh.nkx, msh.nkz)):
54
      self.sethat(hat)
55
    else:
56
      self.setval(np.zeros((msh.nx,msh.nz)))
57
      self.sethat(np.zeros((msh.nkx,msh.nkz),dtype=complex))
58
      self.hasval = False
59
      self.hashat = False
60

61
  def copy(self, name=None, ls=None, mrk=None, clr=None):
62
    new = Var(self.msh, self.name, self.ls, self.mrk, self.clr)
63
    if(self.hasval):
64
      new.val = np.array(self.val)
65
      new.hasval = True
66
    if(self.hashat):
67
      new.hat = np.array(self.hat,dtype=complex)
68
      new.hashat = True
69
    if(name):
70
      new.name = name
71
    if(ls):
72
      new.ls = ls
73
    if(mrk):
74
      new.mrk = mrk
75
    if(clr):
76
      new.clr = clr
77
    return new
78

79
  def setval(self, val_array):
80
    self.val = np.array(val_array)
81
    self.hasval = True
82
    self.hashat = False
83

84
  def sethat(self, hat_array):
85
    self.hat = np.array(hat_array,dtype=complex)
86
    self.hashat = True
87
    self.hasval = False
88

89
  def set_constant(self, const):
90
    self.val = np.array(const*np.ones((self.msh.nx, self.msh.nz)))
91
    self.hasval = True
92
    self.hashat = False
93
    self.fft()
94

95
  def randn(self,sigma=1.):
96
    randns = sigma*np.random.randn(self.msh.nx*self.msh.nz)
97
    self.setval(np.reshape(randns,(self.msh.nx,self.msh.nz)))
98
    self.fft()
99

100
  def fft(self):
101
    if(not self.hashat):
102
      if(self.hasval):
103
        self.hat = np.fft.rfftn(self.val)  / (self.msh.nx*self.msh.nz)
104
        self.hashat = True
105
      else:
106
        raise Exception("Called fft for variable with hasval = False")
107

108

109
  def ifft(self):
110
    if(not self.hasval):
111
      if(self.hashat):
112
        self.val = np.fft.irfftn(self.hat) * (self.msh.nx*self.msh.nz)
113
        self.hasval = True
114
      else:
115
        raise Exception("Called ifft for variable with hashat = False")
116

117

118
  def smult(self, scalar):
119
    if(self.hashat):
120
      return Var(self.msh, hat=scalar*self.hat)
121
    elif(self.hasval):
122
      return Var(self.msh, val=scalar*self.val)
123
    else:
124
      raise Exception("Called smult with neither hashat nor hasval")
125

126
  def mult(self, other):
127
    if(self.hasval and other.hasval):
128
      return Var(self.msh, val=self.val * other.val)
129
    elif(self.hashat and other.hashat):
130
      new = Var(self.msh)
131
      self.ifft()
132
      other.ifft()
133
      new.setval(self.val * other.val)
134
      new.fft()
135
      return new
136
    else:
137
      print('self.hashat = ',  self.hashat)
138
      print('self.hasval = ',  self.hasval)
139
      print('other.hashat = ', other.hashat)
140
      print('other.hasval = ', other.hasval)
141
      raise Exception('mult called with neither hashat nor hasval')
142

143
  def sqrt(self):
144
    if(self.hasval):
145
      return Var(self.msh, val=np.sqrt(np.abs(self.val)))
146
    elif(self.hashat):
147
      self.ifft()
148
      new = Var(self.msh, val=np.sqrt(np.abs(self.val)))
149
      new.fft()
150
      return new
151
    else:
152
      raise Exception('sqrt called with neither hashat nor hasval')
153

154
  def sadd(self, scalar):
155
    if(self.hashat):
156
      new = Var(self.msh, val=scalar*np.ones((self.msh.nx,self.msh.nz)))
157
      new.fft() 
158
      return self.add(new)
159
    elif (self.hasval):
160
      new = Var(self.msh, val=scalar*np.ones((self.msh.nx,self.msh.nz)))
161
      return self.add(new)
162
    else:
163
      raise Exception('sadd called with neither hashat nor hasval')
164

165
  def add(self, other):
166
    if(self.hashat and other.hashat):
167
      return Var(self.msh, hat=self.hat + other.hat)
168
    elif(self.hasval and other.hasval):
169
      return Var(self.msh, val=self.val + other.val)
170
    else:
171
      print('self.hashat = ',  self.hashat)
172
      print('self.hasval = ',  self.hasval)
173
      print('other.hashat = ', other.hashat)
174
      print('other.hasval = ', other.hasval)
175
      raise Exception('add called with neither hashat nor hasval')
176

177
  def diff_x(self):
178
    if(not self.hashat):
179
      warnings.warn("Had to call fft within diff_x")
180
      self.fft()
181
    deriv = Var(self.msh, self.name+'_dx', self.ls, self.mrk, self.clr)
182
    deriv.hat = 1j * self.msh.KX * self.hat
183
    deriv.hashat = True
184
    deriv.filter_kxmax()
185
    return deriv
186

187
  def diff_z(self):
188
    if(not self.hashat):
189
      warnings.warn("Had to call fft within diff_z")
190
      self.fft()
191
    deriv = Var(self.msh, self.name+'_fz', self.ls, self.mrk, self.clr)
192
    deriv.hat = 1j * self.msh.KZ * self.hat
193
    deriv.hashat = True
194
    deriv.filter_kzmax()
195
    return deriv
196

197
  def laplacian(self):
198
    if(not self.hashat):
199
      warnings.warn("Had to call fft within lapl")
200
      self.fft()
201
    lapl = Var(self.msh, self.name+'_lapl', self.ls, self.mrk, self.clr)
202
    lapl.hat = - (self.msh.KX**2 + self.msh.KZ**2) * self.hat 
203
    lapl.hashat = True
204
    lapl.filter_max()
205
    return lapl
206

207
  def viscous_damp(self, nu, dt):
208
    k2 = self.msh.KX**2 + self.msh.KZ**2
209
    self.hat = self.hat * np.exp(-k2**2*nu*dt)
210

211
  def check(self):
212
    if(self.hashat and self.hasval):
213
      val = np.fft.irfftn(self.hat) * (self.msh.nx*self.msh.nz)
214
      if(np.linalg.norm(self.val - val) > check_tol):
215
        raise Exception('Check zero: ' + str(np.linalg.norm(self.val - val)))
216
      hat = np.fft.rfftn(self.val)  / (self.msh.nx*self.msh.nz)
217
      if(np.linalg.norm(self.hat - hat) > check_tol):
218
        raise Exception('Check zero: ' + str(np.linalg.norm(self.hat - hat)))
219
    else:
220
      raise Exception('Check failed: hashat = ' + str(self.hashat) \
221
                                + '; hasval = ' + str(self.hasval))
222

223
  def cp_hat(self,other):
224
    if(self.msh.nkz <= other.msh.nkz and self.msh.nkx <= other.msh.nkx):
225
      # if self is smaller in x and z, spectral cutoff filter
226
      self.hat[:self.msh.nkx//2,:] \
227
              = other.hat[:self.msh.nkx//2,:self.msh.nkz]
228
      self.hat[self.msh.nkx//2:,:] \
229
              = other.hat[(other.msh.nkx-self.msh.nkx//2):,:self.msh.nkz]
230
    elif(self.msh.nkz > other.msh.nkz and self.msh.nkx > other.msh.nkx):
231
      # if self is larger in x and z, pad with zeros
232
      self.hat  = np.zeros((self.msh.nkx,self.msh.nkz),dtype=complex)
233
      self.hat[:other.msh.nkx//2,:other.msh.nkz] \
234
              = other.hat[:other.msh.nkx//2,:]
235
      self.hat[(self.msh.nkx-other.msh.nkx//2):,:other.msh.nkz]\
236
              = other.hat[other.msh.nkx//2:,:]
237
    else:
238
      raise Exception("Not implemented yet: cp_hat()")
239

240
  def filter(self,kx_cut,kz_cut):
241
    k_tol = 1e-9
242
    self.hat[np.abs(self.msh.KX) >= (kx_cut - k_tol)] = 0. + 0.j
243
    self.hat[np.abs(self.msh.KZ) >= (kz_cut - k_tol)] = 0. + 0.j
244
    self.hasval = False
245

246
  def filter_max(self):
247
    self.filter(self.msh.kx_max,self.msh.kz_max)
248

249
  def filter_kxmax(self):
250
    self.filter(self.msh.kx_max,2*self.msh.kz_max)
251

252
  def filter_kzmax(self):
253
    self.filter(2*self.msh.kx_max,self.msh.kz_max)
254

255
  def filter_third(self):
256
    self.filter(2*self.msh.kx_max/3,2*self.msh.kz_max/3)
257

258
  def filter_half(self):
259
    self.filter(self.msh.kx_max/2, self.msh.kz_max/2)
260

261
  def filter_all(self):
262
    self.filter(self.msh.kx[1],self.msh.kz[1])
263

264

265
  def mean(self):
266
    if(self.hashat):
267
      return np.real(self.hat[0,0])
268
    elif(self.hasval):
269
      return np.mean(self.val)
270

271
  def var(self):
272
    if(self.hashat):
273
      spectrum = self.spectrum()
274
      spectrum[0,0] = 0.
275
      return np.sum(spectrum)
276
    elif(self.hasval):
277
      return np.var(self.val)
278

279
  def pdf(self, nbins=100): 
280
    if(self.hashat and not self.hasval):
281
      self.ifft()
282
    (pdf, edges) = np.histogram(self.val,bins=nbins,density=True)
283
    bins = edges[:-1] + np.diff(edges)/2
284
    return (pdf, bins)
285

286
  def spectrum(self):
287
    if(not self.hashat):
288
      self.fft()
289
    return np.abs(self.hat)**2
290

291
  def spectrum_x(self):
292
    spectrum = self.spectrum()
293
    return np.sum(spectrum,axis=1)
294

295
  def spectrum_z(self):
296
    spectrum = self.spectrum()
297
    return np.sum(spectrum,axis=0)
298

299
  def report(self, name='var'):
300
    if(self.hashat and not self.hasval):
301
      self.ifft()
302
    report = (np.min(self.val), np.mean(self.val), np.max(self.val))
303
    print('min/mean/max of '+name+' is:', report)
304

305
def vorticity(u, w):
306
  if(not u.hashat):
307
    warnings.warn("Had to call fft within vorticity()")
308
    u.fft()
309
  if(not w.hashat):
310
    warnings.warn("Had to call fft within vorticity()")
311
    w.fft()
312
  vort = Var(u.msh)
313
  vort.sethat(1.j*u.msh.KX*w.hat - 1.j*u.msh.KZ*u.hat)
314
  vort.filter_max()
315
  return vort
316
 
317
def filter_Gauss(f, ell):
318
  if not f.hashat:
319
    f.fft()
320
  ell2 = ell*ell
321
  kx = f.msh.KX
322
  kz = f.msh.KZ
323
  k2 = kx*kx + kz*kz + 1e-99
324
  g = Var(f.msh)
325
  g.sethat( f.hat * np.exp(-0.5*ell2*k2) )
326
  g.filter_max()
327
  return g
328

329
def filter_Gauss_highpass(f, ell):
330
  if not f.hashat:
331
        f.fft()
332
  ell2 = ell*ell
333
  kx = f.msh.KX
334
  kz = f.msh.KZ
335
  k2 = kx*kx + kz*kz + 1e-99
336
  g = Var(f.msh)
337
  g.sethat( f.hat * (1-np.exp(-0.5*ell2*k2)) )
338
  g.filter_max()
339
  return g
340

341
def filter_highpass(f, kx_cut, kz_cut):
342
  if not f.hashat:
343
    f.fft()
344
  k_tol = 1e-9
345
  g = Var(f.msh, hat=f.hat)
346
  g.hat[np.abs(g.msh.KZ) <= (kz_cut + k_tol)] = 0. + 0.j
347
  return g
348

349
def projection(f_x, f_z):
350
  # (delta_{ij} - k_i k_j/k^2) f_j
351
  if(not f_x.hashat):
352
    f_x.fft()
353
  if(not f_z.hashat):
354
    f_z.fft()
355
  msh = f_x.msh
356
  kx = msh.KX
357
  kz = msh.KZ
358
  k2 = kx*kx + kz*kz + 1e-99 
359
  Pxx = kx*kx/k2
360
  Pxz = kx*kz/k2
361
  Pzz = kz*kz/k2
362
  Pxx[k2==0.] = 0.
363
  Pxz[k2==0.] = 0.
364
  Pzz[k2==0.] = 0.
365
  r_x = Var(msh)
366
  r_z = Var(msh)
367
  r_x.sethat( (1. - Pxx)*f_x.hat - Pxz*f_z.hat )
368
  r_z.sethat( (1. - Pzz)*f_z.hat - Pxz*f_x.hat )
369
  r_x.filter_max()
370
  r_z.filter_max()
371
  return (r_x, r_z)
372

373
def div(f_x, f_z):
374
  # df_x/dx + df_z/fz
375
  if(f_x.hashat and f_z.hashat):
376
    return div_hat(f_x, f_z)
377
  elif(f_x.hasval and f_z.hasval):
378
    return div_val(f_x, f_z)
379
  else:
380
    print('f_x.hashat = ', f_x.hashat)
381
    print('f_x.hasval = ', f_x.hasval)
382
    print('f_z.hashat = ', f_z.hashat)
383
    print('f_z.hasval = ', f_z.hasval)
384
    raise Exception('Called div with neither hashat nor hasval')
385

386
def div_val(f_x, f_z):
387
  if(f_x.hasval and f_z.hasval):
388
    f_x.fft()
389
    f_z.fft()
390
    return div_hat(f_x, f_z)
391
  else:
392
    raise Exception('Called div_val with hasval=False')
393

394
def div_hat(f_x, f_z):
395
  if(f_x.hashat and f_z.hashat):
396
    msh = f_x.msh
397
    kx = msh.KX
398
    kz = msh.KZ
399
    div = Var(msh, hat=1.j*kx*f_x.hat + 1.j*kz*f_z.hat)
400
    div.filter_max()
401
    return div
402
  else:
403
    raise Exception('Called div_hat with hashat=False')
404

405
def friction(u, w, u_avg, pwr):
406
  fx = Var(u.msh)
407
  fz = Var(w.msh)
408
  u.ifft()
409
  w.ifft()
410
  u_mag = np.sqrt(u.val**2 + w.val**2)+1e-99
411
  e_x = u.val / (u_mag)
412
  e_z = w.val / (u_mag)
413
  friction = (u_mag/u_avg)**pwr
414
  fx.setval(friction*e_x)
415
  fz.setval(friction*e_z)
416
  fx.fft()
417
  fz.fft()
418
  fx.filter_third()
419
  fz.filter_third()
420
  return (fx,fz)
421

422

423
Product

Resources

Company
