Surface wave and receiver function - joint inversion (synthetic data)#

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)


What we do in this notebook#

Here we extend the example where CoFI has been used for the inversion of Rayleigh wave phase velocities for a 1D layered earth to a joint inversion where we also account for receiver functions.

Learning outcomes

  • A demonstration of CoFI’s ability to switch between parameter estimation and ensemble methods.

  • An application of CoFI for a joint inversion, here of Rayleigh wave pahse velocity and receiver function data, to a synthetic dataset

Utilities preparation#

In this example, we look at a joint inversion problem of surface wave and receiver function.

We use pysurf96 for computing the forward step of surface wave, and use espresso for receiver function calculations.

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

# !pip install -U cofi
# !pip install git+https://github.com/inlab-geo/pysurf96.git
# !git clone https://github.com/inlab-geo/cofi-examples.git
# %cd cofi-examples/examples/sw_rf_joint
import glob
import numpy as np
import scipy
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

import pysurf96
import bayesbay
import cofi
import espresso
np.seterr(all="ignore");
{'divide': 'ignore', 'over': 'ignore', 'under': 'ignore', 'invalid': 'ignore'}

Model vector

# display theory on the 1D model parameterisation
from IPython.display import display, Markdown

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

display(Markdown(content))
<IPython.core.display.Markdown object>
# layercake model utilities
def form_layercake_model(thicknesses, vs):
    model = np.zeros((len(vs)*2-1,))
    model[1::2] = thicknesses
    model[::2] = vs
    return model

def split_layercake_model(model):
    thicknesses = model[1::2]
    vs = model[::2]
    return thicknesses, vs
# voronoi model utilities
def form_voronoi_model(voronoi_sites, vs):
    return np.hstack((vs, voronoi_sites))

def split_voronoi_model(model):
    voronoi_sites = model[len(model)//2:]
    vs = model[:len(model)//2]
    return voronoi_sites, vs
def voronoi_to_layercake(voronoi_vector: np.ndarray) -> np.ndarray:
    n_layers = len(voronoi_vector) // 2
    velocities = voronoi_vector[:n_layers]
    voronoi_sites = voronoi_vector[n_layers:]
    depths = (voronoi_sites[:-1] + voronoi_sites[1:]) / 2
    thicknesses = depths - np.insert(depths[:-1], 0, 0)
    layercake_vector = np.zeros((2*n_layers-1,))
    layercake_vector[::2] = velocities
    layercake_vector[1::2] = thicknesses
    return layercake_vector

def layercake_to_voronoi(layercake_vector: np.ndarray, first_voronoi_site: float = 0.0) -> np.ndarray:
    n_layers = len(layercake_vector) // 2 + 1
    thicknesses = layercake_vector[1::2]
    velocities = layercake_vector[::2]
    depths = np.cumsum(thicknesses)
    voronoi_sites = np.zeros((n_layers,))
    for i in range(1,len(voronoi_sites)):
        voronoi_sites[i] = 2 * depths[i-1] - voronoi_sites[i-1]
    voronoi_vector = np.hstack((velocities, voronoi_sites))
    return voronoi_vector

Interfacing to pysurf96

# display theory on the using the forward solver
with open("../../theory/geo_surface_wave_dispersion2.md", "r") as f:
    content = f.read()

display(Markdown(content))
<IPython.core.display.Markdown object>
VP_VS = 1.77
RHO_VP_K = 0.32
RHO_VP_B = 0.77
periods = np.geomspace(3, 80, 20)
# forward through pysurf96
def forward_sw(model, periods):
    thicknesses, vs = split_layercake_model(model)
    thicknesses = np.append(thicknesses, 10)
    vp = vs * VP_VS
    rho = RHO_VP_K * vp + RHO_VP_B
    return pysurf96.surf96(
        thicknesses,
        vp,
        vs,
        rho,
        periods,
        wave="rayleigh",
        mode=1,
        velocity="phase",
        flat_earth=False,
    )

Interfacing to Espresso

# obtain the receiver function library
rf_lib = espresso.ReceiverFunctionInversionKnt().rf
t_shift = 5
t_duration = 25
t_sampling_interval = 0.1
gauss = 1
ray_param_s_km = 0.07
# forward through rf in espresso
def forward_rf(
    model,
    t_shift=t_shift,
    t_duration=t_duration,
    t_sampling_interval=t_sampling_interval,
    gauss=gauss,
    ray_param_s_km=ray_param_s_km,
    return_times=False
):
    h, vs = split_layercake_model(model)
    data = rf_lib.rf_calc(ps=0, thik=h, beta=vs, kapa=np.ones((len(vs),))*VP_VS, p=ray_param_s_km,
                      duration=t_duration, dt=t_sampling_interval, shft=t_shift, gauss=gauss)
    if return_times:
        times = np.arange(len(data)) * t_sampling_interval - t_shift
        return data, times
    else:
        return data

Numerical Jacobian

def jacobian(model, data_length, fwd, fwd_kwargs, relative_step=0.01):
    jac = np.zeros((data_length, len(model)))
    original_dpred = fwd(model, **fwd_kwargs)
    for i in range(len(model)):
        perturbed_model = model.copy()
        step = relative_step * model[i]
        perturbed_model[i] += step
        perturbed_dpred = fwd(perturbed_model, **fwd_kwargs)
        derivative = (perturbed_dpred - original_dpred) / step
        jac[:, i] = derivative
    return jac

Plotting

def plot_model(model, ax=None, title="model", **kwargs):
    # process data
    thicknesses = np.append(model[1::2], max(model[1::2]))
    velocities = model[::2]
    y = np.insert(np.cumsum(thicknesses), 0, 0)
    x = np.insert(velocities, 0, velocities[0])

    # plot depth profile
    if ax is None:
        _, ax = plt.subplots()
    plotting_style = {
        "linewidth": kwargs.pop("linewidth", kwargs.pop("lw", 0.5)),
        "alpha": 0.2,
        "color": kwargs.pop("color", kwargs.pop("c", "blue")),
    }
    plotting_style.update(kwargs)
    ax.step(x, y, where="post", **plotting_style)
    if ax.get_ylim()[0] < ax.get_ylim()[1]:
        ax.invert_yaxis()
    ax.set_xlabel("Vs (km/s)")
    ax.set_ylabel("Depth (km)")
    ax.set_title(title)
    return ax
def plot_data(x, y, ax=None, scatter=False, xlabel=None, ylabel=None,
              title="surface wave data", **kwargs):
    if ax is None:
        _, ax = plt.subplots()
    plotting_style = {
        "linewidth": kwargs.pop("linewidth", kwargs.pop("lw", 1)),
        "alpha": 1,
        "color": kwargs.pop("color", kwargs.pop("c", "blue")),
    }
    plotting_style.update(**kwargs)
    if scatter:
        ax.scatter(x, y, **plotting_style)
    else:
        ax.plot(x, y, **plotting_style)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    return ax
def plot_sw_data(rayleigh_phase_velocities, periods, ax=None, scatter=False,
                 title="surface wave data", **kwargs):
    return plot_data(x=periods,
                     y=rayleigh_phase_velocities,
                     ax=ax,
                     scatter=scatter,
                     title=title,
                     xlabel="Periods (s)",
                     ylabel="Rayleigh phase velocities (km/s)",
                     **kwargs)

def plot_rf_data(rf_data, rf_times, ax=None, scatter=False,
                 title="receiver function data", **kwargs):
    return plot_data(x=rf_times,
                     y=rf_data,
                     ax=ax,
                     scatter=scatter,
                     title=title,
                     xlabel="Times (s)",
                     ylabel="Receiver function amplitudes",
                     **kwargs)

Generate synthetic data#

true_thickness = np.array([10, 10, 15, 20, 20, 20, 20, 20])
true_voronoi_positions = np.array([5, 15, 25, 45, 65, 85, 105, 125, 145])
true_vs = np.array([3.38, 3.44, 3.66, 4.25, 4.35, 4.32, 4.315, 4.38, 4.5])
true_model = form_layercake_model(true_thickness, true_vs)
RAYLEIGH_STD = 0.02
RF_STD = 0.015
rayleigh = forward_sw(true_model, periods)
rayleigh_dobs = rayleigh + np.random.normal(0, RAYLEIGH_STD, rayleigh.size)

rf, rf_times = forward_rf(true_model, return_times=True)
rf_dobs = rf + np.random.normal(0, RF_STD, rf.size)
_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 2.5, 2.5]})

# plot true model
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
ax1.grid()

# plot surface wave data
plot_sw_data(rayleigh, periods, ax=ax2, color="orange", label="true rayleigh data (noiseless)")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="brown", s=4,
          label="observed rayleigh data (noisy)")
ax2.grid()

# plot receiver function data
plot_rf_data(rf, rf_times, ax=ax3, color="lightblue", label="true receiver function data (noiseless)")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="darkblue", s=2,
          label="observed receiver function data (noisy)")
ax3.grid()

ax1.legend(loc="lower center", bbox_to_anchor=(0.5, -0.4))
ax2.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax3.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))

plt.tight_layout()
model, surface wave data, receiver function data

Optimisation#

Prepare ``BaseProblem`` for optimisation

n_dims = 17

init_thicknesses = np.ones((n_dims//2,)) * 15
init_vs = np.ones((n_dims//2+1,)) * 4.0
init_model = form_layercake_model(init_thicknesses, init_vs)
_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 2.5, 2.5]})

# plot true model
plot_model(init_model, ax=ax1, alpha=1, color="purple", label="starting model")
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
ax1.grid()

# plot surface wave data
plot_sw_data(forward_sw(init_model, periods), periods, ax=ax2, color="purple",
             label="predicted rayleigh data from starting model")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="orange", s=4,
          label="observed rayleigh data (noisy)")
ax2.grid()

# plot receiver function data
plot_rf_data(forward_rf(init_model), rf_times, ax=ax3, color="purple",
             label="predicted receiver function data from starting model")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="darkblue", s=2,
          label="observed receiver function data (noisy)")
ax3.grid()

ax1.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax2.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax3.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))

