2D MT inversion

import matplotlib.pyplot as plt
import scipy.sparse as sp
import numpy as np
import unittest
from scipy.constants import mu_0
import discretize
import matplotlib.patheffects as pe
from simpeg.electromagnetics import natural_source as nsem
from simpeg.electromagnetics.static import utils as sutils
from simpeg import (
    maps, utils, optimization, objective_function, inversion, inverse_problem, directives,
    data_misfit, regularization, data
)
from discretize import TensorMesh
from pymatsolver import Pardiso
from scipy.spatial import cKDTree
from scipy.stats import norm
# from dask.distributed import Client, LocalCluster
import dill
from geoana.em.fdem import skin_depth
import discretize.utils as dis_utils
import warnings
warnings.filterwarnings("ignore")

def generate_2d_mesh_for_mt(
    rx_locs, frequencies, sigma_background, 
    z_factor_max=5, 
    z_factor_min=5,
    pfz_down = 1.2,
    pfz_up = 1.5,
    npadz_up = 5,
    x_factor_max = 2,
    spacing_factor = 4,
    pfx = 1.5,
    n_max=1000
):
    # Setting the cells in depth dimension
    f_min = frequencies.min()
    f_max = frequencies.max()
    dz_min = np.round(skin_depth(f_max, sigma_background) / z_factor_max) 
    lz = skin_depth(sigma_background, f_min) * z_factor_max
    # Setting the domain length in z-direction
    for nz_down in range(n_max):
        hz_down = dz_min*pfz_down**np.arange(nz_down)[::-1]
        if hz_down.sum()>lz:
            break
    hz_up = [(dz_min, npadz_up, pfz_up)]
    hz_up = dis_utils.unpack_widths(hz_up)
    hz = np.r_[hz_down, hz_up]   
    # Setting the cells in lateral dimension
    d_station = np.diff(rx_locs[:,0]).min()
    dx_min = np.round(d_station/spacing_factor)
    lx = rx_locs[:,0].max() - rx_locs[:,0].min()
    ncx = int(lx / dx_min)
    lx_pad = skin_depth(sigma_background, f_min) * x_factor_max
    for npadx in range(n_max):
        hx_pad = dis_utils.unpack_widths([(dx_min, npadx, -pfx)])
        if hx_pad.sum()>lx_pad:
            break
    hx = [(dx_min, npadx, -pfx), (dx_min, ncx), (dx_min, npadx, pfx)]  
    
    mesh = discretize.TensorMesh([hx, hz])
    mesh.origin = np.r_[-mesh.h[0][:npadx].sum()+rx_locs[:,0].min(), -hz_down.sum()]   
    print (mesh)
    return mesh

input_data = dill.load(open("./synthetic_2d.pik", "rb"))
nskip_freq = 5
nskip_rx = 1
# a decent setup
# nskip_freq = 7
# nskip_rx = 3
relative_error =  0.05
floor_error = 2
alpha_s = 1e-10
alpha_y = 1.
alpha_z = 0.5
rho_0 = 100.
maxIter = 30
use_irls = False
p_s=0
p_y=0
p_z=0

rx_locs = input_data['rx_locs'][::nskip_rx,:]
frequencies = input_data['frequencies'][::nskip_freq]
app_rho_te = input_data['app_rho_te'][::nskip_freq,::nskip_rx]
app_rho_tm = input_data['app_rho_tm'][::nskip_freq,::nskip_rx]
phase_te = input_data['phase_te'][::nskip_freq,::nskip_rx]
phase_tm = input_data['phase_tm'][::nskip_freq,::nskip_rx]


app_rho_err_tm = abs(app_rho_tm) * relative_error
app_rho_err_te = abs (app_rho_te) * relative_error
phase_err_tm = np.ones_like(phase_tm) * floor_error
phase_err_te = np.ones_like(phase_te) * floor_error

mesh = generate_2d_mesh_for_mt(rx_locs, frequencies, 1e-2)
ind_active = mesh.cell_centers[:,1]<0.
print ("Rx")
print (rx_locs)
print ("Freqs")
print (frequencies)

# simulation class for TM mode
rx_list_tm = [
    nsem.receivers.PointNaturalSource(
        rx_locs, orientation="xy", component="apparent_resistivity"
    ),
    nsem.receivers.PointNaturalSource(
        rx_locs, orientation="xy", component="phase"
    ),
]
src_list_tm = [nsem.sources.Planewave(rx_list_tm, frequency=f) for f in frequencies]
survey_tm = nsem.Survey(src_list_tm)

act_map = maps.InjectActiveCells(mesh, ind_active, np.log(1e-8))
exp_map = maps.ExpMap(mesh=mesh)
sigma_map = exp_map * act_map

