More genericity for QDelta coefficients (Parallel-in-Time#584)

* TL: more generic and efficients QDelta coefficients computation + some refactoring

* TL: trying to solve circular imports

* TL: bugfixes

* TL: aaaaaah 😭

* TL: missing this

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/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

@@ -72,9 +72,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) - QIFI(u^k) - QEFE(u^k) + tau
 

pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py

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