1D MT parametric inversion

import numpy as np
from simpeg.electromagnetics import natural_source as nsem
from simpeg import maps
import matplotlib.pyplot as plt
import matplotlib
from simpeg.utils import plot_1d_layer_model
from discretize import TensorMesh
from simpeg import (
    maps,
    data,
    data_misfit,
    regularization,
    optimization,
    inverse_problem,
    inversion,
    directives,
    utils,
)

from mtpy import MTCollection
mc = MTCollection()
mc.open_collection("../../data/transfer_functions/yellowstone_mt_collection.h5")
from ipywidgets import widgets, interact
station_names = mc.dataframe.station.values
def foo(name):
    tf = mc.get_tf(name)
    tf.plot_mt_response()

Q = interact(foo, name=widgets.Select(options=station_names, value='YNP05S'))

name = Q.widget.kwargs['name']
tf = mc.get_tf(name)
fig, ax = plt.subplots(1,1, figsize=(5,5))
mc.dataframe.plot(x='longitude', y='latitude', marker='.', linestyle='None', ax=ax)
ax.plot(tf.longitude, tf.latitude, 'ro')
ax.set_title(name)

frequencies = 1./tf.period
n_layer = 5
wire_map = maps.Wires(("sigma", n_layer), ("t", n_layer - 1))
sigma_map = maps.ExpMap(nP=n_layer) * wire_map.sigma
layer_map = maps.ExpMap(nP=n_layer - 1) * wire_map.t


model_mapping = maps.IdentityMap(nP=n_layer)

receivers_list = [
    nsem.receivers.PointNaturalSource(component="app_res"),
    nsem.receivers.PointNaturalSource(component="phase"),
]

source_list = []
for freq in frequencies:
    source_list.append(nsem.sources.Planewave(receivers_list, freq))

survey = nsem.survey.Survey(source_list)

simulation = nsem.simulation_1d.Simulation1DRecursive(
    survey=survey, 
    sigmaMap=sigma_map,
    thicknessesMap=layer_map,
)

matplotlib.rcParams['font.size'] = 14
fig, axs = plt.subplots(2,1)
ax1, ax2 = axs
ax1.hist(tf.Z.res_error_det/tf.Z.res_det)
ax2.hist(tf.Z.phase_error_det)

app_rho = tf.Z.res_det.copy()
phase = tf.Z.phase_det.copy()
std = np.c_[abs(app_rho)*0.03, np.ones(len(phase))*2].flatten()
noise = np.c_[np.random.randn(app_rho.size)*0.03*abs(app_rho), np.random.randn(app_rho.size)*2].flatten()
dobs = np.c_[app_rho, phase].flatten()
data_object = data.Data(survey, dobs=dobs, standard_deviation=std)

starting_model = np.log(np.r_[np.ones(n_layer)*1./100, np.ones(n_layer-1)*1000])

# 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(simulation=simulation, data=data_object)
mesh = TensorMesh([n_layer])
# Define the regularization on the parameters related to resistivity
mesh_sigma = TensorMesh([mesh.h[0].size])
reg_sigma = regularization.WeightedLeastSquares(mesh_sigma, alpha_s=0.01, alpha_x=0, mapping=wire_map.sigma)

# Define the regularization on the parameters related to layer thickness
mesh_t = TensorMesh([mesh.h[0].size - 1])
reg_t = regularization.WeightedLeastSquares(mesh_t, alpha_s=0.01, alpha_x=0, mapping=wire_map.t)

# Combine to make regularization for the inversion problem
reg = reg_sigma + reg_t

# Define how the optimization problem is solved. Here we will use an inexact
# Gauss-Newton approach that employs the conjugate gradient solver.
opt = optimization.InexactGaussNewton(maxIter=30, maxIterCG=50, maxIterLS=50)

# Define the inverse problem
inv_prob = inverse_problem.BaseInvProblem(dmis, reg, opt)

#######################################################################
# Define Inversion Directives
# ---------------------------
#
# Here we define any directives that are carried out during the inversion. This
# includes the cooling schedule for the trade-off parameter (beta), stopping
# criteria for the inversion and saving inversion results at each iteration.
#

# Apply and update sensitivity weighting as the model updates
update_sensitivity_weights = directives.UpdateSensitivityWeights()

# Defining a starting value for the trade-off parameter (beta) between the data
# misfit and the regularization.
starting_beta = directives.BetaEstimate_ByEig(beta0_ratio=1e2)

# Set the rate of reduction in trade-off parameter (beta) each time the
# the inverse problem is solved. And set the number of Gauss-Newton iterations
# for each trade-off paramter value.
beta_schedule = directives.BetaSchedule(coolingFactor=5.0, coolingRate=3)

# Options for outputting recovered models and predicted data for each beta.
save_iteration = directives.SaveOutputEveryIteration(save_txt=False)
precond = directives.UpdatePreconditioner()
# Setting a stopping criteria for the inversion.
target_misfit = directives.TargetMisfit(chifact=1)

# The directives are defined in a list
directives_list = [
    update_sensitivity_weights,
    precond,
    starting_beta,
    beta_schedule,
    target_misfit,
]

#####################################################################
# Running the Inversion
# ---------------------
#
# To define the inversion object, we need to define the inversion problem and
# the set of directives. We can then run the inversion.
#

# Here we combine the inverse problem and the set of directives
inv = inversion.BaseInversion(inv_prob, directiveList=directives_list)

# Run the inversion
recovered_model = inv.run(starting_model)

import matplotlib.gridspec as gridspec
matplotlib.rcParams['font.size'] = 10
fig = plt.figure(figsize=(16*0.5, 5*0.5))
gs = gridspec.GridSpec(1, 5, figure=fig)

ax0 = fig.add_subplot(gs[0, 0])
x_min = 1
x_max = 1000

plot_1d_layer_model(
    (layer_map * recovered_model)[::-1],
    (1./(sigma_map * recovered_model))[::-1],
    ax=ax0,
    color="k"
)
ax0.set_xlabel("Resistivity ($\Omega$m)")
ax0.set_xlim(1, 1e4)

ax = fig.add_subplot(gs[0, 2:])
ax.loglog(1./frequencies, dobs.reshape((len(frequencies), 2))[:,0], 'x', color='C0', label='Obs.-Amp.')
ax.loglog(1./frequencies, inv_prob.dpred.reshape((len(frequencies), 2))[:,0], color='C0', label='Pred.-Amp.')
ax_1 = ax.twinx()
ax_1.plot(1./frequencies, dobs.reshape((len(frequencies), 2))[:,1], 'x', color='C1', label='Obs.-Pha.')
ax_1.plot(1./frequencies, inv_prob.dpred.reshape((len(frequencies), 2))[:,1], color='C1', label='Pred.-Pha.')
ax.set_xlabel("Period (s)")
ax.grid(True, which='both', alpha=0.5)
ax.set_ylabel("Apparent resistivity ($\Omega$m)")
ax_1.set_ylabel("Phase ($\degree$)")
ax.legend(bbox_to_anchor=(1.1,1))
ax.legend(loc=3)
ax_1.legend(loc=4)
ax.set_ylim(1, 10000)
ax_1.set_ylim(0, 90)    
ax0.set_xlim(1, 10000)
plt.show()