Trial time

src/porepy/models/model_runner.py

@@ -4,57 +4,57 @@
 
 import logging
 import warnings
-from enum import StrEnum
-from typing import TYPE_CHECKING, Optional
+from abc import ABC
+from dataclasses import dataclass
+from typing import Optional
 
 import numpy as np
 
 import porepy as pp
+from porepy.models.solution_strategy import SolutionStrategy
+from porepy.time.time_step_status import TimeStepperStatusFailure
+from porepy.time.time_stepper import TimeStepper
 from porepy.utils.ui_and_logging import DummyProgressBar
 from porepy.utils.ui_and_logging import (
     logging_redirect_tqdm_with_level as logging_redirect_tqdm,
 )
 from porepy.utils.ui_and_logging import progressbar_class
 
-if TYPE_CHECKING:
-    from porepy.numerics.nonlinear.convergence_check import SolverStatus
-
-__all__ = ["SimulationStatus", "ModelRunner"]
+__all__ = [
+    "ModelRunnerStatus",
+    "ModelRunnerStatusSuccess",
+    "ModelRunnerStatusFailure",
+    "ModelRunner",
+]
 
 # Module-wide logger
 logger = logging.getLogger(__name__)
 
 
-class SimulationStatus(StrEnum):
-    """Enumeration of potential simulation statuses."""
+@dataclass
+class ModelRunnerStatus(ABC):
+    """A status object used to indicate the ModelRunner state."""
+
+    def is_success(self) -> bool:
+        """Whether the simulation finished successfully."""
+        return isinstance(self, ModelRunnerStatusSuccess)
 
-    IN_PROGRESS = "in_progress"
-    """Simulation is currently in progress and in a nominal state."""
-    SUCCESSFUL = "successful"
-    """Simulation completed with success."""
-    FAILED = "failed"
-    """Simulation is currently in progress and in a failed state."""
-    STOPPED = "stopped"
-    """Simulation was stopped due to an error."""
+    def is_failure(self) -> bool:
+        """Whether the simulation finished with a failure."""
+        return isinstance(self, ModelRunnerStatusFailure)
 
-    def __str__(self):
-        return self.value
 
-    def is_in_progress(self) -> bool:
-        """Check if the status indicates an ongoing simulation."""
-        return self == SimulationStatus.IN_PROGRESS
+@dataclass
+class ModelRunnerStatusSuccess(ModelRunnerStatus):
+    """A status object that indicates that the simulation finished successfully."""
 
-    def is_successful(self) -> bool:
-        """Check if the status indicates a successful simulation."""
-        return self == SimulationStatus.SUCCESSFUL
 
-    def is_failed(self) -> bool:
-        """Check if the status indicates a failed simulation."""
-        return self == SimulationStatus.FAILED
+@dataclass
+class ModelRunnerStatusFailure(ModelRunnerStatus):
+    """A status object that indicates that the simulation finished with a failure."""
 
-    def is_stopped(self) -> bool:
-        """Check if the status indicates a stopped simulation."""
-        return self == SimulationStatus.STOPPED
+    reason: str
+    "Reason why the model runner failed."
 
 
 def run_stationary_model(model, params: dict) -> None:
@@ -146,7 +146,12 @@ class ModelRunner:
 
     """
 
-    def __init__(self, model: pp.SolutionStrategy, params: dict | None = None) -> None:
+    def __init__(
+        self,
+        model: SolutionStrategy,
+        params: Optional[dict] = None,
+        time_stepper: Optional[TimeStepper] = None,
+    ) -> None:
         self.params = params if isinstance(params, dict) else {}
         """Parameters passed at instantiation."""
 
@@ -156,6 +161,13 @@ def __init__(self, model: pp.SolutionStrategy, params: dict | None = None) -> No
         self.solver: pp.NewtonSolver | pp.LinearSolver
         """Solver instance, set in :meth:`set_solver`."""
 
+        # Construct the default if not provided. This time stepper is constructed even
+        # for a stationary problem, but used only for time-dependent problems.
+        if time_stepper is None:
+            time_stepper = TimeStepper(time_manager=model.time_manager)
+        self.time_stepper: TimeStepper = time_stepper
+        """Responsible for the time stepping logic."""
+
         if self.params.get("prepare_simulation", True):
             self.model.prepare_simulation()
 
@@ -168,7 +180,8 @@ def __init__(self, model: pp.SolutionStrategy, params: dict | None = None) -> No
 
         self.set_solver()
 
-        self.init_time_progressbar()
+        if self._is_time_dependent:
+            self.init_time_progressbar()
 
     def set_solver(self) -> None:
         """Choose between linear and non-linear solver and set :attr:`solver`.
@@ -212,9 +225,6 @@ def init_time_progressbar(self) -> None:
                 " loop will run without progress bars."
             )
 
-        # Save initial time step size; used for progress bar updates.
-        self._dt_0: float = self.model.time_manager.dt
-
         # To display nested ``tqdm`` bars in the correct order, their positions have to
         # be specified. The orders are increasing, i.e., 0 is the lowest level, then 1.
         # Position is passed via '_nl_progress_bar_position' when calling 'NewtonSolver'
