Adding progress bar to forward gravity calculation of prisms

Description of the desired feature:

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:

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?

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.

[mdtanker]

Thanks for the suggestions! I have it working and everything seems alright, except if called with parallel=False then the progress bar only gives 0% and 100%, with no intermediate progress. Not sure why this is happening, but since parallel defaults to True then maybe it’s no big deal.

Here is an example without the progress bar, which took 1m 28.2s :

for k, v in layers.items():
    df_grav[f'{k}_forward_grav'] = v['prisms'].prism_layer.gravity(
        coordinates=(df_grav.x, df_grav.y, df_grav.z),
        field = 'g_z', 
)

And an example with progressbar=True, which took 1min 34.6s:

for k, v in layers.items():
    df_grav[f'{k}_forward_grav'] = v['prisms'].prism_layer.gravity(
        coordinates=(df_grav.x, df_grav.y, df_grav.z),
        field = 'g_z', 
        progressbar=True,
)
100%|██████████| 16641.0/16641 [00:29<00:00, 557.69it/s]
100%|██████████| 16641.0/16641 [00:31<00:00, 527.07it/s]
100%|██████████| 16641.0/16641 [00:33<00:00, 503.12it/s]

In the prism.py file, I’ve added the following before the dispatcher call:

if progressbar is True:
    from numba_progress import ProgressBar
    progress_proxy = ProgressBar(total = coordinates[0].size) 
else: 
    progress_proxy = None

And the following after the dispacher call:

  try:
      progress_proxy.close()
  except:
      pass

And this is the updated function:

def jit_prism_gravity(coordinates, prisms, density, kernel, out, progress_proxy):
    # Iterate over computation points and prisms
    for l in prange(coordinates[0].size):
        # Update progress bar if called
        if progress_proxy is not None:
            progress_proxy.update(1) 
        for m in range(prisms.shape[0]):
            # Iterate over the prism boundaries to compute the result of the
            # integration (see Nagy et al., 2000)
            for i in range(2):
                for j in range(2):
                    for k in range(2):
                        shift_east = prisms[m, 1 - i]
                        shift_north = prisms[m, 3 - j]
                        shift_upward = prisms[m, 5 - k]
                        # If i, j or k is 1, the shift_* will refer to the
                        # lower boundary, meaning the corresponding term should
                        # have a minus sign
                        out[l] += (
                            density[m]
                            * (-1) ** (i + j + k)
                            * kernel(
                                shift_east - coordinates[0][l],
                                shift_north - coordinates[1][l],
                                shift_upward - coordinates[2][l],
                            )
                        )

Is the next step to open a pull request to implement this?

[santisoler]

This looks great! Thanks @mdtanker for the idea! 🥇

Any suggestions for not needing to add the arg progress_proxy to so many locations in the code?

Actually I was thinking that we don’t need to expose the ProgressBar to the user, nor ask them to pass a ProgressBar instance as argument. Instead we can add a progressbar flag to prism_gravity (like the one we have in pooch.retrieve https://www.fatiando.org/pooch/latest/api/generated/pooch.retrieve.html#pooch.retrieve) and initialize a ProgressBar inside the function, which will be passed to the jitted functions down the road.

So, your last example would look like this:

for k, v in layers.items():
    df_grav[f'{k}_forward_grav'] = v['prisms'].prism_layer.gravity(
        coordinates=(df_grav.x, df_grav.y, df_grav.z),
        field = 'g_z',
        progressbar=True,
    )

Since numba-progress will be an optional dependency, the progressbar should be False by default and the jitted functions should be able to work even if no ProgressBar is being passed to them.