EM Depth of Investigation

from emrecharge.app import AppState
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt
mpl.rcParams.update({'font.size':12})
from SimPEG import maps
import SimPEG.electromagnetics.time_domain as tdem
from SimPEG import utils
from ipywidgets import interactive, widgets
from traitlets import HasTraits, Int, Unicode, List, Float, Bool, TraitType, Dict
from IPython.display import display, HTML, JSON
from ipywidgets import (
    Label, VBox, HBox,
    IntSlider, RadioButtons, Checkbox, Output, SelectionSlider, Dropdown, Button,
    AppLayout, Layout, interactive_output, FloatRangeSlider, BoundedIntText, BoundedFloatText
)

from IPython.display import display, clear_output
from traitlets import Int, Float
from scipy.stats import percentileofscore

from SimPEG.electromagnetics.utils import em1d_utils
from discretize.utils import volume_average

import dill
import matplotlib.pyplot as plt

# inverse problem
from discretize import TensorMesh

from SimPEG.utils import mkvc, plot_1d_layer_model
from SimPEG import (
    maps,
    data,
    data_misfit,
    inverse_problem,
    regularization,
    optimization,
    directives,
    inversion,
    utils,
)

import json

fname_aem = './aem_system.json'
fname_ttem = './ttem_system.json'

def get_vertical_discretization(n_layer, minimum_dz, geomtric_factor):
    """
    Creates a list of vertical discretizations generate from a geometric series.

    >>> minimum_dz * geomtric_factor ** np.arange(n_layer)

    Parameters
    ----------
    n_layer : int
        The number of discretizations
    minimum_dz : float
        The smallest discretization
    geometric_factor : float
        The expansion factor of the discretizations

    Returns
    (n_layer) : numpy.ndarray
        The cell widths
    """
    hz = minimum_dz * (geomtric_factor) ** np.arange(n_layer)

    return hz

