Receiver Function Inversion with the Neighbourhood Algorithm#

Open In Colab

If you are running this notebook locally, make sure you’ve followed steps here to set up the environment. (This environment.yml file specifies a list of packages required to run the notebooks)


1. Introduction#

1.1 Receiver functions#

# display theory on receiver functions
from IPython.display import display, Markdown

with open("../../theory/geo_receiver_function.md", "r") as f:
    content = f.read()

display(Markdown(content))
<IPython.core.display.Markdown object>

1.2 Neighbourhood Algorithm#

# display theory on the Neighbourhood Algorithm
from IPython.display import display, Markdown

with open("../../theory/neighbourhood_algorithm.md", "r") as f:
    content = f.read()

display(Markdown(content))
<IPython.core.display.Markdown object>

1.3 CoFI and the NA#

The implementation of the NA that cofi wraps is called `neighpy <inlab-geo/neighpy>`__. This implementation implements both phases of the NA as described in the original papers.

In this notebook we run the two stages of the NA separately using the neighpyI and neighpyII tools in CoFI:

  • Stage I (neighpyI) — Direct search. An ensemble optimiser that minimises a user-defined misfit function to explore the parameter space.

  • Stage II (neighpyII) — Appraisal. An MCMC resampler that resamples the Stage I ensemble according to a user-supplied log posterior probability density (log-PPD). No additional forward model evaluations are required.

The key insight is that the log-PPD used in Stage II does not have to equal the negative of the misfit used in Stage I. The relationship between the two is a modelling choice that depends on the assumed likelihood and prior. In this notebook we use \(\log p(\mathbf{m}|\mathbf{d}) = -\Phi(\mathbf{m},\mathbf{d})\), which corresponds to a Gaussian likelihood with a uniform prior, but we make this choice explicit so the user can modify it.

We apply the NA to two receiver function inversion problems of increasing complexity:

  1. Problem 1 — Invert for S-wave velocities only (4 parameters, layer depths fixed)

  2. Problem 2 — Invert for both layer depths and S-wave velocities (8 parameters)

Note: The hyperparameters in this notebook are intentionally set low for a quick demonstration run. Results may not be fully converged. To improve them, increase n_initial_samples, n_iterations, and n_resample in the cells below.


2. Import modules#

# -------------------------------------------------------- #
#                                                          #
#     Uncomment below to set up environment on "colab"     #
#                                                          #
# -------------------------------------------------------- #

# !pip install -U cofi pyrf96
import math
import numpy as np
import matplotlib.pyplot as plt

from cofi import BaseProblem, InversionOptions, Inversion
import pyrf96

np.random.seed(42)

3. Define the problem#

We use a 4-layer Earth model in pyrf96’s Voronoi nuclei format (mtype=0). Each row is [depth_km, Vs_km/s, Vp/Vs]. The Vp/Vs ratios are held fixed throughout the inversion — only depths and/or velocities are inverted.

We generate synthetic observed data by computing the receiver function for the true model and adding correlated noise.

# True Earth model: 4 layers [depth_km, Vs_km/s, Vp/Vs]
good_model = np.array([
    [ 1.0, 3.0, 1.7],
    [ 8.0, 3.2, 2.0],
    [20.0, 4.0, 1.7],
    [45.0, 4.2, 1.7],
])

# Fixed Vp/Vs ratios
vpvs = good_model[:, 2]

# Generate synthetic data
t, rfunc = pyrf96.rfcalc(good_model)                                  # noise-free
t2, rfunc_noise = pyrf96.rfcalc(good_model, sn=0.5, seed=12345678)   # noisy observed
observed_data = rfunc_noise

# Inverse data covariance matrix for correlated noise
Cdinv = pyrf96.InvDataCov(2.5, 0.01, len(rfunc))
def plot_model_staircase(model, ax, max_depth=60.0, **kwargs):
    """Plot a pyrf96 Voronoi nuclei model as a Vs-depth staircase."""
    order = np.argsort(model[:, 0])
    depths = model[order, 0]
    vs = model[order, 1]
    interfaces = 0.5 * (depths[:-1] + depths[1:])
    plot_depths = np.concatenate([[0.0], np.repeat(interfaces, 2), [max_depth]])
    plot_vs = np.repeat(vs, 2)
    ax.plot(plot_vs, plot_depths, **kwargs)
# Plot the true model and observed data
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 5))

