Repository: fatiando/harmonica (python)

Issue:
Adding progress bar to forward gravity calculation of prisms
**Description of the desired feature:**
<!--
Add as much detail as you can in your description below. If possible, include an
example of how you would like to use this feature (even better if it's a code
example).
-->
During a call to `prism.gravity` (or `prism_layer.gravity`) it would be great to have a progress bar which updates as the calculations advance through the gravity observation points (`coordinates`) .  @LL-Geo suggested the `numba_progress` package, which seems to be a good fit. Looking through `harmonica.forward.prisms.py`, I assume implementing this feature would be done within the `jit_prism_gravity` function, and when calling the function, users would use something along the lines of:


``` python
with ProgressBar(total=max_iter) as numba_progress:
        topo_prisms.prism_layer.gravity(
               coordinates=(df_grav.x, df_grav.y, df_grav.z),
               field = 'g_z',
               progress=numba_progress)
```


**Are you willing to help implement and maintain this feature?** 
<!--
Every feature we add is code that we will have to maintain and keep updated.
Please let us know if you're willing to help maintain this feature in the future.
-->
Yes I'm willing to try, but will likely need some assistance since it's my first code contribution, and I'm unfamiliar with numba_progress.

Pull request description:
Add progressbar to prism forward gravity calculations
<!--
Thank you for contributing a pull request! Please ensure you have taken a look 
at the CONTRIBUTING.md file in this repository (if available) and the general 
guidelines at https://github.com/fatiando/community/blob/main/CONTRIBUTING.md
-->
**Description**
<!--
Please describe changes proposed and WHY you made them.
-->
This PR adds an optional progress bar to the forward gravity calculation of `prisms_gravity`. It uses `numba_progress`, and is called with a flag `progressbar=True`, which defaults to `False`. Currently it only works properly with `parallel=True`.

**Relevant issues/PRs**
<!--
Link to relevant issues/PRs. Example: "Fixes #1234" or "See also #345"
Use keywords (e.g., Fixes, Closes) to create the links and automatically
close issues when this PR is merged. 
See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
Fixes #312 

Patch:
diff --git a/doc/conf.py b/doc/conf.py
index b6fdf3b..ce28632 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -63,6 +63,7 @@ intersphinx_mapping = {
     "matplotlib": ("https://matplotlib.org/", None),
     "pyproj": ("https://pyproj4.github.io/pyproj/stable/", None),
     "pyvista": ("https://docs.pyvista.org", None),
+    "numba_progress": ("https://pypi.org/project/numba-progress/", None),
 }
 
 # Autosummary pages will be generated by sphinx-autogen instead of sphinx-build
diff --git a/doc/install.rst b/doc/install.rst
index 1cb5f44..ece3201 100644
--- a/doc/install.rst
+++ b/doc/install.rst
@@ -40,6 +40,9 @@ Optional:
 * `pyvista <https://www.pyvista.org/>`__ and
   `vtk <https://vtk.org/>`__ (>= 9): for 3D visualizations.
   See :func:`harmonica.prism_to_pyvista`.
+* `numba_progress <https://pypi.org/project/numba-progress/>`__ for
+  printing a progress bar on some forward modelling computations.
+  See :func:`harmonica.prism_gravity`.
 
 The examples in the :ref:`gallery` also use:
 
diff --git a/environment.yml b/environment.yml
index a28daf0..fbaab34 100644
--- a/environment.yml
+++ b/environment.yml
@@ -56,3 +56,6 @@ dependencies:
       # Install flake8-unused-arguments through pip
       # (not available through conda yet)
       - flake8-unused-arguments==0.0.9
+      # Install numba_progress through pip
+      # (not available through conda yet)
+      - numba_progress
diff --git a/harmonica/forward/prism.py b/harmonica/forward/prism.py
index baac487..4d22c29 100644
--- a/harmonica/forward/prism.py
+++ b/harmonica/forward/prism.py
@@ -10,6 +10,12 @@ Forward modelling for prisms
 import numpy as np
 from numba import jit, prange
 
+# Attempt to import numba_progress
+try:
+    from numba_progress import ProgressBar
+except ImportError:
+    ProgressBar = None
+
 from ..constants import GRAVITATIONAL_CONST
 
 
