From 03791a1855fbe81d4a2e8d56655e9103e413a958 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Fri, 8 Aug 2025 13:26:56 +0200 Subject: [PATCH 1/2] Added spectrum computation to 3D RBC --- .../problem_classes/RayleighBenard3D.py | 71 ++++++++++++++++--- .../test_problems/test_RayleighBenard3D.py | 54 ++++++++++++-- 2 files changed, 111 insertions(+), 14 deletions(-) diff --git a/pySDC/implementations/problem_classes/RayleighBenard3D.py b/pySDC/implementations/problem_classes/RayleighBenard3D.py index 97ba6eea79..c3de66771b 100644 --- a/pySDC/implementations/problem_classes/RayleighBenard3D.py +++ b/pySDC/implementations/problem_classes/RayleighBenard3D.py @@ -28,9 +28,9 @@ class RayleighBenard3D(GenericSpectralLinear): The domain, vertical boundary conditions and pressure gauge are - Omega = [0, 8) x (-1, 1) + Omega = [0, Lx) x [0, Ly] x (0, Lz) T(z=+1) = 0 - T(z=-1) = 2 + T(z=-1) = Lz u(z=+-1) = v(z=+-1) = 0 integral over p = 0 @@ -53,16 +53,16 @@ class RayleighBenard3D(GenericSpectralLinear): def __init__( self, Prandtl=1, - Rayleigh=2e6, - nx=256, - ny=256, - nz=64, + Rayleigh=1e6, + nx=64, + ny=64, + nz=32, BCs=None, dealiasing=1.5, comm=None, Lz=1, - Lx=1, - Ly=1, + Lx=4, + Ly=4, useGPU=False, **kwargs, ): @@ -79,7 +79,6 @@ def __init__( comm (mpi4py.Intracomm): Space communicator Lx (float): Horizontal length of the domain """ - # TODO: documentation BCs = {} if BCs is None else BCs BCs = { 'T_top': 0, @@ -328,7 +327,7 @@ def compute_Nusselt_numbers(self, u): _me[0] = u_pad[iw] * u_pad[iT] wT_hat = self.transform(_me, padding=padding)[0] - nusselt_hat = wT_hat - DzT_hat + nusselt_hat = (wT_hat / self.kappa - DzT_hat) * self.axes[-1].L if not hasattr(self, '_zInt'): self._zInt = zAxis.get_integration_matrix() @@ -354,3 +353,55 @@ def compute_Nusselt_numbers(self, u): 't': Nusselt_t, 'b': Nusselt_b, } + + def get_frequency_spectrum(self, u): + xp = self.xp + indices = slice(0, 2) + + # transform the solution to be in frequency space in x and y, but real space in z + if self.spectral_space: + u_hat = self.itransform(u, axes=(-1,)) + else: + u_hat = self.transform( + u, + axes=( + -3, + -2, + ), + ) + u_hat = self.spectral.redistribute(u_hat, axis=2, forward_output=False) + + # compute "energy density" as absolute square of the velocity modes + energy = (u_hat[indices] * xp.conjugate(u_hat[indices])).real / (self.axes[0].N ** 2 * self.axes[1].N ** 2) + + # prepare wave numbers at which to compute the spectrum + abs_kx = xp.abs(self.Kx[:, :, 0]) + abs_ky = xp.abs(self.Ky[:, :, 0]) + + unique_k = xp.unique(xp.append(xp.unique(abs_kx), xp.unique(abs_ky))) + n_k = len(unique_k) + + # compute local spectrum + local_spectrum = self.xp.empty(shape=(2, energy.shape[3], n_k)) + for i, k in zip(range(n_k), unique_k): + mask = xp.logical_or(abs_kx == k, abs_ky == k) + local_spectrum[..., i] = xp.sum(energy[indices, mask, :], axis=1) + + # assemble global spectrum from local spectra + k_all = self.comm.allgather(unique_k) + unique_k_all = [] + for k in k_all: + unique_k_all = xp.unique(xp.append(unique_k_all, xp.unique(k))) + n_k_all = len(unique_k_all) + + spectra = self.comm.allgather(local_spectrum) + spectrum = self.xp.zeros(shape=(2, self.axes[2].N, n_k_all)) + for ks, _spectrum in zip(k_all, spectra): + ks = list(ks) + unique_k_all = list(unique_k_all) + for k in ks: + index_global = unique_k_all.index(k) + index_local = ks.index(k) + spectrum[..., index_global] += _spectrum[..., index_local] + + return xp.array(unique_k_all), spectrum diff --git a/pySDC/tests/test_problems/test_RayleighBenard3D.py b/pySDC/tests/test_problems/test_RayleighBenard3D.py index 5cd3d7a78b..16ab143e7b 100644 --- a/pySDC/tests/test_problems/test_RayleighBenard3D.py +++ b/pySDC/tests/test_problems/test_RayleighBenard3D.py @@ -11,7 +11,7 @@ def test_eval_f(nx, nz, direction, spectral_space): import numpy as np from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D - P = RayleighBenard3D(nx=nx, ny=nx, nz=nz, Rayleigh=1, spectral_space=spectral_space) + P = RayleighBenard3D(nx=nx, ny=nx, nz=nz, Rayleigh=1, spectral_space=spectral_space, Lx=1, Ly=1, Lz=1) iu, iv, iw, ip, iT = P.index(['u', 'v', 'w', 'p', 'T']) X, Y, Z = P.X, P.Y, P.Z cos, sin = np.cos, np.sin @@ -297,7 +297,7 @@ def test_heterogeneous_implementation(): def test_Nusselt_number_computation(w, N=4): from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D - prob = RayleighBenard3D(nx=N, ny=N, nz=N, dealiasing=1.0, spectral_space=False) + prob = RayleighBenard3D(nx=N, ny=N, nz=N, dealiasing=1.0, spectral_space=False, Rayleigh=1.0, Prandtl=1.0) xp = prob.xp iw, iT = prob.index(['w', 'T']) @@ -341,12 +341,58 @@ def test_Nusselt_number_computation(w, N=4): assert xp.isclose(Nu[key], w), f'Expected Nu_{key}={w}, but got {Nu[key]} with constant T and perturbed w!' +@pytest.mark.mpi4py +@pytest.mark.mpi(ranks=[1, 2, 5]) +def test_spectrum_computation(mpi_ranks): + from pySDC.implementations.problem_classes.RayleighBenard3D import RayleighBenard3D + + N = 5 + prob = RayleighBenard3D(nx=N, ny=N, nz=2, dealiasing=1.0, spectral_space=False, Rayleigh=1.0) + xp = prob.xp + iu, iv = prob.index(['u', 'v']) + + num_k = int(N // 2) + 1 + + u = prob.u_exact() * 0 + u[iu] = 1 + ks, spectrum = prob.get_frequency_spectrum(u) + + expect_spectrum = xp.zeros((prob.nz, num_k)) + expect_spectrum[:, 0] = 1 + + assert len(ks) == num_k + assert xp.allclose(spectrum[iv], 0) + assert xp.allclose(spectrum[iu], expect_spectrum) + + assert N >= 5 + u = prob.u_exact() * 0 + u[iu] = xp.sin(prob.Y * 2 * xp.pi / prob.axes[1].L) + ks, spectrum = prob.get_frequency_spectrum(u) + assert xp.allclose(spectrum[iu, :, 2:], 0) + assert not xp.allclose(spectrum[iu, :, :2], 0) + + assert N >= 5 + u = prob.u_exact() * 0 + u[iu] = xp.sin(prob.X * 2 * xp.pi / prob.axes[0].L) * xp.sin(prob.Y * 2 * xp.pi / prob.axes[1].L) + ks, spectrum = prob.get_frequency_spectrum(u) + assert xp.allclose(spectrum[iu, :, 2:], 0) + assert xp.allclose(spectrum[iu, :, 0], 0) + assert not xp.allclose(spectrum[iu, :, 1], 0) + + assert N >= 5 + u = prob.u_exact() * 0 + u[iu] = xp.sin(prob.X * 4 * xp.pi / prob.axes[0].L) * xp.sin(prob.Y * 2 * xp.pi / prob.axes[1].L) + ks, spectrum = prob.get_frequency_spectrum(u) + assert not xp.allclose(spectrum[iu, :, 1:], 0) + assert xp.allclose(spectrum[iu, :, 0], 0) + + if __name__ == '__main__': - # test_eval_f(2**2, 2**1, 'x', False) + test_eval_f(2**2, 2**1, 'x', False) # test_libraries() # test_Poisson_problems(4, 'u') # test_Poisson_problem_w() # test_solver_convergence('bicgstab+ilu', 32, False, True) # test_banded_matrix(False) # test_heterogeneous_implementation() - test_Nusselt_number_computation(N=4, w=3.14) + # test_Nusselt_number_computation(N=4, w=3.14) From 242cbd98849ebeb450eec514086724f75b922c65 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Thu, 14 Aug 2025 16:07:00 +0200 Subject: [PATCH 2/2] Added documentation --- .../problem_classes/RayleighBenard3D.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/pySDC/implementations/problem_classes/RayleighBenard3D.py b/pySDC/implementations/problem_classes/RayleighBenard3D.py index c3de66771b..69bd051e30 100644 --- a/pySDC/implementations/problem_classes/RayleighBenard3D.py +++ b/pySDC/implementations/problem_classes/RayleighBenard3D.py @@ -355,6 +355,21 @@ def compute_Nusselt_numbers(self, u): } def get_frequency_spectrum(self, u): + """ + Compute the frequency spectrum of the velocities in x and y direction in the horizontal plane for every point in + z. If the problem is well resolved, the coefficients will decay quickly with the wave number, and the reverse + indicates that the resolution is too low. + + The returned spectrum has three dimensions. The first is for component (i.e. u or v), the second is for every + point in z and the third is the energy in every wave number. + + Args: + u: The solution you want to compute the spectrum of + + Returns: + RayleighBenard3D.xp.ndarray: wave numbers + RayleighBenard3D.xp.ndarray: spectrum + """ xp = self.xp indices = slice(0, 2)