Update of 2D RBC implementation

pySDC/implementations/hooks/log_solution.py

@@ -208,6 +208,7 @@ class LogToFile(Hooks):
     filename = 'myRun.pySDC'
     time_increment = 0
     allow_overwriting = False
+    counter = 0  # number of stored time points in the file
 
     def __init__(self):
         super().__init__()
@@ -224,8 +225,9 @@ def pre_run(self, step, level_number):
         if os.path.isfile(self.filename) and L.time > 0:
             L.prob.setUpFieldsIO()
             self.outfile = FieldsIO.fromFile(self.filename)
+            self.counter = len(self.outfile.times)
             self.logger.info(
-                f'Set up file {self.filename!r} for writing output. This file already contains solutions up to t={self.outfile.times[-1]:.4f}.'
+                f'Set up file {self.filename!r} for writing output. This file already contains {self.counter} solutions up to t={self.outfile.times[-1]:.4f}.'
             )
         else:
             self.outfile = L.prob.getOutputFile(self.filename)
@@ -236,6 +238,9 @@ def pre_run(self, step, level_number):
                 self.outfile.addField(time=L.time, field=L.prob.processSolutionForOutput(L.u[0]))
                 self.logger.info(f'Written initial conditions at t={L.time:4f} to file')
 
+        type(self).counter = len(self.outfile.times)
+        self.logger.info(f'Will write to disk every {self.time_increment:.4e} time units')
+
     def post_step(self, step, level_number):
         if level_number > 0:
             return None
@@ -252,6 +257,20 @@ def post_step(self, step, level_number):
             self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend))
             self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file')
             self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment
+            type(self).counter = len(self.outfile.times)
+
+    def post_run(self, step, level_number):
+        if level_number > 0:
+            return None
+
+        L = step.levels[level_number]
+
+        value_exists = True in [abs(me - (L.time + L.dt)) < np.finfo(float).eps * 1000 for me in self.outfile.times]
+        if not value_exists:
+            self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend))
+            self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file')
+            self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment
+            type(self).counter = len(self.outfile.times)
 
     @classmethod
     def load(cls, index):

pySDC/implementations/problem_classes/RayleighBenard.py

