Added spectrum computation to 3D RBC

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,70 @@ def compute_Nusselt_numbers(self, u):
             't': Nusselt_t,
             'b': Nusselt_b,
         }
+
+    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)
+
+        # 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

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)