Merge branch 'master' into neuralpint

Author: tlunet

pySDC/core/level.py

@@ -1,3 +1,6 @@
+from pySDC.core.sweeper import Sweeper
+from pySDC.core.problem import Problem
+
 from pySDC.helpers.pysdc_helper import FrozenClass
 
 
@@ -72,8 +75,8 @@ def __init__(self, problem_class, problem_params, sweeper_class, sweeper_params,
         self.status = _Status()
 
         # instantiate sweeper, problem and hooks
-        self.__sweep = sweeper_class(sweeper_params, self)
-        self.__prob = problem_class(**problem_params)
+        self.__sweep: Sweeper = sweeper_class(sweeper_params, self)
+        self.__prob: Problem = problem_class(**problem_params)
 
         # set name
         self.level_index = level_index
@@ -119,7 +122,7 @@ def reset_level(self, reset_status=True):
         self.tau = [None] * self.sweep.coll.num_nodes
 
     @property
-    def sweep(self):
+    def sweep(self) -> Sweeper:
         """
         Getter for the sweeper
 
@@ -129,7 +132,7 @@ def sweep(self):
         return self.__sweep
 
     @property
-    def prob(self):
+    def prob(self) -> Problem:
         """
         Getter for the problem
 

pySDC/core/step.py

@@ -1,6 +1,6 @@
 import logging
 
-from pySDC.core import level as levclass
+from pySDC.core.level import Level
 from pySDC.core.base_transfer import BaseTransfer
 from pySDC.core.errors import ParameterError
 from pySDC.helpers.pysdc_helper import FrozenClass
@@ -73,7 +73,7 @@ def __init__(self, description):
         # empty attributes
         self.__transfer_dict = {}
         self.base_transfer = None
-        self.levels = []
+        self.levels: list[Level] = []
         self.__prev = None
         self.__next = None
 
@@ -149,7 +149,7 @@ def __generate_hierarchy(self, descr):
 
         # generate levels, register and connect if needed
         for l in range(len(descr_list)):
-            L = levclass.Level(
+            L = Level(
                 problem_class=descr_list[l]['problem_class'],
                 problem_params=descr_list[l]['problem_params'],
                 sweeper_class=descr_list[l]['sweeper_class'],

pySDC/core/sweeper.py

@@ -1,13 +1,18 @@
 import logging
 import numpy as np
-from qmat import QDELTA_GENERATORS
+from qmat.qdelta import QDeltaGenerator, QDELTA_GENERATORS
 
 from pySDC.core.errors import ParameterError
-from pySDC.core.level import Level
 from pySDC.core.collocation import CollBase
 from pySDC.helpers.pysdc_helper import FrozenClass
 
 
+# Organize QDeltaGenerator class in dict[type(QDeltaGenerator),set(str)] to retrieve aliases
+QDELTA_GENERATORS_ALIASES = {v: set() for v in set(QDELTA_GENERATORS.values())}
+for k, v in QDELTA_GENERATORS.items():
+    QDELTA_GENERATORS_ALIASES[v].add(k)
+
+
 # short helper class to add params as attributes
 class _Pars(FrozenClass):
     def __init__(self, pars):
@@ -82,27 +87,17 @@ def __init__(self, params, level):
 
         self.__level = level
         self.parallelizable = False
+        for name in ["genQI", "genQE"]:
+            if hasattr(self, name):
+                delattr(self, name)
 
-    def getGenerator(self, qd_type):
-        coll = self.coll
-
-        generator = QDELTA_GENERATORS[qd_type](
-            # for algebraic types (LU, ...)
-            Q=coll.generator.Q,
-            # for MIN in tables, MIN-SR-S ...
-            nNodes=coll.num_nodes,
-            nodeType=coll.node_type,
-            quadType=coll.quad_type,
-            # for time-stepping types, MIN-SR-NS
-            nodes=coll.nodes,
-            tLeft=coll.tleft,
-        )
-
-        return generator
+    def buildGenerator(self, qdType: str) -> QDeltaGenerator:
+        return QDELTA_GENERATORS[qdType](qGen=self.coll.generator, tLeft=self.coll.tleft)
 
     def get_Qdelta_implicit(self, qd_type, k=None):
         QDmat = np.zeros_like(self.coll.Qmat)
-        self.genQI = self.getGenerator(qd_type)
+        if not hasattr(self, "genQI") or qd_type not in QDELTA_GENERATORS_ALIASES[type(self.genQI)]:
+            self.genQI: QDeltaGenerator = self.buildGenerator(qd_type)
         QDmat[1:, 1:] = self.genQI.genCoeffs(k=k)
 
         err_msg = 'Lower triangular matrix expected!'
@@ -114,7 +109,8 @@ def get_Qdelta_implicit(self, qd_type, k=None):
     def get_Qdelta_explicit(self, qd_type, k=None):
         coll = self.coll
         QDmat = np.zeros(coll.Qmat.shape, dtype=float)
-        self.genQE = self.getGenerator(qd_type)
+        if not hasattr(self, "genQE") or qd_type not in QDELTA_GENERATORS_ALIASES[type(self.genQE)]:
+            self.genQE: QDeltaGenerator = self.buildGenerator(qd_type)
         QDmat[1:, 1:], QDmat[1:, 0] = self.genQE.genCoeffs(k=k, dTau=True)
 
         err_msg = 'Strictly lower triangular matrix expected!'
@@ -251,6 +247,8 @@ def level(self, L):
         Args:
             L (pySDC.Level.level): current level
         """
+        from pySDC.core.level import Level
+
         assert isinstance(L, Level)
         self.__level = L
 
@@ -267,5 +265,9 @@ def updateVariableCoeffs(self, k):
         k : int
             Index of the sweep (0 for initial sweep, 1 for the first one, ...).
         """
-        if self.genQI._K_DEPENDENT:
-            self.QI = self.get_Qdelta_implicit(qd_type=self.genQI.__class__.__name__, k=k)
+        if hasattr(self, "genQI") and self.genQI.isKDependent():
+            qdType = type(self.genQI).__name__
+            self.QI = self.get_Qdelta_implicit(qdType, k=k)
+        if hasattr(self, "genQE") and self.genQE.isKDependent():
+            qdType = type(self.genQE).__name__
+            self.QE = self.get_Qdelta_explicit(qdType, k=k)

pySDC/implementations/controller_classes/controller_MPI.py

@@ -28,7 +28,7 @@ def __init__(self, controller_params, description, comm):
         super().__init__(controller_params, description, useMPI=True)
 
         # create single step per processor
-        self.S = Step(description)
+        self.S: Step = Step(description)
 
         # pass communicator for future use
         self.comm = comm
@@ -688,8 +688,11 @@ def it_fine(self, comm, num_procs):
 
             for hook in self.hooks:
                 hook.pre_sweep(step=self.S, level_number=0)
+
+            self.S.levels[0].sweep.updateVariableCoeffs(k + 1)  # update QDelta coefficients if variable preconditioner
             self.S.levels[0].sweep.update_nodes()
             self.S.levels[0].sweep.compute_residual(stage='IT_FINE')
+
             for hook in self.hooks:
                 hook.post_sweep(step=self.S, level_number=0)
 
@@ -718,6 +721,7 @@ def it_down(self, comm, num_procs):
 
                 for hook in self.hooks:
                     hook.pre_sweep(step=self.S, level_number=l)
+
                 self.S.levels[l].sweep.update_nodes()
                 self.S.levels[l].sweep.compute_residual(stage='IT_DOWN')
                 for hook in self.hooks:

pySDC/implementations/controller_classes/controller_nonMPI.py

@@ -4,7 +4,7 @@
 import dill
 
 from pySDC.core.controller import Controller
-from pySDC.core import step as stepclass
+from pySDC.core.step import Step
 from pySDC.core.errors import ControllerError, CommunicationError
 from pySDC.implementations.convergence_controller_classes.basic_restarting import BasicRestarting
 
@@ -32,7 +32,7 @@ def __init__(self, num_procs, controller_params, description):
         # call parent's initialization routine
         super().__init__(controller_params, description, useMPI=False)
 
-        self.MS = [stepclass.Step(description)]
+        self.MS: list[Step] = [Step(description)]
 
         # try to initialize via dill.copy (much faster for many time-steps)
         try:
@@ -42,7 +42,7 @@ def __init__(self, num_procs, controller_params, description):
         except (dill.PicklingError, TypeError, ValueError) as error:
             self.logger.warning(f'Need to initialize steps separately due to pickling error: {error}')
             for _ in range(num_procs - 1):
-                self.MS.append(stepclass.Step(description))
+                self.MS.append(Step(description))
 
         self.base_convergence_controllers += [BasicRestarting.get_implementation(useMPI=False)]
         for convergence_controller in self.base_convergence_controllers:
@@ -542,7 +542,7 @@ def it_check(self, local_MS_running):
         for C in [self.convergence_controllers[i] for i in self.convergence_controller_order]:
             C.reset_buffers_nonMPI(self)
 
-    def it_fine(self, local_MS_running):
+    def it_fine(self, local_MS_running: list[Step]):
         """
         Fine sweeps
 
@@ -567,8 +567,11 @@ def it_fine(self, local_MS_running):
                 # standard sweep workflow: update nodes, compute residual, log progress
                 for hook in self.hooks:
                     hook.pre_sweep(step=S, level_number=0)
+
+                S.levels[0].sweep.updateVariableCoeffs(k + 1)  # update QDelta coefficients if variable preconditioner
                 S.levels[0].sweep.update_nodes()
                 S.levels[0].sweep.compute_residual(stage='IT_FINE')
+
                 for hook in self.hooks:
                     hook.post_sweep(step=S, level_number=0)
 

pySDC/implementations/convergence_controller_classes/adaptive_collocation.py

@@ -2,7 +2,6 @@
 
 from qmat.lagrange import LagrangeApproximation
 from pySDC.core.convergence_controller import ConvergenceController, Status
-from pySDC.core.collocation import CollBase
 
 
 class AdaptiveCollocation(ConvergenceController):

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/implementations/sweeper_classes/generic_implicit.py

@@ -65,9 +65,6 @@ def update_nodes(self):
         # get number of collocation nodes for easier access
         M = self.coll.num_nodes
 
-        # update the MIN-SR-FLEX preconditioner
-        self.updateVariableCoeffs(L.status.sweep)
-
         # gather all terms which are known already (e.g. from the previous iteration)
         # this corresponds to u0 + QF(u^k) - QdF(u^k) + tau
 

pySDC/implementations/sweeper_classes/generic_implicit_MPI.py

@@ -209,9 +209,6 @@ def update_nodes(self):
         # only if the level has been touched before
         assert L.status.unlocked
 
-        # update the MIN-SR-FLEX preconditioner
-        self.updateVariableCoeffs(L.status.sweep)
-
         # gather all terms which are known already (e.g. from the previous iteration)
         # this corresponds to u0 + QF(u^k) - QdF(u^k) + tau
 

pySDC/implementations/sweeper_classes/imex_1st_order.py

@@ -70,9 +70,6 @@ def update_nodes(self):
 
         self.updateVariableCoeffs(L.status.sweep)
 
-        # update the MIN-SR-FLEX preconditioner
-        self.updateVariableCoeffs(L.status.sweep)
-
         # gather all terms which are known already (e.g. from the previous iteration)
         # this corresponds to u0 + QF(u^k) - QIFI(u^k) - QEFE(u^k) + tau