sim_tm= nsem.simulation.Simulation2DElectricField(
    mesh,
    survey=survey_tm,
    sigmaMap=sigma_map,
    solver=Pardiso,
)

# simulation class for TE mode
rx_list_te = [
    nsem.receivers.PointNaturalSource(
        rx_locs, orientation="yx", component="apparent_resistivity"
    ),
    nsem.receivers.PointNaturalSource(
        rx_locs, orientation="yx", component="phase"
    ),
]
src_list_te = [nsem.sources.Planewave(rx_list_te, frequency=f) for f in frequencies]
survey_te = nsem.Survey(src_list_te)

sim_te = nsem.simulation.Simulation2DMagneticField(
    mesh,
    survey=survey_te,
    sigmaMap=sigma_map,
    solver=Pardiso,
)

dobs_te = np.hstack((app_rho_te, phase_te)).flatten()
dobs_tm = np.hstack((app_rho_tm, phase_tm)).flatten()

std_te = np.hstack((app_rho_err_te, phase_err_te)).flatten()
std_tm = np.hstack((app_rho_err_tm, phase_err_tm)).flatten()

m0 = np.ones(ind_active.sum()) * np.log(1./rho_0)

%%time

te_data_object = data.Data(survey_te, dobs=dobs_te, standard_deviation=std_te) 
tm_data_object = data.Data(survey_tm, dobs=dobs_tm, standard_deviation=std_tm) 
dmis_te = data_misfit.L2DataMisfit(data=te_data_object, simulation=sim_te)
dmis_tm = data_misfit.L2DataMisfit(data=tm_data_object, simulation=sim_tm)
dmis = dmis_te + dmis_tm

# Define the regularization (model objective function)
reg = regularization.Sparse(
    mesh,
    active_cells=ind_active,
    reference_model=m0,
    alpha_s=alpha_s,
    alpha_x=alpha_y,
    alpha_y=alpha_z,
    mapping=maps.IdentityMap(nP=int(ind_active.sum()))
)

# Define how the optimization problem is solved. Here we will use an
# Inexact Gauss Newton approach.
opt = optimization.InexactGaussNewton(maxIter=maxIter, maxIterCG=30, tolX=1e-30)

# Here we define the inverse problem that is to be solved
inv_prob = inverse_problem.BaseInvProblem(dmis, reg, opt)

#######################################################################
# Define MT 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.
#

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

# 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=2, coolingRate=1)

save_dictionary = directives.SaveOutputDictEveryIteration()
save_dictionary.outDict = {}

# Setting a stopping criteria for the inversion.
target_misfit = directives.TargetMisfit(chifact=1e0)

if use_irls:
    reg.norms = np.c_[p_s, p_y, p_z]
    # Reach target misfit for L2 solution, then use IRLS until model stops changing.
    IRLS = directives.Update_IRLS(max_irls_iterations=40, minGNiter=1, f_min_change=1e-5)

    # The directives are defined as a list.
    directives_list = [
        IRLS,
        starting_beta,
        save_dictionary,
    ]
else:
    directives_list = [
        starting_beta,
        beta_schedule,
        save_dictionary,
        target_misfit,
    ]

    
#####################################################################
# Running the MT 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
mt_inversion = inversion.BaseInversion(inv_prob, directiveList=directives_list)

# Run inversion
recovered_conductivity_model = mt_inversion.run(m0)

survey_te.nD + survey_tm.nD

output_dict = save_dictionary.outDict
target_misfit = (dobs_te.size+dobs_tm.size)
iterations = list(output_dict.keys())
n_iteration = len(iterations)
phi_ds = np.zeros(n_iteration)
phi_ms = np.zeros(n_iteration)
betas = np.zeros(n_iteration)
for ii, iteration in enumerate(iterations):
    phi_ds[ii] = output_dict[iteration]['phi_d']
    phi_ms[ii] = output_dict[iteration]['phi_m']
    betas[ii] = output_dict[iteration]['beta']

import matplotlib
matplotlib.rcParams['font.size'] = 14
def tikhonov_curve(iteration, scale='log'):
    fig, ax = plt.subplots(1,1, figsize=(5,5))
    ax.plot(phi_ms, phi_ds)
    ax.plot(phi_ms[iteration-1], phi_ds[iteration-1], 'ro')
    ax.set_xlabel("$\phi_m$")
    ax.set_ylabel("$\phi_d$")
    if scale == 'log':
        ax.set_xscale('log')
        ax.set_yscale('log')
    xlim = ax.get_xlim()
    ax.plot(xlim, np.ones(2) * target_misfit, '--')
    ax.set_title("Iteration={:d}, Beta = {:.1e}".format(iteration, betas[iteration-1]))
    ax.set_xlim(xlim)
    plt.show()

