Choclo

Kernel functions for your geophysical models

Choclo is a Python library that hosts optimized forward modelling and kernel functions for running geophysical forward and inverse models, intended to be used by other libraries as the underlying layer of their computation.

“Choclo” is a term used in some countries of South America to refer to corn, originated from the quechua word chuqllu.

See also

Choclo is a part of the Fatiando a Terra project.

Overview#

Choclo provides slim and optimized function to compute the gravitational and magnetic fields of simple geometries like point masses, magnetic dipoles and prisms. It also provides the kernel functions needed to run compute those fields. The goal of Choclo is to provide developers of a simple and efficient way to calculate these fields for a wide range of applications, like forward modellings, sensitivity matrices calculations, equivalent sources implementations and more.

These functions are not designed to be used by final users. Instead they are meant to be part of the underlaying engine of a higher level codebase, like Harmonica.

All Choclo functions rely on Numba for just-in-time compilations, meaning that there’s no need to distribute precompiled code: Choclo provides pure Python code that gets compiled during runtime allowing to run them as fast as they were written in C. Moreover, developers could harness the power of Numba to parallelize processes in a quick and easy way.

Conventions#

Before you jump into Choclo’s functions, it’s worth noting some conventions that will be kept along its codebase:

  • The functions assume a right-handed coordinate system. We avoid using names like “x”, “y” and “z” for the coordinates. Instead we use “easting”, “northing” and “upward” to make extra clear the direction of each axis.

  • We use the first letter of the easting, northing and upward axis to indicate direction of derivatives. For example, a function gravity_e will compute the easting component of the gravitional acceleration, while the gravity_n and gravity_u will compute the northing and upward ones, respectively.

  • The arguments of the functions are always assumed in SI units. And all the functions return results also in SI units. Choclo doesn’t perform unit conversions.

  • The components of the gravitational accelerations and the magnetic fields are computed in the same direction of the easting, northing and upward axis. So gravity_u will compute the upward component of the gravitational acceleration (note the difference with the downward component).

The library#

Choclo is divided in a few different submodules, each with different goals. The three main modules are the ones that host the forward and kernel functions for the different geometries supported by Choclo: point, dipole and prism. Inside each one of these modules we will find forward modelling functions and potentially some kernel functions. The names of the forward modelling functions follow a simple pattern of {field}_{type}. For example, choclo.prism.gravity_e computes the easting component of the gravitational acceleration of a prism, while choclo.prism.gravity_ee computes the easting-easting gravity tensor component.

How to use Choclo#

The simplest case#

Using Choclo is very simple, but it requires some work from our side. Let’s say we need to compute the upward component of the gravitational acceleration that a single rectangular prism produces on a single computation point. To do so we can just call the choclo.prism.gravity_u function:

import numpy as np
from choclo.prism import gravity_u

# Define a single computation point
easting, northing, upward = 0.0, 0.0, 10.0

# Define the boundaries of the prism as a 1d-array
prism = np.array([-10.0, 10.0, -7.0, 7.0, -15.0, -5.0])

# And its density
density = 400.0

# Compute the upward component of the grav. acceleration
g_u = gravity_u(easting, northing, upward, prism, density)
g_u

-1.653643133689388e-07

But this case is very simple: we usually deal with multiple sources and multiple computation points.

Multiple sources and computation points#

In case we have a collection of prisms with certain densities:

prisms = np.array([
    [-10.0, 0.0, -7.0, 0.0, -15.0, -10.0],
    [-10.0, 0.0, 0.0, 7.0, -25.0, -15.0],
    [0.0, 10.0, -7.0, 0.0, -20.0, -13.0],
    [0.0, 10.0, 0.0, 7.0, -12.0, -8.0],
])
densities = np.array([200.0, 300.0, -100.0, 400.0])

And a set of observation points:

easting = np.linspace(-5.0, 5.0, 21)
northing = np.linspace(-4.0, 4.0, 21)
easting, northing = tuple(a.ravel() for a in np.meshgrid(easting, northing))
upward = 10 * np.ones_like(easting)
coordinates = (easting, northing, upward)

And we want to compute the gratitational acceleration that those prisms generate on each observation point, we need to write some kind of loop that computes the effect of each prism on each observation point and adds it to a running result.

A possible solution would be to use Python for loops:

def gravity_upward_slow(coordinates, prisms, densities):
    """
    Compute the upward component of the acceleration of a set of prisms
    """
    # Unpack coordinates of the observation points
    easting, northing, upward = coordinates[:]
    # Initialize a result array full of zeros
    result = np.zeros_like(easting, dtype=np.float64)
    # Compute the upward component that every prism generate on each
    # observation point
    for i in range(len(easting)):
        for j in range(prisms.shape[0]):
            result[i] += gravity_u(
                easting[i], northing[i], upward[i], prisms[j, :], densities[j]
            )
    return result

g_u = gravity_upward_slow(coordinates, prisms, densities)
g_u