@@ -58,7 +58,9 @@ def __init__(
         BCs=None,
         dealiasing=3 / 2,
         comm=None,
-        Lx=8,
+        Lx=4,
+        Lz=1,
+        z0=0,
         **kwargs,
     ):
         """
@@ -73,11 +75,13 @@ def __init__(
             dealiasing (float): Dealiasing for evaluating the non-linear part in real space
             comm (mpi4py.Intracomm): Space communicator
             Lx (float): Horizontal length of the domain
+            Lz (float): Vertical length of the domain
+            z0 (float): Position of lower boundary
         """
         BCs = {} if BCs is None else BCs
         BCs = {
             'T_top': 0,
-            'T_bottom': 2,
+            'T_bottom': 1,
             'v_top': 0,
             'v_bottom': 0,
             'u_top': 0,
@@ -101,11 +105,16 @@ def __init__(
             'dealiasing',
             'comm',
             'Lx',
+            'Lz',
+            'z0',
             localVars=locals(),
             readOnly=True,
         )
 
-        bases = [{'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx}, {'base': 'ultraspherical', 'N': nz}]
+        bases = [
+            {'base': 'fft', 'N': nx, 'x0': 0, 'x1': self.Lx},
+            {'base': 'ultraspherical', 'N': nz, 'x0': self.z0, 'x1': self.Lz},
+        ]
         components = ['u', 'v', 'T', 'p']
         super().__init__(bases, components, comm=comm, **kwargs)
 
@@ -247,8 +256,8 @@ def u_exact(self, t=0, noise_level=1e-3, seed=99):
 
         # linear temperature gradient
         for comp in ['T', 'v', 'u']:
-            a = (self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']) / 2
-            b = (self.BCs[f'{comp}_top'] + self.BCs[f'{comp}_bottom']) / 2
+            a = (self.BCs[f'{comp}_top'] - self.BCs[f'{comp}_bottom']) / self.Lz
+            b = self.BCs[f'{comp}_bottom'] - a * self.z0
             me[self.index(comp)] = a * self.Z + b
 
         # perturb slightly
@@ -257,7 +266,7 @@ def u_exact(self, t=0, noise_level=1e-3, seed=99):
         noise = self.spectral.u_init
         noise[iT] = rng.random(size=me[iT].shape)
 
-        me[iT] += noise[iT].real * noise_level * (self.Z - 1) * (self.Z + 1)
+        me[iT] += noise[iT].real * noise_level * (self.Z - self.z0) * (self.Z - self.z0 + self.Lz)
 
         if self.spectral_space:
             me_hat = self.spectral.u_init_forward
@@ -370,14 +379,41 @@ def compute_vorticity(self, u):
             u_hat = u.copy()
         else:
             u_hat = self.transform(u)
+
         Dz = self.Dz
         Dx = self.Dx
         iu, iv = self.index(['u', 'v'])
 
         vorticity_hat = self.spectral.u_init_forward
-        vorticity_hat[0] = (Dx * u_hat[iv].flatten() + Dz @ u_hat[iu].flatten()).reshape(u[iu].shape)
+        vorticity_hat[0] = (Dx * u_hat[iv].flatten() + Dz @ u_hat[iu].flatten()).reshape(u_hat[iu].shape)
         return self.itransform(vorticity_hat)[0].real
 
+    def getOutputFile(self, fileName):
+        from pySDC.helpers.fieldsIO import Rectilinear
+
+        self.setUpFieldsIO()
+
+        coords = [me.get_1dgrid() for me in self.spectral.axes]
+        assert np.allclose([len(me) for me in coords], self.spectral.global_shape[1:])
+
+        fOut = Rectilinear(np.float64, fileName=fileName)
+        fOut.setHeader(nVar=len(self.components) + 1, coords=coords)
+        fOut.initialize()
+        return fOut
+
+    def processSolutionForOutput(self, u):
+        vorticity = self.compute_vorticity(u)
+
+        if self.spectral_space:
+            u_real = self.itransform(u).real
+        else:
+            u_real = u.real
+
+        me = np.empty(shape=(u_real.shape[0] + 1, *vorticity.shape))
+        me[:-1] = u_real
+        me[-1] = vorticity
+        return me
+
     def compute_Nusselt_numbers(self, u):
         """
         Compute the various versions of the Nusselt number. This reflects the type of heat transport.
@@ -392,52 +428,46 @@ def compute_Nusselt_numbers(self, u):
             dict: Nusselt number averaged over the entire volume and horizontally averaged at the top and bottom.
         """
         iv, iT = self.index(['v', 'T'])
-
-        DzT_hat = self.spectral.u_init_forward
+        zAxis = self.spectral.axes[-1]
 
         if self.spectral_space:
             u_hat = u.copy()
         else:
             u_hat = self.transform(u)
 
-        DzT_hat[iT] = (self.Dz @ u_hat[iT].flatten()).reshape(DzT_hat[iT].shape)
+        DzT_hat = (self.Dz @ u_hat[iT].flatten()).reshape(u_hat[iT].shape)
 
         # compute vT with dealiasing
         padding = (self.dealiasing, self.dealiasing)
         u_pad = self.itransform(u_hat, padding=padding).real
         _me = self.xp.zeros_like(u_pad)
         _me[0] = u_pad[iv] * u_pad[iT]
-        vT_hat = self.transform(_me, padding=padding)
+        vT_hat = self.transform(_me, padding=padding)[0]
 
-        nusselt_hat = (vT_hat[0] - DzT_hat[iT]) / self.nx
-        nusselt_no_v_hat = (-DzT_hat[iT]) / self.nx
+        if not hasattr(self, '_zInt'):
+            self._zInt = zAxis.get_integration_matrix()
 
-        integral_z = self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='integral'), axis=-1).real
-        integral_V = (
-            integral_z[0] * self.axes[0].L
-        )  # only the first Fourier mode has non-zero integral with periodic BCs
-        Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0)
+        nusselt_hat = (vT_hat / self.kappa - DzT_hat) * self.axes[-1].L
 
-        Nusselt_t = self.comm.bcast(
-            self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0
-        )
-        Nusselt_b = self.comm.bcast(
-            self.xp.sum(nusselt_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0], root=0
-        )
-        Nusselt_no_v_t = self.comm.bcast(
-            self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=1), axis=-1).real[0], root=0
-        )
-        Nusselt_no_v_b = self.comm.bcast(
-            self.xp.sum(nusselt_no_v_hat * self.spectral.axes[1].get_BC(kind='Dirichlet', x=-1), axis=-1).real[0],
-            root=0,
-        )
+        # 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]).real
+            integral_z[0] = zAxis.get_integration_constant(integral_z, axis=-1)
+            integral_V = ((top - bot) * integral_z).sum() * self.axes[0].L / self.nx
+
+        Nusselt_V = self.comm.bcast(integral_V / self.spectral.V, root=0)
+        Nusselt_t = self.comm.bcast(self.xp.sum(nusselt_hat.real[0] * top, axis=-1) / self.nx, root=0)
+        Nusselt_b = self.comm.bcast(self.xp.sum(nusselt_hat.real[0] * bot, axis=-1) / self.nx, root=0)
 
         return {
             'V': Nusselt_V,
             't': Nusselt_t,
             'b': Nusselt_b,
-            't_no_v': Nusselt_no_v_t,
-            'b_no_v': Nusselt_no_v_b,
         }
 
     def compute_viscous_dissipation(self, u):