@@ -20,6 +26,7 @@ def prism_gravity(
     field,
     parallel=True,
     dtype="float64",
+    progressbar=False,
     disable_checks=False,
 ):
     """
@@ -73,6 +80,10 @@ def prism_gravity(
     dtype : data-type (optional)
         Data type assigned to the resulting gravitational field. Default to
         ``np.float64``.
+    progressbar : bool (optional)
+        If True, a progress bar of the computation will be printed to standard
+        error (stderr). Requires :mod:`numba_progress` to be installed.
+        Default to ``False``.
     disable_checks : bool (optional)
         Flag that controls whether to perform a sanity check on the model.
         Should be set to ``True`` only when it is certain that the input model
@@ -130,9 +141,23 @@ def prism_gravity(
                 + "mismatch the number of prisms ({})".format(prisms.shape[0])
             )
         _check_prisms(prisms)
+    # Show progress bar for 'jit_prism_gravity' function
+    if progressbar:
+        if ProgressBar is None:
+            raise ImportError(
+                "Missing optional dependency 'numba_progress' required if progressbar=True"
+            )
+        progress_proxy = ProgressBar(total=coordinates[0].size)
+    else:
+        progress_proxy = None
     # Compute gravitational field
-    dispatcher(parallel)(coordinates, prisms, density, kernels[field], result)
+    dispatcher(parallel)(
+        coordinates, prisms, density, kernels[field], result, progress_proxy
+    )
     result *= GRAVITATIONAL_CONST
+    # Close previously created progress bars
+    if progressbar:
+        progress_proxy.close()
     # Convert to more convenient units
     if field == "g_z":
         result *= 1e5  # SI to mGal
@@ -186,7 +211,7 @@ def _check_prisms(prisms):
         raise ValueError(err_msg)
 
 
-def jit_prism_gravity(coordinates, prisms, density, kernel, out):
+def jit_prism_gravity(coordinates, prisms, density, kernel, out, progress_proxy=None):
     """
     Compute gravitational field of prisms on computations points
 
@@ -210,6 +235,8 @@ def jit_prism_gravity(coordinates, prisms, density, kernel, out):
         Array where the resulting field values will be stored.
         Must have the same size as the arrays contained on ``coordinates``.
     """
+    # Check if we need to update the progressbar on each iteration
+    update_progressbar = progress_proxy is not None
     # Iterate over computation points and prisms
     for l in prange(coordinates[0].size):
         for m in range(prisms.shape[0]):
@@ -233,6 +260,9 @@ def jit_prism_gravity(coordinates, prisms, density, kernel, out):
                                 shift_upward - coordinates[2][l],
                             )
                         )
+        # Update progress bar if called
+        if update_progressbar:
+            progress_proxy.update(1)
 
 
 @jit(nopython=True)
diff --git a/harmonica/forward/prism_layer.py b/harmonica/forward/prism_layer.py
index a23a21c..1f3359b 100644
--- a/harmonica/forward/prism_layer.py
+++ b/harmonica/forward/prism_layer.py
@@ -305,7 +305,9 @@ class DatasetAccessorPrismLayer:
         self._obj.coords["top"] = (self.dims, top)
         self._obj.coords["bottom"] = (self.dims, bottom)
 
-    def gravity(self, coordinates, field, density_name="density", **kwargs):
+    def gravity(
+        self, coordinates, field, progressbar=False, density_name="density", **kwargs
+    ):
         """
         Computes the gravity generated by the layer of prisms
 
@@ -358,6 +360,7 @@ class DatasetAccessorPrismLayer:
             prisms=boundaries,
             density=density,
             field=field,
+            progressbar=progressbar,
             **kwargs,
         )
 


Test patch:
diff --git a/env/requirements-tests.txt b/env/requirements-tests.txt
index 2a8058e..ed9de3d 100644
--- a/env/requirements-tests.txt
+++ b/env/requirements-tests.txt
@@ -6,3 +6,4 @@ coverage
 pyvista
 vtk>=9
 netcdf4
+numba_progress
\ No newline at end of file
diff --git a/harmonica/tests/test_prism.py b/harmonica/tests/test_prism.py
index 325072f..d3c5ad5 100644
--- a/harmonica/tests/test_prism.py
+++ b/harmonica/tests/test_prism.py
@@ -7,11 +7,18 @@
 """
 Test forward modelling for prisms.
 """
+from unittest.mock import patch
+
 import numpy as np
 import numpy.testing as npt
 import pytest
 import verde as vd
 
+try:
+    from numba_progress import ProgressBar
+except ImportError:
+    ProgressBar = None
+
 from ..forward.prism import _check_prisms, prism_gravity, safe_atan2, safe_log
 from ..gravity_corrections import bouguer_correction
 from .utils import run_only_with_numba