# Velocity-depth profile
plot_model_staircase(good_model, ax1, color="red", lw=2, label="True model")
ax1.set_xlim(2.5, 5.0)
ax1.set_ylim(60, 0)
ax1.set_xlabel("Vs (km/s)")
ax1.set_ylabel("Depth (km)")
ax1.set_title("True Earth model")
ax1.legend()

# Receiver functions
ax2.plot(t2, rfunc_noise, color="k", label="Observed (noisy)")
ax2.plot(t, rfunc, color="r", ls="--", label="True (noise-free)")
ax2.set_xlabel("Time (s)")
ax2.set_ylabel("Amplitude")
ax2.set_title("Synthetic receiver function")
ax2.legend()

plt.tight_layout()

4. Problem 1: Invert for velocities only#

In this first problem we fix the layer depths at their true values and invert only for the 4 S-wave velocities. This is a simpler 4-dimensional problem.

# Conversion functions for velocity-only inversion
fixed_depths = good_model[:, 0]  # [1, 8, 20, 45] — held fixed

def get_vel_params(fullmodel):
    """Extract velocity parameters from full model."""
    return fullmodel[:, 1].copy()

def make_model_vel(vel_params):
    """Reconstruct full model from velocity parameters (fixed depths and Vp/Vs)."""
    return np.column_stack([fixed_depths, vel_params, vpvs])

# Misfit function for velocity-only inversion (used by Stage I)
def misfit_vel(vel_params):
    model = make_model_vel(vel_params)
    t_pred, predicted_data = pyrf96.rfcalc(model)
    res = observed_data - predicted_data
    misfit_val = np.dot(res, np.transpose(np.dot(Cdinv, res))) / 2.0
    if math.isnan(misfit_val):
        return float("inf")
    return misfit_val

# Log posterior function (used by Stage II)
def log_posterior_vel(vel_params):
    """Log posterior = -misfit (uniform prior, Gaussian likelihood)."""
    return -misfit_vel(vel_params)

# CoFI BaseProblem
inv_problem_vel = BaseProblem()
inv_problem_vel.name = "Receiver Function - Velocity only"
inv_problem_vel.set_objective(misfit_vel)
inv_problem_vel.summary()
=====================================================================
Summary for inversion problem: Receiver Function - Velocity only
=====================================================================
Model shape: Unknown
---------------------------------------------------------------------
List of functions/properties set by you:
['objective']
---------------------------------------------------------------------
List of functions/properties created based on what you have provided:
-- none --
---------------------------------------------------------------------
List of functions/properties that can be further set for the problem:
( not all of these may be relevant to your inversion workflow )
['log_posterior', 'log_posterior_with_blobs', 'log_likelihood', 'log_prior', 'gradient', 'hessian', 'hessian_times_vector', 'residual', 'jacobian', 'jacobian_times_vector', 'data_misfit', 'regularization', 'regularization_matrix', 'forward', 'data', 'data_covariance', 'data_covariance_inv', 'initial_model', 'model_shape', 'blobs_dtype', 'bounds', 'constraints']

4.2 Stage II: Appraisal#

Stage II of the NA resamples the Voronoi cells constructed by Stage I according to a log posterior probability density (log-PPD). Crucially, this stage requires no additional forward model evaluations — it only uses the geometry of the Stage I ensemble and the log-PPD values assigned to each sample.

The user must provide log_ppd values for each sample in the Stage I ensemble. The relationship between the log-PPD and the misfit is a modelling choice:

\[\log p(\mathbf{m}|\mathbf{d}) = -\Phi(\mathbf{m},\mathbf{d})\]

corresponds to a Gaussian likelihood with a uniform prior. More generally, one could incorporate a non-uniform prior or a different likelihood function, which would change the log_ppd values without requiring Stage I to be re-run.

# Compute log-PPD for each Stage I sample.
# Since log_posterior_vel(s) = -misfit_vel(s), this is equivalent to:
#   log_ppd_vel = np.array([log_posterior_vel(s) for s in ds_samples_vel])
# but using the already-computed misfit values avoids redundant forward evaluations.
log_ppd_vel = -ds_objectives_vel
# NA-II hyperparameters
inv_options_vel_app = InversionOptions()
inv_options_vel_app.set_tool("neighpyII")
inv_options_vel_app.set_params(
    bounds=bounds_vel,
    initial_ensemble=ds_samples_vel,
    log_ppd=log_ppd_vel,
    n_resample=20000,
    n_walkers=10,
)

