Merge remote-tracking branch 'origin/master' into neuralpint

Author: tlunet

pySDC/helpers/spectral_helper.py

@@ -535,6 +535,8 @@ def get_BC(self, kind, **kwargs):
             return self.get_integ_BC_row(**kwargs)
         elif kind.lower() == 'dirichlet':
             return self.get_Dirichlet_BC_row(**kwargs)
+        elif kind.lower() == 'neumann':
+            return self.get_Neumann_BC_row(**kwargs)
         else:
             return super().get_BC(kind)
 
@@ -574,6 +576,27 @@ def get_Dirichlet_BC_row(self, x):
         else:
             raise NotImplementedError(f'Don\'t know how to generate Dirichlet BC\'s at {x=}!')
 
+    def get_Neumann_BC_row(self, x):
+        """
+        Get a row for generating Neumann BCs at x with T polynomials.
+
+        Args:
+            x (float): Position of the boundary condition
+
+        Returns:
+            self.xp.ndarray: Row to put into a matrix
+        """
+        n = self.xp.arange(self.N, dtype='D')
+        nn = n**2
+        if x == -1:
+            me = nn
+            me[1:] *= (-1) ** n[:-1]
+            return me
+        elif x == 1:
+            return nn
+        else:
+            raise NotImplementedError(f'Don\'t know how to generate Neumann BC\'s at {x=}!')
+
     def get_Dirichlet_recombination_matrix(self):
         '''
         Get matrix for Dirichlet recombination, which changes the basis to have sparse boundary conditions.

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:]