plt.tight_layout()
model, surface wave data, receiver function data
my_reg = cofi.utils.QuadraticReg(
    weighting_matrix="damping",
    model_shape=(n_dims,),
    reference_model=init_model
)
def my_objective(model, fwd_funcs, d_obs_list, lamda=1.0):
    data_misfit = 0
    for (fwd, fwd_kwargs), d_obs in zip(fwd_funcs, d_obs_list):
        d_pred = fwd(model, **fwd_kwargs)
        data_misfit += np.sum((d_obs - d_pred) ** 2)
    reg = my_reg(model)
    return data_misfit + lamda * reg

def my_objective_gradient(model, fwd_funcs, d_obs_list, lamda=1.0):
    data_misfit_grad = 0
    for (fwd, fwd_kwargs), d_obs in zip(fwd_funcs, d_obs_list):
        d_pred = fwd(model, **fwd_kwargs)
        jac = jacobian(model, len(d_obs), fwd, fwd_kwargs)
        data_misfit_grad += -2 * jac.T @ (d_obs - d_pred)
    reg_grad = my_reg.gradient(model)
    return data_misfit_grad + lamda * reg_grad

def my_objective_hessian(model, fwd_funcs, d_obs_list, lamda=1.0):
    data_misfit_hess = 0
    for (fwd, fwd_kwargs), d_obs in zip(fwd_funcs, d_obs_list):
        jac = jacobian(model, len(d_obs), fwd, fwd_kwargs)
        data_misfit_hess += 2 * jac.T @ jac
    reg_hess = my_reg.hessian(model)
    return data_misfit_hess + lamda * reg_hess
fwd_funcs = [
    (forward_sw, {"periods": periods}),
    (forward_rf, {
        "t_shift": t_shift,
        "t_duration": t_duration,
        "t_sampling_interval": t_sampling_interval,
        "gauss": gauss,
        "ray_param_s_km": ray_param_s_km
    })
]

d_obs_list = [rayleigh_dobs, rf_dobs]

Optimisation with no damping#

lamda = 0

kwargs = {
    "fwd_funcs": fwd_funcs,
    "d_obs_list": d_obs_list,
    "lamda": lamda
}
joint_problem_no_reg = cofi.BaseProblem()
joint_problem_no_reg.set_objective(my_objective, kwargs=kwargs)
joint_problem_no_reg.set_gradient(my_objective_gradient, kwargs=kwargs)
joint_problem_no_reg.set_hessian(my_objective_hessian, kwargs=kwargs)
joint_problem_no_reg.set_initial_model(init_model)

Define ``InversionOptions``

inv_options_optimiser = cofi.InversionOptions()
inv_options_optimiser.set_tool("scipy.optimize.minimize")
inv_options_optimiser.set_params(method="trust-exact")

Define ``Inversion`` and run

inv_optimiser_no_reg = cofi.Inversion(joint_problem_no_reg, inv_options_optimiser)
inv_res_optimiser_no_reg = inv_optimiser_no_reg.run()

Plot results

_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 2.5, 2.5]})

# plot true model
plot_model(inv_res_optimiser_no_reg.model, ax=ax1, alpha=1, color="purple",
           label="inverted model")
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
ax1.grid()

# plot surface wave data
plot_sw_data(forward_sw(inv_res_optimiser_no_reg.model, periods),
             periods, ax=ax2, color="purple",
             label="predicted rayleigh data from inverted model")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="orange", s=4,
          label="observed rayleigh data (noisy)")
ax2.grid()

# plot receiver function data
plot_rf_data(forward_rf(inv_res_optimiser_no_reg.model), rf_times,
             ax=ax3, color="purple",
             label="predicted receiver function data from inverted model")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="darkblue", s=2,
          label="observed receiver function data (noisy)")
ax3.grid()

ax1.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax2.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax3.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))

plt.tight_layout()
model, surface wave data, receiver function data

Optimal damping#

Oh no! The inverted model doesn’t look good, and it even has negative thicknesses.

Maybe adding a damping term to our objective function would help… but how can we determine a good damping factor?

The InversionPool in CoFI can help you answer it.

lambdas = np.logspace(-6, 6, 15)

my_lcurve_problems = []
for lamb in lambdas:
    my_problem = cofi.BaseProblem()
    kwargs = {
        "fwd_funcs": fwd_funcs,
        "d_obs_list": d_obs_list,
        "lamda": lamb
    }
    my_problem.set_objective(my_objective, kwargs=kwargs)
    my_problem.set_gradient(my_objective_gradient, kwargs=kwargs)
    my_problem.set_hessian(my_objective_hessian, kwargs=kwargs)
    my_problem.set_initial_model(init_model)
    my_lcurve_problems.append(my_problem)

def my_callback(inv_result, i):
    m = inv_result.model
    res_norm = 0
    for (fwd, fwd_kwargs), d_obs in zip(fwd_funcs, d_obs_list):
        d_pred = fwd(m, **fwd_kwargs)
        res_norm += np.sum((d_obs - d_pred) ** 2)
    reg_norm = np.sqrt(my_reg(m))
    print(f"Finished inversion with lambda={lambdas[i]}: {res_norm}, {reg_norm}")
    return res_norm, reg_norm

my_inversion_pool = cofi.utils.InversionPool(
    list_of_inv_problems=my_lcurve_problems,
    list_of_inv_options=inv_options_optimiser,
    callback=my_callback,
    parallel=False
)
all_res, all_cb_returns = my_inversion_pool.run()

l_curve_points = list(zip(*all_cb_returns))
Finished inversion with lambda=1e-06: 0.04520739446856946, 6.90707116167134
Finished inversion with lambda=7.196856730011514e-06: 0.04522718798587734, 6.360633520548262
Finished inversion with lambda=5.1794746792312125e-05: 0.04552497636462737, 5.009349361022729
Finished inversion with lambda=0.0003727593720314938: 0.04659449350775741, 4.1587131261737165
Finished inversion with lambda=0.0026826957952797246: 0.05292808006469342, 3.4623278161634596
Finished inversion with lambda=0.019306977288832496: 0.14459133826420748, 1.4196948655389219
Finished inversion with lambda=0.1389495494373136: 0.18703847430755832, 1.01540069556926
Finished inversion with lambda=1.0: 0.4475333284066435, 0.6491293172967686
Finished inversion with lambda=7.196856730011514: 1.366404791820631, 0.30436493811913745
Finished inversion with lambda=51.79474679231202: 2.7231174034586134, 0.06710141486808427
Finished inversion with lambda=372.7593720314938: 3.1382540203146867, 0.010094079042583812
Finished inversion with lambda=2682.6957952797275: 3.204302099415446, 0.0014182699553374812
Finished inversion with lambda=19306.977288832455: 3.2136675737206035, 0.00019739826915063293
Finished inversion with lambda=138949.5494373136: 3.2149867541680957, 2.7434652358831138e-05
Finished inversion with lambda=1000000.0: 3.215168621677044, 3.8121139706656857e-06
# print all the lambdas
lambdas
array([1.00000000e-06, 7.19685673e-06, 5.17947468e-05, 3.72759372e-04,
       2.68269580e-03, 1.93069773e-02, 1.38949549e-01, 1.00000000e+00,
       7.19685673e+00, 5.17947468e+01, 3.72759372e+02, 2.68269580e+03,
       1.93069773e+04, 1.38949549e+05, 1.00000000e+06])

Plot L-curve

# plot the L-curve
res_norm, reg_norm = l_curve_points
plt.plot(reg_norm, res_norm, '.-')
plt.xlabel(r'Norm of regularization term $||Wm||_2$')
plt.ylabel(r'Norm of residual $||g(m)-d||_2$')
for i in range(0, len(lambdas), 2):
    plt.annotate(f'{lambdas[i]:.1e}', (reg_norm[i], res_norm[i]), fontsize=8)
sw rf joint synthetic

Optimisation with damping#

From the L-curve plot above, it seems that a damping factor of around 0.14 would be good.

lamda = 0.14

kwargs = {
    "fwd_funcs": fwd_funcs,
    "d_obs_list": d_obs_list,
    "lamda": lamda
}
joint_problem = cofi.BaseProblem()
joint_problem.set_objective(my_objective, kwargs=kwargs)
joint_problem.set_gradient(my_objective_gradient, kwargs=kwargs)
joint_problem.set_hessian(my_objective_hessian, kwargs=kwargs)
joint_problem.set_initial_model(init_model)

Define ``Inversion`` and run

inv_optimiser = cofi.Inversion(joint_problem, inv_options_optimiser)
inv_res_optimiser = inv_optimiser.run()

Plot results

_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 2.5, 2.5]})

# plot true model
plot_model(inv_res_optimiser.model, ax=ax1, alpha=1, color="purple",
           label="inverted model")
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
ax1.grid()

# plot surface wave data
plot_sw_data(forward_sw(inv_res_optimiser_no_reg.model, periods),
             periods, ax=ax2, color="purple",
             label="predicted rayleigh data from inverted model")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="orange", s=4,
          label="observed rayleigh data (noisy)")
ax2.grid()

# plot receiver function data
plot_rf_data(forward_rf(inv_res_optimiser.model), rf_times,
             ax=ax3, color="purple",
             label="predicted receiver function data from inverted model")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="darkblue", s=2,
          label="observed receiver function data (noisy)")
ax3.grid()

ax1.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax2.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))
ax3.legend(loc="lower center", bbox_to_anchor=(0.5, -0.5))

plt.tight_layout()
model, surface wave data, receiver function data

Fixed-dimensional sampling#

Prepare ``BaseProblem`` for fixed-dimensional sampling

thick_min = 5
thick_max = 30
vs_min = 2
vs_max = 5

def my_log_prior(model):
    thicknesses, vs = split_layercake_model(model)
    thicknesses_out_of_bounds = (thicknesses < thick_min) | (thicknesses > thick_max)
    vs_out_of_bounds = (vs < vs_min) | (vs > vs_max)
    if np.any(thicknesses_out_of_bounds) or np.any(vs_out_of_bounds):
        return float("-inf")
    log_prior = - np.log(thick_max - thick_min) * len(thicknesses) \
                - np.log(vs_max - vs_min) * len(vs)
    return log_prior
