Repository: dynamiqs/dynamiqs

Issue:
Stopping dq.mesolve bugs Jupyter interface
If dq.mesolve computation is stopped (either because of an error, timeout or user action),
when it is launched again, the progress bar does not remain on its line. We have to reboot the kernel every time.

![image](https://github.com/dynamiqs/dynamiqs/assets/146200418/fb3cd022-c9b0-44c4-80eb-8926207f2aa2)

Thanks,


Discussion:
Environment: Ubuntu VM, Jupyter Notebook

Resolving patch:
diff --git a/dynamiqs/solvers/ode/adjoint_ode_solver.py b/dynamiqs/solvers/ode/adjoint_ode_solver.py
--- a/dynamiqs/solvers/ode/adjoint_ode_solver.py
+++ b/dynamiqs/solvers/ode/adjoint_ode_solver.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import warnings
 from abc import abstractmethod
 from typing import Any
 
@@ -8,7 +7,6 @@
 import torch
 from torch import Tensor
 from torch.autograd.function import FunctionCtx, once_differentiable
-from tqdm.std import TqdmWarning
 
 from ..solver import AdjointSolver
 from ..utils.utils import tqdm
@@ -119,33 +117,37 @@ def backward(ctx: FunctionCtx, *grads: Tensor) -> tuple[None, Tensor, Tensor]:
             # initialize progress bar
             solver.pbar = tqdm(total=-t0, disable=not solver.options.verbose)
 
-            # initialize the ODE routine
-            t, y, a, *args = solver.init_augmented(t0, y0, a0)
-
-            # integrate the augmented equation backward between every saved state
-            for i, tnext in enumerate(tstop_bwd[1:]):
-                y, a, g, *args = solver.integrate_augmented(t, tnext, y, a, g, *args)
-
-                if i < len(tstop_bwd) - 2 or saved_ini:
-                    if solver.options.save_states:
-                        # replace y with its checkpointed version
-                        y = ysave[..., -i - 2, :, :]
-                        # update adjoint wrt this time point by adding dL / dy(t)
-                        a += grad_ysave[..., -i - 2, :, :]
-
-                    # update adjoint wrt this time point by adding dL / de(t)
-                    if len(solver.E) > 0:
-                        a += (grad_Esave[..., :, -i - 2, None, None] * solver.E.mH).sum(
-                            dim=-3
-                        )
-
-                # iterate time
-                t = tnext
-
-        # close progress bar
-        with warnings.catch_warnings():  # ignore tqdm precision overflow
-            warnings.simplefilter('ignore', TqdmWarning)
-            solver.pbar.close()
+            try:
+                # initialize the ODE routine
+                t, y, a, *args = solver.init_augmented(t0, y0, a0)
+
+                # integrate the augmented equation backward between every saved state
+                for i, tnext in enumerate(tstop_bwd[1:]):
+                    y, a, g, *args = solver.integrate_augmented(
+                        t, tnext, y, a, g, *args
+                    )
+
+                    if i < len(tstop_bwd) - 2 or saved_ini:
+                        if solver.options.save_states:
+                            # replace y with its checkpointed version
+                            y = ysave[..., -i - 2, :, :]
+                            # update adjoint wrt this time point by adding dL / dy(t)
+                            a += grad_ysave[..., -i - 2, :, :]
+
+                        # update adjoint wrt this time point by adding dL / de(t)
+                        if len(solver.E) > 0:
+                            a += (
+                                grad_Esave[..., :, -i - 2, None, None] * solver.E.mH
+                            ).sum(dim=-3)
+
+                    # iterate time
+                    t = tnext
+            finally:
+                # fix possible precision overflow
+                if solver.pbar.n > -t0:
+                    solver.pbar.n = -t0
+                # close the progress bar
+                solver.pbar.close()
 
         # convert gradients of real-valued parameters to real-valued gradients
         g = tuple(
diff --git a/dynamiqs/solvers/ode/ode_solver.py b/dynamiqs/solvers/ode/ode_solver.py
--- a/dynamiqs/solvers/ode/ode_solver.py
+++ b/dynamiqs/solvers/ode/ode_solver.py
@@ -1,9 +1,7 @@
-import warnings
 from abc import abstractmethod
 from typing import Any
 
 from torch import Tensor
-from tqdm import TqdmWarning
 
 from ..solver import AutogradSolver
 from ..utils.utils import tqdm
@@ -24,13 +22,19 @@ def init_forward(self) -> tuple:
         return self.t0, self.y0
 
     def run_autograd(self):
-        self.run_forward()
+        try:
+            self.run_forward()
+        finally:
+            # fix possible precision overflow
+            if self.pbar.n > self.tstop[-1]:
+                self.pbar.n = self.tstop[-1]
+            # close the progress bar
+            self.pbar.close()
 
     def run_forward(self):
         """Integrates the ODE forward from time `self.t0` to time `self.tstop[-1]`
         starting from initial state `self.y0`, and save the state for each time in
         `self.tstop`."""
-
         # initialize the ODE routine
         t, y, *args = self.init_forward()
 
@@ -40,11 +44,6 @@ def run_forward(self):
             self.save(y)
             t = tnext
 
-        # close the progress bar
-        with warnings.catch_warnings():  # ignore tqdm precision overflow
-            warnings.simplefilter('ignore', TqdmWarning)
-            self.pbar.close()
-
     @abstractmethod
     def integrate(self, t0: float, t1: float, y: Tensor, *args: Any) -> tuple:
         """Integrates the ODE forward from time `t0` to time `t1` with initial state


Test patch: