Progress Bar (Tqdm) (#104)

* implement progress bar

* black, flake8

* autodetect environment

* variable names

* black, flake8

* argument assignment

* inbuild functions, import, ui

* resolve conflict

* revert xr_accessor

* added documentation

* added tests

* added more tests

* update api

* update test section
Test names have been specified, 3rd party test logic removed, tests are
skipped if not run on python3.7, tqdm dependency has been added to pass
travis, minor modifications.

* update dependency

* rework pytest skipif

* rework import

* fix assignment

* rework api

* reinstate progress

* minor fixes

* reworked tests

* black

* include custom description, simulation runtime

* Minor updates
Small phrasing issues, additional testing, tqdm built-in functionality,
update docstring, remove unnecessary code, change logic.

Co-authored-by: Raphael Lange <rlange@sec55-dynip-223.gfz-potsdam.de>
Co-authored-by: Raphael Lange <rlange@macbook-pro-2.localdomain>

ci/requirements/doc.yml

@@ -17,5 +17,6 @@ dependencies:
       - nbconvert==5.6.0
       - sphinx==1.8.5
       - toolz
+      - tqdm==4.40.2
       - xarray==0.13.0
       - zarr==2.4.0

ci/requirements/py37.yml

@@ -12,6 +12,7 @@ dependencies:
   - python-graphviz
   - ipython
   - zarr
+  - tqdm
   - pip:
     - coveralls
     - pytest-cov

doc/api.rst

@@ -170,12 +170,13 @@ listed below. These are defined in ``xsimlab.validators``.
 
 .. _`attrs' validators`: https://www.attrs.org/en/stable/examples.html#validators
 
-Runtime hooks
-=============
+Runtime monitoring
+==================
 
 .. currentmodule:: xsimlab
 .. autosummary::
    :toctree: _api_generated/
 
    runtime_hook
    RuntimeHook
+   ProgressBar

doc/monitor.rst

@@ -6,6 +6,9 @@ Monitor Model Runs
 Models may be complex, built from many processes and may take a while to
 run. xarray-simlab provides functionality to help in monitoring model runs.
 
+This section demonstrates how to use the built-in progress bar. Moreover, 
+it exemplifies how to create your own custom monitoring.
+
 .. ipython:: python
    :suppress:
 
@@ -34,6 +37,51 @@ The following imports are necessary for the examples below.
         },
     )
 
+
+Progress bar
+------------
+
+:class:`~xsimlab.ProgressBar` is based on the `Tqdm`_ package and allows to track
+the progress of simulation runs in ``xarray-simlab``.
+It can be used as a context manager around simulation calls:
+
+.. _Tqdm: https://github.com/tqdm/tqdm/
+
+.. ipython::
+
+   In [2]: with xs.ProgressBar():
+      ...:     out_ds = in_ds.xsimlab.run(model=model2)
+
+Alternatively, you can pass the progress bar via the ``hooks`` argument or use the
+``register`` method (for more information, refer to the :ref:`custom_runtime_hooks` subsection)
+
+``ProgressBar`` and the underlying Tqdm is built to work with different Python
+interfaces. Use the optional argument ``frontend`` according to your
+development environment.
+
+- ``auto``: (default) Automatically detects environment.
+- ``console``: When Python is run from the command line.
+- ``gui``: Tqdm provides a gui version. According to the developers, this is
+  still an experimental feature.
+- ``notebook``: For use in a IPython/Jupyter notebook.
+
+Additionally, you can customize the built-in progress bar, by supplying a
+keyworded argument list to ``ProgressBar``, e.g.:
+
+.. ipython::
+
+   In [4]: with xs.ProgressBar(bar_format="{r_bar}"):
+      ...:     out_ds = in_ds.xsimlab.run(model=model2)
+
+For a full list of customization options, refer to the `Tqdm documentation`_
+
+Note: The ``total`` argument cannot be changed to ensure best performance and
+functionality.
+
+.. _Tqdm documentation: https://tqdm.github.io
+
+.. _custom_runtime_hooks:
+
 Custom runtime hooks
 --------------------
 

doc/whats_new.rst

@@ -71,6 +71,8 @@ Enhancements
 - Added some useful properties and methods to the ``xarray.Dataset.xsimlab``
   extension (:issue:`103`).
 - Save model inputs/outputs using zarr (:issue:`102`).
+- Added :class:`~xsimlab.progress.ProgressBar` to track simulation progress
+  (:issue:`104`).
 
 Bug fixes
 ~~~~~~~~~

xsimlab/__init__.py

@@ -13,6 +13,7 @@
     runtime,
     variable_info,
 )
+from .progress import ProgressBar
 from .variable import variable, index, on_demand, foreign, group
 from .xr_accessor import SimlabAccessor, create_setup
 

xsimlab/drivers.py

@@ -26,6 +26,7 @@ class RuntimeContext(Mapping[str, Any]):
         "sim_start",
         "sim_end",
         "step",
+        "nsteps",
         "step_start",
         "step_end",
         "step_delta",
@@ -270,6 +271,7 @@ def _get_runtime_datasets(self):
 
         init_data_vars = {
             "_sim_start": mclock_coord[0],
+            "_nsteps": self.dataset.xsimlab.nsteps,
             "_sim_end": mclock_coord[-1],
         }
 
@@ -339,7 +341,9 @@ def run_model(self):
         validate_all = self._validate_option is ValidateOption.ALL
 
         runtime_context = RuntimeContext(
-            sim_start=ds_init["_sim_start"].values, sim_end=ds_init["_sim_end"].values,
+            sim_start=ds_init["_sim_start"].values,
+            nsteps=ds_init["_nsteps"].values,
+            sim_end=ds_init["_sim_end"].values,
         )
 
         in_vars = self._get_input_vars(ds_init)

