Add progressbar to prism forward gravity calculations (fatiando#315)

Add optional `progressbar` flag to `prism_gravity` function and to the
`gravity` method of the prism layer accesor to print a progress bar using
`numba_progress`. Add `numba_progress` as optional dependency. Add test
functions for the new feature.

doc/conf.py

@@ -63,6 +63,7 @@
     "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

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:
 

env/requirements-tests.txt

@@ -6,3 +6,4 @@ coverage
 pyvista
 vtk>=9
 netcdf4
+numba_progress

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

harmonica/forward/prism.py

@@ -10,6 +10,12 @@
 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)

harmonica/forward/prism_layer.py

@@ -305,7 +305,9 @@ def update_top_bottom(self, surface, reference):
         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 @@ def gravity(self, coordinates, field, density_name="density", **kwargs):
             prisms=boundaries,
             density=density,
             field=field,
+            progressbar=progressbar,
             **kwargs,
         )
 

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
+        )

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 @@
 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)