array([-3.08024767e-08, -3.10103482e-08, -3.11926215e-08, -3.13484961e-08,
       -3.14772466e-08, -3.15782275e-08, -3.16508775e-08, -3.16947240e-08,
       -3.17093869e-08, -3.16945831e-08, -3.16501302e-08, -3.15759501e-08,
       -3.14720725e-08, -3.13386380e-08, -3.11759003e-08, -3.09842280e-08,
       -3.07641061e-08, -3.05161358e-08, -3.02410339e-08, -2.99396310e-08,
       -2.96128686e-08, -3.11195000e-08, -3.13346703e-08, -3.15240617e-08,
       -3.16868368e-08, -3.18222319e-08, -3.19295618e-08, -3.20082244e-08,
       -3.20577056e-08, -3.20775840e-08, -3.20675354e-08, -3.20273375e-08,
       -3.19568737e-08, -3.18561377e-08, -3.17252367e-08, -3.15643944e-08,
       -3.13739539e-08, -3.11543785e-08, -3.09062534e-08, -3.06302842e-08,
       -3.03272966e-08, -2.99982325e-08, -3.14196998e-08, -3.16420877e-08,
       -3.18385422e-08, -3.20081895e-08, -3.21502273e-08, -3.22639303e-08,
       -3.23486553e-08, -3.24038461e-08, -3.24290393e-08, -3.24238686e-08,
       -3.23880705e-08, -3.23214892e-08, -3.22240807e-08, -3.20959177e-08,
       -3.19371927e-08, -3.17482213e-08, -3.15294445e-08, -3.12814294e-08,
       -3.10048696e-08, -3.07005838e-08, -3.03695135e-08, -3.17021952e-08,
       -3.19316874e-08, -3.21351185e-08, -3.23115780e-08, -3.24602255e-08,
       -3.25802955e-08, -3.26711033e-08, -3.27320505e-08, -3.27626308e-08,
       -3.27624353e-08, -3.27311588e-08, -3.26686047e-08, -3.25746907e-08,
       -3.24494538e-08, -3.22930541e-08, -3.21057787e-08, -3.18880444e-08,
       -3.16403995e-08, -3.13635238e-08, -3.10582283e-08, -3.07254524e-08,
       -3.19661435e-08, -3.22025940e-08, -3.24128819e-08, -3.25960611e-08,
       -3.27512529e-08, -3.28776521e-08, -3.29745324e-08, -3.30412529e-08,
       -3.30772642e-08, -3.30821145e-08, -3.30554558e-08, -3.29970507e-08,
       -3.29067776e-08, -3.27846366e-08, -3.26307546e-08, -3.24453891e-08,
       -3.22289318e-08, -3.19819108e-08, -3.17049911e-08, -3.13989744e-08,
       -3.10647971e-08, -3.22107451e-08, -3.24539745e-08, -3.26709663e-08,
       -3.28607390e-08, -3.30223766e-08, -3.31550343e-08, -3.32579446e-08,
       -3.33304242e-08, -3.33718804e-08, -3.33818184e-08, -3.33598474e-08,
       -3.33056883e-08, -3.32191797e-08, -3.31002846e-08, -3.29490955e-08,
       -3.27658396e-08, -3.25508827e-08, -3.23047313e-08, -3.20280348e-08,
       -3.17215845e-08, -3.13863123e-08, -3.24352494e-08, -3.26850446e-08,
       -3.29085531e-08, -3.31047592e-08, -3.32727101e-08, -3.34115219e-08,
       -3.35203866e-08, -3.35985787e-08, -3.36454629e-08, -3.36605008e-08,
       -3.36432592e-08, -3.35934170e-08, -3.35107730e-08, -3.33952522e-08,
       -3.32469127e-08, -3.30659504e-08, -3.28527044e-08, -3.26076592e-08,
       -3.23314471e-08, -3.20248483e-08, -3.16887888e-08, -3.26389589e-08,
       -3.28950733e-08, -3.31248774e-08, -3.33273224e-08, -3.35014197e-08,
       -3.36462472e-08, -3.37609571e-08, -3.38447825e-08, -3.38970454e-08,
       -3.39171650e-08, -3.39046654e-08, -3.38591840e-08, -3.37804794e-08,
       -3.36684390e-08, -3.35230855e-08, -3.33445839e-08, -3.31332455e-08,
       -3.28895323e-08, -3.26140588e-08, -3.23075927e-08, -3.19710535e-08,
       -3.28212346e-08, -3.30833883e-08, -3.33192332e-08, -3.35276885e-08,
       -3.37077309e-08, -3.38584016e-08, -3.39788135e-08, -3.40681594e-08,
       -3.41257196e-08, -3.41508714e-08, -3.41430967e-08, -3.41019920e-08,
       -3.40272760e-08, -3.39187982e-08, -3.37765467e-08, -3.36006546e-08,
       -3.33914057e-08, -3.31492385e-08, -3.28747492e-08, -3.25686922e-08,
       -3.22319792e-08, -3.29815001e-08, -3.32493807e-08, -3.34909783e-08,
       -3.37051816e-08, -3.38909341e-08, -3.40472413e-08, -3.41731784e-08,
       -3.42678988e-08, -3.43306425e-08, -3.43607454e-08, -3.43576487e-08,
       -3.43209081e-08, -3.42502033e-08, -3.41453465e-08, -3.40062911e-08,
       -3.38331385e-08, -3.36261448e-08, -3.33857248e-08, -3.31124558e-08,
       -3.28070782e-08, -3.24704949e-08, -3.31192457e-08, -3.33925095e-08,
       -3.36395397e-08, -3.38591960e-08, -3.40503903e-08, -3.42120940e-08,
       -3.43433462e-08, -3.44432623e-08, -3.45110433e-08, -3.45459852e-08,
       -3.45474893e-08, -3.45150719e-08, -3.44483743e-08, -3.43471723e-08,
       -3.42113849e-08, -3.40410824e-08, -3.38364930e-08, -3.35980079e-08,
       -3.33261849e-08, -3.30217497e-08, -3.26855958e-08, -3.32340324e-08,
       -3.35123058e-08, -3.37644175e-08, -3.39892003e-08, -3.41855361e-08,
       -3.43523641e-08, -3.44886890e-08, -3.45935899e-08, -3.46662305e-08,
       -3.47058685e-08, -3.47118667e-08, -3.46837034e-08, -3.46209825e-08,
       -3.45234445e-08, -3.43909748e-08, -3.42236131e-08, -3.40215600e-08,
       -3.37851831e-08, -3.35150205e-08, -3.32117829e-08, -3.28763531e-08,
       -3.33254950e-08, -3.36083762e-08, -3.38651893e-08, -3.40947420e-08,
       -3.42958885e-08, -3.44675377e-08, -3.46086619e-08, -3.47183060e-08,
       -3.47955980e-08, -3.48397595e-08, -3.48501163e-08, -3.48261102e-08,
       -3.47673097e-08, -3.46734205e-08, -3.45442961e-08, -3.43799459e-08,
       -3.41805438e-08, -3.39464337e-08, -3.36781342e-08, -3.33763402e-08,
       -3.30419234e-08, -3.33933450e-08, -3.36804060e-08, -3.39415134e-08,
       -3.41754516e-08, -3.43810495e-08, -3.45571878e-08, -3.47028085e-08,
       -3.48169248e-08, -3.48986312e-08, -3.49471148e-08, -3.49616670e-08,
       -3.49416949e-08, -3.48867329e-08, -3.47964541e-08, -3.46706805e-08,
       -3.45093929e-08, -3.43127390e-08, -3.40810394e-08, -3.38147931e-08,
       -3.35146792e-08, -3.31815574e-08, -3.34373730e-08, -3.37281623e-08,
       -3.39931319e-08, -3.42310457e-08, -3.44407092e-08, -3.46209775e-08,
       -3.47707650e-08, -3.48890551e-08, -3.49749115e-08, -3.50274894e-08,
       -3.50460476e-08, -3.50299609e-08, -3.49787315e-08, -3.48920015e-08,
       -3.47695634e-08, -3.46113702e-08, -3.44175443e-08, -3.41883839e-08,
       -3.39243683e-08, -3.36261607e-08, -3.32946084e-08, -3.34574503e-08,
       -3.37514954e-08, -3.40198732e-08, -3.42613296e-08, -3.44746493e-08,
       -3.46586642e-08, -3.48122637e-08, -3.49344043e-08, -3.50241213e-08,
       -3.50805406e-08, -3.51028913e-08, -3.50905176e-08, -3.50428922e-08,
       -3.49596280e-08, -3.48404898e-08, -3.46854042e-08, -3.44944694e-08,
       -3.42679617e-08, -3.40063416e-08, -3.37102559e-08, -3.33805389e-08,
       -3.34535302e-08, -3.37503405e-08, -3.40216536e-08, -3.42661995e-08,
       -3.44827449e-08, -3.46701015e-08, -3.48271362e-08, -3.49527814e-08,
       -3.50460472e-08, -3.51060327e-08, -3.51319398e-08, -3.51230852e-08,
       -3.50789142e-08, -3.49990128e-08, -3.48831198e-08, -3.47311373e-08,
       -3.45431403e-08, -3.43193846e-08, -3.40603115e-08, -3.37665521e-08,
       -3.34389272e-08, -3.34256488e-08, -3.37247189e-08, -3.39984781e-08,
       -3.42456436e-08, -3.44649665e-08, -3.46552412e-08, -3.48153152e-08,
       -3.49440999e-08, -3.50405826e-08, -3.51038392e-08, -3.51330469e-08,
       -3.51274979e-08, -3.50866127e-08, -3.50099526e-08, -3.48972326e-08,
       -3.47483320e-08, -3.45633044e-08, -3.43423853e-08, -3.40859983e-08,
       -3.37947583e-08, -3.34694725e-08, -3.33739245e-08, -3.36747374e-08,
       -3.39504412e-08, -3.41997427e-08, -3.44213805e-08, -3.46141346e-08,
       -3.47768360e-08, -3.49083782e-08, -3.50077294e-08, -3.50739448e-08,
       -3.51061803e-08, -3.51037064e-08, -3.50659214e-08, -3.49923646e-08,
       -3.48827295e-08, -3.47368744e-08, -3.45548329e-08, -3.43368218e-08,
       -3.40832472e-08, -3.37947084e-08, -3.34719985e-08, -3.32985579e-08,
       -3.36005885e-08, -3.38777261e-08, -3.41286699e-08, -3.43521490e-08,
       -3.45469319e-08, -3.47118364e-08, -3.48457414e-08, -3.49475987e-08,
       -3.50164466e-08, -3.50514228e-08, -3.50517789e-08, -3.50168939e-08,
       -3.49462882e-08, -3.48396356e-08, -3.46967757e-08, -3.45177238e-08,
       -3.43026790e-08, -3.40520311e-08, -3.37663636e-08, -3.34464556e-08,
       -3.31998300e-08, -3.35025487e-08, -3.37806038e-08, -3.40326898e-08,
       -3.42575292e-08, -3.44538820e-08, -3.46205563e-08, -3.47564193e-08,
       -3.48604103e-08, -3.49315535e-08, -3.49689717e-08, -3.49719010e-08,
       -3.49397041e-08, -3.48718848e-08, -3.47681001e-08, -3.46281729e-08,
       -3.44521018e-08, -3.42400699e-08, -3.39924510e-08, -3.37098137e-08,
       -3.33929226e-08])