from ipywidgets import interact, widgets
Q_iter = interact(
    tikhonov_curve, 
    iteration=widgets.IntSlider(min=1, max=int(n_iteration), value=n_iteration),
    scale=widgets.RadioButtons(options=['linear', 'log'])
)

from matplotlib.colors import LogNorm
mesh_true = input_data['mesh']
sigma_true = input_data['sigma']
iteration = Q_iter.widget.kwargs['iteration']
m = output_dict[iteration]['m']

def foo_model(iteration):
    fig, axs = plt.subplots(2,1, figsize=(10, 6))
    ax, ax1 = axs
    m = output_dict[iteration]['m']
    sigma = np.ones(mesh.nC) * 1e-8
    sigma[ind_active] = np.exp(m)
    sigma_min = 1e-3
    sigma_max = 10
    grid= False
    out = mesh.plot_image(
        sigma, grid=grid, ax=ax, pcolor_opts={'norm':LogNorm(vmin=sigma_min, vmax=sigma_max), 'cmap':'turbo'},
        range_x=(-60000, 60000), range_y=(-30000, 0)
    )
    cb = plt.colorbar(out[0], fraction=0.01, ax=ax)
    cb.set_label("Conductivity (S/m)")
    ax.set_aspect(1)
    ax.set_xlabel("Easting (m)")
    ax.set_ylabel("Elevation (m)")
    ax.plot(rx_locs[:,0], rx_locs[:,1], 'ro')
    out = mesh_true.plot_image(
        sigma_true, grid=grid, ax=ax1, pcolor_opts={'norm':LogNorm(vmin=sigma_min, vmax=sigma_max), 'cmap':'turbo'},
        range_x=(-60000, 60000), range_y=(-30000, 0)
    )
    cb = plt.colorbar(out[0], fraction=0.01, ax=ax1)
    cb.set_label("Conductivity (S/m)")
    ax1.set_aspect(1)
    ax1.set_xlabel("Easting (m)")
    ax1.set_ylabel("Elevation (m)")

Q_misfit = interact(
    foo_model, 
    iteration=widgets.IntSlider(min=1, max=n_iteration, value=iteration),
)    

pred_te = sim_te.dpred(m)
pred_tm = sim_tm.dpred(m)

n_freq = len(frequencies)
n_rx = rx_locs.shape[0]

PRED_te = pred_te.reshape((n_freq, 2, n_rx))
PRED_tm = pred_tm.reshape((n_freq, 2, n_rx))

rho_app_te_pred = PRED_te[:,0,:]
rho_app_tm_pred = PRED_tm[:,0,:]

phase_te_pred = PRED_te[:,1,:]
phase_tm_pred = PRED_tm[:,1,:]

DOBS_te = dobs_te.reshape((n_freq, 2, n_rx))
DOBS_tm = dobs_tm.reshape((n_freq, 2, n_rx))

rho_app_te_dobs = DOBS_te[:,0,:]
rho_app_tm_dobs = DOBS_tm[:,0,:]

phase_te_dobs = DOBS_te[:,1,:]
phase_tm_dobs = DOBS_tm[:,1,:]

def foo_pred(irx):
    import matplotlib
    matplotlib.rcParams['font.size'] = 10

    fig, axs = plt.subplots(4,1, figsize=(8, 10))
    ax1, ax2, ax3, ax4 = axs
    ax1.loglog(1./frequencies, rho_app_te_pred[:, irx], color='C0')
    ax1.loglog(1./frequencies, rho_app_te_dobs[:, irx], 'x', color='C0')
    ax2.loglog(1./frequencies, rho_app_tm_pred[:, irx], color='C1')
    ax2.loglog(1./frequencies, rho_app_tm_dobs[:, irx], 'x', color='C1')
    for ax in axs[:2]:
        ax.set_ylim(1, 1000)
        ax.set_ylabel("App. Res. (Ohm-m)")
        ax.grid(which='both', alpha=0.3)
    ax3.semilogx(1./frequencies, phase_te_pred[:, irx], color='C0')
    ax3.semilogx(1./frequencies, phase_te_dobs[:, irx], 'x', color='C0')
    ax4.semilogx(1./frequencies, phase_tm_pred[:, irx]+180, color='C1')
    ax4.semilogx(1./frequencies, phase_tm_dobs[:, irx]+180, 'x', color='C1')
    for ax in axs[2:]:
        ax.set_ylim(0, 90)
        ax.set_ylabel("Phase (degree)")
        ax.grid(which='both', alpha=0.3)
    ax4.set_xlabel("Period (s)")

Q_misfit = interact(
    foo_pred, 
    irx=widgets.IntSlider(min=0, max=int(n_rx)-1, value=0),
)