fatiando / harmonica

Forward modeling, inversion, and processing gravity and magnetic data
https://www.fatiando.org/harmonica
BSD 3-Clause "New" or "Revised" License
208 stars 68 forks source link

Check for observartion point inside tesseroid eats too much memory #417

Closed santisoler closed 1 year ago

santisoler commented 1 year ago

Description of the problem:

The tesseroid forward modelling function checks if no observation point falls inside any tesseroid through the _check_points_outside_tesseroids function: https://github.com/fatiando/harmonica/blob/b8f037e6d72b04cb0b4ff62bb12b63dd13e38808/harmonica/_forward/_tesseroid_utils.py#L405-L461

This functions makes use of Numpy arrays to run these checks, but this is super inefficient, since Numpy will try to allocate massive arrays of shapes (n_coords, n_tesseroids).

Solution We should rewrite this function using Numba instead, making sure that we aren't allocating any additional array and the function just checks if observation points don't fall inside any tesseroids. We could also benefit from Numba parallelizations and run the first for loop in parallel (this should follow the user's choice set through the parallel argument in tesseroid_gravity).

As a quick design, we could do something like:


@numba.jit(nopython=True, parallel=True)
def check_points_outside_tesseroids(coordinates, tesseroids):
    """..."""
    longitude, latitude, radius = coordinates[:]
    for i in numba.prange(longitude.size):
        for j in range(tesseroids.shape[0]):
            _check_point_inside_tesseroid(
                longitude[i],
                latitude[i],
                radius[i],
                tesseroids[j, 0],
                tesseroids[j, 1],
                tesseroids[j, 2],
                tesseroids[j, 3],
                tesseroids[j, 4],
                tesseroids[j, 5],
            )

@numba.jit(nopython=True)
def _check_point_inside_tesseroid(
    longitude, latitude, radius, west, east, south, north, bottom, top
):
    """Raise error if observation point falls inside tesseroid"""
    ...
    if ... :
        raise ValueError(...)

Check out https://github.com/fatiando/harmonica/blob/main/harmonica/_forward/prism.py for inspiration on how to code this.

Temporary solution

For the moment, the check can be bypassed by passing disable_checks=True, but it also disables other checks.