# inverse data covariance matrix
Cdinv_rayleigh = np.eye(len(rayleigh_dobs)) / (RAYLEIGH_STD**2)
Cdinv_rf = np.eye(len(rf_dobs)) / (RF_STD**2)
Cdinv_list = [Cdinv_rayleigh, Cdinv_rf]

def my_log_likelihood(
    model,
    fwd_funcs=fwd_funcs,
    d_obs_list=d_obs_list,
    Cdinv_list=Cdinv_list
):
    log_like_sum = 0
    for (fwd, fwd_kwargs), d_obs, Cdinv in zip(fwd_funcs, d_obs_list, Cdinv_list):
        try:
            d_pred = fwd(model, **fwd_kwargs)
        except:
            return float("-inf")
        residual = d_obs - d_pred
        log_like_sum += -0.5 * residual @ (Cdinv @ residual).T
    return log_like_sum
n_walkers = 40

my_walkers_start = np.ones((n_walkers, n_dims)) * inv_res_optimiser.model
for i in range(n_walkers):
    my_walkers_start[i,:] += np.random.normal(0, 0.5, n_dims)
joint_problem.set_log_prior(my_log_prior)
joint_problem.set_log_likelihood(my_log_likelihood)

Define ``InversionOptions``

inv_options_fixed_d_sampling = cofi.InversionOptions()
inv_options_fixed_d_sampling.set_tool("emcee")
inv_options_fixed_d_sampling.set_params(
    nwalkers=n_walkers,
    nsteps=2_000,
    initial_state=my_walkers_start,
    skip_initial_state_check=True,
    progress=True
)

Define ``Inversion`` and run

Sample the prior#

prior_sampling_problem = cofi.BaseProblem()
prior_sampling_problem.set_log_posterior(my_log_prior)
prior_sampling_problem.set_model_shape(init_model.shape)
prior_sampler = cofi.Inversion(prior_sampling_problem, inv_options_fixed_d_sampling)
prior_results = prior_sampler.run()
  0%|          | 0/2000 [00:00<?, ?it/s]
  5%|▌         | 101/2000 [00:00<00:01, 1005.32it/s]
 10%|█         | 203/2000 [00:00<00:01, 1013.63it/s]
 15%|█▌        | 305/2000 [00:00<00:01, 1013.61it/s]
 20%|██        | 407/2000 [00:00<00:01, 1007.48it/s]
 26%|██▌       | 513/2000 [00:00<00:01, 1024.92it/s]
 31%|███       | 616/2000 [00:00<00:01, 1025.36it/s]
 36%|███▌      | 723/2000 [00:00<00:01, 1039.76it/s]
 42%|████▏     | 830/2000 [00:00<00:01, 1049.34it/s]
 47%|████▋     | 936/2000 [00:00<00:01, 1050.95it/s]
 52%|█████▏    | 1042/2000 [00:01<00:00, 1052.62it/s]
 57%|█████▋    | 1148/2000 [00:01<00:00, 1047.93it/s]
 63%|██████▎   | 1253/2000 [00:01<00:00, 1034.44it/s]
 68%|██████▊   | 1357/2000 [00:01<00:00, 1028.20it/s]
 73%|███████▎  | 1461/2000 [00:01<00:00, 1031.27it/s]
 79%|███████▊  | 1571/2000 [00:01<00:00, 1051.34it/s]
 84%|████████▍ | 1678/2000 [00:01<00:00, 1055.65it/s]
 90%|████████▉ | 1790/2000 [00:01<00:00, 1074.38it/s]
 95%|█████████▍| 1899/2000 [00:01<00:00, 1076.20it/s]
100%|██████████| 2000/2000 [00:01<00:00, 1047.42it/s]
import arviz as az

labels = ["v0", "t0", "v1", "t1", "v2", "t2", "v3", "t3", "v4", "t4", "v5", "t5", "v6", "t6", "v7", "t7", "v8"]

prior_results_sampler = prior_results.sampler
az_idata_prior = az.from_emcee(prior_results_sampler, var_names=labels)

axes = az.plot_trace(
    az_idata_prior,
    backend_kwargs={"constrained_layout":True},
    figsize=(10,20),
)

for i, axes_pair in enumerate(axes):
    ax1 = axes_pair[0]
    ax2 = axes_pair[1]
    ax1.set_xlabel("parameter value")
    ax1.set_ylabel("density value")
    ax2.set_xlabel("number of iterations")
    ax2.set_ylabel("parameter value")
t0, t0, t1, t1, t2, t2, t3, t3, t4, t4, t5, t5, t6, t6, t7, t7, v0, v0, v1, v1, v2, v2, v3, v3, v4, v4, v5, v5, v6, v6, v7, v7, v8, v8

Sample the posterior#

