ENH: Choclo for Equivalent Source about simpeg HOT 4 OPEN

Thanks for opening this issue @thibaut-kobold! I think it's fine if we wait for #1321 to be merged to add a minimal example that reproduce the error (kernel shutdown) you faced. No rush for that atm.

from simpeg.

I've wrote a minimal example that (I suspect) reproduces the error that @thibaut-kobold was facing:

# file: eq_sources.py
import numpy as np

import SimPEG
from SimPEG.potential_fields import gravity
import discretize

region = (-10e3, 10e3, -10e3, 10e3)

# Define observation points
rng = np.random.default_rng(seed=42)

n_points = 100
easting = rng.uniform(low=region[0], high=region[1], size=n_points)
northing = rng.uniform(low=region[0], high=region[1], size=n_points)
upward = rng.normal(loc=100, scale=30, size=n_points)
receivers_locations = np.vstack((easting, northing, upward)).T

# Define mesh
cell_size = 1e3
n_cells_easting = int((region[1] - region[0]) / cell_size)
n_cells_northing = int((region[3] - region[2]) / cell_size)
h_easting = [(cell_size, n_cells_easting)]
h_northing = [(cell_size, n_cells_northing)]
h = [h_easting, h_northing]
mesh = discretize.TensorMesh(h, origin="CC")

# Define survey
receivers = gravity.Point(receivers_locations)
source_field = gravity.SourceField([receivers])
survey = gravity.Survey(source_field)

# Define equivalent sources
eq_sources = gravity.SimulationEquivalentSourceLayer(
    mesh,
    cell_z_top=-100,
    cell_z_bottom=-500,
    survey=survey,
    rhoMap=SimPEG.maps.IdentityMap(),
    engine="choclo",
)

# Build sensitivity matrix
sensitivities = eq_sources.G

When trying to run this, we get a bus error:

$ python eq_sources.py
zsh: bus error  python eq_sources.py

To debug this I disabled Numba and we get the following error:

$ NUMBA_DISABLE_JIT=1 python eq_sources.py
Traceback (most recent call last):
  File "/home/santi/git/simpeg/eq_sources.py", line 43, in <module>
    sensitivities = eq_sources.G
                    ^^^^^^^^^^^^
  File "/home/santi/git/simpeg/SimPEG/potential_fields/gravity/simulation.py", line 255, in G
    self._G = self._sensitivity_matrix()
              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/santi/git/simpeg/SimPEG/potential_fields/gravity/simulation.py", line 441, in _sensitivity_matrix
    self._sensitivity_gravity(
  File "/home/santi/git/simpeg/SimPEG/potential_fields/gravity/_numba_functions.py", line 159, in _sensitivity_gravity
    nodes[j, 2],
    ~~~~~^^^^^^
IndexError: index 2 is out of bounds for axis 1 with size 2

Basically, the source of the problem is that the simulation is getting a 2D mesh, while the Numba implementation was written to expect a 3D mesh only: the nodes should be a 2D array containing 3 columns, as the docstring states, instead of a 2D array with 2 columns.

I'll keep digging into this to find a solution. But I wanted to keep a working example here as record.

from simpeg.

I put more thoughts into this and I think the solution to this issue is not very trivial. What's going on is that the equivalent source class takes a 2d mesh and from that it generates the nodes corresponding to the 3d prisms located in a single layer (in a similar way the prism layer is defined in Harmonica). The vertical coordinates of those nodes are defined by the cell_z_top and cell_z_bottom arguments.

In the Numba implementation, the array with the location of the nodes is generated on the fly from the mesh while running the forward: they aren't stored in an attribute so they cannot be overwritten. For this reason, the Numba implementation relies on having a 3D mesh whose neighboring prisms share nodes.

This is not the case for the current equivalent source class, since one can define a layer of prisms that don't share any node. In such scenario, I think that computing the kernels on every node and then computing the additions and subtractions is actually less efficient than iterating through the prisms in the layer.

The current implementation using the nodes does something like (not working example, just something similar to pseudocode):

# Define a result array for the final gravity fields
fields = np.zeros(n_receivers)
# Iterate over receivers
for i in range(n_receivers):
    # Define an empty array for storing the evaluation of kernels on the nodes in the mesh
    kernels = np.empty(n_nodes)
   # Iterate over the nodes in the mesh and evaluate the kernels on them
    for j in range(n_nodes):
        kernels[j] = _evaluate_kernel(receivers[i], nodes[j])
    # Compute the gravity fields from the kernel values for each cell in the mesh
    for k in range(n_cells):
        fields[i] += densities[k] * _kernels_in_nodes_to_cell(kernels, cell_nodes[k])

The advantage of this implementation is that evaluates the kernels (a costly operation) on each node of the mesh only once. This is more efficient than iterating over the prisms in the mesh and forward modelling each of them independently because they share nodes with neighboring cells.

In case they don't share nodes, iterating over prisms is actually faster. It would look something like (not working example):

# Define a result array for the final gravity fields
fields = np.zeros(n_receivers)
# Iterate over receivers
for i in range(n_receivers):
    # Iterate over prisms
    for j in range(n_prisms):
        fields[i] += densities[j] * gravity_field(receivers[i], prisms[j])

Where gravity_field would be a function that forward models a single prism for a single observation point (like the ones in Choclo).

The advantage of this implementation (for a set of arbitrary cases that don't share nodes) is that we can spare the memory required to store the kernels and the nodes arrays, and every array (densities, receivers, prisms, fields) is read or written sequentially, which is more efficient.

Therefore, I think there's room to optimize this here besides just making the equivalent sources work with Choclo.

Possible solutions

Idea 1

One quick and dirty solution we could implement is overwriting the private _get_active_nodes method in SimulationEquivalentSourceLayer that returns the self._nodes and self._unique_inv attributes:

I call it dirty because I don't like the idea of overwriting a parent method, violating Liskov substitution principle. Also, this solution won't take advantage of the optimization bits I mentioned before.

Idea 2

Draft a reimplementation of the equivalent source class that takes advantage of the optimization bits I mentioned before. We could make it clever enough to run a "nodes" implementation in case the bottom and top are constant, or a "prisms" implementation if those are given as arrays.

Moving forward?

I would wait until #1321 is merged, so the BasePFSimulation gets some of the bits I'm moving from the gravity simulation up there, and also we have the magnetic simulation with Choclo implemented. That way it would be easier to apply changes to the equivalent source classes both for the gravity and the magnetic one.

Open to hear ideas

Feel free to share thoughts. I just wanted to capture the ones I had today, but I'm sure there's more to consider here.

from simpeg.

Overall the prism based approach and the node based approach both make the same number of calls to the kernels for the equivalent source layer simulations. The benefit of the prism approach is really only the lack of storing the calls to the kernel.

Your assumption on things being more ordered when working with prisms is only true when working with structured meshes, where the ordering of nodes in memory is related to the ordering of the cells. But even in this case, the nodes are still accessed sequentially in memory when moving between cells. We could probably have an even faster implementation for TensorMesh’s in these cases.

In a TreeMesh, the node ordering is not related to the ordering of cells. The cells actually follow a z-ordering, and the nodes ordering is based on cantor pairing their underlying spatial indices.

from simpeg.