Constructing resistivity distribution for sediment type

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import dill
from emrecharge.colocation import find_closest_locations, find_locations_in_distance, compute_fraction_for_aem_layer
from emrecharge.rockphysics import (
    compute_interval,
    compute_integral,
    from_sigma_to_fraction,
    compute_rho_to_cf_mappings_with_rho
)
from scipy.stats import percentileofscore

from ipywidgets import Output, BoundedIntText, BoundedFloatText, Label, HBox, VBox, widgets, interactive_output
from IPython.display import display, clear_output
from traitlets import Int, Float
from traitlets import HasTraits, Int, Unicode, List, Float, Bool, TraitType, Dict
from emrecharge.app import AppState
import matplotlib as mpl
mpl.rcParams.update({'font.size':12})
df_lith_collar = pd.read_csv("collar.csv")
df_lith_interval = pd.read_csv("interval.csv")
# Code 0: fine-dominated; Code 1: coarse-dominated
df_lith_interval['Code'] = df_lith_interval['Code']-1.
xy_lithology = df_lith_collar[['X', 'Y']].values
df_lith_collar = df_lith_collar.set_index('Hole_ID')
df_lith_interval = df_lith_interval.set_index('Hole_ID')
lith_index = np.array(df_lith_collar.index.to_list())
results = dill.load(open("rho_1d_inversion.pik", "rb"))
xy_em = results['topography'][:,:2]
n_sounding = xy_em.shape[0]
hz = results['hz']
resistivity = results['rho'].reshape((n_sounding, len(hz)))
n_layer = len(hz)

rho_fine_dominated_mean = 15.
rho_coarse_dominated_mean = 35.
log_rho_std = 1.2
rho_fine_dominated = np.exp(np.random.normal(np.log(rho_fine_dominated_mean), np.log(log_rho_std), 1000))
rho_coarse_dominated = np.exp(np.random.normal(np.log(rho_coarse_dominated_mean), np.log(log_rho_std), 1000))

interval_option = 'interval'
colocation_distance = 20

state = AppState()
w_interval_option=widgets.RadioButtons(
    options=["interval", "integral"], 
    value=interval_option, 
    layout={"width":"500px"},
    description='Approach'
)
state.register_stateful_widget(w_interval_option, "interval_option", Unicode(w_interval_option.value))

style = {'description_width': 'initial'}
w_colocation_distance=widgets.FloatSlider(
    min=20, max=600., 
    value=colocation_distance,
    step=10.,
    continuous_update=False, 
    disabled=False,
    description='Colocation distance:',
    style=style,
    layout={"width":"500px"}
)
state.register_stateful_widget(w_colocation_distance, "colocation_distance", Float(w_colocation_distance.value))