inversion_fixed_d_sampler = cofi.Inversion(joint_problem, inv_options_fixed_d_sampling)
inv_result_fixed_d_sampler = inversion_fixed_d_sampler.run()
  0%|          | 0/2000 [00:00<?, ?it/s]
  0%|          | 5/2000 [00:00<00:47, 41.88it/s]
  0%|          | 10/2000 [00:00<00:57, 34.89it/s]
  1%|          | 14/2000 [00:00<01:05, 30.55it/s]
  1%|          | 18/2000 [00:00<01:09, 28.57it/s]
  1%|          | 21/2000 [00:00<01:11, 27.73it/s]
  1%|          | 24/2000 [00:00<01:13, 26.96it/s]
  1%|▏         | 27/2000 [00:00<01:14, 26.53it/s]
  2%|▏         | 30/2000 [00:01<01:15, 26.16it/s]
  2%|▏         | 33/2000 [00:01<01:15, 25.90it/s]
  2%|▏         | 36/2000 [00:01<01:16, 25.72it/s]
  2%|▏         | 39/2000 [00:01<01:16, 25.79it/s]
  2%|▏         | 42/2000 [00:01<01:15, 25.80it/s]
  2%|▏         | 45/2000 [00:01<01:16, 25.65it/s]
  2%|▏         | 48/2000 [00:01<01:15, 25.71it/s]
  3%|▎         | 51/2000 [00:01<01:16, 25.63it/s]
  3%|▎         | 54/2000 [00:01<01:15, 25.71it/s]
  3%|▎         | 57/2000 [00:02<01:15, 25.61it/s]
  3%|▎         | 60/2000 [00:02<01:15, 25.57it/s]
  3%|▎         | 63/2000 [00:02<01:15, 25.76it/s]
  3%|▎         | 66/2000 [00:02<01:15, 25.70it/s]
  3%|▎         | 69/2000 [00:02<01:15, 25.63it/s]
  4%|▎         | 72/2000 [00:02<01:15, 25.62it/s]
  4%|▍         | 75/2000 [00:02<01:15, 25.63it/s]
  4%|▍         | 78/2000 [00:02<01:14, 25.66it/s]
  4%|▍         | 81/2000 [00:03<01:15, 25.35it/s]
  4%|▍         | 84/2000 [00:03<01:16, 25.12it/s]
  4%|▍         | 87/2000 [00:03<01:17, 24.68it/s]
  4%|▍         | 90/2000 [00:03<01:17, 24.55it/s]
  5%|▍         | 93/2000 [00:03<01:16, 24.86it/s]
  5%|▍         | 96/2000 [00:03<01:18, 24.30it/s]
  5%|▍         | 99/2000 [00:03<01:18, 24.33it/s]
  5%|▌         | 102/2000 [00:03<01:17, 24.44it/s]
  5%|▌         | 105/2000 [00:04<01:17, 24.56it/s]
  5%|▌         | 108/2000 [00:04<01:16, 24.64it/s]
  6%|▌         | 111/2000 [00:04<01:16, 24.75it/s]
  6%|▌         | 114/2000 [00:04<01:15, 25.05it/s]
  6%|▌         | 117/2000 [00:04<01:15, 25.05it/s]
  6%|▌         | 120/2000 [00:04<01:16, 24.64it/s]
  6%|▌         | 123/2000 [00:04<01:15, 24.83it/s]
  6%|▋         | 126/2000 [00:04<01:15, 24.78it/s]
  6%|▋         | 129/2000 [00:05<01:15, 24.68it/s]
  7%|▋         | 132/2000 [00:05<01:15, 24.69it/s]
  7%|▋         | 135/2000 [00:05<01:15, 24.77it/s]
  7%|▋         | 138/2000 [00:05<01:15, 24.67it/s]
  7%|▋         | 141/2000 [00:05<01:15, 24.62it/s]
  7%|▋         | 144/2000 [00:05<01:14, 24.90it/s]
  7%|▋         | 147/2000 [00:05<01:14, 24.84it/s]
  8%|▊         | 150/2000 [00:05<01:14, 24.88it/s]
  8%|▊         | 153/2000 [00:05<01:13, 24.99it/s]
  8%|▊         | 156/2000 [00:06<01:13, 25.08it/s]
  8%|▊         | 159/2000 [00:06<01:13, 25.11it/s]
  8%|▊         | 162/2000 [00:06<01:12, 25.22it/s]
  8%|▊         | 165/2000 [00:06<01:12, 25.31it/s]
  8%|▊         | 168/2000 [00:06<01:11, 25.54it/s]
  9%|▊         | 171/2000 [00:06<01:11, 25.54it/s]
  9%|▊         | 174/2000 [00:06<01:11, 25.54it/s]
  9%|▉         | 177/2000 [00:06<01:11, 25.63it/s]
  9%|▉         | 180/2000 [00:07<01:10, 25.65it/s]
  9%|▉         | 183/2000 [00:07<01:10, 25.75it/s]
  9%|▉         | 186/2000 [00:07<01:10, 25.78it/s]
  9%|▉         | 189/2000 [00:07<01:10, 25.85it/s]
 10%|▉         | 192/2000 [00:07<01:09, 25.92it/s]
 10%|▉         | 195/2000 [00:07<01:09, 25.94it/s]
 10%|▉         | 198/2000 [00:07<01:09, 25.93it/s]
 10%|█         | 201/2000 [00:07<01:09, 25.87it/s]
 10%|█         | 204/2000 [00:07<01:09, 25.91it/s]
 10%|█         | 207/2000 [00:08<01:09, 25.77it/s]
 10%|█         | 210/2000 [00:08<01:09, 25.78it/s]
 11%|█         | 213/2000 [00:08<01:09, 25.69it/s]
 11%|█         | 216/2000 [00:08<01:09, 25.71it/s]
 11%|█         | 219/2000 [00:08<01:09, 25.70it/s]
 11%|█         | 222/2000 [00:08<01:09, 25.63it/s]
 11%|█▏        | 225/2000 [00:08<01:09, 25.70it/s]
 11%|█▏        | 228/2000 [00:08<01:08, 25.77it/s]
 12%|█▏        | 231/2000 [00:09<01:08, 25.78it/s]
 12%|█▏        | 234/2000 [00:09<01:08, 25.74it/s]
 12%|█▏        | 237/2000 [00:09<01:08, 25.73it/s]
 12%|█▏        | 240/2000 [00:09<01:08, 25.61it/s]
 12%|█▏        | 243/2000 [00:09<01:08, 25.54it/s]
 12%|█▏        | 246/2000 [00:09<01:08, 25.55it/s]
 12%|█▏        | 249/2000 [00:09<01:08, 25.65it/s]
 13%|█▎        | 252/2000 [00:09<01:08, 25.62it/s]
 13%|█▎        | 255/2000 [00:09<01:08, 25.57it/s]
 13%|█▎        | 258/2000 [00:10<01:08, 25.37it/s]
 13%|█▎        | 261/2000 [00:10<01:08, 25.48it/s]
 13%|█▎        | 264/2000 [00:10<01:07, 25.61it/s]
 13%|█▎        | 267/2000 [00:10<01:07, 25.58it/s]
 14%|█▎        | 270/2000 [00:10<01:07, 25.64it/s]
 14%|█▎        | 273/2000 [00:10<01:07, 25.52it/s]
 14%|█▍        | 276/2000 [00:10<01:07, 25.62it/s]
 14%|█▍        | 279/2000 [00:10<01:07, 25.64it/s]
 14%|█▍        | 282/2000 [00:11<01:07, 25.38it/s]
 14%|█▍        | 285/2000 [00:11<01:07, 25.39it/s]
 14%|█▍        | 288/2000 [00:11<01:08, 25.07it/s]
 15%|█▍        | 291/2000 [00:11<01:08, 25.04it/s]
 15%|█▍        | 294/2000 [00:11<01:08, 24.86it/s]
 15%|█▍        | 297/2000 [00:11<01:08, 24.93it/s]
 15%|█▌        | 300/2000 [00:11<01:07, 25.00it/s]
 15%|█▌        | 303/2000 [00:11<01:08, 24.68it/s]
 15%|█▌        | 306/2000 [00:11<01:08, 24.82it/s]
 15%|█▌        | 309/2000 [00:12<01:08, 24.80it/s]
 16%|█▌        | 312/2000 [00:12<01:07, 25.09it/s]
 16%|█▌        | 315/2000 [00:12<01:07, 24.79it/s]
 16%|█▌        | 318/2000 [00:12<01:08, 24.72it/s]
 16%|█▌        | 321/2000 [00:12<01:07, 24.99it/s]
 16%|█▌        | 324/2000 [00:12<01:06, 25.19it/s]
 16%|█▋        | 327/2000 [00:12<01:06, 25.27it/s]
 16%|█▋        | 330/2000 [00:12<01:06, 25.26it/s]
 17%|█▋        | 333/2000 [00:13<01:05, 25.40it/s]
 17%|█▋        | 336/2000 [00:13<01:05, 25.50it/s]
 17%|█▋        | 339/2000 [00:13<01:04, 25.61it/s]
 17%|█▋        | 342/2000 [00:13<01:04, 25.66it/s]
 17%|█▋        | 345/2000 [00:13<01:05, 25.41it/s]
 17%|█▋        | 348/2000 [00:13<01:04, 25.49it/s]
 18%|█▊        | 351/2000 [00:13<01:04, 25.53it/s]
 18%|█▊        | 354/2000 [00:13<01:04, 25.47it/s]
 18%|█▊        | 357/2000 [00:13<01:04, 25.59it/s]
 18%|█▊        | 360/2000 [00:14<01:04, 25.51it/s]
 18%|█▊        | 363/2000 [00:14<01:04, 25.34it/s]
 18%|█▊        | 366/2000 [00:14<01:04, 25.48it/s]
 18%|█▊        | 369/2000 [00:14<01:04, 25.19it/s]
 19%|█▊        | 372/2000 [00:14<01:04, 25.34it/s]
 19%|█▉        | 375/2000 [00:14<01:04, 25.37it/s]
 19%|█▉        | 378/2000 [00:14<01:03, 25.38it/s]
 19%|█▉        | 381/2000 [00:14<01:03, 25.47it/s]
 19%|█▉        | 384/2000 [00:15<01:03, 25.52it/s]
 19%|█▉        | 387/2000 [00:15<01:02, 25.63it/s]
 20%|█▉        | 390/2000 [00:15<01:02, 25.71it/s]
 20%|█▉        | 393/2000 [00:15<01:02, 25.58it/s]
 20%|█▉        | 396/2000 [00:15<01:03, 25.28it/s]
 20%|█▉        | 399/2000 [00:15<01:04, 25.00it/s]
 20%|██        | 402/2000 [00:15<01:03, 25.13it/s]
 20%|██        | 405/2000 [00:15<01:03, 25.24it/s]
 20%|██        | 408/2000 [00:15<01:02, 25.39it/s]
 21%|██        | 411/2000 [00:16<01:02, 25.39it/s]
 21%|██        | 414/2000 [00:16<01:02, 25.42it/s]
 21%|██        | 417/2000 [00:16<01:01, 25.62it/s]
 21%|██        | 420/2000 [00:16<01:01, 25.75it/s]
 21%|██        | 423/2000 [00:16<01:01, 25.58it/s]
 21%|██▏       | 426/2000 [00:16<01:02, 25.33it/s]
 21%|██▏       | 429/2000 [00:16<01:02, 25.12it/s]
 22%|██▏       | 432/2000 [00:16<01:02, 24.96it/s]
 22%|██▏       | 435/2000 [00:17<01:02, 24.86it/s]
 22%|██▏       | 438/2000 [00:17<01:02, 25.03it/s]
 22%|██▏       | 441/2000 [00:17<01:02, 25.10it/s]
 22%|██▏       | 444/2000 [00:17<01:01, 25.49it/s]
 22%|██▏       | 447/2000 [00:17<01:00, 25.68it/s]
 22%|██▎       | 450/2000 [00:17<00:59, 26.10it/s]
 23%|██▎       | 453/2000 [00:17<00:58, 26.29it/s]
 23%|██▎       | 456/2000 [00:17<00:58, 26.30it/s]
 23%|██▎       | 459/2000 [00:17<00:58, 26.14it/s]
 23%|██▎       | 462/2000 [00:18<00:58, 26.19it/s]
 23%|██▎       | 465/2000 [00:18<00:58, 26.30it/s]
 23%|██▎       | 468/2000 [00:18<00:57, 26.56it/s]
 24%|██▎       | 471/2000 [00:18<00:56, 26.88it/s]
 24%|██▎       | 474/2000 [00:18<00:57, 26.72it/s]
 24%|██▍       | 477/2000 [00:18<00:57, 26.58it/s]
 24%|██▍       | 480/2000 [00:18<00:56, 26.78it/s]
 24%|██▍       | 483/2000 [00:18<00:56, 27.02it/s]
 24%|██▍       | 486/2000 [00:18<00:55, 27.37it/s]
 24%|██▍       | 490/2000 [00:19<00:53, 28.47it/s]
 25%|██▍       | 493/2000 [00:19<00:53, 28.40it/s]
 25%|██▍       | 496/2000 [00:19<00:53, 28.27it/s]
 25%|██▍       | 499/2000 [00:19<00:53, 28.01it/s]
 25%|██▌       | 502/2000 [00:19<00:53, 28.11it/s]
 25%|██▌       | 505/2000 [00:19<00:52, 28.30it/s]
 25%|██▌       | 508/2000 [00:19<00:52, 28.18it/s]
 26%|██▌       | 511/2000 [00:19<00:52, 28.56it/s]
 26%|██▌       | 514/2000 [00:19<00:52, 28.09it/s]
 26%|██▌       | 517/2000 [00:20<00:52, 28.46it/s]
 26%|██▌       | 520/2000 [00:20<00:51, 28.81it/s]
 26%|██▌       | 523/2000 [00:20<00:50, 29.08it/s]
 26%|██▋       | 526/2000 [00:20<00:51, 28.85it/s]
 26%|██▋       | 529/2000 [00:20<00:51, 28.82it/s]
 27%|██▋       | 532/2000 [00:20<00:51, 28.60it/s]
 27%|██▋       | 535/2000 [00:20<00:51, 28.50it/s]
 27%|██▋       | 538/2000 [00:20<00:51, 28.51it/s]
 27%|██▋       | 541/2000 [00:20<00:53, 27.44it/s]
 27%|██▋       | 544/2000 [00:21<00:52, 27.50it/s]
 27%|██▋       | 547/2000 [00:21<00:51, 28.13it/s]
 28%|██▊       | 551/2000 [00:21<00:49, 29.15it/s]
 28%|██▊       | 555/2000 [00:21<00:48, 29.82it/s]
 28%|██▊       | 558/2000 [00:21<00:49, 29.13it/s]
 28%|██▊       | 561/2000 [00:21<00:50, 28.75it/s]
 28%|██▊       | 564/2000 [00:21<00:49, 28.97it/s]
 28%|██▊       | 567/2000 [00:21<00:50, 28.30it/s]
 28%|██▊       | 570/2000 [00:21<00:50, 28.12it/s]
 29%|██▊       | 573/2000 [00:22<00:50, 28.10it/s]
 29%|██▉       | 577/2000 [00:22<00:48, 29.19it/s]
 29%|██▉       | 581/2000 [00:22<00:47, 29.83it/s]
 29%|██▉       | 584/2000 [00:22<00:47, 29.62it/s]
 29%|██▉       | 587/2000 [00:22<00:48, 29.12it/s]
 30%|██▉       | 590/2000 [00:22<00:49, 28.47it/s]
 30%|██▉       | 593/2000 [00:22<00:49, 28.59it/s]
 30%|██▉       | 596/2000 [00:22<00:48, 28.89it/s]
 30%|██▉       | 599/2000 [00:22<00:49, 28.47it/s]
 30%|███       | 602/2000 [00:23<00:49, 28.04it/s]
 30%|███       | 605/2000 [00:23<00:50, 27.43it/s]
 30%|███       | 608/2000 [00:23<00:50, 27.57it/s]
 31%|███       | 611/2000 [00:23<00:49, 27.85it/s]
 31%|███       | 614/2000 [00:23<00:50, 27.44it/s]
 31%|███       | 617/2000 [00:23<00:49, 27.80it/s]
 31%|███       | 620/2000 [00:23<00:50, 27.36it/s]
 31%|███       | 623/2000 [00:23<00:50, 27.04it/s]
 31%|███▏      | 626/2000 [00:23<00:50, 26.96it/s]
 31%|███▏      | 629/2000 [00:24<00:50, 27.19it/s]
 32%|███▏      | 632/2000 [00:24<00:50, 26.97it/s]
 32%|███▏      | 635/2000 [00:24<00:50, 26.98it/s]
 32%|███▏      | 638/2000 [00:24<00:50, 27.01it/s]
 32%|███▏      | 641/2000 [00:24<00:50, 27.07it/s]
 32%|███▏      | 644/2000 [00:24<00:51, 26.31it/s]
 32%|███▏      | 647/2000 [00:24<00:51, 26.21it/s]
 32%|███▎      | 650/2000 [00:24<00:50, 26.60it/s]
 33%|███▎      | 653/2000 [00:24<00:49, 27.17it/s]
 33%|███▎      | 656/2000 [00:25<00:48, 27.73it/s]
 33%|███▎      | 659/2000 [00:25<00:47, 28.19it/s]
 33%|███▎      | 662/2000 [00:25<00:48, 27.86it/s]
 33%|███▎      | 665/2000 [00:25<00:47, 28.06it/s]
 33%|███▎      | 668/2000 [00:25<00:46, 28.36it/s]
 34%|███▎      | 672/2000 [00:25<00:45, 29.15it/s]
 34%|███▍      | 675/2000 [00:25<00:45, 29.24it/s]
 34%|███▍      | 678/2000 [00:25<00:45, 29.05it/s]
 34%|███▍      | 681/2000 [00:25<00:45, 29.05it/s]
 34%|███▍      | 684/2000 [00:25<00:45, 29.09it/s]
 34%|███▍      | 687/2000 [00:26<00:45, 29.10it/s]
 34%|███▍      | 690/2000 [00:26<00:45, 29.11it/s]
 35%|███▍      | 693/2000 [00:26<00:44, 29.29it/s]
 35%|███▍      | 696/2000 [00:26<00:44, 29.43it/s]
 35%|███▍      | 699/2000 [00:26<00:44, 29.15it/s]
 35%|███▌      | 702/2000 [00:26<00:44, 29.25it/s]
 35%|███▌      | 705/2000 [00:26<00:44, 28.84it/s]
 35%|███▌      | 708/2000 [00:26<00:44, 28.83it/s]
 36%|███▌      | 711/2000 [00:26<00:44, 28.89it/s]
 36%|███▌      | 714/2000 [00:27<00:44, 28.88it/s]
 36%|███▌      | 718/2000 [00:27<00:43, 29.56it/s]
 36%|███▌      | 721/2000 [00:27<00:44, 28.67it/s]
 36%|███▌      | 724/2000 [00:27<00:44, 28.89it/s]
 36%|███▋      | 727/2000 [00:27<00:44, 28.75it/s]
 36%|███▋      | 730/2000 [00:27<00:44, 28.57it/s]
 37%|███▋      | 733/2000 [00:27<00:44, 28.75it/s]
 37%|███▋      | 736/2000 [00:27<00:44, 28.21it/s]
 37%|███▋      | 739/2000 [00:27<00:44, 28.48it/s]
 37%|███▋      | 742/2000 [00:28<00:44, 28.10it/s]
 37%|███▋      | 745/2000 [00:28<00:45, 27.45it/s]
 37%|███▋      | 748/2000 [00:28<00:45, 27.25it/s]
 38%|███▊      | 751/2000 [00:28<00:44, 27.83it/s]
 38%|███▊      | 754/2000 [00:28<00:45, 27.21it/s]
 38%|███▊      | 758/2000 [00:28<00:44, 28.03it/s]
 38%|███▊      | 762/2000 [00:28<00:43, 28.77it/s]
 38%|███▊      | 765/2000 [00:28<00:42, 29.03it/s]
 38%|███▊      | 768/2000 [00:28<00:42, 29.18it/s]
 39%|███▊      | 771/2000 [00:29<00:41, 29.38it/s]
 39%|███▉      | 775/2000 [00:29<00:40, 30.09it/s]
 39%|███▉      | 779/2000 [00:29<00:40, 30.17it/s]
 39%|███▉      | 783/2000 [00:29<00:41, 29.50it/s]
 39%|███▉      | 786/2000 [00:29<00:41, 29.26it/s]
 39%|███▉      | 789/2000 [00:29<00:41, 29.39it/s]
 40%|███▉      | 793/2000 [00:29<00:40, 29.80it/s]
 40%|███▉      | 796/2000 [00:29<00:41, 29.29it/s]
 40%|███▉      | 799/2000 [00:29<00:40, 29.34it/s]
 40%|████      | 802/2000 [00:30<00:41, 28.53it/s]
 40%|████      | 805/2000 [00:30<00:41, 28.78it/s]
 40%|████      | 808/2000 [00:30<00:42, 27.88it/s]
 41%|████      | 811/2000 [00:30<00:42, 28.09it/s]
 41%|████      | 814/2000 [00:30<00:41, 28.61it/s]
 41%|████      | 817/2000 [00:30<00:41, 28.84it/s]
 41%|████      | 821/2000 [00:30<00:39, 29.65it/s]
 41%|████▏     | 825/2000 [00:30<00:39, 30.00it/s]
 41%|████▏     | 829/2000 [00:30<00:38, 30.33it/s]
 42%|████▏     | 833/2000 [00:31<00:38, 30.18it/s]
 42%|████▏     | 837/2000 [00:31<00:39, 29.55it/s]
 42%|████▏     | 841/2000 [00:31<00:38, 30.25it/s]
 42%|████▏     | 845/2000 [00:31<00:37, 30.83it/s]
 42%|████▏     | 849/2000 [00:31<00:36, 31.27it/s]
 43%|████▎     | 853/2000 [00:31<00:36, 31.79it/s]
 43%|████▎     | 857/2000 [00:31<00:35, 31.91it/s]
 43%|████▎     | 861/2000 [00:32<00:36, 31.26it/s]
 43%|████▎     | 865/2000 [00:32<00:36, 31.07it/s]
 43%|████▎     | 869/2000 [00:32<00:35, 31.46it/s]
 44%|████▎     | 873/2000 [00:32<00:35, 32.00it/s]
 44%|████▍     | 877/2000 [00:32<00:35, 31.67it/s]
 44%|████▍     | 881/2000 [00:32<00:35, 31.42it/s]
 44%|████▍     | 885/2000 [00:32<00:35, 31.16it/s]
 44%|████▍     | 889/2000 [00:32<00:35, 31.53it/s]
 45%|████▍     | 893/2000 [00:33<00:35, 31.37it/s]
 45%|████▍     | 897/2000 [00:33<00:35, 31.35it/s]
 45%|████▌     | 901/2000 [00:33<00:35, 31.35it/s]
 45%|████▌     | 905/2000 [00:33<00:34, 32.02it/s]
 45%|████▌     | 909/2000 [00:33<00:34, 31.72it/s]
 46%|████▌     | 913/2000 [00:33<00:33, 32.01it/s]
 46%|████▌     | 917/2000 [00:33<00:33, 31.96it/s]
 46%|████▌     | 921/2000 [00:33<00:33, 32.26it/s]
 46%|████▋     | 925/2000 [00:34<00:33, 31.68it/s]
 46%|████▋     | 929/2000 [00:34<00:34, 30.68it/s]
 47%|████▋     | 933/2000 [00:34<00:34, 31.03it/s]
 47%|████▋     | 937/2000 [00:34<00:34, 31.09it/s]
 47%|████▋     | 941/2000 [00:34<00:33, 31.68it/s]
 47%|████▋     | 945/2000 [00:34<00:33, 31.07it/s]
 47%|████▋     | 949/2000 [00:34<00:33, 31.45it/s]
 48%|████▊     | 953/2000 [00:34<00:32, 32.41it/s]
 48%|████▊     | 957/2000 [00:35<00:32, 31.80it/s]
 48%|████▊     | 961/2000 [00:35<00:32, 31.82it/s]
 48%|████▊     | 965/2000 [00:35<00:32, 31.83it/s]
 48%|████▊     | 969/2000 [00:35<00:32, 31.75it/s]
 49%|████▊     | 973/2000 [00:35<00:32, 31.36it/s]
 49%|████▉     | 977/2000 [00:35<00:32, 31.75it/s]
 49%|████▉     | 981/2000 [00:35<00:31, 32.27it/s]
 49%|████▉     | 985/2000 [00:35<00:31, 31.74it/s]
 49%|████▉     | 989/2000 [00:36<00:31, 31.83it/s]
 50%|████▉     | 993/2000 [00:36<00:31, 31.56it/s]
 50%|████▉     | 997/2000 [00:36<00:31, 31.56it/s]
 50%|█████     | 1001/2000 [00:36<00:31, 31.74it/s]
 50%|█████     | 1005/2000 [00:36<00:32, 30.59it/s]
 50%|█████     | 1009/2000 [00:36<00:32, 30.59it/s]
 51%|█████     | 1013/2000 [00:36<00:33, 29.54it/s]
 51%|█████     | 1017/2000 [00:36<00:32, 30.31it/s]
 51%|█████     | 1021/2000 [00:37<00:31, 31.03it/s]
 51%|█████▏    | 1025/2000 [00:37<00:30, 31.68it/s]
 51%|█████▏    | 1029/2000 [00:37<00:30, 31.69it/s]
 52%|█████▏    | 1033/2000 [00:37<00:30, 32.07it/s]
 52%|█████▏    | 1037/2000 [00:37<00:29, 32.18it/s]
 52%|█████▏    | 1041/2000 [00:37<00:30, 31.91it/s]
 52%|█████▏    | 1045/2000 [00:37<00:30, 31.42it/s]
 52%|█████▏    | 1049/2000 [00:37<00:30, 31.17it/s]
 53%|█████▎    | 1053/2000 [00:38<00:29, 31.83it/s]
 53%|█████▎    | 1057/2000 [00:38<00:29, 31.95it/s]
 53%|█████▎    | 1061/2000 [00:38<00:28, 32.66it/s]
 53%|█████▎    | 1065/2000 [00:38<00:28, 32.52it/s]
 53%|█████▎    | 1069/2000 [00:38<00:28, 32.66it/s]
 54%|█████▎    | 1073/2000 [00:38<00:28, 32.38it/s]
 54%|█████▍    | 1077/2000 [00:38<00:28, 32.06it/s]
 54%|█████▍    | 1081/2000 [00:38<00:28, 32.31it/s]
 54%|█████▍    | 1085/2000 [00:39<00:27, 33.09it/s]
 54%|█████▍    | 1089/2000 [00:39<00:27, 33.70it/s]
 55%|█████▍    | 1093/2000 [00:39<00:26, 33.68it/s]
 55%|█████▍    | 1097/2000 [00:39<00:26, 33.99it/s]
 55%|█████▌    | 1101/2000 [00:39<00:25, 35.01it/s]
 55%|█████▌    | 1105/2000 [00:39<00:25, 35.12it/s]
 55%|█████▌    | 1109/2000 [00:39<00:25, 35.34it/s]
 56%|█████▌    | 1113/2000 [00:39<00:24, 35.98it/s]
 56%|█████▌    | 1117/2000 [00:39<00:24, 35.66it/s]
 56%|█████▌    | 1121/2000 [00:40<00:24, 35.91it/s]
 56%|█████▋    | 1125/2000 [00:40<00:24, 35.78it/s]
 56%|█████▋    | 1129/2000 [00:40<00:24, 35.63it/s]
 57%|█████▋    | 1133/2000 [00:40<00:24, 34.82it/s]
 57%|█████▋    | 1137/2000 [00:40<00:24, 35.17it/s]
 57%|█████▋    | 1141/2000 [00:40<00:24, 35.73it/s]
 57%|█████▋    | 1145/2000 [00:40<00:23, 36.03it/s]
 57%|█████▋    | 1149/2000 [00:40<00:23, 36.87it/s]
 58%|█████▊    | 1154/2000 [00:40<00:22, 37.86it/s]
 58%|█████▊    | 1159/2000 [00:41<00:21, 38.47it/s]
 58%|█████▊    | 1163/2000 [00:41<00:21, 38.71it/s]
 58%|█████▊    | 1168/2000 [00:41<00:21, 39.31it/s]
 59%|█████▊    | 1172/2000 [00:41<00:21, 38.81it/s]
 59%|█████▉    | 1176/2000 [00:41<00:21, 38.52it/s]
 59%|█████▉    | 1180/2000 [00:41<00:21, 37.49it/s]
 59%|█████▉    | 1184/2000 [00:41<00:22, 36.66it/s]
 59%|█████▉    | 1188/2000 [00:41<00:21, 36.93it/s]
 60%|█████▉    | 1192/2000 [00:42<00:22, 35.69it/s]
 60%|█████▉    | 1196/2000 [00:42<00:22, 35.18it/s]
 60%|██████    | 1200/2000 [00:42<00:22, 34.92it/s]
 60%|██████    | 1204/2000 [00:42<00:22, 34.97it/s]
 60%|██████    | 1208/2000 [00:42<00:21, 36.12it/s]
 61%|██████    | 1212/2000 [00:42<00:21, 36.53it/s]
 61%|██████    | 1216/2000 [00:42<00:21, 36.84it/s]
 61%|██████    | 1220/2000 [00:42<00:21, 36.03it/s]
 61%|██████    | 1224/2000 [00:42<00:20, 37.05it/s]
 61%|██████▏   | 1228/2000 [00:43<00:20, 36.87it/s]
 62%|██████▏   | 1232/2000 [00:43<00:20, 36.81it/s]
 62%|██████▏   | 1236/2000 [00:43<00:20, 36.53it/s]
 62%|██████▏   | 1240/2000 [00:43<00:20, 36.68it/s]
 62%|██████▏   | 1244/2000 [00:43<00:21, 35.97it/s]
 62%|██████▏   | 1248/2000 [00:43<00:20, 36.03it/s]
 63%|██████▎   | 1252/2000 [00:43<00:20, 36.09it/s]
 63%|██████▎   | 1256/2000 [00:43<00:20, 36.61it/s]
 63%|██████▎   | 1260/2000 [00:43<00:20, 36.50it/s]
 63%|██████▎   | 1264/2000 [00:43<00:19, 37.12it/s]
 63%|██████▎   | 1268/2000 [00:44<00:19, 37.27it/s]
 64%|██████▎   | 1272/2000 [00:44<00:19, 37.46it/s]
 64%|██████▍   | 1276/2000 [00:44<00:19, 37.19it/s]
 64%|██████▍   | 1280/2000 [00:44<00:19, 37.70it/s]
 64%|██████▍   | 1284/2000 [00:44<00:18, 38.12it/s]
 64%|██████▍   | 1288/2000 [00:44<00:18, 37.89it/s]
 65%|██████▍   | 1293/2000 [00:44<00:18, 39.17it/s]
 65%|██████▍   | 1297/2000 [00:44<00:18, 38.55it/s]
 65%|██████▌   | 1301/2000 [00:44<00:18, 37.39it/s]
 65%|██████▌   | 1305/2000 [00:45<00:18, 38.09it/s]
 66%|██████▌   | 1310/2000 [00:45<00:17, 38.57it/s]
 66%|██████▌   | 1314/2000 [00:45<00:17, 38.49it/s]
 66%|██████▌   | 1318/2000 [00:45<00:18, 37.84it/s]
 66%|██████▌   | 1322/2000 [00:45<00:18, 37.29it/s]
 66%|██████▋   | 1326/2000 [00:45<00:17, 37.89it/s]
 66%|██████▋   | 1330/2000 [00:45<00:17, 38.35it/s]
 67%|██████▋   | 1334/2000 [00:45<00:17, 38.06it/s]
 67%|██████▋   | 1338/2000 [00:45<00:17, 37.96it/s]
 67%|██████▋   | 1342/2000 [00:46<00:17, 37.37it/s]
 67%|██████▋   | 1346/2000 [00:46<00:17, 36.87it/s]
 68%|██████▊   | 1350/2000 [00:46<00:17, 37.21it/s]
 68%|██████▊   | 1354/2000 [00:46<00:17, 37.45it/s]
 68%|██████▊   | 1358/2000 [00:46<00:17, 37.24it/s]
 68%|██████▊   | 1362/2000 [00:46<00:17, 37.16it/s]
 68%|██████▊   | 1366/2000 [00:46<00:17, 36.15it/s]
 68%|██████▊   | 1370/2000 [00:46<00:17, 35.10it/s]
 69%|██████▊   | 1374/2000 [00:46<00:17, 35.77it/s]
 69%|██████▉   | 1378/2000 [00:47<00:16, 36.81it/s]
 69%|██████▉   | 1382/2000 [00:47<00:16, 36.46it/s]
 69%|██████▉   | 1386/2000 [00:47<00:16, 36.44it/s]
 70%|██████▉   | 1390/2000 [00:47<00:16, 36.41it/s]
 70%|██████▉   | 1394/2000 [00:47<00:16, 36.69it/s]
 70%|██████▉   | 1398/2000 [00:47<00:16, 37.02it/s]
 70%|███████   | 1403/2000 [00:47<00:15, 37.76it/s]
 70%|███████   | 1407/2000 [00:47<00:15, 37.56it/s]
 71%|███████   | 1411/2000 [00:47<00:15, 37.95it/s]
 71%|███████   | 1415/2000 [00:48<00:15, 37.34it/s]
 71%|███████   | 1419/2000 [00:48<00:15, 37.38it/s]
 71%|███████   | 1424/2000 [00:48<00:15, 38.19it/s]
 71%|███████▏  | 1428/2000 [00:48<00:15, 37.63it/s]
 72%|███████▏  | 1433/2000 [00:48<00:14, 38.96it/s]
 72%|███████▏  | 1437/2000 [00:48<00:14, 38.74it/s]
 72%|███████▏  | 1442/2000 [00:48<00:14, 39.45it/s]
 72%|███████▏  | 1446/2000 [00:48<00:14, 39.18it/s]
 72%|███████▎  | 1450/2000 [00:48<00:14, 39.17it/s]
 73%|███████▎  | 1454/2000 [00:49<00:14, 38.95it/s]
 73%|███████▎  | 1459/2000 [00:49<00:13, 39.22it/s]
 73%|███████▎  | 1463/2000 [00:49<00:13, 39.15it/s]
 73%|███████▎  | 1467/2000 [00:49<00:13, 39.21it/s]
 74%|███████▎  | 1472/2000 [00:49<00:13, 39.91it/s]
 74%|███████▍  | 1477/2000 [00:49<00:12, 40.39it/s]
 74%|███████▍  | 1482/2000 [00:49<00:12, 40.51it/s]
 74%|███████▍  | 1487/2000 [00:49<00:12, 40.44it/s]
 75%|███████▍  | 1492/2000 [00:49<00:12, 41.89it/s]
 75%|███████▍  | 1497/2000 [00:50<00:12, 41.63it/s]
 75%|███████▌  | 1502/2000 [00:50<00:12, 41.34it/s]
 75%|███████▌  | 1507/2000 [00:50<00:12, 40.91it/s]
 76%|███████▌  | 1512/2000 [00:50<00:11, 41.27it/s]
 76%|███████▌  | 1517/2000 [00:50<00:11, 40.27it/s]
 76%|███████▌  | 1522/2000 [00:50<00:11, 39.93it/s]
 76%|███████▋  | 1527/2000 [00:50<00:11, 40.08it/s]
 77%|███████▋  | 1532/2000 [00:50<00:11, 39.58it/s]
 77%|███████▋  | 1536/2000 [00:51<00:11, 39.35it/s]
 77%|███████▋  | 1541/2000 [00:51<00:11, 39.51it/s]
 77%|███████▋  | 1545/2000 [00:51<00:11, 38.34it/s]
 77%|███████▋  | 1549/2000 [00:51<00:11, 37.93it/s]
 78%|███████▊  | 1553/2000 [00:51<00:11, 38.27it/s]
 78%|███████▊  | 1557/2000 [00:51<00:11, 38.53it/s]
 78%|███████▊  | 1561/2000 [00:51<00:11, 38.63it/s]
 78%|███████▊  | 1565/2000 [00:51<00:11, 38.71it/s]
 78%|███████▊  | 1569/2000 [00:51<00:11, 38.99it/s]
 79%|███████▊  | 1573/2000 [00:52<00:11, 38.15it/s]
 79%|███████▉  | 1577/2000 [00:52<00:11, 37.97it/s]
 79%|███████▉  | 1581/2000 [00:52<00:10, 38.32it/s]
 79%|███████▉  | 1586/2000 [00:52<00:10, 39.96it/s]
 80%|███████▉  | 1591/2000 [00:52<00:09, 41.71it/s]
 80%|███████▉  | 1596/2000 [00:52<00:09, 42.95it/s]
 80%|████████  | 1601/2000 [00:52<00:09, 42.72it/s]
 80%|████████  | 1606/2000 [00:52<00:08, 43.85it/s]
 81%|████████  | 1611/2000 [00:52<00:09, 42.83it/s]
 81%|████████  | 1616/2000 [00:53<00:09, 42.49it/s]
 81%|████████  | 1621/2000 [00:53<00:09, 41.44it/s]
 81%|████████▏ | 1626/2000 [00:53<00:09, 40.97it/s]
 82%|████████▏ | 1631/2000 [00:53<00:09, 40.21it/s]
 82%|████████▏ | 1636/2000 [00:53<00:09, 40.27it/s]
 82%|████████▏ | 1641/2000 [00:53<00:08, 40.67it/s]
 82%|████████▏ | 1646/2000 [00:53<00:08, 41.15it/s]
 83%|████████▎ | 1651/2000 [00:53<00:08, 40.60it/s]
 83%|████████▎ | 1656/2000 [00:54<00:08, 40.92it/s]
 83%|████████▎ | 1661/2000 [00:54<00:08, 41.53it/s]
 83%|████████▎ | 1666/2000 [00:54<00:08, 41.71it/s]
 84%|████████▎ | 1671/2000 [00:54<00:07, 41.66it/s]
 84%|████████▍ | 1676/2000 [00:54<00:08, 40.26it/s]
 84%|████████▍ | 1681/2000 [00:54<00:07, 41.44it/s]
 84%|████████▍ | 1686/2000 [00:54<00:07, 41.34it/s]
 85%|████████▍ | 1691/2000 [00:54<00:07, 42.20it/s]
 85%|████████▍ | 1696/2000 [00:54<00:07, 41.30it/s]
 85%|████████▌ | 1701/2000 [00:55<00:07, 41.04it/s]
 85%|████████▌ | 1706/2000 [00:55<00:07, 41.55it/s]
 86%|████████▌ | 1711/2000 [00:55<00:07, 40.67it/s]
 86%|████████▌ | 1716/2000 [00:55<00:07, 40.53it/s]
 86%|████████▌ | 1721/2000 [00:55<00:06, 40.34it/s]
 86%|████████▋ | 1726/2000 [00:55<00:06, 40.02it/s]
 87%|████████▋ | 1731/2000 [00:55<00:07, 38.29it/s]
 87%|████████▋ | 1735/2000 [00:55<00:06, 38.39it/s]
 87%|████████▋ | 1740/2000 [00:56<00:06, 39.04it/s]
 87%|████████▋ | 1745/2000 [00:56<00:06, 39.43it/s]
 87%|████████▋ | 1749/2000 [00:56<00:06, 38.73it/s]
 88%|████████▊ | 1753/2000 [00:56<00:06, 38.80it/s]
 88%|████████▊ | 1758/2000 [00:56<00:06, 40.31it/s]
 88%|████████▊ | 1763/2000 [00:56<00:05, 41.38it/s]
 88%|████████▊ | 1768/2000 [00:56<00:05, 40.44it/s]
 89%|████████▊ | 1773/2000 [00:56<00:05, 41.37it/s]
 89%|████████▉ | 1778/2000 [00:57<00:05, 41.50it/s]
 89%|████████▉ | 1783/2000 [00:57<00:05, 40.65it/s]
 89%|████████▉ | 1788/2000 [00:57<00:05, 40.83it/s]
 90%|████████▉ | 1793/2000 [00:57<00:05, 41.29it/s]
 90%|████████▉ | 1798/2000 [00:57<00:04, 41.44it/s]
 90%|█████████ | 1803/2000 [00:57<00:04, 42.40it/s]
 90%|█████████ | 1808/2000 [00:57<00:04, 42.41it/s]
 91%|█████████ | 1813/2000 [00:57<00:04, 42.27it/s]
 91%|█████████ | 1818/2000 [00:57<00:04, 43.04it/s]
 91%|█████████ | 1823/2000 [00:58<00:04, 42.54it/s]
 91%|█████████▏| 1828/2000 [00:58<00:04, 42.14it/s]
 92%|█████████▏| 1833/2000 [00:58<00:03, 43.17it/s]
 92%|█████████▏| 1838/2000 [00:58<00:03, 42.49it/s]
 92%|█████████▏| 1843/2000 [00:58<00:03, 42.84it/s]
 92%|█████████▏| 1848/2000 [00:58<00:03, 43.32it/s]
 93%|█████████▎| 1853/2000 [00:58<00:03, 43.42it/s]
 93%|█████████▎| 1858/2000 [00:58<00:03, 44.38it/s]
 93%|█████████▎| 1863/2000 [00:59<00:03, 43.69it/s]
 93%|█████████▎| 1868/2000 [00:59<00:03, 42.36it/s]
 94%|█████████▎| 1873/2000 [00:59<00:03, 41.28it/s]
 94%|█████████▍| 1878/2000 [00:59<00:02, 41.73it/s]
 94%|█████████▍| 1883/2000 [00:59<00:02, 41.09it/s]
 94%|█████████▍| 1888/2000 [00:59<00:02, 40.17it/s]
 95%|█████████▍| 1893/2000 [00:59<00:02, 39.44it/s]
 95%|█████████▍| 1898/2000 [00:59<00:02, 40.58it/s]
 95%|█████████▌| 1903/2000 [01:00<00:02, 41.65it/s]
 95%|█████████▌| 1908/2000 [01:00<00:02, 42.12it/s]
 96%|█████████▌| 1913/2000 [01:00<00:02, 41.75it/s]
 96%|█████████▌| 1918/2000 [01:00<00:02, 40.33it/s]
 96%|█████████▌| 1923/2000 [01:00<00:01, 40.42it/s]
 96%|█████████▋| 1928/2000 [01:00<00:01, 40.74it/s]
 97%|█████████▋| 1933/2000 [01:00<00:01, 41.31it/s]
 97%|█████████▋| 1938/2000 [01:00<00:01, 41.50it/s]
 97%|█████████▋| 1943/2000 [01:00<00:01, 42.08it/s]
 97%|█████████▋| 1949/2000 [01:01<00:01, 44.47it/s]
 98%|█████████▊| 1954/2000 [01:01<00:01, 44.18it/s]
 98%|█████████▊| 1959/2000 [01:01<00:00, 43.45it/s]
 98%|█████████▊| 1965/2000 [01:01<00:00, 44.84it/s]
 98%|█████████▊| 1970/2000 [01:01<00:00, 44.43it/s]
 99%|█████████▉| 1975/2000 [01:01<00:00, 44.05it/s]
 99%|█████████▉| 1980/2000 [01:01<00:00, 43.77it/s]
 99%|█████████▉| 1985/2000 [01:01<00:00, 43.18it/s]
