Added Nusselt number computation to 3D RBC

pySDC/implementations/problem_classes/RayleighBenard3D.py

@@ -298,62 +298,59 @@ def u_exact(self, t=0, noise_level=1e-3, seed=99):
         else:
             return me
 
-    def get_fig(self):  # pragma: no cover
+    def compute_Nusselt_numbers(self, u):
         """
-        Get a figure suitable to plot the solution of this problem
+        Compute the various versions of the Nusselt number. This reflects the type of heat transport.
+        If the Nusselt number is equal to one, it indicates heat transport due to conduction. If it is larger,
+        advection is present.
+        Computing the Nusselt number at various places can be used to check the code.
 
-        Returns
-        -------
-        self.fig : matplotlib.pyplot.figure.Figure
-        """
-        import matplotlib.pyplot as plt
-        from mpl_toolkits.axes_grid1 import make_axes_locatable
-
-        plt.rcParams['figure.constrained_layout.use'] = True
-        self.fig, axs = plt.subplots(2, 1, sharex=True, sharey=True, figsize=((10, 5)))
-        self.cax = []
-        divider = make_axes_locatable(axs[0])
-        self.cax += [divider.append_axes('right', size='3%', pad=0.03)]
-        divider2 = make_axes_locatable(axs[1])
-        self.cax += [divider2.append_axes('right', size='3%', pad=0.03)]
-        return self.fig
-
-    def plot(self, u, t=None, fig=None, quantity='T'):  # pragma: no cover
-        r"""
-        Plot the solution.
-
-        Parameters
-        ----------
-        u : dtype_u
-            Solution to be plotted
-        t : float
-            Time to display at the top of the figure
-        fig : matplotlib.pyplot.figure.Figure
-            Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated.
-        quantity : (str)
-            quantity you want to plot
-
-        Returns
-        -------
-        None
-        """
-        fig = self.get_fig() if fig is None else fig
-        axs = fig.axes
+        Args:
+            u: The solution you want to compute the Nusselt numbers of
 