class EMSIMULATION:
    """EMSIMULATION"""
    def __init__(self):        
        self.n_layer = 3
        self._initialize()
        self.dobs = None        
        # Widget variables
        
    def _initialize(self):         
        
        # AEM 
        with open(fname_aem, 'r') as infile:
            system_aem = json.load(infile)
        area = system_aem['tx_area']

        i_start_hm = 10
        i_start_lm = 10
        i_end_hm = 10 + 27
        i_end_lm = 10 + 20 -2 

        waveform_hm = np.array(system_aem['hm']['waveform'])
        waveform_lm = np.array(system_aem['lm']['waveform'])
        time_input_currents_hm = waveform_hm[:,0] 
        input_currents_hm = waveform_hm[:,1]
        time_input_currents_lm = waveform_lm[:,0] 
        input_currents_lm = waveform_lm[:,1]

        time_gates = np.array(system_aem['time_gates'])
        GateTimeShift=system_aem['lm']['GateTimeShift']
        MeaTimeDelay=system_aem['lm']['MeaTimeDelay']
        NoGates=int(system_aem['lm']['NoGates'])
        t0_lm = waveform_lm[:,0].max()
        times_lm = (time_gates[:NoGates,0] + GateTimeShift + MeaTimeDelay)[i_start_lm:i_end_lm]

        GateTimeShift=system_aem['hm']['GateTimeShift']
        MeaTimeDelay=system_aem['hm']['MeaTimeDelay']
        NoGates=int(system_aem['hm']['NoGates'])
        t0_hm = waveform_hm[:,0].max()
        times_hm = (time_gates[:NoGates,0] + GateTimeShift + MeaTimeDelay)[i_start_hm:i_end_hm]    
        
        receiver_location = np.array([13.25, 0., 30.])
        receiver_orientation = "z"                    # "x", "y" or "z"
        
        source_location = np.array([0., 0., 30.])  
        source_orientation = 'z'
        source_list = []
        waveform_lm = tdem.sources.PiecewiseLinearWaveform(times=time_input_currents_lm, currents=input_currents_lm)
        waveform_hm = tdem.sources.PiecewiseLinearWaveform(times=time_input_currents_hm, currents=input_currents_hm)
        # Define receivers at each location.
        dbzdt_receiver_lm = tdem.receivers.PointMagneticFluxTimeDerivative(
            receiver_location, times_lm, receiver_orientation,
        )
        dbzdt_receiver_hm = tdem.receivers.PointMagneticFluxTimeDerivative(
            receiver_location, times_hm, receiver_orientation,
        )

        # Must define the transmitter properties and associated receivers
        src_0 = tdem.sources.MagDipole(
            [dbzdt_receiver_lm],
            location=source_location,
            waveform=waveform_lm,
            orientation=source_orientation,
        )
        src_1 = tdem.sources.MagDipole(
            [dbzdt_receiver_hm],
            location=source_location,
            waveform=waveform_hm,
            orientation=source_orientation,
        )
        # this needs to be fixed later
        source_list.append(src_0)
        source_list.append(src_1)

        survey = tdem.Survey(source_list)
        wire_map = maps.Wires(("sigma", self.n_layer), ("thickness", self.n_layer-1))
        
        self._simulation_aem = tdem.Simulation1DLayered(
            survey=survey, thicknessesMap=wire_map.thickness, sigmaMap=wire_map.sigma,
        )

        # tTEM 
        with open(fname_ttem, 'r') as infile:
            system_ttem = json.load(infile)
        area = system_ttem['tx_area']
        i_start_hm = int(system_ttem['hm']['RemoveInitialGates'])
        i_start_lm = int(system_ttem['lm']['RemoveInitialGates'])
        i_end_lm = int(system_ttem['lm']['RemoveGatesFrom'])
        waveform_hm = np.array(system_ttem['hm']['waveform'])
        waveform_lm = np.array(system_ttem['lm']['waveform'])
        time_input_currents_hm = waveform_hm[:,0] - waveform_hm[:,0].max()
        input_currents_hm = waveform_hm[:,1]
        time_input_currents_lm = waveform_lm[:,0] - waveform_lm[:,0].max()
        input_currents_lm = waveform_lm[:,1]

        time_gates = np.array(system_ttem['time_gates'])
        GateTimeShift=system_ttem['lm']['GateTimeShift']
        MeaTimeDelay=system_ttem['lm']['MeaTimeDelay']
        NoGates=int(system_ttem['lm']['NoGates'])
        t0_lm = waveform_lm[:,0].max()
        times_lm = (time_gates[:NoGates,0] + GateTimeShift + MeaTimeDelay)[i_start_lm:i_end_lm]
        GateTimeShift=system_ttem['hm']['GateTimeShift']
        MeaTimeDelay=system_ttem['hm']['MeaTimeDelay']
        NoGates=int(system_ttem['hm']['NoGates'])
        t0_hm = waveform_hm[:,0].max()
        times_hm = (time_gates[:NoGates,0] + GateTimeShift + MeaTimeDelay)[i_start_hm:]
        
        receiver_location = np.array([9.28, 0., 0.45])
        receiver_orientation = "z"                    # "x", "y" or "z"        
        source_location = np.array([0., 0., 0.5])  
        source_orientation = 'z'
        source_list = []
        waveform_lm = tdem.sources.PiecewiseLinearWaveform(times=time_input_currents_lm, currents=input_currents_lm)
        waveform_hm = tdem.sources.PiecewiseLinearWaveform(times=time_input_currents_hm, currents=input_currents_hm)
        # Define receivers at each location.
        dbzdt_receiver_lm = tdem.receivers.PointMagneticFluxTimeDerivative(
            receiver_location, times_lm, receiver_orientation,
        )
        dbzdt_receiver_hm = tdem.receivers.PointMagneticFluxTimeDerivative(
            receiver_location, times_hm, receiver_orientation,
        )

        # Must define the transmitter properties and associated receivers
        src_0 = tdem.sources.MagDipole(
            [dbzdt_receiver_lm],
            location=source_location,
            waveform=waveform_lm,
            orientation=source_orientation,
        )
        src_1 = tdem.sources.MagDipole(
            [dbzdt_receiver_hm],
            location=source_location,
            waveform=waveform_hm,
            orientation=source_orientation,
        )
        # this needs to be fixed later
        source_list.append(src_0)
        source_list.append(src_1)

        survey_ttem = tdem.Survey(source_list)
        
        self._simulation_ttem = tdem.Simulation1DLayered(
            survey=survey_ttem, thicknessesMap=wire_map.thickness, sigmaMap=wire_map.sigma,
        )
        
        thicknesses_aem_doi = get_vertical_discretization(30, 3, 1.1)
        
        self._simulation_aem_doi = tdem.Simulation1DLayered(
            survey=survey, thicknesses=thicknesses_aem_doi, sigmaMap=maps.IdentityMap()
        )
        
        thicknesses_ttem_doi = get_vertical_discretization(30, 1, 1.1)                
        
        self._simulation_ttem_doi = tdem.Simulation1DLayered(
            survey=survey_ttem, thicknesses=thicknesses_ttem_doi, sigmaMap=maps.IdentityMap()
        )        
        
    @property
    def times_lm(self):
        src = self._simulation.survey.source_list[0]
        return src.receiver_list[0].times
    @property
    def times_hm(self):
        src = self._simulation.survey.source_list[1]
        return src.receiver_list[0].times
        
    def simulate(self, em_system, z, h, rho_background, rho_layer):
        
        thicknesses = np.array([z, h])        
        rho = np.array([rho_background, rho_layer, rho_background])
        sigma = 1./rho
        model = np.r_[sigma, thicknesses]
        if em_system == 'AEM':
            self._simulation = self._simulation_aem
            self._simulation_doi = self._simulation_aem_doi
        elif em_system == 'tTEM':
            self._simulation = self._simulation_ttem
            self._simulation_doi = self._simulation_ttem_doi
        
        dpred = self._simulation.dpred(model)
        mesh_1d = TensorMesh([[z, h, 1e3]])
        hz = np.r_[self._simulation_doi.thicknesses, self._simulation_doi.thicknesses[-1]]
        mesh_1d_doi = TensorMesh([hz])
        sigma_doi = volume_average(mesh_1d, mesh_1d_doi, values=sigma)
        dpred_doi = self._simulation_doi.dpred(sigma_doi)
        
        return dpred

    def depth_of_investigation_christiansen_2012(self, dobs, relative_error, floor, threshold):
        delta_d = relative_error * abs(dobs) + floor
        J = self._simulation_doi.getJ(self._simulation_doi.model)
        J_sum = abs(np.diag(1./delta_d) @ J['ds'] @ np.diag(self._simulation_doi.sigma)).sum(axis=0)
        hz = np.r_[self._simulation_doi.thicknesses, self._simulation_doi.thicknesses[-1]]
        S = np.cumsum(J_sum[::-1]/hz[::-1])[::-1]
        active = np.argmax(np.argwhere((S - threshold) > 0.0))+1
        if active == len(self._simulation_doi.sigma):
            doi = abs(self._simulation_doi.depth[-1])
        else:
            doi = abs(self._simulation_doi.depth[active])
        return S, doi
    