100%|█████████▉| 1990/2000 [01:02<00:00, 42.16it/s]
100%|█████████▉| 1995/2000 [01:02<00:00, 42.36it/s]
100%|██████████| 2000/2000 [01:02<00:00, 43.67it/s]
100%|██████████| 2000/2000 [01:02<00:00, 32.12it/s]
sampler = inv_result_fixed_d_sampler.sampler
az_idata = az.from_emcee(sampler, var_names=labels)
az_idata.get("posterior")
<xarray.Dataset> Size: 11MB
Dimensions:  (chain: 40, draw: 2000)
Coordinates:
  * chain    (chain) int64 320B 0 1 2 3 4 5 6 7 8 ... 31 32 33 34 35 36 37 38 39
  * draw     (draw) int64 16kB 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
Data variables: (12/17)
    t0       (chain, draw) float64 640kB 14.98 14.98 14.98 ... 19.64 19.64 19.64
    t1       (chain, draw) float64 640kB 15.58 15.58 15.58 ... 15.32 15.32 15.32
    t2       (chain, draw) float64 640kB 14.37 14.37 14.37 ... 7.025 7.025 7.025
    t3       (chain, draw) float64 640kB 15.57 15.57 15.57 ... 16.38 16.38 16.38
    t4       (chain, draw) float64 640kB 14.97 14.97 14.97 ... 17.39 17.39 17.39
    t5       (chain, draw) float64 640kB 14.84 14.84 14.84 ... 21.26 21.26 21.26
    ...       ...
    v3       (chain, draw) float64 640kB 4.327 4.327 4.327 ... 4.326 4.326 4.326
    v4       (chain, draw) float64 640kB 4.081 4.081 4.081 ... 4.392 4.392 4.392
    v5       (chain, draw) float64 640kB 4.884 4.884 4.884 ... 4.264 4.264 4.264
    v6       (chain, draw) float64 640kB 4.007 4.007 4.007 ... 4.299 4.299 4.299
    v7       (chain, draw) float64 640kB 4.69 4.69 4.69 ... 4.343 4.343 4.343
    v8       (chain, draw) float64 640kB 4.741 4.741 4.741 ... 4.562 4.562 4.562