@@ -509,8 +539,8 @@ def compute_max_step_size(P, u):
         grid_spacing_x = P.X[1, 0] - P.X[0, 0]
 
         cell_wallz = P.xp.zeros(P.nz + 1)
-        cell_wallz[0] = 1
-        cell_wallz[-1] = -1
+        cell_wallz[0] = P.Lz
+        cell_wallz[-1] = 0
         cell_wallz[1:-1] = (P.Z[0, :-1] + P.Z[0, 1:]) / 2
         grid_spacing_z = cell_wallz[:-1] - cell_wallz[1:]
 

pySDC/projects/GPU/configs/RBC_configs.py

@@ -29,54 +29,6 @@ class RayleighBenardRegular(Config):
     sweeper_type = 'IMEX'
     Tend = 50
 
-    def get_LogToFile(self, ranks=None):
-        import numpy as np
-        from pySDC.implementations.hooks.log_solution import LogToPickleFileAfterXS as LogToFile
-
-        LogToFile.path = f'{self.base_path}/data/'
-        LogToFile.file_name = f'{self.get_path(ranks=ranks)}-solution'
-        LogToFile.time_increment = 1e-1
-
-        def process_solution(L):
-            P = L.prob
-
-            if P.spectral_space:
-                uend = P.itransform(L.uend)
-
-            data = {
-                't': L.time + L.dt,
-                'local_slice': P.local_slice,
-                'shape': P.global_shape,
-            }
-
-            if P.useGPU:
-                data['u'] = uend.get().view(np.ndarray)
-                data['vorticity'] = L.prob.compute_vorticity(L.uend).get().view(np.ndarray)
-                if L.time == 0:
-                    data['X'] = L.prob.X.get()
-                    data['Z'] = L.prob.Z.get()
-            else:
-                data['u'] = uend.view(np.ndarray)
-                data['vorticity'] = L.prob.compute_vorticity(L.uend).view(np.ndarray)
-                if L.time == 0:
-                    data['X'] = L.prob.X
-                    data['Z'] = L.prob.Z
-            return data
-
-        def logging_condition(L):
-            sweep = L.sweep
-            if hasattr(sweep, 'comm'):
-                if sweep.comm.rank == sweep.comm.size - 1:
-                    return True
-                else:
-                    return False
-            else:
-                return True
-
-        LogToFile.process_solution = process_solution
-        LogToFile.logging_condition = logging_condition
-        return LogToFile
-
     def get_controller_params(self, *args, **kwargs):
         from pySDC.implementations.hooks.log_step_size import LogStepSize
 
@@ -125,12 +77,12 @@ def get_description(self, *args, MPIsweeper=False, res=-1, **kwargs):
         return desc
 
     def get_initial_condition(self, P, *args, restart_idx=0, **kwargs):