inv_vel_app = Inversion(inv_problem_vel, inv_options_vel_app)
inv_result_vel_app = inv_vel_app.run()
inv_result_vel_app.summary()
============================
Summary for inversion result
============================
SUCCESS
----------------------------
new_samples: [[3.08936785 3.14424749 4.04079552 4.44101329]
 [3.11039626 3.09111185 4.03866865 4.43234406]
 [3.10455271 3.17238427 4.02898751 4.49612852]
 ...
 [3.11025929 3.1584892  3.99314115 4.33662172]
 [3.09754188 3.15536398 4.03084395 4.48797973]
 [3.12008656 3.08642429 4.02306098 4.48868035]]
appraisal_samples_vel = inv_result_vel_app.new_samples
appraisal_mean_vel = appraisal_samples_vel.mean(axis=0)

fig, ax = plt.subplots(figsize=(5, 8))
for i in range(0, len(appraisal_samples_vel), 10):
    m = make_model_vel(appraisal_samples_vel[i])
    plot_model_staircase(m, ax, color="k", alpha=0.01, lw=0.5)

plot_model_staircase(good_model, ax, color="b", lw=2, label="True model")
plot_model_staircase(make_model_vel(best_vel), ax, color="g", lw=2, ls="--",
                     label="Best model (NA-I)")
plot_model_staircase(make_model_vel(appraisal_mean_vel), ax, color="r", lw=2,
                     ls="--", label="Mean model (NA-II)")
ax.set_xlim(2.5, 5.0)
ax.set_ylim(60, 0)
ax.set_xlabel("Vs (km/s)")
ax.set_ylabel("Depth (km)")
ax.set_title("Problem 1: NA-II appraisal ensemble")
ax.legend(loc="lower left")
<matplotlib.legend.Legend object at 0x7efcbc64c910>
# 4x4 corner plot for velocity-only problem
true_vel = get_vel_params(good_model)
n_params_vel = 4
var_names_vel = [f"Vs{i+1}" for i in range(n_params_vel)]

fig, axes = plt.subplots(n_params_vel, n_params_vel, figsize=(8, 8),
                         tight_layout=True,
                         gridspec_kw={"hspace": 0, "wspace": 0})
for i in range(n_params_vel):
    for j in range(n_params_vel):
        axes[i, j].set_xlim(bounds_vel[j])
        axes[i, j].set_xticks([])
        axes[i, j].set_yticks([])
        if i == j:
            axes[i, j].hist(appraisal_samples_vel[:, j], bins=100,
                            histtype="step", color="k")
            axes[i, j].axvline(best_vel[j], c="g", ls="--")
            axes[i, j].axvline(true_vel[j], c="b", ls="--")
            axes[i, j].axvline(appraisal_mean_vel[j], c="r", ls="--")
        elif j < i:
            axes[i, j].scatter(appraisal_samples_vel[:, j],
                               appraisal_samples_vel[:, i],
                               s=1, c="k", alpha=0.01)
            axes[i, j].scatter(best_vel[j], best_vel[i], c="g", s=10, zorder=10)
            axes[i, j].scatter(true_vel[j], true_vel[i], c="b", s=10, zorder=10)
            axes[i, j].scatter(appraisal_mean_vel[j], appraisal_mean_vel[i],
                               c="r", s=10, zorder=10)
            axes[i, j].set_ylim(bounds_vel[i])
        else:
            axes[i, j].axis("off")
        if i == n_params_vel - 1 and j <= i:
            axes[i, j].set_xlabel(var_names_vel[j])
            axes[i, j].set_xticks(np.linspace(*bounds_vel[j], 3))
        if j == 0 and i > 0:
            axes[i, j].set_ylabel(var_names_vel[i])

fig.suptitle("Problem 1: Posterior corner plot (velocity-only)")
Text(0.5, 0.98, 'Problem 1: Posterior corner plot (velocity-only)')

5. Problem 2: Invert for depths and velocities#

Now we invert for both layer depths and S-wave velocities simultaneously. This is a harder 8-dimensional problem with interleaved parameters: [d1, vs1, d2, vs2, d3, vs3, d4, vs4].

# Conversion functions for depth+velocity inversion
def get_inversion_parameters(fullmodel):
    """Extract [depth, Vs] pairs as flat 8-vector."""
    return fullmodel[:, :2].flatten()