Attributes:
    created_at:                 2024-04-19T05:03:16.440166+00:00
    arviz_version:              0.18.0
    inference_library:          emcee
    inference_library_version:  3.1.5


flat_samples = sampler.get_chain(discard=500, thin=500, flat=True)
rand_indices = np.random.randint(len(flat_samples), size=100)

_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 3, 3]})

ax1.set_ylim(170)

# plot samples and data predictions from samples
for idx in rand_indices:
    sample = flat_samples[idx]
    plot_model(sample, ax=ax1, alpha=0.2, lw=0.5, color="gray")
    plot_sw_data(forward_sw(sample, periods), periods,
                 ax=ax2, alpha=0.2, lw=0.5, color="gray")
    plot_rf_data(forward_rf(sample), rf_times,
                 ax=ax3, alpha=0.2, lw=0.5, color="gray")

# add labels to samples
sample_0 = flat_samples[rand_indices[0]]
plot_model(sample_0, ax=ax1, alpha=0.5, lw=0.5, color="gray", label="samples")
plot_sw_data(forward_sw(sample_0, periods), periods, ax=ax2,
             alpha=0.5, lw=0.5, color="gray", label="rayleigh_dpred from samples")
plot_rf_data(forward_rf(sample_0), rf_times, ax=ax3,
             alpha=0.5, lw=0.5, color="gray", label="rf_dpred from samples")