-        if restart_idx > 0:
-            return super().get_initial_condition(P, *args, restart_idx=restart_idx, **kwargs)
-        else:
+        if restart_idx == 0:
             u0 = P.u_exact(t=0, seed=P.comm.rank, noise_level=1e-3)
             u0_with_pressure = P.solve_system(u0, 1e-9, u0)
             return u0_with_pressure, 0
+        else:
+            return super().get_initial_condition(P, *args, restart_idx=restart_idx, **kwargs)
 
     def prepare_caches(self, prob):
         """
@@ -139,6 +91,7 @@ def prepare_caches(self, prob):
         prob.eval_f(prob.u_init)
 
     def plot(self, P, idx, n_procs_list, quantitiy='T', quantitiy2='vorticity'):
+        from pySDC.helpers.fieldsIO import FieldsIO
         import numpy as np
 
         cmaps = {'vorticity': 'bwr', 'p': 'bwr'}
@@ -147,48 +100,30 @@ def plot(self, P, idx, n_procs_list, quantitiy='T', quantitiy2='vorticity'):
         cax = P.cax
         axs = fig.get_axes()
 
-        buffer = {}
-        vmin = {quantitiy: np.inf, quantitiy2: np.inf}
-        vmax = {quantitiy: -np.inf, quantitiy2: -np.inf}
-
-        X = {}
-        Z = {}
-
-        for rank in range(n_procs_list[2]):
-            ranks = self.ranks[:-1] + [rank]
-            LogToFile = self.get_LogToFile(ranks=ranks)
-
-            buffer[f'u-{rank}'] = LogToFile.load(idx)
-
-            vmin[quantitiy2] = min([vmin[quantitiy2], buffer[f'u-{rank}'][quantitiy2].real.min()])
-            vmax[quantitiy2] = max([vmax[quantitiy2], buffer[f'u-{rank}'][quantitiy2].real.max()])
-            vmin[quantitiy] = min([vmin[quantitiy], buffer[f'u-{rank}']['u'][P.index(quantitiy)].real.min()])
-            vmax[quantitiy] = max([vmax[quantitiy], buffer[f'u-{rank}']['u'][P.index(quantitiy)].real.max()])
-
-            first_frame = LogToFile.load(0)
-            X[rank] = first_frame['X']
-            Z[rank] = first_frame['Z']
-
-        for rank in range(n_procs_list[2]):
-            im2 = axs[1].pcolormesh(
-                X[rank],
-                Z[rank],
-                buffer[f'u-{rank}'][quantitiy2].real,
-                vmin=-vmax[quantitiy2] if cmaps.get(quantitiy2, None) in ['bwr'] else vmin[quantitiy2],
-                vmax=vmax[quantitiy2],
-                cmap=cmaps.get(quantitiy2, None),
-            )
-            im = axs[0].pcolormesh(
-                X[rank],
-                Z[rank],
-                buffer[f'u-{rank}']['u'][P.index(quantitiy)].real,
-                vmin=vmin[quantitiy],
-                vmax=-vmin[quantitiy] if cmaps.get(quantitiy, None) in ['bwr', 'seismic'] else vmax[quantitiy],
-                cmap=cmaps.get(quantitiy, 'plasma'),
-            )
+        outfile = FieldsIO.fromFile(self.get_file_name())
+
+        x = outfile.header['coords'][0]
+        z = outfile.header['coords'][1]
+        X, Z = np.meshgrid(x, z, indexing='ij')
+
+        t, data = outfile.readField(idx)
+        im = axs[0].pcolormesh(
+            X,
+            Z,
+            data[P.index(quantitiy)].real,
+            cmap=cmaps.get(quantitiy, 'plasma'),
+        )
+
+        im2 = axs[1].pcolormesh(
+            X,
+            Z,
+            data[-1].real,
+            cmap=cmaps.get(quantitiy2, None),
+        )
+
         fig.colorbar(im2, cax[1])
         fig.colorbar(im, cax[0])
-        axs[0].set_title(f't={buffer[f"u-{rank}"]["t"]:.2f}')
+        axs[0].set_title(f't={t:.2f}')
         axs[1].set_xlabel('x')
         axs[1].set_ylabel('z')
         axs[0].set_aspect(1.0)