Add progressbar to tesseroid forward modelling

harmonica/_forward/prism.py

@@ -29,6 +29,8 @@
 )
 from numba import jit, prange
 
+from .utils import initialize_progressbar
+
 # Define dictionary with available gravity fields for prisms
 FIELDS = {
  "potential": gravity_pot,
@@ -43,12 +45,6 @@
  "g_nz": gravity_nu,
 }
 
-# Attempt to import numba_progress
-try:
- from numba_progress import ProgressBar
-except ImportError:
- ProgressBar = None
-
 
 def prism_gravity(
  coordinates,
@@ -216,28 +212,16 @@ def prism_gravity(
  _check_singular_points(coordinates, prisms, field)
  # Discard null prisms (zero volume or zero density)
  prisms, density = _discard_null_prisms(prisms, density)
- # 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
  # Choose parallelized or serialized forward function
  if parallel:
  gravity_prism_func = jit_prism_gravity_parallel
  else:
  gravity_prism_func = jit_prism_gravity_serial
  # Compute gravitational field
- gravity_prism_func(
- coordinates, prisms, density, FIELDS[field], result, progress_proxy
- )
- # Close previously created progress bars
- if progress_proxy:
- progress_proxy.close()
+ with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy:
+ gravity_prism_func(
+ coordinates, prisms, density, FIELDS[field], result, progress_proxy
+ )
  # Invert sign of gravity_u, gravity_eu, gravity_nu
  if field in ("g_z", "g_ez", "g_nz"):
  result *= -1
@@ -559,6 +543,8 @@ def jit_prism_gravity(
  out : 1d-array
  Array where the resulting field values will be stored.
  Must have the same size as the arrays contained on ``coordinates``.
+ progress_proxy : :class:`numba_progress.ProgressBar` or None
+ Instance of :class:`numba_progress.ProgressBar` or None.
  """
  # Unpack coordinates
  easting, northing, upward = coordinates
@@ -579,7 +565,7 @@ def jit_prism_gravity(
  prisms[m, 5],
  density[m],
  )
- # Update progress bar if called
+ # Update progress bar
  if update_progressbar:
  progress_proxy.update(1)
 

harmonica/_forward/prism_layer.py

@@ -342,6 +342,10 @@ def gravity(
  - Downward acceleration: ``g_z``
  - Diagonal tensor components: ``g_ee``, ``g_nn``, ``g_zz``
  - Non-diagonal tensor components: ``g_en``, ``g_ez``, ``g_nz``
+ 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``.
  density_name : str (optional)
  Name of the property layer (or ``data_var`` of the
  :class:`xarray.Dataset`) that will be used for the density of each

harmonica/_forward/tesseroid.py

@@ -24,6 +24,7 @@
  gauss_legendre_quadrature_variable_density,
 )
 from .point import kernel_g_z_spherical, kernel_potential_spherical
+from .utils import initialize_progressbar
 
 STACK_SIZE = 100
 MAX_DISCRETIZATIONS = 100000
@@ -39,6 +40,7 @@ def tesseroid_gravity(
  parallel=True,
  radial_adaptive_discretization=False,
  dtype=np.float64,
+ progressbar=False,
  disable_checks=False,
 ):
  r"""
@@ -99,6 +101,10 @@ def tesseroid_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
@@ -181,18 +187,20 @@ def tesseroid_gravity(
  # tesseroid
  glq_nodes, glq_weights = glq_nodes_weights(GLQ_DEGREES)
  # Compute gravitational field
- dispatcher(parallel, density)(
- coordinates,
- tesseroids,
- density,
- result,
- DISTANCE_SIZE_RATII[field],
- radial_adaptive_discretization,
- glq_nodes,
- glq_weights,
- kernels[field],
- dtype,
- )
+ with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy:
+ dispatcher(parallel, density)(
+ coordinates,
+ tesseroids,
+ density,
+ result,
+ DISTANCE_SIZE_RATII[field],
+ radial_adaptive_discretization,
+ glq_nodes,
+ glq_weights,
+ kernels[field],
+ dtype,
+ progress_proxy,
+ )
  result *= GRAVITATIONAL_CONST
  # Convert to more convenient units
  if field == "g_z":
@@ -220,7 +228,7 @@ def dispatcher(parallel, density):
  return dispatchers[parallel]
 
 
-def jit_tesseroid_gravity(
+def jit_tesseroid_gravity( # noqa: CFQ002
  coordinates,
  tesseroids,
  density,
@@ -231,6 +239,7 @@ def jit_tesseroid_gravity(
  glq_weights,
  kernel,
  dtype,
+ progress_proxy,
 ):
  """
  Compute gravitational field of tesseroids on computations points
@@ -276,7 +285,11 @@ def jit_tesseroid_gravity(
  Kernel function for the gravitational field of point masses.
  dtype : data-type
  Data type assigned to the resulting gravitational field.
+ progress_proxy : :class:`numba_progress.ProgressBar` or None
+ Instance of :class:`numba_progress.ProgressBar` or None.
  """
+ # Check if we need to update the progressbar on each iteration
+ update_progressbar = progress_proxy is not None
  # Get coordinates of the observation points
  # and precompute trigonometric functions
  longitude, latitude, radius = coordinates[:]
@@ -313,9 +326,12 @@ def jit_tesseroid_gravity(
  glq_weights,
  kernel,
  )
+ # Update progress bar
+ if update_progressbar:
+ progress_proxy.update(1)
 
 
-def jit_tesseroid_gravity_variable_density(
+def jit_tesseroid_gravity_variable_density( # noqa: CFQ002
  coordinates,
  tesseroids,
  density,
@@ -326,6 +342,7 @@ def jit_tesseroid_gravity_variable_density(
  glq_weights,
  kernel,
  dtype,
+ progress_proxy,
 ):
  """
  Compute gravitational field of tesseroids on computations points
@@ -371,7 +388,11 @@ def jit_tesseroid_gravity_variable_density(
  Kernel function for the gravitational field of point masses.
  dtype : data-type
  Data type assigned to the resulting gravitational field.
+ progress_proxy : :class:`numba_progress.ProgressBar` or None
+ Instance of :class:`numba_progress.ProgressBar` or None.
  """
+ # Check if we need to update the progressbar on each iteration
+ update_progressbar = progress_proxy is not None
  # Get coordinates of the observation points
  # and precompute trigonometric functions
  longitude, latitude, radius = coordinates[:]
@@ -408,6 +429,9 @@ def jit_tesseroid_gravity_variable_density(
  glq_weights,
  kernel,
  )
+ # Update progress bar
+ if update_progressbar:
+ progress_proxy.update(1)
 
 
 # Define jitted versions of the forward modelling function

harmonica/_forward/tesseroid_layer.py

@@ -250,7 +250,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 tesseroids
 
@@ -266,6 +268,10 @@ def gravity(self, coordinates, field, density_name="density", **kwargs):
  The variable fields are:
  - Gravitational potential: ``potential``
  - Downward acceleration: ``g_z``
+ 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``.
  density_name : str (optional)
  Name of the property layer (or ``data_var`` of the
  :class:`xarray.Dataset`) that will be used for the density of each
@@ -275,7 +281,8 @@ def gravity(self, coordinates, field, density_name="density", **kwargs):
  -------
  result : array
  Gravitational field generated by the tesseroid on the computation
- point in mGal
+ point. Gravitational potential will be returned in SI units, while
+ acceleration components will be in mGal.
 
  See also
  --------
@@ -296,6 +303,7 @@ def gravity(self, coordinates, field, density_name="density", **kwargs):
  tesseroids=boundaries,
  density=density,
  field=field,
+ progressbar=progressbar,
  **kwargs,
  )
 

harmonica/_forward/utils.py

@@ -7,9 +7,17 @@
 """
 Utilities for forward modelling
 """
+import contextlib
+
 import numpy as np
 from numba import jit
 
+# Attempt to import numba_progress
+try:
+ from numba_progress import ProgressBar
+except ImportError:
+ ProgressBar = None
+
 
 def distance(point_p, point_q, coordinate_system="cartesian", ellipsoid=None):
  """
@@ -325,3 +333,60 @@
  )
  )
  return dist
+
+
+def initialize_progressbar(total, use_progressbar):
+ """
+ Generate an instance of ``numba_progress.ProgressBar``
+
+ This function is meant to be used within other forward modelling functions
+ that have a ``progressbar`` optional argument.
+
+ Parameters
+ ----------
+ total : int
+ Number of iterations that will be assumed as the total for the
+ :class:`numba_progress.ProgressBar`.
+ use_progressbar : bool
+ Weather to initialize a progressbar or not. If True, then the function
+ will return None.
+
+ Returns
+ -------
+ :class:`numba_progress.ProgressBar` or :class:`contextlib.nullcontext`
+ Instance of :class:`numba_progress.ProgressBar` if ``use_progressbar``
+ is True. Instance of :class:`contextlib.nullcontext` if
+ ``use_progressbar`` is False.
+
+ Raises
+ ------
+ ImportError
+ If ``numba_progress`` is missing and ``use_progress`` is set to True.
+
+ See also
+ --------
+ :func:`harmonica.prism_gravity`
+
+ Examples
+ --------
+ This function is meant to be used inside a context manager, like ``with``:
+
+ >>> with initialize_progressbar(3, True) as progress_proxy:
+ ... for _ in range(3):
+ ... progress_proxy.update(3) # doctest: +SKIP
+
+ >>> with initialize_progressbar(3, False) as progress_proxy:
+ ... print(progress_proxy is None) # doctest: +SKIP
+ True
+
+ """
+ # Return None if progressbar is not desired
+ if not use_progressbar:
+ return contextlib.nullcontext()
+ # Raise error if numba_progress is not installed
+ if ProgressBar is None:
+ raise ImportError(
+ "Missing optional dependency 'numba_progress' required "
+ "if progressbar=True"
+ )
+ return ProgressBar(total=total)

harmonica/tests/test_forward_utils.py

@@ -7,12 +7,19 @@
 """
 Test utils functions for forward modelling
 """
+from unittest.mock import patch
+
 import boule as bl
 import numpy as np
 import numpy.testing as npt
 import pytest
 
-from .._forward.utils import check_coordinate_system, distance
+from .._forward.utils import check_coordinate_system, distance, initialize_progressbar
+
+try:
+ from numba_progress import ProgressBar
+except ImportError:
+ ProgressBar = None
 
 
 @pytest.mark.use_numba
@@ -104,3 +111,27 @@ def test_geodetic_distance_equator_poles():
  ellipsoid=ellipsoid,
  ),
  )
+
+
+@pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress")
+@pytest.mark.parametrize("use_progressbar", (True, False))
+def test_initialize_progressbar(use_progressbar):
+ """Test initialize_progressbar."""
+ with initialize_progressbar(3, use_progressbar) as progress_proxy:
+ if use_progressbar:
+ assert isinstance(progress_proxy, ProgressBar)
+ else:
+ assert progress_proxy is None
+
+
+@patch("harmonica._forward.utils.ProgressBar", None)
+@pytest.mark.parametrize("use_progressbar", (True, False))
+def test_initialize_progressbar_import_error(use_progressbar):
+ """Test if it raises ImportError when numba_progress is missing."""
+ if use_progressbar:
+ with pytest.raises(ImportError): # noqa: SIM117
+ with initialize_progressbar(3, use_progressbar) as progress_proxy:
+ pass
+ else:
+ with initialize_progressbar(3, use_progressbar) as progress_proxy:
+ assert progress_proxy is None