em_sim = EMSIMULATION()

def get_cmap_range_value(vmin, vmax, vrange):
    return [vmin + 0.0*vrange, vmax]

res_min, res_max = 1, 1000
res_range = res_max-res_min
depth_min, depth_max = 0, 500
depth_range = depth_max-depth_min
doi_min, doi_max = 0, 100
doi_range = doi_max-doi_min
state = AppState()
np.random.seed(1)

w_em_system=widgets.RadioButtons(options=["AEM", "tTEM"], value='AEM', layout={"width":"200px"})
state.register_stateful_widget(w_em_system, "em_system", Unicode(w_em_system.value))

w_rho_range=widgets.IntRangeSlider(
    min=res_min, max=res_max, 
    value=(3, 300),
    step=1,
    continuous_update=False, 
    disabled=False
)
state.register_stateful_widget(w_rho_range, "rho_range", List(get_cmap_range_value(res_min, res_max, res_range)))
w_depth_range=widgets.IntRangeSlider(
    min=depth_min, max=depth_max, 
    value=(0, 500),
    step=1,
    continuous_update=False, 
    disabled=False
)
state.register_stateful_widget(w_depth_range, "depth_range", List(get_cmap_range_value(depth_min, depth_max, depth_range)),)
w_doi_range=widgets.IntRangeSlider(
    min=doi_min, max=doi_max, 
    value=(0, 10),
    step=1,
    continuous_update=False, 
    disabled=False
)
state.register_stateful_widget(w_doi_range, "doi_range", List(get_cmap_range_value(doi_min, doi_max, doi_range)),)

