Minor optimization in prism forward modelling

[santisoler]

In the prism forward modelling we need to evaluate the kernel functions on each
of the vertices of the prism by passing shifted coordinates, i.e. the
coordinates of the vertices from a coordinate system located in the observation
points. Now we avoid recomputing shifted coordinates by moving their definition
to outer for loops.

[leouieda]

💯 Any idea how much an impact this has?

[santisoler]

I imagined it wasn't significant, but run this example to check it:

import time
import numpy as np
import verde as vd

import harmonica as hm

# Create a layer of prisms
region = (0, 100e3, -40e3, 40e3)
spacing = 2e3
(easting, northing) = vd.grid_coordinates(region=region, spacing=spacing)
surface = 100 * np.exp(-((easting - 50e3) ** 2 + northing**2) / 1e9)
density = 2670.0 * np.ones_like(surface)
prisms = hm.prism_layer(
    coordinates=(easting[0, :], northing[:, 0]),
    surface=surface,
    reference=0,
    properties={"density": density},
)

# Compute gravity field of prisms on a regular grid of observation points
coordinates = vd.grid_coordinates(region, spacing=spacing, extra_coords=1e3)

# Run a first time to compile
gravity = prisms.prism_layer.gravity(coordinates, field="g_z")

# Run and benchmark
elapsed_time = []
for _ in range(100):
    start = time.time()
    gravity = prisms.prism_layer.gravity(coordinates, field="g_z")
    end = time.time()
    elapsed_time.append(end - start)

mean = np.mean(elapsed_time)
std = np.std(elapsed_time)
print(f"Elapsed time: {mean:.4f} +/- {std:.4f} s")

I got the following results:

	Mean [s]	Std [s]
Without optimization	0.2599	0.0042
With optimization	0.2513	0.0077

So it speed it up a little bit, but not significantly.

[leouieda]

I suspect this is numba being smart and putting those out of the loop automatically anyway. Still, worth having it in the code to make it explicit.