@@ -381,3 +388,50 @@ def test_prisms_parallel_vs_serial():
             coordinates, prisms, densities, field=field, parallel=False
         )
         npt.assert_allclose(result_parallel, result_serial)
+
+
+@pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress")
+@pytest.mark.use_numba
+def test_progress_bar():
+    """
+    Check if forward gravity results with and without progress bar match
+    """
+    prisms = [
+        [-100, 0, -100, 0, -10, 0],
+        [0, 100, -100, 0, -10, 0],
+        [-100, 0, 0, 100, -10, 0],
+        [0, 100, 0, 100, -10, 0],
+    ]
+    densities = [2000, 3000, 4000, 5000]
+    coordinates = vd.grid_coordinates(
+        region=(-100, 100, -100, 100), spacing=20, extra_coords=10
+    )
+    for field in ("potential", "g_z"):
+        result_progress_true = prism_gravity(
+            coordinates, prisms, densities, field=field, progressbar=True
+        )
+        result_progress_false = prism_gravity(
+            coordinates, prisms, densities, field=field, progressbar=False
+        )
+        npt.assert_allclose(result_progress_true, result_progress_false)
+
+
+@patch("harmonica.forward.prism.ProgressBar", None)
+def test_numba_progress_missing_error():
+    """
+    Check if error is raised when progresbar=True and numba_progress package
+    is not installed.
+    """
+    prisms = [
+        [-100, 0, -100, 0, -10, 0],
+        [0, 100, -100, 0, -10, 0],
+        [-100, 0, 0, 100, -10, 0],
+        [0, 100, 0, 100, -10, 0],
+    ]
+    densities = [2000, 3000, 4000, 5000]
+    coordinates = [0, 0, 0]
+    # Check if error is raised
+    with pytest.raises(ImportError):
+        prism_gravity(
+            coordinates, prisms, densities, field="potential", progressbar=True
+        )
diff --git a/harmonica/tests/test_prism_layer.py b/harmonica/tests/test_prism_layer.py
index ddcbe57..634e563 100644
--- a/harmonica/tests/test_prism_layer.py
+++ b/harmonica/tests/test_prism_layer.py
@@ -8,6 +8,7 @@
 Test prisms layer
 """
 import warnings
+from unittest.mock import patch
 
 import numpy as np
 import numpy.testing as npt
@@ -22,6 +23,11 @@ try:
 except ImportError:
     pyvista = None
 
+try:
+    from numba_progress import ProgressBar
+except ImportError:
+    ProgressBar = None
+
 
 @pytest.fixture(params=("numpy", "xarray"))
 def dummy_layer(request):
@@ -422,3 +428,40 @@ def test_to_pyvista(dummy_layer, properties):
         assert pv_grid.array_names == ["density"]
         assert pv_grid.get_array("density").ndim == 1
         npt.assert_allclose(pv_grid.get_array("density"), layer.density.values.ravel())
+
+
+@pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress")
+@pytest.mark.use_numba
+def test_progress_bar(dummy_layer):
+    """
+    Check if forward gravity results with and without progress bar match
+    """
+    coordinates = vd.grid_coordinates((1, 3, 7, 10), spacing=1, extra_coords=30.0)
+    (easting, northing), surface, reference, density = dummy_layer
+    layer = prism_layer(
+        (easting, northing), surface, reference, properties={"density": density}
+    )
+    result_progress_true = layer.prism_layer.gravity(
+        coordinates, field="g_z", progressbar=True
+    )
+
+    result_progress_false = layer.prism_layer.gravity(
+        coordinates, field="g_z", progressbar=False
+    )
+    npt.assert_allclose(result_progress_true, result_progress_false)
+
+
+@patch("harmonica.forward.prism.ProgressBar", None)
+def test_numba_progress_missing_error(dummy_layer):
+    """
+    Check if error is raised when progressbar=True and numba_progress package
+    is not installed.
+    """
+    coordinates = vd.grid_coordinates((1, 3, 7, 10), spacing=1, extra_coords=30.0)
+    (easting, northing), surface, reference, density = dummy_layer
+    layer = prism_layer(
+        (easting, northing), surface, reference, properties={"density": density}
+    )
+    # Check if error is raised
+    with pytest.raises(ImportError):
+        layer.prism_layer.gravity(coordinates, field="g_z", progressbar=True)