# plot true model and data observations
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="r", s=4,
          label="rayleigh_dobs")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="r", s=2,
          label="rf_dobs")

# plot damped optimisation result
plot_model(inv_res_optimiser.model, ax=ax1, alpha=1, color="green",
           label="damped solution")
plot_sw_data(forward_sw(inv_res_optimiser_no_reg.model, periods),
             periods, ax=ax2, color="green",
             label="rayleigh_dpred from damped solution")
plot_rf_data(forward_rf(inv_res_optimiser.model), rf_times,
             ax=ax3, color="green",
             label="rf_dpred from damped solution")

# plot initial model for dampied optimsiation
plot_model(init_model, ax=ax1, alpha=1, color="purple",
           label="initial model for damped solution")
plot_sw_data(forward_sw(init_model, periods), periods, ax=ax2, color="purple",
             label="rayleigh_dpred from initial model for damped solution")
plot_rf_data(forward_rf(init_model), rf_times, ax=ax3, color="purple",
             label="rf_dpred from initial model for damped solution")

ax1.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax2.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax3.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax1.grid()
ax2.grid()
ax3.grid()

plt.tight_layout()
model, surface wave data, receiver function data
axes = az.plot_trace(
    az_idata,
    backend_kwargs={"constrained_layout":True},
    figsize=(10,20)
)