def plot(**kwargs):
    np.random.seed(1)
    colocation_distance = kwargs['colocation_distance']
    interval_option = kwargs['interval_option']    
    loc, inds = find_locations_in_distance(xy_em, xy_lithology, distance=colocation_distance)
    if inds is not None:
        if len(inds)>1:
            lith_index_colocated = lith_index[inds]
            _, inds_em = find_closest_locations(xy_lithology[inds, :], xy_em)
            df_lith_collar_colocated = df_lith_collar.loc[lith_index_colocated]

            well_names = np.array(df_lith_collar_colocated.index.tolist())


            rock_physics_data = []
            rock_physics_well_names = []
            for i_sounding in range(len(inds_em)):
                well_name = well_names[i_sounding]
                rho_tmp = resistivity[inds_em[i_sounding],:][:n_layer]
                df_driller_tmp = df_lith_interval.loc[well_name]
                inds_sort = np.argsort(df_driller_tmp['From'].values)
                z_top = df_driller_tmp['From'].values[inds_sort]
                z_bot = df_driller_tmp['To'].values[inds_sort]
                code_tmp = df_driller_tmp['Code'].values.astype(float)[inds_sort]
                lith_data = pd.DataFrame(data=np.c_[z_top, z_bot, code_tmp], columns=['From', 'To', 'Code'])
                fraction = compute_fraction_for_aem_layer(hz[:n_layer], lith_data)
                depth = np.cumsum(np.r_[0, hz[:n_layer]])
                top = depth[:-1]
                bottom = depth[1:]
                rock_physics_data.append(np.c_[top, bottom, fraction, rho_tmp])
                rock_physics_well_names.append(np.array([well_name] * n_layer))
            rock_physics_data = np.vstack(rock_physics_data)
            rock_physics_well_names = np.hstack(rock_physics_well_names)
            df_rock_physics_data = pd.DataFrame(data=rock_physics_data, columns=['top', 'bottom', 'f_fine', 'f_coarse', 'rho_aem'])
            df_rock_physics_data['well_names'] = rock_physics_well_names
            df_rock_physics_data['gse_wse'] = 1e5 * np.ones(df_rock_physics_data.shape[0])
            df_rock_physics_data = df_rock_physics_data.dropna()
            n_bootstrap = 1000


            hist_fig, hist_ax = plt.subplots(1,1, figsize=(5*1.5, 2.5*1.5))
            plt.figure(hist_fig.number)

            if interval_option == 'interval':
                rho_above = compute_interval(df_rock_physics_data, n_bootstrap=n_bootstrap)
            elif interval_option == 'integral':
                rho_above = compute_integral(df_rock_physics_data, n_bootstrap=n_bootstrap)  

            np.clip(rho_above, 1, 1000)

            rho_min = 13
            rho_max = 25

            rho_fine_above_tmp = np.percentile(rho_above[:,0], 50)
            rho_coarse_above_tmp = np.percentile(rho_above[:,1], 50)


            n_bins = 60

            p_fine = round(percentileofscore(np.sort(rho_above[:,0]), rho_fine_above_tmp))
            p_coarse = round(percentileofscore(np.sort(rho_above[:,1]), rho_coarse_above_tmp))
            bins_above = np.linspace(np.log10(rho_min), np.log10(rho_max), n_bins)
            hist_ax.hist(np.log10(rho_above[:,0]), bins=bins_above, alpha=0.5, label='fine-dominated', color='blue', edgecolor='k')
            hist_ax.hist(np.log10(rho_above[:,1]), bins=bins_above, alpha=0.5, label='coarse-dominated', color='red', edgecolor='k')
            hist_ax.hist(np.log10(resistivity.flatten()), bins=bins_above, alpha=0.2, label='EM resistivity', color='C0', edgecolor='k', zorder=0)
            hist_ax.axvline(x=np.log10(rho_fine_above_tmp), color='blue', lw=2, linestyle='--')
            hist_ax.axvline(x=np.log10(rho_coarse_above_tmp), color='red', lw=2, linestyle='--')

            xticks = hist_ax.get_xticks()
            hist_ax.set_xticks(xticks)
            hist_ax.set_xticklabels(["{:.0f}".format(10**x) for x in xticks])
            hist_ax.set_xlim((np.log10(rho_min), np.log10(rho_max)))

            hist_ax.set_xlabel("Resistivity ($\Omega$m)")  
            hist_ax.set_ylabel("Counts")
            hist_ax.set_title(f"The number of colocated wells: {len(inds):d}\n rho_FD: {rho_fine_above_tmp:.1f} (std={np.std(rho_above[:,0]):.2f}), rho_CD: {rho_coarse_above_tmp:.1f} (std={np.std(rho_above[:,1]):.2f})")
            hist_ax.legend(prop={'size': 12}, loc=1)
            plt.tight_layout()
            plt.show(hist_fig)
        else:
            print ("There is only one well")
    else:
        print ("There is no wells")
        
w_output = interactive_output(
    plot,
    { k: state.widgets[k] for k in state.widgets.keys()}
);
panel = VBox([w_colocation_distance, w_interval_option, w_output])
display(panel)