Coverage for pySDC/implementations/problem_classes/RayleighBenard3D.py: 85%
126 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-11 11:36 +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, 8) x (-1, 1)
32 T(z=+1) = 0
33 T(z=-1) = 2
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=2e6,
57 nx=256,
58 ny=256,
59 nz=64,
60 BCs=None,
61 dealiasing=1.5,
62 comm=None,
63 Lz=1,
64 Lx=1,
65 Ly=1,
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 # TODO: documentation
83 BCs = {} if BCs is None else BCs
84 BCs = {
85 'T_top': 0,
86 'T_bottom': Lz,
87 'w_top': 0,
88 'w_bottom': 0,
89 'v_top': 0,
90 'v_bottom': 0,
91 'u_top': 0,
92 'u_bottom': 0,
93 'p_integral': 0,
94 **BCs,
95 }
96 if comm is None:
97 try:
98 from mpi4py import MPI
100 comm = MPI.COMM_WORLD
101 except ModuleNotFoundError:
102 pass
103 self._makeAttributeAndRegister(
104 'Prandtl',
105 'Rayleigh',
106 'nx',
107 'ny',
108 'nz',
109 'BCs',
110 'dealiasing',
111 'comm',
112 'Lx',
113 'Ly',
114 'Lz',
115 localVars=locals(),
116 readOnly=True,
117 )
119 bases = [
120 {'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx, 'useFFTW': not useGPU},
121 {'base': 'fft', 'N': ny, 'x0': 0, 'x1': self.Ly, 'useFFTW': not useGPU},
122 {'base': 'ultraspherical', 'N': nz, 'x0': 0, 'x1': self.Lz},
123 ]
124 components = ['u', 'v', 'w', 'T', 'p']
125 super().__init__(bases, components, comm=comm, useGPU=useGPU, **kwargs)
127 self.X, self.Y, self.Z = self.get_grid()
128 self.Kx, self.Ky, self.Kz = self.get_wavenumbers()
130 # construct 3D matrices
131 Dzz = self.get_differentiation_matrix(axes=(2,), p=2)
132 Dz = self.get_differentiation_matrix(axes=(2,))
133 Dy = self.get_differentiation_matrix(axes=(1,))
134 Dyy = self.get_differentiation_matrix(axes=(1,), p=2)
135 Dx = self.get_differentiation_matrix(axes=(0,))
136 Dxx = self.get_differentiation_matrix(axes=(0,), p=2)
137 Id = self.get_Id()
139 S1 = self.get_basis_change_matrix(p_out=0, p_in=1)
140 S2 = self.get_basis_change_matrix(p_out=0, p_in=2)
142 U01 = self.get_basis_change_matrix(p_in=0, p_out=1)
143 U12 = self.get_basis_change_matrix(p_in=1, p_out=2)
144 U02 = self.get_basis_change_matrix(p_in=0, p_out=2)
146 self.Dx = Dx
147 self.Dxx = Dxx
148 self.Dy = Dy
149 self.Dyy = Dyy
150 self.Dz = S1 @ Dz
151 self.Dzz = S2 @ Dzz
152 self.S2 = S2
153 self.S1 = S1
155 # compute rescaled Rayleigh number to extract viscosity and thermal diffusivity
156 Ra = Rayleigh / (max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * self.axes[2].L ** 3)
157 self.kappa = (Ra * Prandtl) ** (-1 / 2.0)
158 self.nu = (Ra / Prandtl) ** (-1 / 2.0)
160 # construct operators
161 _D = U02 @ (Dxx + Dyy) + Dzz
162 L_lhs = {
163 'p': {'u': U01 @ Dx, 'v': U01 @ Dy, 'w': Dz}, # divergence free constraint
164 'u': {'p': U02 @ Dx, 'u': -self.nu * _D},
165 'v': {'p': U02 @ Dy, 'v': -self.nu * _D},
166 'w': {'p': U12 @ Dz, 'w': -self.nu * _D, 'T': -U02 @ Id},
167 'T': {'T': -self.kappa * _D},
168 }
169 self.setup_L(L_lhs)
171 # mass matrix
172 _U02 = U02 @ Id
173 M_lhs = {i: {i: _U02} for i in ['u', 'v', 'w', 'T']}
174 self.setup_M(M_lhs)
176 # BCs
177 self.add_BC(
178 component='p', equation='p', axis=2, v=self.BCs['p_integral'], kind='integral', line=-1, scalar=True
179 )
180 self.add_BC(component='T', equation='T', axis=2, x=-1, v=self.BCs['T_bottom'], kind='Dirichlet', line=-1)
181 self.add_BC(component='T', equation='T', axis=2, x=1, v=self.BCs['T_top'], kind='Dirichlet', line=-2)
182 self.add_BC(component='w', equation='w', axis=2, x=1, v=self.BCs['w_top'], kind='Dirichlet', line=-1)
183 self.add_BC(component='w', equation='w', axis=2, x=-1, v=self.BCs['w_bottom'], kind='Dirichlet', line=-2)
184 self.remove_BC(component='w', equation='w', axis=2, x=-1, kind='Dirichlet', line=-2, scalar=True)
185 for comp in ['u', 'v']:
186 self.add_BC(
187 component=comp, equation=comp, axis=2, v=self.BCs[f'{comp}_top'], x=1, kind='Dirichlet', line=-2
188 )
189 self.add_BC(
190 component=comp,
191 equation=comp,
192 axis=2,
193 v=self.BCs[f'{comp}_bottom'],
194 x=-1,
195 kind='Dirichlet',
196 line=-1,
197 )
199 # eliminate Nyquist mode if needed
200 if nx % 2 == 0:
201 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index()
202 for component in self.components:
203 self.add_BC(
204 component=component, equation=component, axis=0, kind='Nyquist', line=int(Nyquist_mode_index), v=0
205 )
206 if ny % 2 == 0:
207 Nyquist_mode_index = self.axes[0].get_Nyquist_mode_index()
208 for component in self.components:
209 self.add_BC(
210 component=component, equation=component, axis=1, kind='Nyquist', line=int(Nyquist_mode_index), v=0
211 )
212 self.setup_BCs()
214 self.work_counters['rhs'] = WorkCounter()
216 def eval_f(self, u, *args, **kwargs):
217 f = self.f_init
219 if self.spectral_space:
220 u_hat = u.copy()
221 else:
222 u_hat = self.transform(u)
224 f_impl_hat = self.u_init_forward
226 iu, iv, iw, iT, ip = self.index(['u', 'v', 'w', 'T', 'p'])
227 derivative_indices = [iu, iv, iw, iT]
229 # evaluate implicit terms
230 f_impl_hat = -(self.L @ u_hat.flatten()).reshape(u_hat.shape)
231 for i in derivative_indices:
232 f_impl_hat[i] = (self.S2 @ f_impl_hat[i].flatten()).reshape(f_impl_hat[i].shape)
233 f_impl_hat[ip] = (self.S1 @ f_impl_hat[ip].flatten()).reshape(f_impl_hat[ip].shape)
235 if self.spectral_space:
236 self.xp.copyto(f.impl, f_impl_hat)
237 else:
238 f.impl[:] = self.itransform(f_impl_hat).real
240 # -------------------------------------------
241 # treat convection explicitly with dealiasing
243 # start by computing derivatives
244 padding = (self.dealiasing,) * self.ndim
245 derivatives = []
246 u_hat_flat = [u_hat[i].flatten() for i in derivative_indices]
248 _D_u_hat = self.u_init_forward
249 for D in [self.Dx, self.Dy, self.Dz]:
250 _D_u_hat *= 0
251 for i in derivative_indices:
252 self.xp.copyto(_D_u_hat[i], (D @ u_hat_flat[i]).reshape(_D_u_hat[i].shape))
253 derivatives.append(self.itransform(_D_u_hat, padding=padding).real)
255 u_pad = self.itransform(u_hat, padding=padding).real
257 fexpl_pad = self.xp.zeros_like(u_pad)
258 for i in derivative_indices:
259 for i_vel, iD in zip([iu, iv, iw], range(self.ndim)):
260 fexpl_pad[i] -= u_pad[i_vel] * derivatives[iD][i]
262 if self.spectral_space:
263 self.xp.copyto(f.expl, self.transform(fexpl_pad, padding=padding))
264 else:
265 f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real
267 self.work_counters['rhs']()
268 return f
270 def u_exact(self, t=0, noise_level=1e-3, seed=99):
271 assert t == 0
272 assert (
273 self.BCs['v_top'] == self.BCs['v_bottom']
274 ), 'Initial conditions are only implemented for zero velocity gradient'
276 me = self.spectral.u_init
277 iu, iw, iT, ip = self.index(['u', 'w', 'T', 'p'])
279 # linear temperature gradient
280 assert self.Lz == 1
281 for comp in ['T', 'u', 'v', 'w']:
282 a = self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']
283 b = self.BCs[f'{comp}_bottom']
284 me[self.index(comp)] = a * self.Z + b
286 # perturb slightly
287 rng = self.xp.random.default_rng(seed=seed)
289 noise = self.spectral.u_init
290 noise[iT] = rng.random(size=me[iT].shape)
292 me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1)
294 if self.spectral_space:
295 me_hat = self.spectral.u_init_forward
296 me_hat[:] = self.transform(me)
297 return me_hat
298 else:
299 return me
301 def get_fig(self): # pragma: no cover
302 """
303 Get a figure suitable to plot the solution of this problem
305 Returns
306 -------
307 self.fig : matplotlib.pyplot.figure.Figure
308 """
309 import matplotlib.pyplot as plt
310 from mpl_toolkits.axes_grid1 import make_axes_locatable
312 plt.rcParams['figure.constrained_layout.use'] = True
313 self.fig, axs = plt.subplots(2, 1, sharex=True, sharey=True, figsize=((10, 5)))
314 self.cax = []
315 divider = make_axes_locatable(axs[0])
316 self.cax += [divider.append_axes('right', size='3%', pad=0.03)]
317 divider2 = make_axes_locatable(axs[1])
318 self.cax += [divider2.append_axes('right', size='3%', pad=0.03)]
319 return self.fig
321 def plot(self, u, t=None, fig=None, quantity='T'): # pragma: no cover
322 r"""
323 Plot the solution.
325 Parameters
326 ----------
327 u : dtype_u
328 Solution to be plotted
329 t : float
330 Time to display at the top of the figure
331 fig : matplotlib.pyplot.figure.Figure
332 Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated.
333 quantity : (str)
334 quantity you want to plot
336 Returns
337 -------
338 None
339 """
340 fig = self.get_fig() if fig is None else fig
341 axs = fig.axes
343 imV = axs[1].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real)
345 if self.spectral_space:
346 u = self.itransform(u)
348 imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real)
350 for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity']):
351 axs[i].set_aspect(1)
352 axs[i].set_title(label)
354 if t is not None:
355 fig.suptitle(f't = {t:.2f}')
356 axs[1].set_xlabel(r'$x$')
357 axs[1].set_ylabel(r'$z$')
358 fig.colorbar(imT, self.cax[0])
359 fig.colorbar(imV, self.cax[1])