w_z_1=widgets.FloatText(min=1, max=300, step=1, value=20, description='$z$', continuous_update=False, layout={"width":"300px"})
state.register_stateful_widget(w_z_1, "z_1", Float(w_z_1.value))
w_h_1=widgets.FloatText(min=1, max=300, step=1, value=20, description='$h$', continuous_update=False, layout={"width":"300px"})
state.register_stateful_widget(w_h_1, "h_1", Float(w_h_1.value))
w_rho_background_1=widgets.FloatText(min=1, max=1000, step=1, value=20, description='$\\rho_b$', continuous_update=False, layout={"width":"300px"})
state.register_stateful_widget(w_rho_background_1, "rho_background_1", Float(w_rho_background_1.value))
w_rho_layer_1=widgets.FloatText(min=1, max=1000, step=1, value=50, description='$\\rho$', continuous_update=False, layout={"width":"300px"})
state.register_stateful_widget(w_rho_layer_1, "rho_layer_1", Float(w_rho_layer_1.value))

w_relative_error=widgets.FloatText(min=1, max=100, step=1, value=3, description='% error', continuous_update=False, layout={"width":"300px"})
state.register_stateful_widget(w_relative_error, "relative_error", Float(w_relative_error.value))

style = {'description_width': 'initial'}
w_halfspace=widgets.Checkbox(value=False, description='Assume halfspace', layout={"width":"200px", "left": "20px"}, style=style)
state.register_stateful_widget(w_halfspace, "halfspace", Bool(w_halfspace.value))


 
# w_noise_floor=widgets.FloatText(value=1e-15, description='floor', continuous_update=False, layout={"width":"300px"})
# state.register_stateful_widget(w_noise_floor, "noise_floor", Float(w_noise_floor.value))
w_doi_threshold = widgets.FloatText(min=0, max=300, value=0.8, step=0.1, continuous_update=False, description='threshold:', layout={"width":"300px"})
state.register_stateful_widget(w_doi_threshold, "doi_threshold", Float(w_doi_threshold.value))

def apply_change_system(change):
    if change["new"]=='AEM':
        w_depth_range.value = (0, 500)
        w_noise_floor.value = 1e-15
    else:
        w_depth_range.value = (0, 150)
        w_noise_floor.value = 1e-11
        
state.widgets["em_system"].observe(apply_change_system, names=['value'])

def apply_change(change):
    if change["new"]:
        w_rho_layer_1.disabled = True      
        w_h_1.disabled = True  
        w_z_1.disabled = True  
    else:
        w_rho_layer_1.disabled = False      
        w_h_1.disabled = False  
        w_z_1.disabled = False  
        
state.widgets["halfspace"].observe(apply_change, names=['value'])

