Source code for core.level

from typing import Any, Dict, List, Optional, Type
from pySDC.core.sweeper import Sweeper
from pySDC.core.problem import Problem

from pySDC.helpers.pysdc_helper import FrozenClass


# short helper class to add params as attributes
class _Pars(FrozenClass):
    def __init__(self, params: Dict[str, Any]) -> None:
        self.dt: Optional[float] = None
        self.dt_initial: Optional[float] = None
        self.restol: float = -1.0
        self.nsweeps: int = 1
        self.residual_type: str = 'full_abs'
        for k, v in params.items():
            setattr(self, k, v)
        # freeze class, no further attributes allowed from this point
        self._freeze()

        self.dt_initial = self.dt * 1.0 if self.dt is not None else None


# short helper class to bundle all status variables
class _Status(FrozenClass):
    """
    This class carries the status of the level. All variables that the core SDC / PFASST functionality depend on are
    initialized here.
    """

    def __init__(self) -> None:
        self.residual: Optional[float] = None
        self.unlocked: bool = False
        self.updated: bool = False
        self.time: Optional[float] = None
        self.dt_new: Optional[float] = None
        self.sweep: Optional[int] = None
        # freeze class, no further attributes allowed from this point
        self._freeze()




class Level(FrozenClass):
    """
    Level class containing all management functionality for a single level

    A level contains all data structures, types and objects to perform sweeps on this particular level. It does not
    know about other levels.

    Attributes:
        params (__Pars): parameter object containing the custom parameters passed by the user
        status (__Status): status object
        level_index (int): custom string naming this level
        uend: dof values at the right end point of the interval
        u (list of dtype_u): dof values at the nodes
        uold (list of dtype_u): copy of dof values for saving data during restriction)
        f (list of dtype_f): RHS values at the nodes
        fold (list of dtype_f): copy of RHS values for saving data during restriction
        tau (list of dtype_u): FAS correction, allocated via step class if necessary
    """

    def __init__(
        self,
        problem_class: Type[Problem],
        problem_params: Dict[str, Any],
        sweeper_class: Type[Sweeper],
        sweeper_params: Dict[str, Any],
        level_params: Dict[str, Any],
        level_index: int,
    ) -> None:
        """
        Initialization routine

        Args:
            problem_class: problem class
            problem_params (dict): parameters for the problem to be initialized
            sweeper_class: sweeper class
            sweeper_params (dict): parameters for the sweeper (contains collocation)
            level_params (dict): parameters given by the user, will be added as attributes
            level_index (int): custom name for this level
        """

        # set level parameters and status
        self.params: _Pars = _Pars(level_params)
        self.status: _Status = _Status()

        # instantiate sweeper, problem and hooks
        self.__sweep: Sweeper = sweeper_class(sweeper_params, self)
        self.__prob: Problem = problem_class(**problem_params)

        # set name
        self.level_index: int = level_index

        # empty data at the nodes, the right end point and tau
        self.uend: Optional[Any] = None
        self.u: List[Optional[Any]] = [None] * (self.sweep.coll.num_nodes + 1)
        self.uold: List[Optional[Any]] = [None] * (self.sweep.coll.num_nodes + 1)
        self.u_avg: List[Optional[Any]] = [None] * self.sweep.coll.num_nodes
        self.residual: List[Optional[Any]] = [None] * self.sweep.coll.num_nodes
        self.increment: List[Optional[Any]] = [None] * self.sweep.coll.num_nodes
        self.f: List[Optional[Any]] = [None] * (self.sweep.coll.num_nodes + 1)
        self.fold: List[Optional[Any]] = [None] * (self.sweep.coll.num_nodes + 1)

        self.tau: List[Optional[Any]] = [None] * self.sweep.coll.num_nodes

        self.__tag: Optional[Any] = None

        # freeze class, no further attributes allowed from this point
        self._freeze()



    def reset_level(self, reset_status: bool = True) -> None:
        """
        Routine to clean-up the level for the next time step

        Args:
            reset_status (bool): Reset the status or only the solution

        Returns:
            None
        """

        # reset status
        if reset_status:
            self.status = _Status()

        # all data back to None
        self.uend = None
        self.u = [None] * (self.sweep.coll.num_nodes + 1)
        self.uold = [None] * (self.sweep.coll.num_nodes + 1)
        self.f = [None] * (self.sweep.coll.num_nodes + 1)
        self.fold = [None] * (self.sweep.coll.num_nodes + 1)
        self.tau = [None] * self.sweep.coll.num_nodes


    @property
    def sweep(self) -> Sweeper:
        """
        Getter for the sweeper

        Returns:
            pySDC.Sweeper.sweeper: the sweeper associated to this level
        """
        return self.__sweep

    @property
    def prob(self) -> Problem:
        """
        Getter for the problem

        Returns:
            pySDC.Problem.ptype: the problem associated to this level
        """
        return self.__prob

    @property
    def time(self) -> Optional[float]:
        """
        Meta-getter for the current time

        Returns:
            float: referencing status time for convenience
        """
        return self.status.time

    @property
    def dt(self) -> Optional[float]:
        """
        Meta-getter for the time-step size

        Returns:
            float: referencing dt from parameters for convenience
        """
        return self.params.dt

    @property
    def tag(self) -> Optional[Any]:
        """
        Getter for tag

        Returns:
            tag for sending/receiving
        """
        return self.__tag

    @tag.setter
    def tag(self, t: Any) -> None:
        """
        Setter for tag

        Args:
            t: new tag for sending/receiving
        """
        self.__tag = t