Fixed Rayleigh number definition in RBC (#522)

Author: brownbaerchen

pySDC/helpers/spectral_helper.py

@@ -882,6 +882,7 @@ def __init__(self, comm=None, useGPU=False, debug=False):
         self.BCs = None
 
         self.fft_cache = {}
+        self.fft_dealias_shape_cache = {}
 
     @property
     def u_init(self):
@@ -1470,7 +1471,9 @@ def _transform_dct(self, u, axes, padding=None, **kwargs):
 
             if padding is not None:
                 shape = list(v.shape)
-                if self.comm:
+                if ('forward', *padding) in self.fft_dealias_shape_cache.keys():
+                    shape[0] = self.fft_dealias_shape_cache[('forward', *padding)]
+                elif self.comm:
                     send_buf = np.array(v.shape[0])
                     recv_buf = np.array(v.shape[0])
                     self.comm.Allreduce(send_buf, recv_buf)
@@ -1645,7 +1648,9 @@ def _transform_idct(self, u, axes, padding=None, **kwargs):
             if padding is not None:
                 if padding[axis] != 1:
                     shape = list(v.shape)
-                    if self.comm:
+                    if ('backward', *padding) in self.fft_dealias_shape_cache.keys():
+                        shape[0] = self.fft_dealias_shape_cache[('backward', *padding)]
+                    elif self.comm:
                         send_buf = np.array(v.shape[0])
                         recv_buf = np.array(v.shape[0])
                         self.comm.Allreduce(send_buf, recv_buf)
@@ -1754,8 +1759,6 @@ def get_aligned(self, u, axis_in, axis_out, fft=None, forward=False, **kwargs):
         if self.comm.size == 1:
             return u.copy()
 
-        fft = self.get_fft(**kwargs) if fft is None else fft
-
         global_fft = self.get_fft(**kwargs)
         axisA = [me.axisA for me in global_fft.transfer]
         axisB = [me.axisB for me in global_fft.transfer]
@@ -1793,6 +1796,8 @@ def get_aligned(self, u, axis_in, axis_out, fft=None, forward=False, **kwargs):
         else:  # go the potentially slower route of not reusing transfer classes
             from mpi4py_fft import newDistArray
 
+            fft = self.get_fft(**kwargs) if fft is None else fft
+
             _in = newDistArray(fft, forward).redistribute(axis_in)
             _in[...] = u
 

pySDC/implementations/problem_classes/RayleighBenard.py

@@ -6,6 +6,7 @@
 from pySDC.core.convergence_controller import ConvergenceController
 from pySDC.core.hooks import Hooks
 from pySDC.implementations.convergence_controller_classes.check_convergence import CheckConvergence
+from pySDC.core.problem import WorkCounter
 
 
 class RayleighBenard(GenericSpectralLinear):
@@ -20,7 +21,7 @@ class RayleighBenard(GenericSpectralLinear):
         v_t - nu (v_xx + v_zz) + p_z - T = -uv_x - vv_z
 
     with u the horizontal velocity, v the vertical velocity (in z-direction), T the temperature, p the pressure, indices
-    denoting derivatives, kappa=(Rayleigh * Prandl)**(-1/2) and nu = (Rayleigh / Prandl)**(-1/2). Everything on the left
+    denoting derivatives, kappa=(Rayleigh * Prandtl)**(-1/2) and nu = (Rayleigh / Prandtl)**(-1/2). Everything on the left
     hand side, that is the viscous part, the pressure gradient and the buoyancy due to temperature are treated
     implicitly, while the non-linear convection part on the right hand side is integrated explicitly.
 
@@ -36,7 +37,7 @@ class RayleighBenard(GenericSpectralLinear):
     facilitate the Dirichlet BCs.
 
     Parameters:
-        Prandl (float): Prandl number
+        Prandtl (float): Prandtl number
         Rayleigh (float): Rayleigh number
         nx (int): Horizontal resolution
         nz (int): Vertical resolution
@@ -50,26 +51,28 @@ class RayleighBenard(GenericSpectralLinear):
 
     def __init__(
         self,
-        Prandl=1,
+        Prandtl=1,
         Rayleigh=2e6,
         nx=256,
         nz=64,
         BCs=None,
         dealiasing=3 / 2,
         comm=None,
+        Lx=8,
         **kwargs,
     ):
         """
         Constructor. `kwargs` are forwarded to parent class constructor.
 
         Args:
-            Prandl (float): Prandtl number
+            Prandtl (float): Prandtl number
             Rayleigh (float): Rayleigh number
             nx (int): Resolution in x-direction
             nz (int): Resolution in z direction
             BCs (dict): Vertical boundary conditions
             dealiasing (float): Dealiasing for evaluating the non-linear part in real space
             comm (mpi4py.Intracomm): Space communicator
+            Lx (float): Horizontal length of the domain
         """
         BCs = {} if BCs is None else BCs
         BCs = {
@@ -90,18 +93,19 @@ def __init__(
             except ModuleNotFoundError:
                 pass
         self._makeAttributeAndRegister(
-            'Prandl',
+            'Prandtl',
             'Rayleigh',
             'nx',
             'nz',
             'BCs',
             'dealiasing',
             'comm',
+            'Lx',
             localVars=locals(),
             readOnly=True,
         )
 
-        bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': 8}, {'base': 'ultraspherical', 'N': nz}]
+        bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx}, {'base': 'ultraspherical', 'N': nz}]
         components = ['u', 'v', 'T', 'p']
         super().__init__(bases, components, comm=comm, **kwargs)
 
@@ -127,15 +131,17 @@ def __init__(
         self.Dz = S1 @ Dz
         self.Dzz = S2 @ Dzz
 
-        kappa = (Rayleigh * Prandl) ** (-1 / 2.0)
-        nu = (Rayleigh / Prandl) ** (-1 / 2.0)
+        # compute rescaled Rayleigh number to extract viscosity and thermal diffusivity
+        Ra = Rayleigh / (max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * self.axes[1].L ** 3)
+        self.kappa = (Ra * Prandtl) ** (-1 / 2.0)
+        self.nu = (Ra / Prandtl) ** (-1 / 2.0)
 
         # construct operators
         L_lhs = {
             'p': {'u': U01 @ Dx, 'v': Dz},  # divergence free constraint
-            'u': {'p': U02 @ Dx, 'u': -nu * (U02 @ Dxx + Dzz)},
-            'v': {'p': U12 @ Dz, 'v': -nu * (U02 @ Dxx + Dzz), 'T': -U02 @ Id},
-            'T': {'T': -kappa * (U02 @ Dxx + Dzz)},
+            'u': {'p': U02 @ Dx, 'u': -self.nu * (U02 @ Dxx + Dzz)},
+            'v': {'p': U12 @ Dz, 'v': -self.nu * (U02 @ Dxx + Dzz), 'T': -U02 @ Id},
+            'T': {'T': -self.kappa * (U02 @ Dxx + Dzz)},
         }
         self.setup_L(L_lhs)
 
@@ -175,6 +181,8 @@ def __init__(
                 )
         self.setup_BCs()
 
+        self.work_counters['rhs'] = WorkCounter()
+
     def eval_f(self, u, *args, **kwargs):
         f = self.f_init
 
@@ -225,6 +233,7 @@ def eval_f(self, u, *args, **kwargs):
         else:
             f.expl[:] = self.itransform(self.transform(fexpl_pad, padding=padding)).real
 
+        self.work_counters['rhs']()
         return f
 
     def u_exact(self, t=0, noise_level=1e-3, seed=99):

pySDC/implementations/problem_classes/generic_spectral.py

@@ -28,21 +28,26 @@ class attributes of the spectral helper. This class will automatically switch th
         Pr (sparse matrix): Right preconditioner
     """
 
-    @classmethod
-    def setup_GPU(cls):
+    def setup_GPU(self):
         """switch to GPU modules"""
         import cupy as cp
         from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh
         from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
 
-        cls.dtype_u = cupy_mesh
+        self.dtype_u = cupy_mesh
 
         GPU_versions = {
             mesh: cupy_mesh,
             imex_mesh: imex_cupy_mesh,
         }
 
-        cls.dtype_f = GPU_versions[cls.dtype_f]
+        self.dtype_f = GPU_versions[self.dtype_f]
+
+        if self.comm is not None:
+            from pySDC.helpers.NCCL_communicator import NCCLComm
+
+            if not isinstance(self.comm, NCCLComm):
+                self.__dict__['comm'] = NCCLComm(self.comm)
 
     def __init__(
         self,
@@ -94,6 +99,9 @@ def __init__(
 
         if useGPU:
             self.setup_GPU()
+            if self.solver_args is not None:
+                if 'rtol' in self.solver_args.keys():
+                    self.solver_args['tol'] = self.solver_args.pop('rtol')
 
         for base in bases:
             self.spectral.add_axis(**base)
@@ -218,14 +226,21 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
 
         if self.spectral_space:
             rhs_hat = rhs.copy()
+            if u0 is not None:
+                u0_hat = self.Pr.T @ u0.copy().flatten()
         else:
             rhs_hat = self.spectral.transform(rhs)
+            if u0 is not None:
+                u0_hat = self.Pr.T @ self.spectral.transform(u0).flatten()
+
+        if self.useGPU:
+            self.xp.cuda.Device().synchronize()
 
         rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape)
         rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat)
         rhs_hat = self.Pl @ rhs_hat.flatten()
 
-        if dt not in self.cached_factorizations.keys():
+        if dt not in self.cached_factorizations.keys() or not self.solver_type.lower() == 'cached_direct':
             A = self.M + dt * self.L
             A = self.Pl @ self.spectral.put_BCs_in_matrix(A) @ self.Pr
 
@@ -255,25 +270,58 @@ def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs)
 
         elif self.solver_type.lower() == 'direct':
             _sol_hat = sp.linalg.spsolve(A, rhs_hat)
+        elif self.solver_type.lower() == 'lsqr':
+            lsqr = sp.linalg.lsqr(
+                A,
+                rhs_hat,
+                x0=u0_hat,
+                **self.solver_args,
+            )
+            _sol_hat = lsqr[0]
         elif self.solver_type.lower() == 'gmres':
             _sol_hat, _ = sp.linalg.gmres(
                 A,
                 rhs_hat,
-                x0=u0.flatten(),
+                x0=u0_hat,
                 **self.solver_args,
                 callback=self.work_counters[self.solver_type],
-                callback_type='legacy',
+                callback_type='pr_norm',
+            )
+        elif self.solver_type.lower() == 'gmres+ilu':
+            linalg = self.spectral.linalg
+
+            if dt not in self.cached_factorizations.keys():
+                if len(self.cached_factorizations) >= self.max_cached_factorizations:
+                    to_evict = list(self.cached_factorizations.keys())[0]
+                    self.cached_factorizations.pop(to_evict)
+                    self.logger.debug(f'Evicted matrix factorization for {to_evict=:.6f} from cache')
+                iLU = linalg.spilu(A, drop_tol=dt * 1e-4, fill_factor=100)
+                self.cached_factorizations[dt] = linalg.LinearOperator(A.shape, iLU.solve)
+                self.logger.debug(f'Cached matrix factorization for {dt=:.6f}')
+                self.work_counters['factorizations']()
+
+            _sol_hat, _ = linalg.gmres(
+                A,
+                rhs_hat,
+                x0=u0_hat,
+                **self.solver_args,
+                callback=self.work_counters[self.solver_type],
+                callback_type='pr_norm',
+                M=self.cached_factorizations[dt],
             )
         elif self.solver_type.lower() == 'cg':
             _sol_hat, _ = sp.linalg.cg(
-                A, rhs_hat, x0=u0.flatten(), **self.solver_args, callback=self.work_counters[self.solver_type]
+                A, rhs_hat, x0=u0_hat, **self.solver_args, callback=self.work_counters[self.solver_type]
             )
         else:
-            raise NotImplementedError(f'Solver {self.solver_type:!} not implemented in {type(self).__name__}!')
+            raise NotImplementedError(f'Solver {self.solver_type=} not implemented in {type(self).__name__}!')
 
         sol_hat = self.spectral.u_init_forward
         sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape)
 
+        if self.useGPU:
+            self.xp.cuda.Device().synchronize()
+
         if self.spectral_space:
             return sol_hat
         else:
@@ -319,7 +367,6 @@ def compute_residual_DAE(self, stage=''):
             res[m] += L.tau[m]
         # use abs function from data type here
         res_norm.append(abs(res[m]))
-        # print(m, [abs(me) for me in res[m]], [abs(me) for me in L.u[0] - L.u[m + 1]])
 
     # find maximal residual over the nodes
     if L.params.residual_type == 'full_abs':

pySDC/tests/test_problems/test_RayleighBenard.py

@@ -15,8 +15,8 @@ def test_eval_f(nx, nz, direction, spectral_space):
     X, Z = P.X, P.Z
     cos, sin = np.cos, np.sin
 
-    kappa = (P.Rayleigh * P.Prandl) ** (-1 / 2)
-    nu = (P.Rayleigh / P.Prandl) ** (-1 / 2)
+    kappa = P.kappa
+    nu = P.nu
 
     if direction == 'x':
         y = sin(X * np.pi)
@@ -181,7 +181,9 @@ def test_Poisson_problems(nx, component):
         'T_top': 0,
         'T_bottom': 0,
     }
-    P = RayleighBenard(nx=nx, nz=6, BCs=BCs, Rayleigh=1.0)
+    P = RayleighBenard(
+        nx=nx, nz=6, BCs=BCs, Rayleigh=(max([abs(BCs['T_top'] - BCs['T_bottom']), np.finfo(float).eps]) * 2**3)
+    )
     rhs = P.u_init
 
     idx = P.index(f'{component}')