diff --git a/pySDC/core/level.py b/pySDC/core/level.py index d9f99630d6..b68135ac8e 100644 --- a/pySDC/core/level.py +++ b/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 diff --git a/pySDC/core/step.py b/pySDC/core/step.py index 4eae64c4a5..93da015734 100644 --- a/pySDC/core/step.py +++ b/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'], diff --git a/pySDC/core/sweeper.py b/pySDC/core/sweeper.py index 1db9eaf356..ab3af710e1 100644 --- a/pySDC/core/sweeper.py +++ b/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) diff --git a/pySDC/implementations/controller_classes/controller_MPI.py b/pySDC/implementations/controller_classes/controller_MPI.py index 829d384b23..1bb07b8581 100644 --- a/pySDC/implementations/controller_classes/controller_MPI.py +++ b/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: diff --git a/pySDC/implementations/controller_classes/controller_nonMPI.py b/pySDC/implementations/controller_classes/controller_nonMPI.py index 924a6a705f..334d34c3ea 100644 --- a/pySDC/implementations/controller_classes/controller_nonMPI.py +++ b/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) diff --git a/pySDC/implementations/convergence_controller_classes/adaptive_collocation.py b/pySDC/implementations/convergence_controller_classes/adaptive_collocation.py index 65ce0280b6..534265f7b8 100644 --- a/pySDC/implementations/convergence_controller_classes/adaptive_collocation.py +++ b/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): diff --git a/pySDC/implementations/sweeper_classes/generic_implicit.py b/pySDC/implementations/sweeper_classes/generic_implicit.py index cf33d5bd35..8762bfaf19 100644 --- a/pySDC/implementations/sweeper_classes/generic_implicit.py +++ b/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 diff --git a/pySDC/implementations/sweeper_classes/generic_implicit_MPI.py b/pySDC/implementations/sweeper_classes/generic_implicit_MPI.py index cf3cba2087..8d76a014b2 100644 --- a/pySDC/implementations/sweeper_classes/generic_implicit_MPI.py +++ b/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 diff --git a/pySDC/implementations/sweeper_classes/imex_1st_order.py b/pySDC/implementations/sweeper_classes/imex_1st_order.py index bd20e19953..da21c02e5e 100644 --- a/pySDC/implementations/sweeper_classes/imex_1st_order.py +++ b/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 diff --git a/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py b/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py index 4c16ca8b61..6dffb3cf42 100644 --- a/pySDC/implementations/sweeper_classes/imex_1st_order_MPI.py +++ b/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() diff --git a/pySDC/tests/test_sweepers/test_preconditioners.py b/pySDC/tests/test_sweepers/test_preconditioners.py index a313a0ad56..e20eb228a9 100644 --- a/pySDC/tests/test_sweepers/test_preconditioners.py +++ b/pySDC/tests/test_sweepers/test_preconditioners.py @@ -112,7 +112,7 @@ def test_FLEX_preconditioner_in_sweepers(imex, num_nodes, MPI=False): for k in range(1, level_params['nsweeps'] + 1): lvl.status.sweep = k - sweep.update_nodes() + sweep.updateVariableCoeffs(k) assert np.allclose( sweep.QI, sweep.get_Qdelta_implicit(sweeper_params['QI'], k) ), f'Got incorrect FLEX preconditioner in sweep {k}' @@ -216,4 +216,4 @@ def test_PIC(node_type, quad_type, M): test_LU('LEGENDRE', 'RADAU-RIGHT', 4) test_LU('EQUID', 'LOBATTO', 5) - test_FLEX_preconditioner_in_sweepers(True) + test_FLEX_preconditioner_in_sweepers(True, 4) diff --git a/pyproject.toml b/pyproject.toml index 37ad344377..47e393f676 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ 'sympy>=1.0', 'numba>=0.35', 'dill>=0.2.6', - 'qmat>=0.1.8', + 'qmat>=0.1.19', ] [project.urls]