@@ -232,7 +242,7 @@ def init_time_progressbar(self) -> None:
                         self.model.time_manager.schedule[-1]
                         - self.model.time_manager.schedule[0]
                     )
-                    / self._dt_0
+                    / self.model.time_manager.dt
                 )
             )
 
@@ -246,177 +256,73 @@ def init_time_progressbar(self) -> None:
         else:
             self.time_progressbar = DummyProgressBar()
 
-    def run(self, *args, **kwargs) -> None:
-        """Runs the model as specified."""
+        self.use_progress_bar = use_progress_bar
 
+    def run(self) -> ModelRunnerStatus:
+        """Run the model (stationary or time-dependent)."""
+        # Run simulation.
         if self._is_time_dependent:
-            # Redirect the root logger, to avoid logger-progressbars interference.
-            with logging_redirect_tqdm([logging.root]):
-                # Time loop.
-                while not self.model.time_manager.final_time_reached():
-                    self.before_time_step()
-                    solver_status = self.solver.solve(self.model)
-                    simulation_status = self.after_time_step(solver_status)
-                    if (
-                        simulation_status.is_successful()
-                        or simulation_status.is_stopped()
-                    ):
-                        break
+            simulation_status = self._run_time_dependent()
         else:
-            solver_status = self.solver.solve(self.model)
-            simulation_status = self.after_stationary_solve(solver_status)
+            simulation_status = self._run_stationary()
 
+        # Clean up model after simulation.
         self.model.after_simulation()
 
-    def before_time_step(self) -> None:
-        """Method to be executed at the beginning of each time step.
-
-        Increases the time and sets the model's AD time step value.
-        Executes :meth:`~porepy.models.solution_strategy.ModelSolverInterface.
-        before_time_step` and logs the progress.
-
-        """
-        # Increase the simulation time.
-        self.model.time_manager.increase_time()
-        self.model.time_manager.increase_time_index()
-        # Prepare model.
-        self.model.before_time_step()
-
-        # Logging and progressbar update.
-        logger.info(
-            f"\nTime step {self.model.time_manager.time_index} at time"
-            + f" {self.model.time_manager.time:.1e}"
-            + f" of {self.model.time_manager.time_final:.1e}"
-            + f" with time step {self.model.time_manager.dt:.1e}"
-        )
-        self.time_progressbar.set_postfix_str(
-            f"Time step size {self.model.time_manager.dt:.2e}"
-        )
-
-    def after_time_step(self, solver_status: SolverStatus) -> SimulationStatus:
-        """Method to be executed at the end of each time step.
-
-        React to solver status, updates the time step size and logs the progress.
-
-        Parameters:
-            solver_status: Status of the time step, as returned by the solver.
-
-        Raises:
-            RuntimeError: If the simulation is stopped due to failures in solver and
-                time step recomputation.
-
-        Returns:
-            A simulation-status object of the time step, which can be used to determine
-            whether to continue the simulation or not.
-
-        """
-        simulation_status: SimulationStatus
-
-        if solver_status.is_successful():
-            # Conclude simulation status if final time reached.
-            simulation_status = (
-                SimulationStatus.SUCCESSFUL
-                if self.model.time_manager.final_time_reached()
-                else SimulationStatus.IN_PROGRESS
-            )
-            self.model.after_time_step_convergence()
-
-            # Need to log before updating the time step size.
-            self.logging(simulation_status)
-
-            # Update the time step magnitude if the dynamic scheme is used.
-            if not self.model.time_manager.is_constant:
-                assert isinstance(
-                    self.model.nonlinear_solver_statistics, pp.NonlinearSolverStatistics
-                )  # For type checking, to ensure the method is available.
-                self.model.time_manager.compute_time_step(
-                    iterations=self.model.nonlinear_solver_statistics.num_iterations
-                )
-
-            # Update progressbar length.
-            self.time_progressbar.update(n=self.model.time_manager.dt / self._dt_0)
-
-        elif solver_status.is_failed() or solver_status.is_stopped():
-            # If solver failed or stopped, base notion to propagate is that the
-            # simulation failed in the current time step.
-            simulation_status = SimulationStatus.FAILED
-
-            # If constant time step, simulation will be stopped.
-            if self.model.time_manager.is_constant:
-                logger.warning(
-                    "Solver failed to converge but time step size is constant and "
-                    "cannot be reduced."
-                )
-                simulation_status = SimulationStatus.STOPPED
-                self.logging(simulation_status)
-
-            # Else recompute time step and attempt to solve again.
-            else:
-                # This calls
-                # ``time_manager._adaptation_based_on_recomputation``, which substracts
-                # the current ``dt`` from the simulation time, computes a shorter
-                # ``dt``, and adds the updated ``dt`` to the simulation time again.
-                # It will also raise a TimeSteppingError if the minimal time step
-                # is reached.
-                try:
-                    self.model.after_time_step_failure()
-                    # Need to log before updating the time step size since the failed
-                    # time step is part of the log.
-                    self.logging(simulation_status)
-                    # Update the time step size for the next attempt.
-                    self.model.time_manager.compute_time_step(recompute_solution=True)
-                # If time step recomputations fails for any reason, stop the simulation.
-                except Exception as e:
-                    # Redirect the exception as a warning, and give the control to
-                    # the ModelRunner to stop the simulation.
-                    logger.warning(str(e))
-                    simulation_status = SimulationStatus.STOPPED
-                    self.logging(simulation_status)
-
-        else:
-            raise ValueError(f"Unrecognized solver stats {solver_status}.")
-
-        if simulation_status.is_stopped():
-            logger.warning("Simulation stopped.")
-            raise RuntimeError("Simulation stopped.")
+        if simulation_status.is_failure():
+            raise RuntimeError(simulation_status)
 
         return simulation_status
 