def get_model_parameters(invmodel):
    """Reconstruct full [4,3] model from 8-vector, restoring fixed Vp/Vs."""
    return np.append(invmodel.reshape(len(vpvs), -1), vpvs[:, None], axis=1)

# Misfit function for depth+velocity inversion (used by Stage I)
def misfit_dv(imodel):
    model = get_model_parameters(imodel)
    t_pred, predicted_data = pyrf96.rfcalc(model)
    res = observed_data - predicted_data
    misfit_val = np.dot(res, np.transpose(np.dot(Cdinv, res))) / 2.0
    if math.isnan(misfit_val):
        return float("inf")
    return misfit_val

# Log posterior function (used by Stage II)
def log_posterior_dv(imodel):
    """Log posterior = -misfit (uniform prior, Gaussian likelihood)."""
    return -misfit_dv(imodel)

# CoFI BaseProblem
inv_problem_dv = BaseProblem()
inv_problem_dv.name = "Receiver Function - Depth + Velocity"
inv_problem_dv.set_objective(misfit_dv)
inv_problem_dv.summary()
=====================================================================
Summary for inversion problem: Receiver Function - Depth + Velocity
=====================================================================
Model shape: Unknown
---------------------------------------------------------------------
List of functions/properties set by you:
['objective']
---------------------------------------------------------------------
List of functions/properties created based on what you have provided:
-- none --
---------------------------------------------------------------------
List of functions/properties that can be further set for the problem:
( not all of these may be relevant to your inversion workflow )
['log_posterior', 'log_posterior_with_blobs', 'log_likelihood', 'log_prior', 'gradient', 'hessian', 'hessian_times_vector', 'residual', 'jacobian', 'jacobian_times_vector', 'data_misfit', 'regularization', 'regularization_matrix', 'forward', 'data', 'data_covariance', 'data_covariance_inv', 'initial_model', 'model_shape', 'blobs_dtype', 'bounds', 'constraints']

5.1 Stage I: Direct Search#

# NA-I hyperparameters for 8-D problem (larger than Problem 1)
bounds_dv = [(0.0, 60.0), (2.0, 4.5)] * 4

inv_options_dv = InversionOptions()
inv_options_dv.set_tool("neighpyI")
inv_options_dv.set_params(
    bounds=bounds_dv,
    n_initial_samples=2000,
    n_samples_per_iteration=100,
    n_cells_to_resample=20,
    n_iterations=50,
)

inv_dv = Inversion(inv_problem_dv, inv_options_dv)
inv_result_dv = inv_dv.run()
inv_result_dv.summary()
NAI - Initial Random Search

