Adding Progress Bars (#1703)

* Customize bar format

* Display bars again after reset when full

* Get data from configuration file

* Track all packets at once and don't reset the progress bar

* Close last container and add docstrings

* Add option to enable/disble the progress bar

* Display container only in notebooks

* Add a progress bar to track iterations

* Align bars

* Refactor code

* Format using black

* Add message regarding progress bars in the quickstart notebook

* Demonstrate progress bars in the documentation

* Change show_progress_bar to show_progress_bars

* [build docs]

docs/io/output/index.rst

@@ -11,4 +11,5 @@ In addition to the widgets, TARDIS can output information in several other forms
     access_iterations
     to_hdf
     callback
-    vpacket_logging
+    vpacket_logging
+    progress_bars

docs/io/output/progress_bars.ipynb

@@ -0,0 +1,63 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cc613a48",
+   "metadata": {},
+   "source": [
+    "### Progress Bars for Simulation Run ###\n",
+    "TARDIS displays progress bars by default to track the simulation. The progress bars are not displayed in the documentation but show up when you run the notebook.\n",
+    "\n",
+    "![TARDIS Progress Bars](../images/progress_bars_demo.gif)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6e6b8b8f",
+   "metadata": {},
+   "source": [
+    "You can disable the progress bars as well."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "319759fd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tardis import run_tardis\n",
+    "sim = run_tardis('tardis_example.yml', show_progress_bars = False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e51d94c9",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/quickstart/quickstart.ipynb

@@ -70,6 +70,19 @@
     "#### Running the simulation (long output) ####"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<div class=\"alert alert-info\">\n",
+    "\n",
+    "Note\n",
+    "    \n",
+    "The progress of the simulation can be tracked using progress bars which are displayed when the notebook is run, but are not displayed in the documentation. \n",
+    "    \n",
+    "</div>"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

tardis/base.py

@@ -15,6 +15,7 @@ def run_tardis(
     show_convergence_plots=True,
     log_level=None,
     specific_log_level=None,
+    show_progress_bars=True,
     **kwargs,
 ):
     """
@@ -52,6 +53,8 @@ def run_tardis(
         The default value None means that the `specific_log_level` specified in the configuration file will be used.
     show_convergence_plots : bool, default: True, optional
         Option to enable tardis convergence plots.
+    show_progress_bars : bool, default: True, optional
+        Option to enable the progress bar.
     **kwargs : dict, optional
         Optional keyword arguments including those
         supported by :obj:`tardis.visualization.tools.convergence_plot.ConvergencePlots`.
@@ -101,6 +104,7 @@ def run_tardis(
         atom_data=atom_data,
         virtual_packet_logging=virtual_packet_logging,
         show_convergence_plots=show_convergence_plots,
+        show_progress_bars=show_progress_bars,
         **kwargs,
     )
     for cb in simulation_callbacks:

tardis/montecarlo/base.py

@@ -268,6 +268,8 @@ def run(
         nthreads=1,
         last_run=False,
         iteration=0,
+        total_iterations=0,
+        show_progress_bars=True,
     ):
         """
         Run the montecarlo calculation
@@ -280,6 +282,8 @@ def run(
         no_of_virtual_packets : int
         nthreads : int
         last_run : bool
+        total_iterations : int
+            The total number of iterations in the simulation.
 
         Returns
         -------
@@ -301,7 +305,15 @@ def run(
         )
 
         configuration_initialize(self, no_of_virtual_packets)
-        montecarlo_radial1d(model, plasma, self)
+        montecarlo_radial1d(
+            model,
+            plasma,
+            iteration,
+            no_of_packets,
+            total_iterations,
+            show_progress_bars,
+            self,
+        )
         self._integrator = FormalIntegrator(model, plasma, self)
         # montecarlo.montecarlo_radial1d(
         #    model, plasma, self,

tardis/montecarlo/montecarlo_numba/base.py

@@ -1,4 +1,4 @@
-from numba import prange, njit, jit
+from numba import prange, njit, jit, objmode
 import logging
 import numpy as np
 
@@ -26,9 +26,18 @@
 )
 from tardis.montecarlo.montecarlo_numba import njit_dict
 from numba.typed import List
+from tardis.util.base import update_iterations_pbar, update_packet_pbar
 
 
-def montecarlo_radial1d(model, plasma, runner):
+def montecarlo_radial1d(
+    model,
+    plasma,
+    iteration,
+    no_of_packets,
+    total_iterations,
+    show_progress_bars,
+    runner,
+):
     packet_collection = PacketCollection(
         runner.input_r,
         runner.input_nu,
@@ -76,6 +85,10 @@ def montecarlo_radial1d(model, plasma, runner):
         runner.spectrum_frequency.value,
         number_of_vpackets,
         packet_seeds,
+        iteration=iteration,
+        show_progress_bars=show_progress_bars,
+        no_of_packets=no_of_packets,
+        total_iterations=total_iterations,
     )
 
     runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist
@@ -109,6 +122,7 @@ def montecarlo_radial1d(model, plasma, runner):
         runner.virt_packet_last_line_interaction_out_id = np.concatenate(
             np.array(virt_packet_last_line_interaction_out_id)
         ).ravel()
+    update_iterations_pbar(1)
 
 
 @njit(**njit_dict)
@@ -120,6 +134,10 @@ def montecarlo_main_loop(
     spectrum_frequency,
     number_of_vpackets,
     packet_seeds,
+    iteration,
+    show_progress_bars,
+    no_of_packets,
+    total_iterations,
 ):
     """
     This is the main loop of the MonteCarlo routine that generates packets
@@ -178,6 +196,16 @@ def montecarlo_main_loop(
     virt_packet_last_line_interaction_out_id = []
 
     for i in prange(len(output_nus)):
+        if show_progress_bars:
+            with objmode:
+                update_amount  = 1
+                update_packet_pbar(
+                    update_amount,
+                    current_iteration=iteration,
+                    no_of_packets=no_of_packets,
+                    total_iterations=total_iterations,
+                )
+
         if montecarlo_configuration.single_packet_seed != -1:
             seed = packet_seeds[montecarlo_configuration.single_packet_seed]
             np.random.seed(seed)
@@ -214,8 +242,12 @@ def montecarlo_main_loop(
 
         vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
         vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx]
-        vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx]
-        vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx]
+        vpackets_initial_mu = vpacket_collection.initial_mus[
+            : vpacket_collection.idx
+        ]
+        vpackets_initial_r = vpacket_collection.initial_rs[
+            : vpacket_collection.idx
+        ]
 
         v_packets_idx = np.floor(
             (vpackets_nu - spectrum_frequency[0]) / delta_nu
@@ -233,17 +265,29 @@ def montecarlo_main_loop(
     if montecarlo_configuration.VPACKET_LOGGING:
         for vpacket_collection in vpacket_collections:
             vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx]
-            vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx]
-            vpackets_initial_mu = vpacket_collection.initial_mus[: vpacket_collection.idx]
-            vpackets_initial_r = vpacket_collection.initial_rs[: vpacket_collection.idx]
+            vpackets_energy = vpacket_collection.energies[
+                : vpacket_collection.idx
+            ]
+            vpackets_initial_mu = vpacket_collection.initial_mus[
+                : vpacket_collection.idx
+            ]
+            vpackets_initial_r = vpacket_collection.initial_rs[
+                : vpacket_collection.idx
+            ]
             virt_packet_nus.append(np.ascontiguousarray(vpackets_nu))
             virt_packet_energies.append(np.ascontiguousarray(vpackets_energy))
-            virt_packet_initial_mus.append(np.ascontiguousarray(vpackets_initial_mu))
-            virt_packet_initial_rs.append(np.ascontiguousarray(vpackets_initial_r))
-            virt_packet_last_interaction_in_nu.append(np.ascontiguousarray(
-                vpacket_collection.last_interaction_in_nu[
-                    : vpacket_collection.idx
-                ])
+            virt_packet_initial_mus.append(
+                np.ascontiguousarray(vpackets_initial_mu)
+            )
+            virt_packet_initial_rs.append(
+                np.ascontiguousarray(vpackets_initial_r)
+            )
+            virt_packet_last_interaction_in_nu.append(
+                np.ascontiguousarray(
+                    vpacket_collection.last_interaction_in_nu[
+                        : vpacket_collection.idx
+                    ]
+                )
             )
             virt_packet_last_interaction_type.append(
                 np.ascontiguousarray(

tardis/simulation/base.py

@@ -134,6 +134,7 @@ def __init__(
         nthreads,
         show_convergence_plots,
         convergence_plots_kwargs,
+        show_progress_bars,
     ):
 
         super(Simulation, self).__init__(iterations, model.no_of_shells)
@@ -151,6 +152,7 @@ def __init__(
         self.luminosity_nu_end = luminosity_nu_end
         self.luminosity_requested = luminosity_requested
         self.nthreads = nthreads
+        self.show_progress_bars = show_progress_bars
 
         if convergence_strategy.type in ("damped"):
             self.convergence_strategy = convergence_strategy
@@ -372,6 +374,8 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0, last_run=False):
             nthreads=self.nthreads,
             last_run=last_run,
             iteration=self.iterations_executed,
+            total_iterations=self.iterations,
+            show_progress_bars=self.show_progress_bars,
         )
         output_energy = self.runner.output_energy
         if np.sum(output_energy < 0) == len(output_energy):
@@ -592,6 +596,7 @@ def from_config(
         packet_source=None,
         virtual_packet_logging=False,
         show_convergence_plots=True,
+        show_progress_bars=True,
         **kwargs,
     ):
         """
@@ -684,4 +689,5 @@ def from_config(
             convergence_strategy=config.montecarlo.convergence_strategy,
             nthreads=config.montecarlo.nthreads,
             convergence_plots_kwargs=convergence_plots_kwargs,
+            show_progress_bars=show_progress_bars,
         )