-    def after_stationary_solve(self, solver_status: SolverStatus) -> SimulationStatus:
-        """Method to be executed at the end of a stationary solve.
-
-        React to solver status and logs the progress.
-
-        Parameters:
-            solver_status: Status of the solve, as returned by the solver.
-
-        Returns:
-            A simulation-status enum which can be used to determine whether the
-            simulation was successful or not.
+    def _run_stationary(self) -> ModelRunnerStatus:
+        """Run a stationary model."""
+        # Perform stationary solve.
+        convergence_status = self.solver.solve(self.model)
 
-        """
-        if solver_status.is_successful():
+        # Conclude the simulation status based on the solver status.
+        if convergence_status.is_converged():
             # NOTE: time_step_convergence can be considered a misnomer.
             # But technically this is the only time we solve for. Thus we reuse the
             # method to set the solution and save data.
             self.model.after_time_step_convergence()
-            simulation_status = SimulationStatus.SUCCESSFUL
+            return ModelRunnerStatusSuccess()
         else:
             self.model.after_time_step_failure()
-            simulation_status = SimulationStatus.STOPPED
+            return ModelRunnerStatusFailure("Solver did not converge.")
 
-        return simulation_status
+    def _run_time_dependent(self) -> ModelRunnerStatus:
+        """Run a time-dependent model with trial-based time stepping."""
+
+        with logging_redirect_tqdm([logging.root]):
+            while not self.model.time_manager.final_time_reached():
+                # Perform the time step.
+                time_step_status = self.time_stepper.perform_time_step(
+                    self.model, self.solver
+                )
+
+                # Progressbar update.
+                self.update_time_progressbar()
+
+                # Abort simulation if time step was stopped.
+                match time_step_status:
+                    case TimeStepperStatusFailure(nonlinear_solver_status, reason):
+                        logger.error(f"Time stepping failed: {reason}")
+                        return ModelRunnerStatusFailure(reason=reason)
 
-    def logging(self, simulation_status: SimulationStatus) -> None:
-        self.model.nonlinear_solver_statistics.log_simulation_status(simulation_status)
-        self.model.nonlinear_solver_statistics.log_mesh_information(
-            self.model.mdg.subdomains()
+            # Conclude the simulation status.
+            if self.model.time_manager.final_time_reached():
+                return ModelRunnerStatusSuccess()
+            return ModelRunnerStatusFailure("Final time was not reached.")
+
+    def update_time_progressbar(self) -> None:
+        """Update the time progressbar with the current time and time step size."""
+        self.time_progressbar.set_postfix_str(
+            f"Time step size {self.model.time_manager.dt:.2e}"
         )
-        if self._is_time_dependent:
-            assert isinstance(self.model.nonlinear_solver_statistics, pp.TimeStatistics)
-            self.model.nonlinear_solver_statistics.log_time_information(
-                self.model.time_manager.time_index,
-                self.model.time_manager.time,
-                self.model.time_manager.dt,
-                self.model.time_manager.final_time_reached(),
+        self.time_progressbar.update(
+            n=np.round(
+                self.model.time_manager.time
+                / self.model.time_manager.time_final
+                * self.time_progressbar.total
             )
+        )

src/porepy/models/protocol.py

@@ -666,6 +666,9 @@ def _is_reference_component_eliminated(self) -> bool:
             """Returns True if ``params['eliminate_reference_component'] == True`.
             Defaults to True."""
 
+        def before_time_step(self) -> None:
+            """Called at the start of each time step by model runners."""
+
         def before_nonlinear_loop(self) -> None:
             """Method to be called before the non-linear loop.
 
@@ -689,6 +692,12 @@ def after_nonlinear_convergence(self) -> None:
 
             """
 
+        def after_time_step_convergence(self) -> None:
+            """Called after a new time step solution has been achieved."""
+
+        def after_time_step_failure(self) -> None:
+            """Called after a time step has failed to converge."""
+
         def set_nonlinear_discretizations(self) -> None:
             """Set the list of nonlinear discretizations.