NAI - Optimisation Loop:   0%|          | 0/50 [00:00<?, ?it/s]
NAI - Optimisation Loop:   2%|▏         | 1/50 [00:01<00:55,  1.13s/it]
NAI - Optimisation Loop:   4%|▍         | 2/50 [00:01<00:26,  1.79it/s]
NAI - Optimisation Loop:   6%|▌         | 3/50 [00:01<00:17,  2.67it/s]
NAI - Optimisation Loop:   8%|▊         | 4/50 [00:01<00:13,  3.49it/s]
NAI - Optimisation Loop:  10%|█         | 5/50 [00:01<00:10,  4.17it/s]
NAI - Optimisation Loop:  12%|█▏        | 6/50 [00:01<00:09,  4.74it/s]
NAI - Optimisation Loop:  14%|█▍        | 7/50 [00:02<00:08,  5.17it/s]
NAI - Optimisation Loop:  16%|█▌        | 8/50 [00:02<00:07,  5.49it/s]
NAI - Optimisation Loop:  18%|█▊        | 9/50 [00:02<00:07,  5.72it/s]
NAI - Optimisation Loop:  20%|██        | 10/50 [00:02<00:06,  5.92it/s]
NAI - Optimisation Loop:  22%|██▏       | 11/50 [00:02<00:06,  6.02it/s]
NAI - Optimisation Loop:  24%|██▍       | 12/50 [00:02<00:06,  6.12it/s]
NAI - Optimisation Loop:  26%|██▌       | 13/50 [00:03<00:05,  6.29it/s]
NAI - Optimisation Loop:  28%|██▊       | 14/50 [00:03<00:05,  6.39it/s]
NAI - Optimisation Loop:  30%|███       | 15/50 [00:03<00:05,  6.49it/s]
NAI - Optimisation Loop:  32%|███▏      | 16/50 [00:03<00:05,  6.47it/s]
NAI - Optimisation Loop:  34%|███▍      | 17/50 [00:03<00:05,  6.46it/s]
NAI - Optimisation Loop:  36%|███▌      | 18/50 [00:03<00:04,  6.54it/s]
NAI - Optimisation Loop:  38%|███▊      | 19/50 [00:03<00:04,  6.58it/s]
NAI - Optimisation Loop:  40%|████      | 20/50 [00:04<00:04,  6.50it/s]
NAI - Optimisation Loop:  42%|████▏     | 21/50 [00:04<00:04,  6.56it/s]
NAI - Optimisation Loop:  44%|████▍     | 22/50 [00:04<00:04,  6.50it/s]
NAI - Optimisation Loop:  46%|████▌     | 23/50 [00:04<00:04,  6.47it/s]
NAI - Optimisation Loop:  48%|████▊     | 24/50 [00:04<00:04,  6.45it/s]
NAI - Optimisation Loop:  50%|█████     | 25/50 [00:04<00:03,  6.51it/s]
NAI - Optimisation Loop:  52%|█████▏    | 26/50 [00:04<00:03,  6.59it/s]
NAI - Optimisation Loop:  54%|█████▍    | 27/50 [00:05<00:03,  6.65it/s]
NAI - Optimisation Loop:  56%|█████▌    | 28/50 [00:05<00:03,  6.68it/s]
NAI - Optimisation Loop:  58%|█████▊    | 29/50 [00:05<00:03,  6.59it/s]
NAI - Optimisation Loop:  60%|██████    | 30/50 [00:05<00:03,  6.64it/s]
NAI - Optimisation Loop:  62%|██████▏   | 31/50 [00:05<00:02,  6.64it/s]
NAI - Optimisation Loop:  64%|██████▍   | 32/50 [00:05<00:02,  6.57it/s]
NAI - Optimisation Loop:  66%|██████▌   | 33/50 [00:06<00:02,  6.59it/s]
NAI - Optimisation Loop:  68%|██████▊   | 34/50 [00:06<00:02,  6.56it/s]
NAI - Optimisation Loop:  70%|███████   | 35/50 [00:06<00:02,  6.53it/s]
NAI - Optimisation Loop:  72%|███████▏  | 36/50 [00:06<00:02,  6.56it/s]
NAI - Optimisation Loop:  74%|███████▍  | 37/50 [00:06<00:01,  6.53it/s]
NAI - Optimisation Loop:  76%|███████▌  | 38/50 [00:06<00:01,  6.60it/s]
NAI - Optimisation Loop:  78%|███████▊  | 39/50 [00:06<00:01,  6.52it/s]
NAI - Optimisation Loop:  80%|████████  | 40/50 [00:07<00:01,  6.48it/s]
NAI - Optimisation Loop:  82%|████████▏ | 41/50 [00:07<00:01,  6.46it/s]
NAI - Optimisation Loop:  84%|████████▍ | 42/50 [00:07<00:01,  6.52it/s]
NAI - Optimisation Loop:  86%|████████▌ | 43/50 [00:07<00:01,  6.49it/s]
NAI - Optimisation Loop:  88%|████████▊ | 44/50 [00:07<00:00,  6.56it/s]
NAI - Optimisation Loop:  90%|█████████ | 45/50 [00:07<00:00,  6.62it/s]
NAI - Optimisation Loop:  92%|█████████▏| 46/50 [00:08<00:00,  6.56it/s]
NAI - Optimisation Loop:  94%|█████████▍| 47/50 [00:08<00:00,  6.62it/s]
NAI - Optimisation Loop:  96%|█████████▌| 48/50 [00:08<00:00,  6.68it/s]
NAI - Optimisation Loop:  98%|█████████▊| 49/50 [00:08<00:00,  6.72it/s]
NAI - Optimisation Loop: 100%|██████████| 50/50 [00:08<00:00,  6.66it/s]
NAI - Optimisation Loop: 100%|██████████| 50/50 [00:08<00:00,  5.79it/s]
============================
Summary for inversion result
============================
SUCCESS
----------------------------
model: [ 3.65446741  3.22530365  5.33884335  3.43241498 25.54124344  4.16511517
 41.47990624  4.49602258]
