Coverage for pySDC/implementations/problem_classes/RayleighBenard3D.py: 82%
183 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-04 15:08 +0000
1import numpy as np
2from mpi4py import MPI
4from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear
5from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
6from pySDC.core.convergence_controller import ConvergenceController
7from pySDC.core.hooks import Hooks
8from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence
9from pySDC.core.problem import WorkCounter
12class RayleighBenard3D(GenericSpectralLinear):
13 """
14 Rayleigh-Benard Convection is a variation of incompressible Navier-Stokes.
16 The equations we solve are
18 u_x + v_y + w_z = 0
19 T_t - kappa (T_xx + T_yy + T_zz) = -uT_x - vT_y - wT_z
20 u_t - nu (u_xx + u_yy + u_zz) + p_x = -uu_x - vu_y - wu_z
21 v_t - nu (v_xx + v_yy + v_zz) + p_y = -uv_x - vv_y - wv_z
22 w_t - nu (w_xx + w_yy + w_zz) + p_z - T = -uw_x - vw_y - ww_z
24 with u the horizontal velocity, v the vertical velocity (in z-direction), T the temperature, p the pressure, indices
25 denoting derivatives, kappa=(Rayleigh * Prandtl)**(-1/2) and nu = (Rayleigh / Prandtl)**(-1/2). Everything on the left
26 hand side, that is the viscous part, the pressure gradient and the buoyancy due to temperature are treated
27 implicitly, while the non-linear convection part on the right hand side is integrated explicitly.
29 The domain, vertical boundary conditions and pressure gauge are
31 Omega = [0, Lx) x [0, Ly] x (0, Lz)
32 T(z=+1) = 0
33 T(z=-1) = Lz
34 u(z=+-1) = v(z=+-1) = 0
35 integral over p = 0
37 The spectral discretization uses FFT horizontally, implying periodic BCs, and an ultraspherical method vertically to
38 facilitate the Dirichlet BCs.
40 Parameters:
41 Prandtl (float): Prandtl number
42 Rayleigh (float): Rayleigh number
43 nx (int): Horizontal resolution
44 nz (int): Vertical resolution
45 BCs (dict): Can specify boundary conditions here
46 dealiasing (float): Dealiasing factor for evaluating the non-linear part
47 comm (mpi4py.Intracomm): Space communicator
48 """
50 dtype_u = mesh
51 dtype_f = imex_mesh
53 def __init__(
54 self,
55 Prandtl=1,
56 Rayleigh=1e6,
57 nx=64,
58 ny=64,
59 nz=32,
60 BCs=None,
61 dealiasing=1.5,
62 comm=None,
63 Lz=1,
64 Lx=4,
65 Ly=4,
66 useGPU=False,
67 **kwargs,
68 ):
69 """
70 Constructor. `kwargs` are forwarded to parent class constructor.
72 Args:
73 Prandtl (float): Prandtl number
74 Rayleigh (float): Rayleigh number
75 nx (int): Resolution in x-direction
76 nz (int): Resolution in z direction
77 BCs (dict): Vertical boundary conditions
78 dealiasing (float): Dealiasing for evaluating the non-linear part in real space
79 comm (mpi4py.Intracomm): Space communicator
80 Lx (float): Horizontal length of the domain
81 """
82 BCs = {} if BCs is None else BCs
83 BCs = {
84 'T_top': 0,
85 'T_bottom': Lz,
86 'w_top': 0,
87 'w_bottom': 0,
88 'v_top': 0,
89 'v_bottom': 0,
90 'u_top': 0,
91 'u_bottom': 0,
92 'p_integral': 0,
93 **BCs,
94 }
95 if comm is None:
96 try:
97 from mpi4py import MPI
99 comm = MPI.COMM_WORLD
100 except ModuleNotFoundError:
101 pass
102 self._makeAttributeAndRegister(
103 'Prandtl',
104 'Rayleigh',
105 'nx',
106 'ny',
107 'nz',
108 'BCs',
109 'dealiasing',
110 'comm',
111 'Lx',
112 'Ly',
113 'Lz',
114 localVars=locals(),
115 readOnly=True,
116 )
118 bases = [
119 {'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx, 'useFFTW': not useGPU},
120 {'base': 'fft', 'N': ny, 'x0': 0, 'x1': self.Ly, 'useFFTW': not useGPU},
121 {'base': 'ultraspherical', 'N': nz, 'x0': 0, 'x1': self.Lz},
122 ]
123 components = ['u', 'v', 'w', 'T', 'p']
124 super().__init__(bases, components, comm=comm, useGPU=useGPU, **kwargs)
126 self.X, self.Y, self.Z = self.get_grid()
127 self.Kx, self.Ky, self.Kz = self.get_wavenumbers()
129 # construct 3D matrices
130 Dzz = self.get_differentiation_matrix(axes=(2,), p=2)
131 Dz = self.get_differentiation_matrix(axes=(2,))
132 Dy = self.get_differentiation_matrix(axes=(1,))
133 Dyy = self.get_differentiation_matrix(axes=(1,), p=2)
134 Dx = self.get_differentiation_matrix(axes=(0,))
135 Dxx = self.get_differentiation_matrix(axes=(0,), p=2)
136 Id = self.get_Id()
138 S1 = self.get_basis_change_matrix(p_out=0, p_in=1)
139 S2 = self.get_basis_change_matrix(p_out=0, p_in=2)
141 U01 = self.get_basis_change_matrix(p_in=0, p_out=1)
142 U12 = self.get_basis_change_matrix(p_in=1, p_out=2)
143 U02 = self.get_basis_change_matrix(p_in=0, p_out=2)
145 self.Dx = Dx
146 self.Dxx = Dxx
147 self.Dy = Dy
148 self.Dyy = Dyy
149 self.Dz = S1 @ Dz
150 self.Dzz = S2 @ Dzz
151 self.S2 = S2
152 self.S1 = S1
154 # compute rescaled Rayleigh number to extract viscosity and thermal diffusivity
155 Ra = Rayleigh / (max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * self.axes[2].L ** 3)
156 self.kappa = (Ra * Prandtl) ** (-1 / 2.0)
157 self.nu = (Ra / Prandtl) ** (-1 / 2.0)
159 # construct operators
160 _D = U02 @ (Dxx + Dyy) + Dzz
161 L_lhs = {
162 'p': {'u': U01 @ Dx, 'v': U01 @ Dy, 'w': Dz}, # divergence free constraint
163 'u': {'p': U02 @ Dx, 'u': -self.nu * _D},
164 'v': {'p': U02 @ Dy, 'v': -self.nu * _D},
165 'w': {'p': U12 @ Dz, 'w': -self.nu * _D, 'T': -U02 @ Id},
166 'T': {'T': -self.kappa * _D},
167 }
168 self.setup_L(L_lhs)
170 # mass matrix
171 _U02 = U02 @ Id
172 M_lhs = {i: {i: _U02} for i in ['u', 'v', 'w', 'T']}
173 self.setup_M(M_lhs)
175 # BCs
176 self.add_BC(
177 component='p', equation='p', axis=2, v=self.BCs['p_integral'], kind='integral', line=-1, scalar=True
178 )
179 self.add_BC(component='T', equation='T', axis=2, x=-1, v=self.BCs['T_bottom'], kind='Dirichlet', line=-1)
180 self.add_BC(component='T', equation='T', axis=2, x=1, v=self.BCs['T_top'], kind='Dirichlet', line=-2)
181 self.add_BC(component='w', equation='w', axis=2, x=1, v=self.BCs['w_top'], kind='Dirichlet', line=-1)
182 self.add_BC(component='w', equation='w', axis=2, x=-1, v=self.BCs['w_bottom'], kind='Dirichlet', line=-2)
183 self.remove_BC(component='w', equation='w', axis=2, x=-1, kind='Dirichlet', line=-2, scalar=True)
184 for comp in ['u', 'v']:
185 self.add_BC(
186 component=comp, equation=comp, axis=2, v=self.BCs[f'{comp}_top'], x=1, kind='Dirichlet', line=-2
187 )
188 self.add_BC(
189 component=comp,
190 equation=comp,
191 axis=2,
192 v=self.BCs[f'{comp}_bottom'],
193 x=-1,
194 kind='Dirichlet',
195 line=-1,
196 )
198 # eliminate Nyquist mode if needed
199 if nx % 2 == 0:
200 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index()
201 for component in self.components:
202 self.add_BC(
203 component=component, equation=component, axis=0, kind='Nyquist', line=int(Nyquist_mode_index), v=0
204 )
205 if ny % 2 == 0:
206 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index()
207 for component in self.components:
208 self.add_BC(
209 component=component, equation=component, axis=1, kind='Nyquist', line=int(Nyquist_mode_index), v=0
210 )
211 self.setup_BCs()
213 self.work_counters['rhs'] = WorkCounter()
215 def eval_f(self, u, *args, **kwargs):
216 f = self.f_init
218 if self.spectral_space:
219 u_hat = u.copy()
220 else:
221 u_hat = self.transform(u)
223 f_impl_hat = self.u_init_forward
225 iu, iv, iw, iT, ip = self.index(['u', 'v', 'w', 'T', 'p'])
226 derivative_indices = [iu, iv, iw, iT]
228 # evaluate implicit terms
229 f_impl_hat = -(self.L @ u_hat.flatten()).reshape(u_hat.shape)
230 for i in derivative_indices:
231 f_impl_hat[i] = (self.S2 @ f_impl_hat[i].flatten()).reshape(f_impl_hat[i].shape)
232 f_impl_hat[ip] = (self.S1 @ f_impl_hat[ip].flatten()).reshape(f_impl_hat[ip].shape)
234 if self.spectral_space:
235 self.xp.copyto(f.impl, f_impl_hat)
236 else:
237 f.impl[:] = self.itransform(f_impl_hat).real
239 # -------------------------------------------
240 # treat convection explicitly with dealiasing
242 # start by computing derivatives
243 padding = (self.dealiasing,) * self.ndim
244 derivatives = []
245 u_hat_flat = [u_hat[i].flatten() for i in derivative_indices]
247 _D_u_hat = self.u_init_forward
248 for D in [self.Dx, self.Dy, self.Dz]:
249 _D_u_hat *= 0
250 for i in derivative_indices:
251 self.xp.copyto(_D_u_hat[i], (D @ u_hat_flat[i]).reshape(_D_u_hat[i].shape))
252 derivatives.append(self.itransform(_D_u_hat, padding=padding).real)
254 u_pad = self.itransform(u_hat, padding=padding).real
256 fexpl_pad = self.xp.zeros_like(u_pad)
257 for i in derivative_indices:
258 for i_vel, iD in zip([iu, iv, iw], range(self.ndim)):
259 fexpl_pad[i] -= u_pad[i_vel] * derivatives[iD][i]
261 if self.spectral_space:
262 self.xp.copyto(f.expl, self.transform(fexpl_pad, padding=padding))
263 else:
264 f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real
266 self.work_counters['rhs']()
267 return f
269 def u_exact(self, t=0, noise_level=1e-3, seed=99):
270 assert t == 0
271 assert (
272 self.BCs['v_top'] == self.BCs['v_bottom']
273 ), 'Initial conditions are only implemented for zero velocity gradient'
275 me = self.spectral.u_init
276 iu, iw, iT, ip = self.index(['u', 'w', 'T', 'p'])
278 # linear temperature gradient
279 assert self.Lz == 1
280 for comp in ['T', 'u', 'v', 'w']:
281 a = self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']
282 b = self.BCs[f'{comp}_bottom']
283 me[self.index(comp)] = a * self.Z + b
285 # perturb slightly
286 rng = self.xp.random.default_rng(seed=seed)
288 noise = self.spectral.u_init
289 noise[iT] = rng.random(size=me[iT].shape)
291 me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1)
293 if self.spectral_space:
294 me_hat = self.spectral.u_init_forward
295 me_hat[:] = self.transform(me)
296 return me_hat
297 else:
298 return me
300 def compute_Nusselt_numbers(self, u):
301 """
302 Compute the various versions of the Nusselt number. This reflects the type of heat transport.
303 If the Nusselt number is equal to one, it indicates heat transport due to conduction. If it is larger,
304 advection is present.
305 Computing the Nusselt number at various places can be used to check the code.
307 Args:
308 u: The solution you want to compute the Nusselt numbers of
310 Returns:
311 dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom.
312 """
313 iw, iT = self.index(['w', 'T'])
314 zAxis = self.spectral.axes[-1]
316 if self.spectral_space:
317 u_hat = u.copy()
318 else:
319 u_hat = self.transform(u)
321 DzT_hat = (self.Dz @ u_hat[iT].flatten()).reshape(u_hat[iT].shape)
323 # compute wT with dealiasing
324 padding = (self.dealiasing,) * self.ndim
325 u_pad = self.itransform(u_hat, padding=padding).real
326 _me = self.xp.zeros_like(u_pad)
327 _me[0] = u_pad[iw] * u_pad[iT]
328 wT_hat = self.transform(_me, padding=padding)[0]
330 nusselt_hat = (wT_hat / self.kappa - DzT_hat) * self.axes[-1].L
332 if not hasattr(self, '_zInt'):
333 self._zInt = zAxis.get_integration_matrix()
335 # get coefficients for evaluation on the boundary
336 top = zAxis.get_BC(kind='Dirichlet', x=1)
337 bot = zAxis.get_BC(kind='Dirichlet', x=-1)
339 integral_V = 0
340 if self.comm.rank == 0:
342 integral_z = (self._zInt @ nusselt_hat[0, 0]).real
343 integral_z[0] = zAxis.get_integration_constant(integral_z, axis=-1)
344 integral_V = ((top - bot) * integral_z).sum() * self.axes[0].L * self.axes[1].L / self.nx / self.ny
346 Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0)
348 Nusselt_t = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0, :] * top, axis=-1) / self.nx / self.ny, root=0)
349 Nusselt_b = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0] * bot / self.nx / self.ny, axis=-1), root=0)
351 return {
352 'V': Nusselt_V,
353 't': Nusselt_t,
354 'b': Nusselt_b,
355 }
357 def get_frequency_spectrum(self, u):
358 """
359 Compute the frequency spectrum of the velocities in x and y direction in the horizontal plane for every point in
360 z. If the problem is well resolved, the coefficients will decay quickly with the wave number, and the reverse
361 indicates that the resolution is too low.
363 The returned spectrum has three dimensions. The first is for component (i.e. u or v), the second is for every
364 point in z and the third is the energy in every wave number.
366 Args:
367 u: The solution you want to compute the spectrum of
369 Returns:
370 RayleighBenard3D.xp.ndarray: wave numbers
371 RayleighBenard3D.xp.ndarray: spectrum
372 """
373 xp = self.xp
374 indices = slice(0, 2)
376 # transform the solution to be in frequency space in x and y, but real space in z
377 if self.spectral_space:
378 u_hat = self.itransform(u, axes=(-1,))
379 else:
380 u_hat = self.transform(
381 u,
382 axes=(
383 -3,
384 -2,
385 ),
386 )
387 u_hat = self.spectral.redistribute(u_hat, axis=2, forward_output=False)
389 # compute "energy density" as absolute square of the velocity modes
390 energy = (u_hat[indices] * xp.conjugate(u_hat[indices])).real / (self.axes[0].N ** 2 * self.axes[1].N ** 2)
392 # prepare wave numbers at which to compute the spectrum
393 abs_kx = xp.abs(self.Kx[:, :, 0])
394 abs_ky = xp.abs(self.Ky[:, :, 0])
396 unique_k = xp.unique(xp.append(xp.unique(abs_kx), xp.unique(abs_ky)))
397 n_k = len(unique_k)
399 # compute local spectrum
400 local_spectrum = self.xp.empty(shape=(2, energy.shape[3], n_k))
401 for i, k in zip(range(n_k), unique_k):
402 mask = xp.logical_or(abs_kx == k, abs_ky == k)
403 local_spectrum[..., i] = xp.sum(energy[indices, mask, :], axis=1)
405 # assemble global spectrum from local spectra
406 k_all = self.comm.allgather(unique_k)
407 unique_k_all = []
408 for k in k_all:
409 unique_k_all = xp.unique(xp.append(unique_k_all, xp.unique(k)))
410 n_k_all = len(unique_k_all)
412 spectra = self.comm.allgather(local_spectrum)
413 spectrum = self.xp.zeros(shape=(2, self.axes[2].N, n_k_all))
414 for ks, _spectrum in zip(k_all, spectra):
415 ks = list(ks)
416 unique_k_all = list(unique_k_all)
417 for k in ks:
418 index_global = unique_k_all.index(k)
419 index_local = ks.index(k)
420 spectrum[..., index_global] += _spectrum[..., index_local]
422 return xp.array(unique_k_all), spectrum