Repository: simpeg/simpeg
Pull request #1310 (2023-12-20 01:43:02)

Re-implement distance weighting and add a strategy comparison
#### Squash and Merge summary

Add a new `distance_weighting` function with a new Numba-based implementation that makes it more memory efficient that using Scipy's `spatial_distance.cdist`. Add more tests and a tutorial comparing depth, distance, and sensitivity weighting.

---

#### Summary
distance weighting was sitting in an old location, with old naming convention, and quite inefficient to compute.
This PR brings distance weighting at the same level of maturity as depth weighting, with a 10x speed up.
This PR also adds a new tutorial comparing depth, distance, and sensitivity weighting.

#### PR Checklist
* [x] If this is a work in progress PR, set as a Draft PR
* [x] Linted my code according to the [style guides](https://docs.simpeg.xyz/content/getting_started/contributing/code-style.html).
* [ ] Added [tests](https://docs.simpeg.xyz/content/getting_started/practices.html#testing) to verify changes to the code.
* [x] Added necessary documentation to any new functions/classes following the
      expect [style](https://docs.simpeg.xyz/content/getting_started/practices.html#documentation).
* [ ] Marked as ready for review (if this is was a draft PR), and converted 
      to a Pull Request
* [ ] Tagged ``@simpeg/simpeg-developers`` when ready for review.

#### Reference issue
<!--Example: write "Closes #NNNN" to automatically close that issue on merge.-->

#### What does this implement/fix?
<!--Please explain your changes.-->

#### Additional information
<!--Any additional information you think is important.-->


<!--
Once all tests pass and the code has been reviewed and approved, it will be merged into main
-->

**Visualizing weights**
<img width="985" alt="image" src="https://github.com/simpeg/simpeg/assets/97514898/9e6eb0c0-b184-4f10-b749-43779b23b292">

**efficiency comparison with former `get_dist_wgt` on a small example**

<img width="913" alt="Screenshot 2023-12-19 at 6 26 22 PM" src="https://github.com/simpeg/simpeg/assets/97514898/20e7edd2-caa7-49cd-ade6-801847bf3754">

**comparison with former `get_dist_wgt` on a small example**
<img width="802" alt="image" src="https://github.com/simpeg/simpeg/assets/97514898/92bac62c-8fd3-4934-bf1d-2a3a717ea8e1">
<img width="820" alt="image" src="https://github.com/simpeg/simpeg/assets/97514898/4f87cfde-3a6a-428f-a2a4-9e7422d8288e">


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
Attention: `17 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`102f323`)](https://app.codecov.io/gh/simpeg/simpeg/commit/102f32356260d15bbd90f748948155fbfd6676bd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 82.35% compared to head [(`a925fad`)](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 47.88%.

| [Files](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) | Patch % | Lines |
|---|---|---|
| [SimPEG/utils/model\_utils.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL3V0aWxzL21vZGVsX3V0aWxzLnB5) | 22.72% | [17 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1310       +/-   ##
===========================================
- Coverage   82.35%   47.88%   -34.47%     
===========================================
  Files         165      165               
  Lines       25239    25255       +16     
===========================================
- Hits        20785    12094     -8691     
- Misses       4454    13161     +8707     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/simpeg/simpeg/pull/1310?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).

jcapriot: Where was distance weighting previously? I don’t see any removals here in my quick glance.

The other thing is that this creates multiple large dense matrices of (n_rx, n_cells), which will likely put a lot of pressure on memory.
jcapriot: Also, I don’t think you need to put in the deprecated keyword argument check here since it’s a new function. No need to create something that is automatically deprecated.
thibaut-kobold: Hi @jcapriot : this is where distance weighting is located at the moment: https://docs.simpeg.xyz/content/api/generated/SimPEG.potential_fields.get_dist_wgt.html#SimPEG.potential_fields.get_dist_wgt

I have not removed or depreciated it yet. It is not used anywhere (and is not even tested). What do you think? Can we simply remove it, or depreciate it but still functional, or depreciate it and throw an error pointing to the new function?
thibaut-kobold: > The other thing is that this creates multiple large dense matrices of (n_rx, n_cells), which will likely put a lot of pressure on memory.

There is one big matrix `distance` indeed. For potential field, this is the same size as the sensitivity matrix that most keep in memory. I don't think the code is keeping it once we return the final (n_cells) array.
jcapriot: Thanks for pointing to that. I would say we should just mark that function as deprecated and point to the new one.

Just so that everything is clear, the old function computes the radial distance to 8 points surrounding a cell center, and uses the average weight, as opposed to this implementation which only computes to the cell center. This likely accounts for much of the speedup you're seeing (8 calculations in 3D vs 1 calculation here). Let me track through the history to see who originally wrote the function so we can understand the reasoning
jcapriot: Looks like this is the oldest bit I can find before the history is lost: https://github.com/simpeg/simpeg/blame/487d178b706e0d5549a7f9a2494d5b737f18474e/SimPEG/potential_fields/magnetics/simulation.py#L1068

In there it says that @domfournier created it in 2015.
domfournier: I probably got this from a paper or the old Fortran code. I don't have an opinion to go whichever way, other than intuitively the nodal approach would be closer to the sensitivity calc.
I can do a review too now.
thibaut-kobold: I am not familiar with `numba` but I took a very quick shot at it, proposing the option to either use `numba` or `numpy` (similar in idea with the choice of `geoana` or `choclo` for PF)
dccowan: It's great that you made a tutorial. Please note that new tutorials are being developed within github.com/simpeg/user-tutorials ; see https://simpeg.xyz/user-tutorials/ .

Another option would be add the discussion of weighting to the current tutorials for magnetic and gravity inversion. 
jcapriot: > Since each element of the unnormalized distance weights array is independent from each other, it sounds better to me to parallelize the loop over the active cells in the mesh, so each thread will be writing to a single element of the distance_weights array. This way we avoid racing conditions in which two threads try to write the same memory address, and therefore will have to wait for the first one to finish.

Just fyi, from what I understand the llvm compiler should recognize that there is a reduction operation here when in parallel. They don’t write to the same location during parallel execution, but parallelizing over the active cells is likelier faster as well, because the operation over the data is more aligned in memory.
thibaut-kobold: i am only coming back to this PR now: thank you for the in-depth comment, and proposing an actual working implementation of an alternative; this is great stuff.
Anyway, I swapped my function with yours; things looks good, but tests are failing for things that seems unrelated to that PR

Patch:
diff --git a/simpeg/potential_fields/base.py b/simpeg/potential_fields/base.py
index 524a4e541..f99757533 100644
--- a/simpeg/potential_fields/base.py
+++ b/simpeg/potential_fields/base.py
@@ -544,6 +544,12 @@ def get_dist_wgt(mesh, receiver_locations, actv, R, R0):
     wr : (n_cell) numpy.ndarray
         Distance weighting model; 0 for all inactive cells
     """
+    warnings.warn(
+        "The get_dist_wgt function has been deprecated, please import "
+        "simpeg.utils.distance_weighting. This will be removed in SimPEG 0.24.0",
+        FutureWarning,
+        stacklevel=2,
+    )
     # Find non-zero cells
     if actv.dtype == "bool":
         inds = (
diff --git a/simpeg/utils/__init__.py b/simpeg/utils/__init__.py
index de6cb066b..6e0895fd4 100644
--- a/simpeg/utils/__init__.py
+++ b/simpeg/utils/__init__.py
@@ -76,6 +76,7 @@ Model Utility Functions
   :toctree: generated/
 
   depth_weighting
+  distance_weighting
   model_builder.add_block
   model_builder.create_2_layer_model
   model_builder.create_block_in_wholespace
@@ -235,7 +236,7 @@ from .coord_utils import (
     rotation_matrix_from_normals,
     rotate_points_from_normals,
 )
-from .model_utils import depth_weighting
+from .model_utils import depth_weighting, distance_weighting
 from .plot_utils import plot2Ddata, plotLayer, plot_1d_layer_model
 from .io_utils import download
 from .pgi_utils import (
diff --git a/simpeg/utils/model_utils.py b/simpeg/utils/model_utils.py
index 91df15da7..dc48a0e17 100644
--- a/simpeg/utils/model_utils.py
+++ b/simpeg/utils/model_utils.py
@@ -1,8 +1,27 @@
-from .mat_utils import mkvc
+import warnings
+from typing import Literal, Optional
+
+import discretize
 import numpy as np
+import scipy.sparse as sp
 from scipy.interpolate import griddata
 from scipy.spatial import cKDTree
-import scipy.sparse as sp
+from scipy.spatial.distance import cdist
+
+from .mat_utils import mkvc
+
+try:
+    import numba
+    from numba import njit, prange
+except ImportError:
+    numba = None
+
+    # Define dummy njit decorator
+    def njit(*args, **kwargs):
+        return lambda f: f
+
+    # Define dummy prange function
+    prange = range
 
 
 def surface_layer_index(mesh, topo, index=0):
@@ -150,3 +169,174 @@ def depth_weighting(
         wz = wz[active_cells]
 
     return wz / np.nanmax(wz)
+
+
+@njit(parallel=True)
+def _distance_weighting_numba(
+    cell_centers: np.ndarray,
+    reference_locs: np.ndarray,
+    threshold: float,
+    exponent: float = 2.0,
+) -> np.ndarray:
+    r"""
+    distance weighting kernel in numba.
+
+    If numba is not installed, this will work as a regular for loop.
+
+    Parameters
+    ----------
+    cell_centers : np.ndarray
+        cell centers of the mesh.
+    reference_locs : (n, ndim) numpy.ndarray
+        The coordinate of the reference location, usually the receiver locations,
+        for the distance weighting.
+        A 2d array, with multiple reference locations, where each row should
+        contain the coordinates of a single location point in the following
+        order: _x_, _y_, _z_ (for 3D meshes) or _x_, _z_ (for 2D meshes).
+    threshold : float
+        Threshold parameters used in the distance weighting.
+    exponent : float, optional
+        Exponent parameter for distance weighting.
+        The exponent should match the natural decay power of the potential
+        field. For example, for gravity acceleration, set it to 2; for magnetic
+        fields, to 3.
+
+    Returns
+    -------
+    (n_active) numpy.ndarray
+        Normalized distance weights for the mesh at every active cell as
+        a 1d-array.
+    """
+    n_active_cells = cell_centers.shape[0]
+    n_reference_locs = len(reference_locs)
+
+    distance_weights = np.zeros(n_active_cells)
+    for j in prange(n_active_cells):
+        cell_center = cell_centers[j]
+        for i in range(n_reference_locs):
+            reference_loc = reference_locs[i]
+            distance = np.sqrt(((cell_center - reference_loc) ** 2).sum())
+            distance_weights[j] += (distance + threshold) ** (-2 * exponent)
+
+    distance_weights = np.sqrt(distance_weights)
+    distance_weights /= np.nanmax(distance_weights)
+    return distance_weights
+
+
+def distance_weighting(
+    mesh: discretize.base.BaseMesh,
+    reference_locs: np.ndarray,
+    active_cells: Optional[np.ndarray] = None,
+    exponent: Optional[float] = 2.0,
+    threshold: Optional[float] = None,
+    engine: Literal["numba", "scipy"] = "numba",
+    cdist_opts: Optional[dict] = None,
+):
+    r"""
+    Construct diagonal elements of a distance weighting matrix
+
+    Builds the model weights following the distance weighting strategy, a method
+    to generate weights based on the distance between mesh cell centers and some
+    reference location(s).
+    Use these weights in regularizations to counteract the natural decay of
+    potential field data with distance.
+
+    Parameters
+    ----------
+    mesh : discretize.base.BaseMesh
+        Discretized model space.
+    reference_locs : (n, ndim) numpy.ndarray
+        The coordinate of the reference location, usually the receiver locations,
+        for the distance weighting.
+        A 2d array, with multiple reference locations, where each row should
+        contain the coordinates of a single location point in the following
+        order: _x_, _y_, _z_ (for 3D meshes) or _x_, _z_ (for 2D meshes).
+    active_cells : (mesh.n_cells) numpy.ndarray of bool, optional
+        Index vector for the active cells on the mesh.
+        If ``None``, every cell will be assumed to be active.
+    exponent : float or None, optional
+        Exponent parameter for distance weighting.
+        The exponent should match the natural decay power of the potential
+        field. For example, for gravity acceleration, set it to 2; for magnetic
+        fields, to 3.
+    threshold : float or None, optional
+        Threshold parameters used in the distance weighting.
+        If ``None``, it will be set to half of the smallest cell width.
+    engine: str, 'numba' or 'scipy'
+        Pick between a ``scipy.spatial.distance.cdist`` computation (memory
+        intensive) or `for` loop implementation, parallelized with numba if
+        available. Default to ``"numba"``.
+    cdist_opts: dict, optional
+        Only valid with ``engine=="scipy"``. Options to pass to
+        ``scipy.spatial.distance.cdist``. Default to None.
+
+    Returns
+    -------
+    (n_active) numpy.ndarray
+        Normalized distance weights for the mesh at every active cell as
+        a 1d-array.
+    """
+
+    active_cells = (
+        np.ones(mesh.n_cells, dtype=bool) if active_cells is None else active_cells
+    )
+
+    # Default threshold value
+    if threshold is None:
+        threshold = 0.5 * mesh.h_gridded.min()
+
+    reference_locs = np.atleast_2d(reference_locs)
+    cell_centers = mesh.cell_centers[active_cells]
+
+    # address 1D case
+    if mesh.dim == 1:
+        cell_centers = cell_centers.reshape(-1, 1)
+        reference_locs = reference_locs.reshape(-1, 1)
+
+    if reference_locs.shape[1] != mesh.dim:
+        raise ValueError(
+            f"Invalid 'reference_locs' with shape '{reference_locs.shape}'. "
+            "The number of columns of the reference_locs array should match "
+            f"the dimensions of the mesh ({mesh.dim})."
+        )
+
+    if engine == "numba" and cdist_opts is not None:
+        raise TypeError(
+            "The `cdist_opts` is valid only when engine is 'scipy'."
+            "The current engine is 'numba'."
+        )
+
+    if engine == "numba":
+        if numba is None:
+            warnings.warn(
+                "Numba is not installed. Distance computations will be slower.",
+                stacklevel=2,
+            )
+        distance_weights = _distance_weighting_numba(
+            cell_centers,
+            reference_locs,
+            exponent=exponent,
+            threshold=threshold,
+        )
+
+    elif engine == "scipy":
+        warnings.warn(
+            "``scipy.spatial.distance.cdist`` computations can be memory intensive. "
+            "Consider switching to `engine='numba'` "
+            "if you run into memory overflow issues.",
+            stacklevel=2,
+        )
+        cdist_opts = cdist_opts or dict()
+        distance = cdist(cell_centers, reference_locs, **cdist_opts)
+
+        distance_weights = (((distance + threshold) ** exponent) ** -2).sum(axis=1)
+
+        distance_weights = distance_weights**0.5
+        distance_weights /= np.nanmax(distance_weights)
+
+    else:
+        raise ValueError(
+            f"Invalid engine '{engine}'. Engine should be either 'scipy' or 'numba'."
+        )
+
+    return distance_weights
diff --git a/tutorials/03-gravity/plot_inv_1c_gravity_anomaly_irls_compare_weighting.py b/tutorials/03-gravity/plot_inv_1c_gravity_anomaly_irls_compare_weighting.py
new file mode 100644
index 000000000..2fc2d9210
--- /dev/null
+++ b/tutorials/03-gravity/plot_inv_1c_gravity_anomaly_irls_compare_weighting.py
@@ -0,0 +1,576 @@
+"""
+Compare weighting strategy with Inversion of surface Gravity Anomaly Data
+=========================================================================
+
+Here we invert gravity anomaly data to recover a density contrast model. We formulate the inverse problem as an iteratively
+re-weighted least-squares (IRLS) optimization problem. For this tutorial, we
+focus on the following:
+
+    - Setting regularization weights
+    - Defining the survey from xyz formatted data
+    - Generating a mesh based on survey geometry
+    - Including surface topography
+    - Defining the inverse problem (data misfit, regularization, optimization)
+    - Specifying directives for the inversion
+    - Setting sparse and blocky norms
+    - Plotting the recovered model and data misfit
+
+Although we consider gravity anomaly data in this tutorial, the same approach
+can be used to invert gradiometry and other types of geophysical data.
+"""
+
+#########################################################################
+# Import modules
+# --------------
+#
+
+import os
+import tarfile
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+from discretize import TensorMesh
+from discretize.utils import active_from_xyz
+
+from simpeg import (
+    data,
+    data_misfit,
+    directives,
+    inverse_problem,
+    inversion,
+    maps,
+    optimization,
+    regularization,
+    utils,
+)
+from simpeg.potential_fields import gravity
+from simpeg.utils import model_builder, plot2Ddata
+
+# sphinx_gallery_thumbnail_number = 3
+
+#############################################
+# Define File Names
+# -----------------
+#
+# File paths for assets we are loading. To set up the inversion, we require
+# topography and field observations. The true model defined on the whole mesh
+# is loaded to compare with the inversion result. These files are stored as a
+# tar-file on our google cloud bucket:
+# "https://storage.googleapis.com/simpeg/doc-assets/gravity.tar.gz"
+#
+
+# storage bucket where we have the data
+data_source = "https://storage.googleapis.com/simpeg/doc-assets/gravity.tar.gz"
+
+# download the data
+downloaded_data = utils.download(data_source, overwrite=True)
+
+# unzip the tarfile
+tar = tarfile.open(downloaded_data, "r")
+tar.extractall()
+tar.close()
+
+# path to the directory containing our data
+dir_path = downloaded_data.split(".")[0] + os.path.sep
+
+# files to work with
+topo_filename = dir_path + "gravity_topo.txt"
+data_filename = dir_path + "gravity_data.obs"
+
+
+#############################################
+# Load Data and Plot
+# ------------------
+#
+# Here we load and plot synthetic gravity anomaly data. Topography is generally
+# defined as an (N, 3) array. Gravity data is generally defined with 4 columns:
+# x, y, z and data.
+#
+
+# Load topography
+xyz_topo = np.loadtxt(str(topo_filename))
+
+# Load field data
+dobs = np.loadtxt(str(data_filename))
+
+# Define receiver locations and observed data
+receiver_locations = dobs[:, 0:3]
+dobs = dobs[:, -1]
+
+# Plot
+mpl.rcParams.update({"font.size": 12})
+fig = plt.figure(figsize=(7, 5))
+
+ax1 = fig.add_axes([0.1, 0.1, 0.73, 0.85])
+plot2Ddata(
+    receiver_locations,
+    dobs,
+    ax=ax1,
+    contourOpts={"cmap": "bwr"},
+    shade=True,
+    nx=20,
+    ny=20,
+    dataloc=True,
+)
+ax1.set_title("Gravity Anomaly")
+ax1.set_xlabel("x (m)")
+ax1.set_ylabel("y (m)")
+
+ax2 = fig.add_axes([0.8, 0.1, 0.03, 0.85])
+norm = mpl.colors.Normalize(vmin=-np.max(np.abs(dobs)), vmax=np.max(np.abs(dobs)))
+cbar = mpl.colorbar.ColorbarBase(
+    ax2, norm=norm, orientation="vertical", cmap=mpl.cm.bwr, format="%.1e"
+)
+cbar.set_label("$mGal$", rotation=270, labelpad=15, size=12)
+
+plt.show()
+
+#############################################
+# Assign Uncertainties
+# --------------------
+#
+# Inversion with simpeg requires that we define the standard deviation of our data.
+# This represents our estimate of the noise in our data. For a gravity inversion,
+# a constant floor value is generally applied to all data. For this tutorial,
+# the standard deviation on each datum will be 1% of the maximum observed
+# gravity anomaly value.
+#
+
+maximum_anomaly = np.max(np.abs(dobs))
+
+uncertainties = 0.01 * maximum_anomaly * np.ones(np.shape(dobs))
+
+#############################################
+# Defining the Survey
+# -------------------
+#
+# Here, we define the survey that will be used for this tutorial. Gravity
+# surveys are simple to create. The user only needs an (N, 3) array to define
+# the xyz locations of the observation locations. From this, the user can
+# define the receivers and the source field.
+#
+
+# Define the receivers. The data consists of vertical gravity anomaly measurements.
+# The set of receivers must be defined as a list.
+receiver_list = gravity.receivers.Point(receiver_locations, components="gz")
+
+receiver_list = [receiver_list]
+
+# Define the source field
+source_field = gravity.sources.SourceField(receiver_list=receiver_list)
+
+# Define the survey
+survey = gravity.survey.Survey(source_field)
+
+#############################################
+# Defining the Data
+# -----------------
+#
+# Here is where we define the data that is inverted. The data is defined by
+# the survey, the observation values and the standard deviation.
+#
+
+data_object = data.Data(survey, dobs=dobs, standard_deviation=uncertainties)
+
+
+#############################################
+# Defining a Tensor Mesh
+# ----------------------
+#
+# Here, we create the tensor mesh that will be used to invert gravity anomaly
+# data. If desired, we could define an OcTree mesh.
+#
+
+dh = 5.0
+hx = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+hy = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+hz = [(dh, 5, -1.3), (dh, 15)]
+mesh = TensorMesh([hx, hy, hz], "CCN")
+
+########################################################
+# Starting/Reference Model and Mapping on Tensor Mesh
+# ---------------------------------------------------
+#
+# Here, we create starting and/or reference models for the inversion as
+# well as the mapping from the model space to the active cells. Starting and
+# reference models can be a constant background value or contain a-priori
+# structures.
+#
+
+# Find the indices of the active cells in forward model (ones below surface)
+ind_active = active_from_xyz(mesh, xyz_topo)
+
+# Define mapping from model to active cells
+nC = int(ind_active.sum())
+model_map = maps.IdentityMap(nP=nC)  # model consists of a value for each active cell
+
+# Define and plot starting model
+starting_model = np.zeros(nC)
+
+
+##############################################
+# Define the Physics and data misfit
+# ----------------------------------
+#
+# Here, we define the physics of the gravity problem by using the simulation
+# class.
+#
+
+simulation = gravity.simulation.Simulation3DIntegral(
+    survey=survey, mesh=mesh, rhoMap=model_map, ind_active=ind_active
+)
+
+# Define the data misfit. Here the data misfit is the L2 norm of the weighted
+# residual between the observed data and the data predicted for a given model.
+# Within the data misfit, the residual between predicted and observed data are
+# normalized by the data's standard deviation.
+dmis = data_misfit.L2DataMisfit(data=data_object, simulation=simulation)
+
+
+#######################################################################
+# Running the Depth Weighted inversion
+# ------------------------------------
+#
+# Here we define the directives, weights, regularization, and optimization
+# for a depth-weighted inversion
+#
+
+# inversion directives
+# Defining a starting value for the trade-off parameter (beta) between the data
+# misfit and the regularization.
+starting_beta = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
+
+# Defines the directives for the IRLS regularization. This includes setting
+# the cooling schedule for the trade-off parameter.
+update_IRLS = directives.Update_IRLS(
+    f_min_change=1e-4,
+    max_irls_iterations=30,
+    coolEpsFact=1.5,
+    beta_tol=1e-2,
+)
+
+# Options for outputting recovered models and predicted data for each beta.
+save_iteration = directives.SaveOutputEveryIteration(save_txt=False)
+
+# Updating the preconditionner if it is model dependent.
+update_jacobi = directives.UpdatePreconditioner()
+
+# The directives are defined as a list
+directives_list = [
+    update_IRLS,
+    starting_beta,
+    save_iteration,
+    update_jacobi,
+]
+
+# Define the regularization (model objective function) with depth weighting.
+reg_dpth = regularization.Sparse(mesh, active_cells=ind_active, mapping=model_map)
+reg_dpth.norms = [0, 2, 2, 2]
+depth_weights = utils.depth_weighting(
+    mesh, receiver_locations, active_cells=ind_active, exponent=2
+)
+reg_dpth.set_weights(depth_weights=depth_weights)
+
+# Define how the optimization problem is solved. Here we will use a projected
+# Gauss-Newton approach that employs the conjugate gradient solver.
+opt = optimization.ProjectedGNCG(
+    maxIter=100, lower=-1.0, upper=1.0, maxIterLS=20, maxIterCG=10, tolCG=1e-3
+)
+
+# Here we define the inverse problem that is to be solved
+inv_prob = inverse_problem.BaseInvProblem(dmis, reg_dpth, opt)
+
+# Here we combine the inverse problem and the set of directives
+inv = inversion.BaseInversion(inv_prob, directives_list)
+
+# Run inversion
+recovered_model_dpth = inv.run(starting_model)
+
+#######################################################################
+# Running the Distance Weighted inversion
+# ---------------------------------------
+#
+# Here we define the directives, weights, regularization, and optimization
+# for a distance-weighted inversion
+#
+
+# inversion directives
+# Defining a starting value for the trade-off parameter (beta) between the data
+# misfit and the regularization.
+starting_beta = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
+
+# Defines the directives for the IRLS regularization. This includes setting
+# the cooling schedule for the trade-off parameter.
+update_IRLS = directives.Update_IRLS(
+    f_min_change=1e-4,
+    max_irls_iterations=30,
+    coolEpsFact=1.5,
+    beta_tol=1e-2,
+)
+
+# Options for outputting recovered models and predicted data for each beta.
+save_iteration = directives.SaveOutputEveryIteration(save_txt=False)
+
+# Updating the preconditionner if it is model dependent.
+update_jacobi = directives.UpdatePreconditioner()
+
+# The directives are defined as a list
+directives_list = [
+    update_IRLS,
+    starting_beta,
+    save_iteration,
+    update_jacobi,
+]
+
+# Define the regularization (model objective function) with distance weighting.
+reg_dist = regularization.Sparse(mesh, active_cells=ind_active, mapping=model_map)
+reg_dist.norms = [0, 2, 2, 2]
+distance_weights = utils.distance_weighting(
+    mesh, receiver_locations, active_cells=ind_active, exponent=2
+)
+reg_dist.set_weights(distance_weights=distance_weights)
+
+# Define how the optimization problem is solved. Here we will use a projected
+# Gauss-Newton approach that employs the conjugate gradient solver.
+opt = optimization.ProjectedGNCG(
+    maxIter=100, lower=-1.0, upper=1.0, maxIterLS=20, maxIterCG=10, tolCG=1e-3
+)
+
+# Here we define the inverse problem that is to be solved
+inv_prob = inverse_problem.BaseInvProblem(dmis, reg_dist, opt)
+
+# Here we combine the inverse problem and the set of directives
+inv = inversion.BaseInversion(inv_prob, directives_list)
+
+# Run inversion
+recovered_model_dist = inv.run(starting_model)
+
+#######################################################################
+# Running the Distance Weighted inversion
+# ---------------------------------------
+#
+# Here we define the directives, weights, regularization, and optimization
+# for a sensitivity weighted inversion
+#
+
+# inversion directives
+# Defining a starting value for the trade-off parameter (beta) between the data
+# misfit and the regularization.
+starting_beta = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
+
+# Defines the directives for the IRLS regularization. This includes setting
+# the cooling schedule for the trade-off parameter.
+update_IRLS = directives.Update_IRLS(
+    f_min_change=1e-4,
+    max_irls_iterations=30,
+    coolEpsFact=1.5,
+    beta_tol=1e-2,
+)
+
+# Options for outputting recovered models and predicted data for each beta.
+save_iteration = directives.SaveOutputEveryIteration(save_txt=False)
+
+# Updating the preconditionner if it is model dependent.
+update_jacobi = directives.UpdatePreconditioner()
+
+# Add sensitivity weights
+sensitivity_weights = directives.UpdateSensitivityWeights(every_iteration=False)
+
+# The directives are defined as a list
+directives_list = [
+    update_IRLS,
+    sensitivity_weights,
+    starting_beta,
+    save_iteration,
+    update_jacobi,
+]
+
+# Define the regularization (model objective function) for sensitivity weighting.
+reg_sensw = regularization.Sparse(mesh, active_cells=ind_active, mapping=model_map)
+reg_sensw.norms = [0, 2, 2, 2]
+
+# Define how the optimization problem is solved. Here we will use a projected
+# Gauss-Newton approach that employs the conjugate gradient solver.
+opt = optimization.ProjectedGNCG(
+    maxIter=100, lower=-1.0, upper=1.0, maxIterLS=20, maxIterCG=10, tolCG=1e-3
+)
+
+# Here we define the inverse problem that is to be solved
+inv_prob = inverse_problem.BaseInvProblem(dmis, reg_sensw, opt)
+
+# Here we combine the inverse problem and the set of directives
+inv = inversion.BaseInversion(inv_prob, directives_list)
+
+# Run inversion
+recovered_model_sensw = inv.run(starting_model)
+
+############################################################
+# Recreate True Model
+# -------------------
+#
+
+# Define density contrast values for each unit in g/cc
+background_density = 0.0
+block_density = -0.2
+sphere_density = 0.2
+
+# Define model. Models in simpeg are vector arrays.
+true_model = background_density * np.ones(nC)
+
+# You could find the indicies of specific cells within the model and change their
+# value to add structures.
+ind_block = (
+    (mesh.gridCC[ind_active, 0] > -50.0)
+    & (mesh.gridCC[ind_active, 0] < -20.0)
+    & (mesh.gridCC[ind_active, 1] > -15.0)
+    & (mesh.gridCC[ind_active, 1] < 15.0)
+    & (mesh.gridCC[ind_active, 2] > -50.0)
+    & (mesh.gridCC[ind_active, 2] < -30.0)
+)
+true_model[ind_block] = block_density
+
+# You can also use simpeg utilities to add structures to the model more concisely
+ind_sphere = model_builder.get_indices_sphere(
+    np.r_[35.0, 0.0, -40.0], 15.0, mesh.gridCC
+)
+ind_sphere = ind_sphere[ind_active]
+true_model[ind_sphere] = sphere_density
+
+
+############################################################
+# Plotting True Model and Recovered Models
+# ----------------------------------------
+#
+
+# Plot Models
+fig, ax = plt.subplots(2, 2, figsize=(20, 10), sharex=True, sharey=True)
+ax = ax.flatten()
+plotting_map = maps.InjectActiveCells(mesh, ind_active, np.nan)
+cmap = "coolwarm"
+slice_y_loc = 0.0
+
+mm = mesh.plot_slice(
+    plotting_map * true_model,
+    normal="Y",
+    ax=ax[0],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap, "norm": norm},
+)
+ax[0].set_title(f"True model slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="$g/cm^3$", ax=ax[0])
+
+# plot depth weighting result
+vmax = np.abs(recovered_model_dpth).max()
+norm = mpl.colors.TwoSlopeNorm(vcenter=0, vmin=-vmax, vmax=vmax)
+mm = mesh.plot_slice(
+    plotting_map * recovered_model_dpth,
+    normal="Y",
+    ax=ax[1],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap, "norm": norm},
+)
+ax[1].set_title(f"Depth weighting Model slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="$g/cm^3$", ax=ax[1])
+
+# plot distance weighting result
+vmax = np.abs(recovered_model_dist).max()
+norm = mpl.colors.TwoSlopeNorm(vcenter=0, vmin=-vmax, vmax=vmax)
+mm = mesh.plot_slice(
+    plotting_map * recovered_model_dist,
+    normal="Y",
+    ax=ax[2],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap, "norm": norm},
+)
+ax[2].set_title(f"Distance weighting Model slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="$g/cm^3$", ax=ax[2])
+
+# plot sensitivity weighting result
+vmax = np.abs(recovered_model_sensw).max()
+norm = mpl.colors.TwoSlopeNorm(vcenter=0, vmin=-vmax, vmax=vmax)
+mm = mesh.plot_slice(
+    plotting_map * recovered_model_sensw,
+    normal="Y",
+    ax=ax[3],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap, "norm": norm},
+)
+ax[3].set_title(f"Sensitivity weighting Model slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="$g/cm^3$", ax=ax[3])
+
+# shared plotting
+plotting_map = maps.InjectActiveCells(mesh, ind_active, 0.0)
+slice_y_ind = (
+    mesh.cell_centers[:, 1] == np.abs(mesh.cell_centers[:, 1] - slice_y_loc).min()
+)
+for axx in ax:
+    utils.plot2Ddata(
+        mesh.cell_centers[slice_y_ind][:, [0, 2]],
+        (plotting_map * true_model)[slice_y_ind],
+        contourOpts={"alpha": 0},
+        level=True,
+        ncontour=2,
+        levelOpts={"colors": "grey", "linewidths": 2, "linestyles": "--"},
+        method="nearest",
+        ax=axx,
+    )
+    axx.set_aspect(1)
+
+plt.tight_layout()
+
+############################################################
+# Visualize weights
+# -----------------
+#
+# Plot Weights
+fig, ax = plt.subplots(1, 3, figsize=(20, 4), sharex=True, sharey=True)
+plotting_map = maps.InjectActiveCells(mesh, ind_active, np.nan)
+cmap = "magma"
+slice_y_loc = 0.0
+
+# plot depth weights
+mm = mesh.plot_slice(
+    plotting_map * np.log10(depth_weights),
+    normal="Y",
+    ax=ax[0],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap},
+)
+ax[0].set_title(f"log10(depth weights) slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="log10(depth weights)", ax=ax[0])
+
+# plot distance weights
+mm = mesh.plot_slice(
+    plotting_map * np.log10(distance_weights),
+    normal="Y",
+    ax=ax[1],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap},
+)
+ax[1].set_title(f"log10(distance weights) slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="log10(distance weights)", ax=ax[1])
+
+# plot sensitivity weights
+mm = mesh.plot_slice(
+    plotting_map * np.log10(reg_sensw.objfcts[0].get_weights(key="sensitivity")),
+    normal="Y",
+    ax=ax[2],
+    grid=False,
+    slice_loc=slice_y_loc,
+    pcolor_opts={"cmap": cmap},
+)
+ax[2].set_title(f"log10(sensitivity weights) slice at y = {slice_y_loc} m")
+plt.colorbar(mm[0], label="log10(sensitivity weights)", ax=ax[2])
+
+# shared plotting
+for axx in ax:
+    axx.set_aspect(1)
+
+plt.tight_layout()


Test patch:
diff --git a/tests/base/test_model_utils.py b/tests/base/test_model_utils.py
index 97802fbf2..4cc34beaf 100644
--- a/tests/base/test_model_utils.py
+++ b/tests/base/test_model_utils.py
@@ -1,13 +1,10 @@
-import pytest
 import unittest
 
 import numpy as np
-
+import pytest
 from discretize import TensorMesh
 
-from simpeg import (
-    utils,
-)
+from simpeg import utils
 
 
 class DepthWeightingTest(unittest.TestCase):
@@ -77,6 +74,109 @@ class DepthWeightingTest(unittest.TestCase):
         np.testing.assert_allclose(wz, wz2)
 
 
+class TestDistancehWeighting:
+    def test_distance_weighting_3D(self):
+        # Mesh
+        dh = 5.0
+        hx = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+        hy = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+        hz = [(dh, 15)]
+        mesh = TensorMesh([hx, hy, hz], "CCN")
+
+        rng = np.random.default_rng(seed=42)
+        actv = rng.integers(low=0, high=2, size=mesh.n_cells, dtype=bool)
+
+        # Define reference locs at random locations
+        reference_locs = rng.uniform(
+            low=mesh.nodes.min(axis=0), high=mesh.nodes.max(axis=0), size=(1000, 3)
+        )
+
+        # distance weighting
+        with pytest.warns():
+            wz_scipy = utils.distance_weighting(
+                mesh, reference_locs, active_cells=actv, exponent=3, engine="scipy"
+            )
+        wz_numba = utils.distance_weighting(
+            mesh, reference_locs, active_cells=actv, exponent=3, engine="numba"
+        )
+        np.testing.assert_allclose(wz_scipy, wz_numba)
+
+        with pytest.raises(ValueError):
+            utils.distance_weighting(
+                mesh, reference_locs, active_cells=actv, exponent=3, engine="test"
+            )
+
+    def test_distance_weighting_2D(self):
+        # Mesh
+        dh = 5.0
+        hx = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+        hz = [(dh, 15)]
+        mesh = TensorMesh([hx, hz], "CN")
+
+        rng = np.random.default_rng(seed=42)
+        actv = rng.integers(low=0, high=2, size=mesh.n_cells, dtype=bool)
+
+        # Define reference locs at random locations
+        reference_locs = rng.uniform(
+            low=mesh.nodes.min(axis=0), high=mesh.nodes.max(axis=0), size=(1000, 2)
+        )
+
+        # distance weighting
+        with pytest.warns():
+            wz_scipy = utils.distance_weighting(
+                mesh, reference_locs, active_cells=actv, exponent=3, engine="scipy"
+            )
+        wz_numba = utils.distance_weighting(
+            mesh, reference_locs, active_cells=actv, exponent=3, engine="numba"
+        )
+        np.testing.assert_allclose(wz_scipy, wz_numba)
+
+    def test_distance_weighting_1D(self):
+        # Mesh
+        dh = 5.0
+        hx = [(dh, 5, -1.3), (dh, 40), (dh, 5, 1.3)]
+        mesh = TensorMesh([hx], "C")
+
+        rng = np.random.default_rng(seed=42)
+        actv = rng.integers(low=0, high=2, size=mesh.n_cells, dtype=bool)
+
+        # Define reference locs at random locations
+        reference_locs = rng.uniform(
+            low=mesh.nodes.min(axis=0), high=mesh.nodes.max(axis=0), size=(1000, 1)
+        )
+
+        # distance weighting
+        with pytest.warns():
+            wz_scipy = utils.distance_weighting(
+                mesh, reference_locs, active_cells=actv, exponent=3, engine="scipy"
+            )
+        wz_numba = utils.distance_weighting(
+            mesh, reference_locs, active_cells=actv, exponent=3, engine="numba"
+        )
+        np.testing.assert_allclose(wz_scipy, wz_numba)
+
+    @pytest.mark.parametrize("ndim", (2, 3))
+    def test_invalid_reference_locs(self, ndim):
+        """
+        Test if errors are raised when invalid reference_locs are passed.
+        """
+        hx = [5.0, 10]
+        h = [hx] * ndim
+        origin = "CCN" if ndim == 3 else "CC"
+        reference_locs = [1.0, 2.0] if ndim == 3 else [1.0]
+        mesh = TensorMesh(h, origin)
+        with pytest.raises(ValueError):
+            utils.distance_weighting(mesh, reference_locs)
+
+    def test_numba_and_cdist_opts_error(self):
+        """Test error when passing numba and cdist_opts."""
+        hx = [5.0, 10]
+        mesh = TensorMesh([hx, hx, hx])
+        msg = "The `cdist_opts` is valid only when engine is 'scipy'."
+        with pytest.raises(TypeError, match=msg):
+            utils.distance_weighting(mesh, [1.0, 2.0, 3.0], cdist_opts={"foo": "bar"})
+
+
 @pytest.fixture
 def mesh():
     """Sample mesh."""