samples: [[12.35888073  4.17986845 14.53017432 ...  3.29953543 27.92913793
   2.36780655]
 [54.77833827  3.25919094 54.25336229 ...  2.56894737 22.78114425
   3.13115686]
 [45.3030486   3.66195155 11.13520233 ...  3.19308944 55.49065324
   3.20351945]
 ...
 [ 3.65377737  3.22624301  5.32845836 ...  4.15493385 41.49363447
   4.49990418]
 [ 3.65387412  3.22655375  5.32909702 ...  4.15635468 41.49323695
   4.49834524]
 [ 3.65340715  3.22533198  5.32865265 ...  4.15686466 41.48756262
   4.49832081]]
objectives: [4494.49996292 1122.94609476 1054.46175399 ...  494.62492545  494.87238816
  494.59962133]
best_dv = inv_result_dv.model
ds_samples_dv = inv_result_dv.samples
ds_objectives_dv = inv_result_dv.objectives

print("Best model (depth + velocity):")
print(get_model_parameters(best_dv))
print("\nTrue model:")
print(good_model)
Best model (depth + velocity):
[[ 3.65446741  3.22530365  1.7       ]
 [ 5.33884335  3.43241498  2.        ]
 [25.54124344  4.16511517  1.7       ]
 [41.47990624  4.49602258  1.7       ]]

True model:
[[ 1.   3.   1.7]
 [ 8.   3.2  2. ]
 [20.   4.   1.7]
 [45.   4.2  1.7]]

Convergence#

best_i_dv = np.argmin(ds_objectives_dv)

fig, ax = plt.subplots(figsize=(10, 4))
ax.plot(ds_objectives_dv, marker=".", linestyle="", markersize=2, color="black")
ax.scatter(best_i_dv, ds_objectives_dv[best_i_dv], c="g", s=30, zorder=10,
           label="Best model")
ax.axvline(2000, c="grey", ls="--", lw=0.8)
ax.set_yscale("log")
ax.set_xlabel("Sample number")
ax.set_ylabel("Misfit")
ax.set_title("Problem 2: NA-I convergence (depth + velocity)")
ax.text(0.15, 0.95, "Initial random search", transform=ax.transAxes,
        ha="center", va="top")
ax.text(0.65, 0.95, "Neighbourhood search", transform=ax.transAxes,
        ha="center", va="top")
ax.legend()
<matplotlib.legend.Legend object at 0x7efcc1193390>

5.2 Stage II: Appraisal#

# Compute log-PPD for each Stage I sample (same choice as Problem 1)
log_ppd_dv = -ds_objectives_dv
# NA-II hyperparameters
inv_options_dv_app = InversionOptions()
inv_options_dv_app.set_tool("neighpyII")
inv_options_dv_app.set_params(
    bounds=bounds_dv,
    initial_ensemble=ds_samples_dv,
    log_ppd=log_ppd_dv,
    n_resample=20000,
    n_walkers=10,
)

inv_dv_app = Inversion(inv_problem_dv, inv_options_dv_app)
inv_result_dv_app = inv_dv_app.run()
inv_result_dv_app.summary()
============================
Summary for inversion result
============================
SUCCESS
----------------------------
new_samples: [[ 3.79248677  3.22706369  9.79108    ...  4.18531913 57.20266037
   4.46681503]
 [ 8.3232985   3.22053145 14.96308353 ...  4.18659007 58.11078973
   4.45000291]
 [10.62230822  3.20394371  6.736847   ...  4.2095668  43.51925628
   4.45956822]
 ...
 [11.66904023  3.16648035 15.01326392 ...  4.37806961 48.76035271
   4.49182749]
 [16.96376464  3.29605069 12.21207934 ...  4.43404615 42.62371863
   4.3759516 ]
 [14.44362114  3.12602231  2.08165964 ...  4.48057032 55.05805663
   4.4281363 ]]
appraisal_samples_dv = inv_result_dv_app.new_samples
appraisal_mean_dv = appraisal_samples_dv.mean(axis=0)

fig, ax = plt.subplots(figsize=(5, 8))
for i in range(0, len(appraisal_samples_dv), 10):
    m = get_model_parameters(appraisal_samples_dv[i])
    plot_model_staircase(m, ax, color="k", alpha=0.01, lw=0.5)

plot_model_staircase(good_model, ax, color="b", lw=2, label="True model")
plot_model_staircase(get_model_parameters(best_dv), ax, color="g", lw=2,
                     ls="--", label="Best model (NA-I)")