For loops are known to be slow, and in case we are working with very large models and a large number of computation points these calculations could take too long. So this solution is not recommended.

Important

Using Python for loops to run Choclo’s functions is not advisable!

We can write a much faster and efficient solution relying on numba. Since every function in Choclo is being JIT compiled, we can safely include calls to these functions inside other JIT compiled functions. So we can write an alterantive function by adding a @numba.jit decorator:

import numba

@numba.jit(nopython=True)
def gravity_upward_jit(coordinates, prisms, densities):
    """
    Compute the upward component of the acceleration of a set of prisms
    """
    # Unpack coordinates of the observation points
    easting, northing, upward = coordinates[:]
    # Initialize a result array full of zeros
    result = np.zeros_like(easting, dtype=np.float64)
    # Compute the upward component that every prism generate on each
    # observation point
    for i in range(len(easting)):
        for j in range(prisms.shape[0]):
            result[i] += gravity_u(
                easting[i], northing[i], upward[i], prisms[j, :], densities[j]
            )
    return result

g_u = gravity_upward_jit(coordinates, prisms, densities)
g_u

array([-3.08024767e-08, -3.10103482e-08, -3.11926215e-08, -3.13484961e-08,
       -3.14772466e-08, -3.15782275e-08, -3.16508775e-08, -3.16947240e-08,
       -3.17093869e-08, -3.16945831e-08, -3.16501302e-08, -3.15759501e-08,
       -3.14720725e-08, -3.13386380e-08, -3.11759003e-08, -3.09842280e-08,
       -3.07641061e-08, -3.05161358e-08, -3.02410339e-08, -2.99396310e-08,
       -2.96128686e-08, -3.11195000e-08, -3.13346703e-08, -3.15240617e-08,
       -3.16868368e-08, -3.18222319e-08, -3.19295618e-08, -3.20082244e-08,
       -3.20577056e-08, -3.20775840e-08, -3.20675354e-08, -3.20273375e-08,
       -3.19568737e-08, -3.18561377e-08, -3.17252367e-08, -3.15643944e-08,
       -3.13739539e-08, -3.11543785e-08, -3.09062534e-08, -3.06302842e-08,
       -3.03272966e-08, -2.99982325e-08, -3.14196998e-08, -3.16420877e-08,
       -3.18385422e-08, -3.20081895e-08, -3.21502273e-08, -3.22639303e-08,
       -3.23486553e-08, -3.24038461e-08, -3.24290393e-08, -3.24238686e-08,
       -3.23880705e-08, -3.23214892e-08, -3.22240807e-08, -3.20959177e-08,
       -3.19371927e-08, -3.17482213e-08, -3.15294445e-08, -3.12814294e-08,
       -3.10048696e-08, -3.07005838e-08, -3.03695135e-08, -3.17021952e-08,
       -3.19316874e-08, -3.21351185e-08, -3.23115780e-08, -3.24602255e-08,
       -3.25802955e-08, -3.26711033e-08, -3.27320505e-08, -3.27626308e-08,
       -3.27624353e-08, -3.27311588e-08, -3.26686047e-08, -3.25746907e-08,
       -3.24494538e-08, -3.22930541e-08, -3.21057787e-08, -3.18880444e-08,
       -3.16403995e-08, -3.13635238e-08, -3.10582283e-08, -3.07254524e-08,
       -3.19661435e-08, -3.22025940e-08, -3.24128819e-08, -3.25960611e-08,
       -3.27512529e-08, -3.28776521e-08, -3.29745324e-08, -3.30412529e-08,
       -3.30772642e-08, -3.30821145e-08, -3.30554558e-08, -3.29970507e-08,
       -3.29067776e-08, -3.27846366e-08, -3.26307546e-08, -3.24453891e-08,
       -3.22289318e-08, -3.19819108e-08, -3.17049911e-08, -3.13989744e-08,
       -3.10647971e-08, -3.22107451e-08, -3.24539745e-08, -3.26709663e-08,
       -3.28607390e-08, -3.30223766e-08, -3.31550343e-08, -3.32579446e-08,
       -3.33304242e-08, -3.33718804e-08, -3.33818184e-08, -3.33598474e-08,
       -3.33056883e-08, -3.32191797e-08, -3.31002846e-08, -3.29490955e-08,
       -3.27658396e-08, -3.25508827e-08, -3.23047313e-08, -3.20280348e-08,
       -3.17215845e-08, -3.13863123e-08, -3.24352494e-08, -3.26850446e-08,
       -3.29085531e-08, -3.31047592e-08, -3.32727101e-08, -3.34115219e-08,
       -3.35203866e-08, -3.35985787e-08, -3.36454629e-08, -3.36605008e-08,
       -3.36432592e-08, -3.35934170e-08, -3.35107730e-08, -3.33952522e-08,
       -3.32469127e-08, -3.30659504e-08, -3.28527044e-08, -3.26076592e-08,
       -3.23314471e-08, -3.20248483e-08, -3.16887888e-08, -3.26389589e-08,
       -3.28950733e-08, -3.31248774e-08, -3.33273224e-08, -3.35014197e-08,
       -3.36462472e-08, -3.37609571e-08, -3.38447825e-08, -3.38970454e-08,
       -3.39171650e-08, -3.39046654e-08, -3.38591840e-08, -3.37804794e-08,
       -3.36684390e-08, -3.35230855e-08, -3.33445839e-08, -3.31332455e-08,
       -3.28895323e-08, -3.26140588e-08, -3.23075927e-08, -3.19710535e-08,
       -3.28212346e-08, -3.30833883e-08, -3.33192332e-08, -3.35276885e-08,
       -3.37077309e-08, -3.38584016e-08, -3.39788135e-08, -3.40681594e-08,
       -3.41257196e-08, -3.41508714e-08, -3.41430967e-08, -3.41019920e-08,
       -3.40272760e-08, -3.39187982e-08, -3.37765467e-08, -3.36006546e-08,
       -3.33914057e-08, -3.31492385e-08, -3.28747492e-08, -3.25686922e-08,
       -3.22319792e-08, -3.29815001e-08, -3.32493807e-08, -3.34909783e-08,
       -3.37051816e-08, -3.38909341e-08, -3.40472413e-08, -3.41731784e-08,
       -3.42678988e-08, -3.43306425e-08, -3.43607454e-08, -3.43576487e-08,
       -3.43209081e-08, -3.42502033e-08, -3.41453465e-08, -3.40062911e-08,
       -3.38331385e-08, -3.36261448e-08, -3.33857248e-08, -3.31124558e-08,
       -3.28070782e-08, -3.24704949e-08, -3.31192457e-08, -3.33925095e-08,
       -3.36395397e-08, -3.38591960e-08, -3.40503903e-08, -3.42120940e-08,
       -3.43433462e-08, -3.44432623e-08, -3.45110433e-08, -3.45459852e-08,
       -3.45474893e-08, -3.45150719e-08, -3.44483743e-08, -3.43471723e-08,
       -3.42113849e-08, -3.40410824e-08, -3.38364930e-08, -3.35980079e-08,
       -3.33261849e-08, -3.30217497e-08, -3.26855958e-08, -3.32340324e-08,
       -3.35123058e-08, -3.37644175e-08, -3.39892003e-08, -3.41855361e-08,
       -3.43523641e-08, -3.44886890e-08, -3.45935899e-08, -3.46662305e-08,
       -3.47058685e-08, -3.47118667e-08, -3.46837034e-08, -3.46209825e-08,
       -3.45234445e-08, -3.43909748e-08, -3.42236131e-08, -3.40215600e-08,
       -3.37851831e-08, -3.35150205e-08, -3.32117829e-08, -3.28763531e-08,
       -3.33254950e-08, -3.36083762e-08, -3.38651893e-08, -3.40947420e-08,
       -3.42958885e-08, -3.44675377e-08, -3.46086619e-08, -3.47183060e-08,
       -3.47955980e-08, -3.48397595e-08, -3.48501163e-08, -3.48261102e-08,
       -3.47673097e-08, -3.46734205e-08, -3.45442961e-08, -3.43799459e-08,
       -3.41805438e-08, -3.39464337e-08, -3.36781342e-08, -3.33763402e-08,
       -3.30419234e-08, -3.33933450e-08, -3.36804060e-08, -3.39415134e-08,
       -3.41754516e-08, -3.43810495e-08, -3.45571878e-08, -3.47028085e-08,
       -3.48169248e-08, -3.48986312e-08, -3.49471148e-08, -3.49616670e-08,
       -3.49416949e-08, -3.48867329e-08, -3.47964541e-08, -3.46706805e-08,
       -3.45093929e-08, -3.43127390e-08, -3.40810394e-08, -3.38147931e-08,
       -3.35146792e-08, -3.31815574e-08, -3.34373730e-08, -3.37281623e-08,
       -3.39931319e-08, -3.42310457e-08, -3.44407092e-08, -3.46209775e-08,
       -3.47707650e-08, -3.48890551e-08, -3.49749115e-08, -3.50274894e-08,
       -3.50460476e-08, -3.50299609e-08, -3.49787315e-08, -3.48920015e-08,
       -3.47695634e-08, -3.46113702e-08, -3.44175443e-08, -3.41883839e-08,
       -3.39243683e-08, -3.36261607e-08, -3.32946084e-08, -3.34574503e-08,
       -3.37514954e-08, -3.40198732e-08, -3.42613296e-08, -3.44746493e-08,
       -3.46586642e-08, -3.48122637e-08, -3.49344043e-08, -3.50241213e-08,
       -3.50805406e-08, -3.51028913e-08, -3.50905176e-08, -3.50428922e-08,
       -3.49596280e-08, -3.48404898e-08, -3.46854042e-08, -3.44944694e-08,
       -3.42679617e-08, -3.40063416e-08, -3.37102559e-08, -3.33805389e-08,
       -3.34535302e-08, -3.37503405e-08, -3.40216536e-08, -3.42661995e-08,
       -3.44827449e-08, -3.46701015e-08, -3.48271362e-08, -3.49527814e-08,
       -3.50460472e-08, -3.51060327e-08, -3.51319398e-08, -3.51230852e-08,
       -3.50789142e-08, -3.49990128e-08, -3.48831198e-08, -3.47311373e-08,
       -3.45431403e-08, -3.43193846e-08, -3.40603115e-08, -3.37665521e-08,
       -3.34389272e-08, -3.34256488e-08, -3.37247189e-08, -3.39984781e-08,
       -3.42456436e-08, -3.44649665e-08, -3.46552412e-08, -3.48153152e-08,
       -3.49440999e-08, -3.50405826e-08, -3.51038392e-08, -3.51330469e-08,
       -3.51274979e-08, -3.50866127e-08, -3.50099526e-08, -3.48972326e-08,
       -3.47483320e-08, -3.45633044e-08, -3.43423853e-08, -3.40859983e-08,
       -3.37947583e-08, -3.34694725e-08, -3.33739245e-08, -3.36747374e-08,
       -3.39504412e-08, -3.41997427e-08, -3.44213805e-08, -3.46141346e-08,
       -3.47768360e-08, -3.49083782e-08, -3.50077294e-08, -3.50739448e-08,
       -3.51061803e-08, -3.51037064e-08, -3.50659214e-08, -3.49923646e-08,
       -3.48827295e-08, -3.47368744e-08, -3.45548329e-08, -3.43368218e-08,
       -3.40832472e-08, -3.37947084e-08, -3.34719985e-08, -3.32985579e-08,
       -3.36005885e-08, -3.38777261e-08, -3.41286699e-08, -3.43521490e-08,
       -3.45469319e-08, -3.47118364e-08, -3.48457414e-08, -3.49475987e-08,
       -3.50164466e-08, -3.50514228e-08, -3.50517789e-08, -3.50168939e-08,
       -3.49462882e-08, -3.48396356e-08, -3.46967757e-08, -3.45177238e-08,
       -3.43026790e-08, -3.40520311e-08, -3.37663636e-08, -3.34464556e-08,
       -3.31998300e-08, -3.35025487e-08, -3.37806038e-08, -3.40326898e-08,
       -3.42575292e-08, -3.44538820e-08, -3.46205563e-08, -3.47564193e-08,
       -3.48604103e-08, -3.49315535e-08, -3.49689717e-08, -3.49719010e-08,
       -3.49397041e-08, -3.48718848e-08, -3.47681001e-08, -3.46281729e-08,
       -3.44521018e-08, -3.42400699e-08, -3.39924510e-08, -3.37098137e-08,
       -3.33929226e-08])