-        imV = axs[1].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real)
+        Returns:
+            dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom.
+        """
+        iw, iT = self.index(['w', 'T'])
+        zAxis = self.spectral.axes[-1]
 
         if self.spectral_space:
-            u = self.itransform(u)
+            u_hat = u.copy()
+        else:
+            u_hat = self.transform(u)
 
-        imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real)
+        DzT_hat = (self.Dz @ u_hat[iT].flatten()).reshape(u_hat[iT].shape)
 
-        for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity']):
-            axs[i].set_aspect(1)
-            axs[i].set_title(label)
+        # compute wT with dealiasing
+        padding = (self.dealiasing,) * self.ndim
+        u_pad = self.itransform(u_hat, padding=padding).real
+        _me = self.xp.zeros_like(u_pad)
+        _me[0] = u_pad[iw] * u_pad[iT]
+        wT_hat = self.transform(_me, padding=padding)[0]
 
-        if t is not None:
-            fig.suptitle(f't = {t:.2f}')
-        axs[1].set_xlabel(r'$x$')
-        axs[1].set_ylabel(r'$z$')
-        fig.colorbar(imT, self.cax[0])
-        fig.colorbar(imV, self.cax[1])
+        nusselt_hat = wT_hat - DzT_hat
+
+        if not hasattr(self, '_zInt'):
+            self._zInt = zAxis.get_integration_matrix()
+
+        # get coefficients for evaluation on the boundary
+        top = zAxis.get_BC(kind='Dirichlet', x=1)
+        bot = zAxis.get_BC(kind='Dirichlet', x=-1)
+
+        integral_V = 0
+        if self.comm.rank == 0:
+
+            integral_z = (self._zInt @ nusselt_hat[0, 0]).real
+            integral_z[0] = zAxis.get_integration_constant(integral_z, axis=-1)
+            integral_V = ((top - bot) * integral_z).sum() * self.axes[0].L * self.axes[1].L / self.nx / self.ny
+
+        Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0)
+
+        Nusselt_t = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0, :] * top, axis=-1) / self.nx / self.ny, root=0)
+        Nusselt_b = self.comm.bcast(self.xp.sum(nusselt_hat.real[0, 0] * bot / self.nx / self.ny, axis=-1), root=0)
+
+        return {
+            'V': Nusselt_V,
+            't': Nusselt_t,
+            'b': Nusselt_b,
+        }

pySDC/tests/test_helpers/test_spectral_helper_1d_ultraspherical.py

@@ -71,6 +71,42 @@ def test_integration(N):
     assert np.allclose(poly.integ(lbnd=-1).coef[:-1], U_hat)
 
 
+@pytest.mark.base
+@pytest.mark.parametrize('N', [4, 7, 32])
+@pytest.mark.parametrize('x0', [-1, 0])
+def test_integration_rescaled_domain(N, x0, x1=1):
+    import numpy as np
+    from pySDC.helpers.spectral_helper import UltrasphericalHelper
+
+    helper = UltrasphericalHelper(N, x0=x0, x1=x1)
+
+    x = helper.get_1dgrid()
+    y = N - 2
+    u = x**y * (y + 1)
+    u_hat = helper.transform(u)
+
+    S = helper.get_integration_matrix()
+    U_hat = S @ u_hat
+    U_hat[0] = helper.get_integration_constant(U_hat, axis=-1)
+
+    int_expect = x ** (y + 1)
+    int_hat_expect = helper.transform(int_expect)
+
+    # compare indefinite integral
+    assert np.allclose(int_hat_expect[1:], U_hat[1:])
+    if x0 == 0:
+        assert np.allclose(int_hat_expect[0], U_hat[0]), 'Integration constant is wrong!'
+
+    # compare definite integral
+    top = helper.get_BC(kind='dirichlet', x=1)
+    bot = helper.get_BC(kind='dirichlet', x=-1)
+    res = ((top - bot) * U_hat).sum()
+    if x0 == 0:
+        assert np.isclose(res, 1)
+    elif x0 == -1:
+        assert np.isclose(res, 0 if y % 2 == 1 else 2)
+
+
 @pytest.mark.base
 @pytest.mark.parametrize('N', [6, 33])
 @pytest.mark.parametrize('deg', [1, 3])

pySDC/tests/test_problems/test_RayleighBenard3D.py

@@ -292,11 +292,61 @@ def test_heterogeneous_implementation():
     assert xp.allclose(*un)
 
 
+@pytest.mark.mpi4py
+@pytest.mark.parametrize('w', [0, 1, 3.14])
+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)
+    xp = prob.xp
+    iw, iT = prob.index(['w', 'T'])
+
+    # constant temperature and perturbed velocity
+    u = prob.u_init
+    u[iT, ...] = 3 * prob.Z**2 + 1
+    u[iw] = w * (1 + xp.sin(prob.Y / prob.axes[1].L * 2 * xp.pi))
+    Nu = prob.compute_Nusselt_numbers(u)
+
+    for key, expect in zip(['t', 'b', 'V'], [prob.Lz * (3 + 1) * w - 6, w, w * (1 + 1) - 3]):
+        assert xp.isclose(Nu[key], expect), f'Expected Nu_{key}={expect}, but got {Nu[key]}'
+
+    # zero
+    u = prob.u_init
+    Nu = prob.compute_Nusselt_numbers(u)
+    assert xp.allclose(list(Nu.values()), 0)
+
+    # constant gradient
+    u = prob.u_init
+    u[iT] = prob.Z**2 + 1
+    Nu = prob.compute_Nusselt_numbers(u)
+
+    for key, expect in zip(['t', 'b', 'V'], [-prob.Lz * 2, 0, -1]):
+        assert xp.isclose(Nu[key], expect), f'Expected Nu_{key}={expect}, but got {Nu[key]} with T=z**2!'
+
+    # gradient plus fluctuations
+    u = prob.u_init
+    u[iT] = prob.Z * 1 + (xp.sin(prob.X / prob.axes[0].L * 2 * xp.pi) * xp.sin(prob.Y / prob.axes[1].L * 2 * xp.pi))
+    Nu = prob.compute_Nusselt_numbers(u)
+
+    for key in ['t', 'b', 'V']:
+        assert xp.isclose(Nu[key], -1), f'Expected Nu_{key}=-1, but got {Nu[key]} with T=z*(1+sin(x)+sin(y))!'
+
+    # constant temperature and perturbed velocity
+    u = prob.u_init
+    u[iT, ...] = 1
+    u[iw] = w * (1 + xp.sin(prob.Y / prob.axes[1].L * 2 * xp.pi))
+    Nu = prob.compute_Nusselt_numbers(u)
+
+    for key in ['t', 'b', 'V']:
+        assert xp.isclose(Nu[key], w), f'Expected Nu_{key}={w}, but got {Nu[key]} with constant T and perturbed w!'
+
+
 if __name__ == '__main__':
     # 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_heterogeneous_implementation()
+    test_Nusselt_number_computation(N=4, w=3.14)