xsimlab/progress.py

@@ -0,0 +1,91 @@
+from xsimlab.hook import RuntimeHook, runtime_hook
+
+
+class ProgressBar(RuntimeHook):
+    """
+    Progress bar implementation using the tqdm package.
+
+    Parameters
+    ----------
+    frontend : {"auto", "console", "gui", "notebook"}, optional
+        Selects a frontend for displaying the progress bar. By default ("auto"),
+        the frontend is chosen by guessing in which environment the simulation
+        is run. The "console" frontend displays an ascii progress bar, while the
+        "gui" frontend is based on matplotlib and the "notebook" frontend is based
+        on ipywidgets.
+    **kwargs : dict, optional
+        Arbitrary keyword arguments for progress bar customization.
+
+    Examples
+    --------
+    :class:`ProgressBar` takes full advantage of :class:`RuntimeHook`.
+
+    Call it as part of :func:`run`:
+    >>> out_ds = in_ds.xsimlab.run(model=model, hooks=[xs.ProgressBar()])
+
+    In a context manager using the `with` statement`:
+    >>> with xs.ProgressBar():
+    ...    out_ds = in_ds.xsimlab.run(model=model)
+
+    Globally with `register` method:
+    >>> pbar = xs.ProgressBar()
+    >>> pbar.register()
+    >>> out_ds = in_ds.xsimlab.run(model=model)
+    >>> pbar.unregister()
+
+    For additional customization, see: https://tqdm.github.io/docs/tqdm/
+    """
+
+    def __init__(self, frontend="auto", **kwargs):
+        if frontend == "auto":
+            from tqdm.auto import tqdm
+        elif frontend == "console":
+            from tqdm import tqdm
+        elif frontend == "gui":
+            from tqdm.gui import tqdm
+        elif frontend == "notebook":
+            from tqdm.notebook import tqdm
+        else:
+            raise ValueError(
+                f"Frontend argument {frontend!r} not supported. Please select one of the following: {', '.join(['auto', 'console', 'gui', 'notebook'])}"
+            )
+
+        self.custom_description = False
+        if "desc" in kwargs.keys():
+            self.custom_description = True
+
+        self.tqdm = tqdm
+        self.tqdm_kwargs = {"bar_format": "{bar} {percentage:3.0f}% | {desc} "}
+        self.tqdm_kwargs.update(kwargs)
+
+    @runtime_hook("initialize", trigger="pre")
+    def init_bar(self, model, context, state):
+        if self.custom_description:
+            self.tqdm_kwargs.update(total=context["nsteps"] + 2)
+        else:
+            self.tqdm_kwargs.update(total=context["nsteps"] + 2, desc="initialize")
+        self.pbar_model = self.tqdm(**self.tqdm_kwargs)
+
+    @runtime_hook("initialize", trigger="post")
+    def update_init(self, mode, context, state):
+        self.pbar_model.update(1)
+
+    @runtime_hook("run_step", trigger="post")
+    def update_runstep(self, mode, context, state):
+        if not self.custom_description:
+            self.pbar_model.set_description_str(
+                f"run step {context['step']}/{context['nsteps']}"
+            )
+        self.pbar_model.update(1)
+
+    @runtime_hook("finalize", trigger="pre")
+    def update_finalize(self, model, context, state):
+        if not self.custom_description:
+            self.pbar_model.set_description_str("finalize")
+
+    @runtime_hook("finalize", trigger="post")
+    def close_bar(self, model, context, state):
+        self.pbar_model.update(1)
+        elapsed_time = self.tqdm.format_interval(self.pbar_model.format_dict["elapsed"])
+        self.pbar_model.set_description_str(f"Simulation finished in {elapsed_time}")
+        self.pbar_model.close()

xsimlab/tests/__init__.py

@@ -0,0 +1,15 @@
+import importlib
+import pytest
+
+
+def _importorskip(modname):
+    try:
+        mod = importlib.import_module(modname)
+        has = True
+    except ImportError:
+        has = False
+    func = pytest.mark.skipif(not has, reason=f"requires {modname}")
+    return has, func
+
+
+has_tqdm, requires_tqdm = _importorskip("tqdm")

xsimlab/tests/test_progress.py

@@ -0,0 +1,39 @@
+import importlib
+import pytest
+
+from . import has_tqdm
+import xsimlab as xs
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+@pytest.mark.parametrize(
+    "frontend,tqdm_module",
+    [
+        ("auto", "tqdm"),  # assume tests are run in a terminal evironment
+        ("console", "tqdm"),
+        ("gui", "tqdm.gui"),
+        ("notebook", "tqdm.notebook"),
+    ],
+)
+def test_progress_bar_init(frontend, tqdm_module):
+    pbar = xs.ProgressBar(frontend=frontend)
+    tqdm = importlib.import_module(tqdm_module)
+
+    assert pbar.tqdm is tqdm.tqdm
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+@pytest.mark.parametrize("kw", [{}, {"bar_format": "{bar}"}])
+def test_progress_bar_init_kwargs(kw):
+    pbar = xs.ProgressBar(**kw)
+
+    assert "bar_format" in pbar.tqdm_kwargs
+
+    if "bar_format" in kw:
+        assert pbar.tqdm_kwargs["bar_format"] == kw["bar_format"]
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_init_error(in_dataset, model):
+    with pytest.raises(ValueError, match=r".*not supported.*"):
+        pbar = xs.ProgressBar(frontend="invalid_frontend")