Let’s benchmark these two functions to see how much faster the decorated function runs:

%timeit gravity_upward_slow(coordinates, prisms, densities)

3.56 ms ± 22.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit gravity_upward_jit(coordinates, prisms, densities)

699 µs ± 10.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

From these numbers we can see that we have significantly reduced the computation time by several factors by just decorating our function.

Note

The benchmarked times may vary if you run them in your system.

See also

Check How to measure the performance of Numba? to learn more about how to properly benchmark jitted functions.

Parallelizing our runs#

We have already shown how we can reduce the computation times of our forward modelling by decorating our functions with @numba.jit(nopython=True). But this is just the first step: all the computations were being run in serial in a single CPU. We can harness the full power of our modern multiprocessors CPUs by parallelizing our runs. To do so we need to use the numba.prange instead of the regular Python range function and slightly change the decorator of our function by adding a parallel=True:

import numba

@numba.jit(nopython=True, parallel=True)
def gravity_upward_parallel(coordinates, prisms, densities):
    """
    Compute the upward component of the acceleration of a set of prisms
    """
    # Unpack coordinates of the observation points
    easting, northing, upward = coordinates[:]
    # Initialize a result array full of zeros
    result = np.zeros_like(easting, dtype=np.float64)
    # Compute the upward component that every prism generate on each
    # observation point
    for i in numba.prange(len(easting)):
        for j in range(prisms.shape[0]):
            result[i] += gravity_u(
                easting[i], northing[i], upward[i], prisms[j, :], densities[j]
            )
    return result