def plot(**kwargs):
    rho_layer_1 = kwargs['rho_layer_1']
    rho_background_1 = kwargs['rho_background_1']
    z_1 = kwargs['z_1']
    h_1 = kwargs['h_1']
    em_system = kwargs['em_system']
    rho_range = kwargs['rho_range']
    depth_range = kwargs['depth_range']
    doi_range = kwargs['doi_range']
    relative_error = kwargs['relative_error']
    # noise_floor = kwargs['noise_floor']
    doi_threshold = kwargs['doi_threshold']
    halfspace = kwargs['halfspace']
    noise_floor = 0.
    thickness_1 = np.r_[z_1, h_1]
    if halfspace:
        rho = np.r_[rho_background_1, rho_background_1, rho_background_1]
    else:
        rho = np.r_[rho_background_1, rho_layer_1, rho_background_1]
    rng = np.random.default_rng(42)
    fig, axs = plt.subplots(1,2, figsize=(10, 5))
    ax1, ax2 = axs
    dpred = em_sim.simulate(em_system, z_1, h_1, rho_background_1, rho_layer_1)  
    em_sim.dobs = dpred + relative_error*0.01*abs(dpred)*rng.standard_normal(len(dpred))
    S, doi = em_sim.depth_of_investigation_christiansen_2012(em_sim.dobs, relative_error*0.01, noise_floor, doi_threshold)
    
    utils.plot_1d_layer_model(thickness_1, rho, ax=ax1, **{'color':'C0', 'label': 'True'})
    utils.plot_1d_layer_model(em_sim._simulation_doi.thicknesses, S, ax=ax2, scale='linear', **{'color':'k', 'linestyle':'-', 'lw':3})
    ax2.hlines(doi, doi_range[0], doi_range[1], color='crimson', lw=2, linestyle='--')
    ax2.vlines(doi_threshold, depth_range[1], depth_range[0], color='crimson', lw=2, linestyle='--')    
    ax1.set_xlabel("Resistivity (Ohm-m)")
    ax1.set_title("Vertical resistivity profile")
    ax1.set_xlim(rho_range[0], rho_range[1])
    ax1.set_ylim(depth_range[1], depth_range[0])
    ax2.set_xlabel("DOI index")
    ax2.set_title("DOI estimation (DOI={:.0f} m)".format(doi))
    ax2.set_xlim(doi_range[0], doi_range[1])
    ax2.set_ylim(depth_range[1], depth_range[0])
    ax2.set_ylabel("")

    plt.tight_layout()
    plt.show()
    return   

# Some texts
top_label_1 = Label(value="EM system:", layout={"width":"100px"})
# top_right_label_1 = Label(value="Simulation or Inversion:", layout={"width":"150px"})
top_label_2 = Label(value="Resistivity range", layout={"width":"100px"})
top_label_3 = Label(value="Depth range", layout={"width":"100px"})
top_label_4 = Label(value="DOI index range", layout={"width":"100px"})
edit_plot_limit_label = Label(value="Edit plot access limit:", layout={"width":"150px", "padding":"0px"})
top_panel_1 = HBox([top_label_1, w_em_system], layout={"align_items":"center"})
top_panel_2 = HBox([top_label_2, w_rho_range], layout={"align_items":"center"})
top_panel_3 = HBox([top_label_3, w_depth_range], layout={"align_items":"center"})
top_panel_4 = HBox([top_label_4, w_doi_range], layout={"align_items":"center"})
# top_panel_4 = HBox([run_simulation_label, w_run_simulation])
left_label = Label(value="Parameters:" ,layout={"margin":"0px"})
left_panel = VBox(
    [w_halfspace, w_z_1, w_h_1, w_rho_background_1, w_rho_layer_1, w_relative_error, w_doi_threshold], layout={"align_items":"center"}
)

# right_label = Label(value="Other parameters:")
# right_panel = VBox(
#     [right_label, w_relative_error, w_doi_threshold], layout={"align_items":"center"}
# )


w_output = interactive_output(
    plot,
    { k: state.widgets[k] for k in state.widgets.keys()}
);

bottom_panel = HBox([left_panel], layout={"align_items":"center"})
# edit_panel = HBox([edit_plot_limit_label, w_edit_plot], layout={"align_items":"center"})
panel = VBox([top_panel_1, left_label, bottom_panel, w_output, edit_plot_limit_label, top_panel_2, top_panel_3, top_panel_4])
display(panel)