for i, axes_pair in enumerate(axes):
    ax1 = axes_pair[0]
    ax2 = axes_pair[1]
    ax1.axvline(true_model[i], linestyle='dotted', color='red')
    ax1.set_xlabel("parameter value")
    ax1.set_ylabel("density value")
    ax2.set_xlabel("number of iterations")
    ax2.set_ylabel("parameter value")
t0, t0, t1, t1, t2, t2, t3, t3, t4, t4, t5, t5, t6, t6, t7, t7, v0, v0, v1, v1, v2, v2, v3, v3, v4, v4, v5, v5, v6, v6, v7, v7, v8, v8

Trans-dimensional sampling#

Prepare utilities for trans-d sampling

def fwd_rayleigh_for_bayesbay(state):
    vs = state["voronoi"]["vs"]
    voronoi_sites = state["voronoi"]["discretization"]
    depths = (voronoi_sites[:-1] + voronoi_sites[1:]) / 2
    thicknesses = depths - np.insert(depths[:-1], 0, 0)
    model = form_layercake_model(thicknesses, vs)
    return forward_sw(model, periods)

def fwd_rf_for_bayesbay(state):
    vs = state["voronoi"]["vs"]
    voronoi_sites = state["voronoi"]["discretization"]
    depths = (voronoi_sites[:-1] + voronoi_sites[1:]) / 2
    thicknesses = depths - np.insert(depths[:-1], 0, 0)
    model = form_layercake_model(thicknesses, vs)
    return forward_rf(model)
targets = [
    bayesbay.Target("rayleigh", rayleigh_dobs, covariance_mat_inv=1 / RAYLEIGH_STD**2),
    bayesbay.Target("rf", rf_dobs, covariance_mat_inv=1 / RF_STD**2),
]
forward_funcs = [fwd_rayleigh_for_bayesbay, fwd_rf_for_bayesbay]

my_log_likelihood = bayesbay.LogLikelihood(targets, forward_funcs)
param_vs = bayesbay.prior.UniformPrior(
    name="vs",
    vmin=[2.7, 3.2, 3.75],
    vmax=[4, 4.75, 5],
    position=[0, 40, 80],
    perturb_std=0.15
)

def param_vs_initialize(param, positions):
    vmin, vmax = param.get_vmin_vmax(positions)
    sorted_vals = np.sort(np.random.uniform(vmin, vmax, positions.size))
    for i in range (len(sorted_vals)):
        val = sorted_vals[i]
        vmin_i = vmin if np.isscalar(vmin) else vmin[i]
        vmax_i = vmax if np.isscalar(vmax) else vmax[i]
        if val < vmin_i or val > vmax_i:
            if val > vmax_i: sorted_vals[i] = vmax_i
            if val < vmin_i: sorted_vals[i] = vmin_i
    return sorted_vals

param_vs.set_custom_initialize(param_vs_initialize)
parameterization = bayesbay.parameterization.Parameterization(
    bayesbay.discretization.Voronoi1D(
        name="voronoi",
        vmin=0,
        vmax=150,
        perturb_std=10,
        n_dimensions=None,
        n_dimensions_min=4,
        n_dimensions_max=15,
        parameters=[param_vs],
    )
)
my_perturbation_funcs = parameterization.perturbation_functions
n_chains=12
walkers_start = []
for i in range(n_chains):
    walkers_start.append(parameterization.initialize())

Define ``InversionOptions``

inv_options_trans_d = cofi.InversionOptions()
inv_options_trans_d.set_tool("bayesbay")
inv_options_trans_d.set_params(
    walkers_starting_states=walkers_start,
    perturbation_funcs=my_perturbation_funcs,
    log_like_ratio_func=my_log_likelihood,
    n_chains=n_chains,
    n_iterations=2_000,
    burnin_iterations=1_000,
    verbose=False,
    save_every=200,
)

Define ``Inversion`` and run

inversion_trans_d_sampler = cofi.Inversion(joint_problem, inv_options_trans_d)
inv_res_trans_d_sampler = inversion_trans_d_sampler.run()
saved_states = inv_res_trans_d_sampler.models
samples_voronoi = saved_states["voronoi.discretization"]
samples_vs = saved_states["voronoi.vs"]
interp_depths = np.arange(150, dtype=float)
rand_indices = np.random.randint(len(samples_voronoi), size=100)

_, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4), gridspec_kw={'width_ratios': [1, 3, 3]})

ax1.set_ylim(170)

# plot samples and data predictions from samples
for idx in rand_indices:
    sample_voronoi = form_voronoi_model(samples_voronoi[idx], samples_vs[idx])
    sample = voronoi_to_layercake(sample_voronoi)
    plot_model(sample, ax=ax1, alpha=0.2, lw=0.5, color="gray")
    plot_sw_data(forward_sw(sample, periods), periods,
                 ax=ax2, alpha=0.2, lw=0.5, color="gray")
    plot_rf_data(forward_rf(sample), rf_times,
                 ax=ax3, alpha=0.2, lw=0.5, color="gray")

# add labels to samples
sample_0_voronoi = form_voronoi_model(samples_voronoi[0], samples_vs[0])
sample_0 = voronoi_to_layercake(sample_0_voronoi)
plot_model(sample_0, ax=ax1, alpha=0.2, lw=0.5, color="gray", label="samples")
plot_sw_data(forward_sw(sample_0, periods), periods, ax=ax2,
             alpha=0.2, lw=0.5, color="gray", label="rayleigh_dpred from samples")
plot_rf_data(forward_rf(sample_0), rf_times, ax=ax3,
             alpha=0.2, lw=0.5, color="gray", label="rf_dpred from samples")

# plot true model and data observations
plot_model(true_model, ax=ax1, alpha=1, color="r", label="true model")
plot_sw_data(rayleigh_dobs, periods, ax=ax2, scatter=True, color="r", s=4,
          label="rayleigh_dobs")
plot_rf_data(rf_dobs, rf_times, ax=ax3, scatter=True, color="r", s=2,
          label="rf_dobs")

# plot damped optimisation result
plot_model(inv_res_optimiser.model, ax=ax1, alpha=1, color="green",
           label="damped solution")
plot_sw_data(forward_sw(inv_res_optimiser_no_reg.model, periods),
             periods, ax=ax2, color="green",
             label="rayleigh_dpred from damped solution")
plot_rf_data(forward_rf(inv_res_optimiser.model), rf_times,
             ax=ax3, color="green",
             label="rf_dpred from damped solution")

# plot initial model for dampied optimsiation
plot_model(init_model, ax=ax1, alpha=1, color="purple",
           label="initial model for damped solution")
plot_sw_data(forward_sw(init_model, periods), periods, ax=ax2, color="purple",
             label="rayleigh_dpred from initial model for damped solution")
plot_rf_data(forward_rf(init_model), rf_times, ax=ax3, color="purple",
             label="rf_dpred from initial model for damped solution")

ax1.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax2.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax3.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18))
ax1.grid()
ax2.grid()
ax3.grid()

plt.tight_layout()
model, surface wave data, receiver function data

Watermark#

watermark_list = ["cofi", "espresso", "numpy", "matplotlib", "scipy", "bayesbay"]
for pkg in watermark_list:
    pkg_var = __import__(pkg)
    print(pkg, getattr(pkg_var, "__version__"))
cofi 0.2.8
espresso 0.3.13
numpy 1.26.4
matplotlib 3.8.3
scipy 1.12.0
bayesbay 0.3.0

sphinx_gallery_thumbnail_number = -1

Total running time of the script: (1 minutes 29.331 seconds)

Gallery generated by Sphinx-Gallery