g_u = gravity_upward_parallel(coordinates, prisms, densities)
g_u

array([-3.08024767e-08, -3.10103482e-08, -3.11926215e-08, -3.13484961e-08,
       -3.14772466e-08, -3.15782275e-08, -3.16508775e-08, -3.16947240e-08,
       -3.17093869e-08, -3.16945831e-08, -3.16501302e-08, -3.15759501e-08,
       -3.14720725e-08, -3.13386380e-08, -3.11759003e-08, -3.09842280e-08,
       -3.07641061e-08, -3.05161358e-08, -3.02410339e-08, -2.99396310e-08,
       -2.96128686e-08, -3.11195000e-08, -3.13346703e-08, -3.15240617e-08,
       -3.16868368e-08, -3.18222319e-08, -3.19295618e-08, -3.20082244e-08,
       -3.20577056e-08, -3.20775840e-08, -3.20675354e-08, -3.20273375e-08,
       -3.19568737e-08, -3.18561377e-08, -3.17252367e-08, -3.15643944e-08,
       -3.13739539e-08, -3.11543785e-08, -3.09062534e-08, -3.06302842e-08,
       -3.03272966e-08, -2.99982325e-08, -3.14196998e-08, -3.16420877e-08,
       -3.18385422e-08, -3.20081895e-08, -3.21502273e-08, -3.22639303e-08,
       -3.23486553e-08, -3.24038461e-08, -3.24290393e-08, -3.24238686e-08,
       -3.23880705e-08, -3.23214892e-08, -3.22240807e-08, -3.20959177e-08,
       -3.19371927e-08, -3.17482213e-08, -3.15294445e-08, -3.12814294e-08,
       -3.10048696e-08, -3.07005838e-08, -3.03695135e-08, -3.17021952e-08,
       -3.19316874e-08, -3.21351185e-08, -3.23115780e-08, -3.24602255e-08,
       -3.25802955e-08, -3.26711033e-08, -3.27320505e-08, -3.27626308e-08,
       -3.27624353e-08, -3.27311588e-08, -3.26686047e-08, -3.25746907e-08,
       -3.24494538e-08, -3.22930541e-08, -3.21057787e-08, -3.18880444e-08,
       -3.16403995e-08, -3.13635238e-08, -3.10582283e-08, -3.07254524e-08,
       -3.19661435e-08, -3.22025940e-08, -3.24128819e-08, -3.25960611e-08,
       -3.27512529e-08, -3.28776521e-08, -3.29745324e-08, -3.30412529e-08,
       -3.30772642e-08, -3.30821145e-08, -3.30554558e-08, -3.29970507e-08,
       -3.29067776e-08, -3.27846366e-08, -3.26307546e-08, -3.24453891e-08,
       -3.22289318e-08, -3.19819108e-08, -3.17049911e-08, -3.13989744e-08,
       -3.10647971e-08, -3.22107451e-08, -3.24539745e-08, -3.26709663e-08,
       -3.28607390e-08, -3.30223766e-08, -3.31550343e-08, -3.32579446e-08,
       -3.33304242e-08, -3.33718804e-08, -3.33818184e-08, -3.33598474e-08,
       -3.33056883e-08, -3.32191797e-08, -3.31002846e-08, -3.29490955e-08,
       -3.27658396e-08, -3.25508827e-08, -3.23047313e-08, -3.20280348e-08,
       -3.17215845e-08, -3.13863123e-08, -3.24352494e-08, -3.26850446e-08,
       -3.29085531e-08, -3.31047592e-08, -3.32727101e-08, -3.34115219e-08,
       -3.35203866e-08, -3.35985787e-08, -3.36454629e-08, -3.36605008e-08,
       -3.36432592e-08, -3.35934170e-08, -3.35107730e-08, -3.33952522e-08,
       -3.32469127e-08, -3.30659504e-08, -3.28527044e-08, -3.26076592e-08,
       -3.23314471e-08, -3.20248483e-08, -3.16887888e-08, -3.26389589e-08,
       -3.28950733e-08, -3.31248774e-08, -3.33273224e-08, -3.35014197e-08,
       -3.36462472e-08, -3.37609571e-08, -3.38447825e-08, -3.38970454e-08,
       -3.39171650e-08, -3.39046654e-08, -3.38591840e-08, -3.37804794e-08,
       -3.36684390e-08, -3.35230855e-08, -3.33445839e-08, -3.31332455e-08,
       -3.28895323e-08, -3.26140588e-08, -3.23075927e-08, -3.19710535e-08,
       -3.28212346e-08, -3.30833883e-08, -3.33192332e-08, -3.35276885e-08,
       -3.37077309e-08, -3.38584016e-08, -3.39788135e-08, -3.40681594e-08,
       -3.41257196e-08, -3.41508714e-08, -3.41430967e-08, -3.41019920e-08,
       -3.40272760e-08, -3.39187982e-08, -3.37765467e-08, -3.36006546e-08,
       -3.33914057e-08, -3.31492385e-08, -3.28747492e-08, -3.25686922e-08,
       -3.22319792e-08, -3.29815001e-08, -3.32493807e-08, -3.34909783e-08,
       -3.37051816e-08, -3.38909341e-08, -3.40472413e-08, -3.41731784e-08,
       -3.42678988e-08, -3.43306425e-08, -3.43607454e-08, -3.43576487e-08,
       -3.43209081e-08, -3.42502033e-08, -3.41453465e-08, -3.40062911e-08,
       -3.38331385e-08, -3.36261448e-08, -3.33857248e-08, -3.31124558e-08,
       -3.28070782e-08, -3.24704949e-08, -3.31192457e-08, -3.33925095e-08,
       -3.36395397e-08, -3.38591960e-08, -3.40503903e-08, -3.42120940e-08,
       -3.43433462e-08, -3.44432623e-08, -3.45110433e-08, -3.45459852e-08,
       -3.45474893e-08, -3.45150719e-08, -3.44483743e-08, -3.43471723e-08,
       -3.42113849e-08, -3.40410824e-08, -3.38364930e-08, -3.35980079e-08,
       -3.33261849e-08, -3.30217497e-08, -3.26855958e-08, -3.32340324e-08,
       -3.35123058e-08, -3.37644175e-08, -3.39892003e-08, -3.41855361e-08,
       -3.43523641e-08, -3.44886890e-08, -3.45935899e-08, -3.46662305e-08,
       -3.47058685e-08, -3.47118667e-08, -3.46837034e-08, -3.46209825e-08,
       -3.45234445e-08, -3.43909748e-08, -3.42236131e-08, -3.40215600e-08,
       -3.37851831e-08, -3.35150205e-08, -3.32117829e-08, -3.28763531e-08,
       -3.33254950e-08, -3.36083762e-08, -3.38651893e-08, -3.40947420e-08,
       -3.42958885e-08, -3.44675377e-08, -3.46086619e-08, -3.47183060e-08,
       -3.47955980e-08, -3.48397595e-08, -3.48501163e-08, -3.48261102e-08,
       -3.47673097e-08, -3.46734205e-08, -3.45442961e-08, -3.43799459e-08,
       -3.41805438e-08, -3.39464337e-08, -3.36781342e-08, -3.33763402e-08,
       -3.30419234e-08, -3.33933450e-08, -3.36804060e-08, -3.39415134e-08,
       -3.41754516e-08, -3.43810495e-08, -3.45571878e-08, -3.47028085e-08,
       -3.48169248e-08, -3.48986312e-08, -3.49471148e-08, -3.49616670e-08,
       -3.49416949e-08, -3.48867329e-08, -3.47964541e-08, -3.46706805e-08,
       -3.45093929e-08, -3.43127390e-08, -3.40810394e-08, -3.38147931e-08,
       -3.35146792e-08, -3.31815574e-08, -3.34373730e-08, -3.37281623e-08,
       -3.39931319e-08, -3.42310457e-08, -3.44407092e-08, -3.46209775e-08,
       -3.47707650e-08, -3.48890551e-08, -3.49749115e-08, -3.50274894e-08,
       -3.50460476e-08, -3.50299609e-08, -3.49787315e-08, -3.48920015e-08,
       -3.47695634e-08, -3.46113702e-08, -3.44175443e-08, -3.41883839e-08,
       -3.39243683e-08, -3.36261607e-08, -3.32946084e-08, -3.34574503e-08,
       -3.37514954e-08, -3.40198732e-08, -3.42613296e-08, -3.44746493e-08,
       -3.46586642e-08, -3.48122637e-08, -3.49344043e-08, -3.50241213e-08,
       -3.50805406e-08, -3.51028913e-08, -3.50905176e-08, -3.50428922e-08,
       -3.49596280e-08, -3.48404898e-08, -3.46854042e-08, -3.44944694e-08,
       -3.42679617e-08, -3.40063416e-08, -3.37102559e-08, -3.33805389e-08,
       -3.34535302e-08, -3.37503405e-08, -3.40216536e-08, -3.42661995e-08,
       -3.44827449e-08, -3.46701015e-08, -3.48271362e-08, -3.49527814e-08,
       -3.50460472e-08, -3.51060327e-08, -3.51319398e-08, -3.51230852e-08,
       -3.50789142e-08, -3.49990128e-08, -3.48831198e-08, -3.47311373e-08,
       -3.45431403e-08, -3.43193846e-08, -3.40603115e-08, -3.37665521e-08,
       -3.34389272e-08, -3.34256488e-08, -3.37247189e-08, -3.39984781e-08,
       -3.42456436e-08, -3.44649665e-08, -3.46552412e-08, -3.48153152e-08,
       -3.49440999e-08, -3.50405826e-08, -3.51038392e-08, -3.51330469e-08,
       -3.51274979e-08, -3.50866127e-08, -3.50099526e-08, -3.48972326e-08,
       -3.47483320e-08, -3.45633044e-08, -3.43423853e-08, -3.40859983e-08,
       -3.37947583e-08, -3.34694725e-08, -3.33739245e-08, -3.36747374e-08,
       -3.39504412e-08, -3.41997427e-08, -3.44213805e-08, -3.46141346e-08,
       -3.47768360e-08, -3.49083782e-08, -3.50077294e-08, -3.50739448e-08,
       -3.51061803e-08, -3.51037064e-08, -3.50659214e-08, -3.49923646e-08,
       -3.48827295e-08, -3.47368744e-08, -3.45548329e-08, -3.43368218e-08,
       -3.40832472e-08, -3.37947084e-08, -3.34719985e-08, -3.32985579e-08,
       -3.36005885e-08, -3.38777261e-08, -3.41286699e-08, -3.43521490e-08,
       -3.45469319e-08, -3.47118364e-08, -3.48457414e-08, -3.49475987e-08,
       -3.50164466e-08, -3.50514228e-08, -3.50517789e-08, -3.50168939e-08,
       -3.49462882e-08, -3.48396356e-08, -3.46967757e-08, -3.45177238e-08,
       -3.43026790e-08, -3.40520311e-08, -3.37663636e-08, -3.34464556e-08,
       -3.31998300e-08, -3.35025487e-08, -3.37806038e-08, -3.40326898e-08,
       -3.42575292e-08, -3.44538820e-08, -3.46205563e-08, -3.47564193e-08,
       -3.48604103e-08, -3.49315535e-08, -3.49689717e-08, -3.49719010e-08,
       -3.49397041e-08, -3.48718848e-08, -3.47681001e-08, -3.46281729e-08,
       -3.44521018e-08, -3.42400699e-08, -3.39924510e-08, -3.37098137e-08,
       -3.33929226e-08])

With numba.prange we can specify which loop we want to run in parallel. Since we are updating the values of results on each iteration, it’s advisable to parallelize the loop over the observation points. By setting parallel=True in the decorator we are telling Numba to pararellize this function, otherwise Numba will reinterpret the numba.prange function as a regular range and run this loop in serial.

Note

In some applications it’s desirable that our forward models are run in serial. For example, if they are part of larger problem that gets parallelized at a higher level. The parallel parameter in the numba.jit decorator allows us to change this behaviour at will without having to modify the function code.

Let’s benchmark this function against the non-parallelized gravity_upward_jit:

%timeit gravity_upward_jit(coordinates, prisms, densities)

694 µs ± 2.84 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

%timeit gravity_upward_parallel(coordinates, prisms, densities)

348 µs ± 3.42 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

By distributing the load between multiple processors we were capable of lowering the computation time by a few more factors.