diff --git a/doc/conf.py b/doc/conf.py
index b6fdf3bee..ce28632a0 100644
--- a/doc/conf.py
+++ b/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
diff --git a/doc/install.rst b/doc/install.rst
index 1cb5f442d..ece3201a2 100644
--- a/doc/install.rst
+++ b/doc/install.rst
@@ -40,6 +40,9 @@ Optional:
* `pyvista `__ and
`vtk `__ (>= 9): for 3D visualizations.
See :func:`harmonica.prism_to_pyvista`.
+* `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/env/requirements-tests.txt b/env/requirements-tests.txt
index 2a8058e81..ed9de3da8 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/environment.yml b/environment.yml
index a28daf01a..fbaab3449 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 baac487dc..4d22c29c9 100644
--- a/harmonica/forward/prism.py
+++ b/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)
diff --git a/harmonica/forward/prism_layer.py b/harmonica/forward/prism_layer.py
index a23a21cb1..1f3359b0c 100644
--- a/harmonica/forward/prism_layer.py
+++ b/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,
)
diff --git a/harmonica/tests/test_prism.py b/harmonica/tests/test_prism.py
index 325072f38..d3c5ad58d 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 ddcbe57c4..634e5630d 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 @@
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)