plot_model_staircase(get_model_parameters(appraisal_mean_dv), ax, color="r",
                     lw=2, ls="--", label="Mean model (NA-II)")
ax.set_xlim(2.5, 5.0)
ax.set_ylim(60, 0)
ax.set_xlabel("Vs (km/s)")
ax.set_ylabel("Depth (km)")
ax.set_title("Problem 2: NA-II appraisal ensemble")
ax.legend(loc="lower left")
<matplotlib.legend.Legend object at 0x7efcb64f8410>

Corner plot#

# 8x8 corner plot for depth+velocity problem
true_dv = get_inversion_parameters(good_model)
n_params_dv = 8
var_names_dv = ["d1", "Vs1", "d2", "Vs2", "d3", "Vs3", "d4", "Vs4"]

fig, axes = plt.subplots(n_params_dv, n_params_dv, figsize=(12, 12),
                         tight_layout=True,
                         gridspec_kw={"hspace": 0, "wspace": 0})
for i in range(n_params_dv):
    for j in range(n_params_dv):
        axes[i, j].set_xlim(bounds_dv[j])
        axes[i, j].set_xticks([])
        axes[i, j].set_yticks([])
        if i == j:
            axes[i, j].hist(appraisal_samples_dv[::10, j], bins=100,
                            histtype="step", color="k")
            axes[i, j].axvline(best_dv[j], c="g", ls="--")
            axes[i, j].axvline(true_dv[j], c="b", ls="--")
            axes[i, j].axvline(appraisal_mean_dv[j], c="r", ls="--")
        elif j < i:
            axes[i, j].scatter(appraisal_samples_dv[::10, j],
                               appraisal_samples_dv[::10, i],
                               s=1, c="k", alpha=0.01)
            axes[i, j].scatter(best_dv[j], best_dv[i], c="g", s=10, zorder=10)
            axes[i, j].scatter(true_dv[j], true_dv[i], c="b", s=10, zorder=10)
            axes[i, j].scatter(appraisal_mean_dv[j], appraisal_mean_dv[i],
                               c="r", s=10, zorder=10)
            axes[i, j].set_ylim(bounds_dv[i])
        else:
            axes[i, j].axis("off")

for ax_, xlabel in zip(axes[-1, :], var_names_dv):
    ax_.set_xlabel(xlabel)
for ax_, ylabel in zip(axes[:, 0], var_names_dv):
    if ylabel != var_names_dv[0]:
        ax_.set_ylabel(ylabel)

fig.suptitle("Problem 2: Posterior corner plot (depth + velocity)")
Text(0.5, 0.98, 'Problem 2: Posterior corner plot (depth + velocity)')

Predicted data ensemble#

fig, ax = plt.subplots(figsize=(10, 4))

for i in range(0, len(appraisal_samples_dv), 15):
    m = get_model_parameters(appraisal_samples_dv[i])
    try:
        t_pred, rf_pred = pyrf96.rfcalc(m)
        ax.plot(t_pred, rf_pred, color="k", alpha=0.01, lw=0.5)
    except Exception:
        pass

ax.plot(t2, observed_data, color="k", lw=1.5, label="Observed", zorder=10)
t_best, rf_best = pyrf96.rfcalc(get_model_parameters(best_dv))
ax.plot(t_best, rf_best, color="g", ls="--", lw=1.5, label="Best model (NA-I)",
        zorder=11)
t_mean, rf_mean = pyrf96.rfcalc(get_model_parameters(appraisal_mean_dv))
ax.plot(t_mean, rf_mean, color="r", ls="--", lw=1.5, label="Mean model (NA-II)",
        zorder=11)
ax.set_xlabel("Time (s)")
ax.set_ylabel("Amplitude")
ax.set_title("Problem 2: Predicted receiver functions from NA-II ensemble")
ax.legend()
<matplotlib.legend.Legend object at 0x7efcca251bd0>

Watermark#

watermark_list = ["cofi", "numpy", "scipy", "matplotlib"]
for pkg in watermark_list:
    pkg_var = __import__(pkg)
    print(pkg, getattr(pkg_var, "__version__"))
cofi 0.2.11+68.gc319bc9
numpy 2.2.6
scipy 1.17.1
matplotlib 3.10.9

sphinx_gallery_thumbnail_number = -1

Total running time of the script: (2 minutes 21.252 seconds)

Gallery generated by Sphinx-Gallery