# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from scipy.stats import laplace, norm, uniform

param_files = ["exp_params.txt", "pair_params.txt", "peptide_params.txt"]
params = {}
for pf in param_files:
    pfp = os.path.join("/gd/bayesian_RT/Alignments/SQC_20180815_2", pf)
    #pfp = os.path.join("/Users/albert/git/RTLib/Alignments/NCE_20180520_4", pf)
    if os.path.exists(pfp):
      with open(pfp, "rb") as f:
        try:
          params[pf.split("_")[0]] = pd.read_csv(pfp, sep="\t")
        except:
          print("some error")
params

dfa = pd.read_csv('/gd/bayesian_RT/Alignments/SQC_20180815_2/ev_updated.txt', sep='\t', low_memory=False)
dfa

# +
dff = dfa[-(dfa["remove"])].reset_index(drop=True)

# factorize experiments into exp_id
#eidx = dff["raw_file"].unique()
#print(len(eidx))
#dff["exp_id"] = dff["raw_file"].map({ind: val for val, ind in enumerate(dff["raw_file"].unique())})

# factorize sequence into peptide_id
#dff["stan_peptide_id"] = dff["sequence"].map({ind: val for val, ind in enumerate(dff["sequence"].unique())})

num_experiments = dff["exp_id"].max() + 1
num_observations = dff.shape[0]
num_peptides = dff["peptide_id"].max() + 1
retention_times = dff["Retention time"]

# cap PEP at 1
dff['PEP'].loc[dff['PEP'] > 1] = 1
dff['pep_new'].loc[dff['pep_new'] > 1] = 1
dff['pep_updated'].loc[dff['pep_updated'] > 1] = 1

pep = dff["PEP"]
stan_peptide_id = dff["peptide_id"]
exp_id = dff["exp_id"]
exp_names = dff['Raw file'].unique()
mean_log_rt = np.mean(np.log(retention_times))
sd_log_rt = np.std(np.log(retention_times))
max_rt = retention_times.max()

pep_exp_all = dff["peptide_id"].map(str) + " - " + dff["exp_id"].map(str)
pep_exp_pairs = pep_exp_all.unique()
num_pep_exp_pairs = len(pep_exp_pairs)
print(len(pep_exp_all))
print(len(pep_exp_all.unique()))
muij_map = pep_exp_all.map({ind: val for val, ind in enumerate(pep_exp_pairs)})
#muij_map
#dff

# +
pep_col_code = pd.cut(dff["PEP"], 10)

for exp in range(1,2):
    print("Generating Summary for Experiment", exp, "|", exp_names[exp])

    exp_params = params["exp"].iloc[exp]
    exp_inds = (dff['exp_id'] == exp) & (~pd.isnull(dff['pep_new']))

    predicted = dff['muij'][exp_inds].values
    predicted_sd = dff['sigmaij'][exp_inds].values
    mus = dff['mu'][exp_inds].values

    # observed values
    observed = dff['Retention time'][exp_inds].values
    obs_peps = dff['PEP'][exp_inds].values
    obs_code = pep_col_code[exp_inds].values
    residual = observed - predicted

    plt.subplot(121)
    plt.scatter(mus, observed, s=1, color="black")
    #plt.scatter(predicted, observed)
    plt.plot([0, exp_params["split_point"]],
             [exp_params["beta_0"], (exp_params["split_point"] * exp_params["beta_1"]) + exp_params["beta_0"]],
            color="red")
    plt.plot([exp_params["split_point"], 300], 
             [(exp_params["split_point"] * exp_params["beta_1"]) + exp_params["beta_0"], (exp_params["split_point"] * exp_params["beta_1"]) + ((300-exp_params["split_point"]) * exp_params["beta_2"]) + exp_params["beta_0"]],
            color="green")
    plt.plot(np.repeat(exp_params["split_point"], 2), [-100, 300], color="blue", linestyle="dashed")
    plt.axis([0, mus.max() + 10, exp_params["beta_0"]-10, observed.max() + 10])
    plt.title(exp_names[exp])
    plt.xlabel("Reference RT (min)")
    plt.ylabel("Observed RT (min)")

    plt.subplot(122)
    plt.scatter(predicted, residual, s=4, c=pep_col_code.cat.codes.values[exp_inds], alpha=0.5)
    plt.plot([0, exp_params["split_point"]], [0, 0], color="red")
    plt.plot([exp_params["split_point"], 300], [0, 0], color="green")
    plt.plot(np.repeat(exp_params["split_point"], 2), [-100, 300], color="blue", linestyle="dashed")
    
    # confidence intervals, 2.5% and 97.5%
    conf_x = predicted[np.argsort(predicted)]
    conf_2p5 = laplace.ppf(0.025, loc=0, scale=predicted_sd)[np.argsort(predicted)]
    conf_97p5 = laplace.ppf(0.975, loc=0, scale=predicted_sd)[np.argsort(predicted)]
    
    plt.plot(conf_x, conf_2p5, color="red")
    plt.plot(conf_x, conf_97p5, color="red")
    plt.axis([predicted.min()-5, predicted.max()+5, residual.min()-5, residual.max()+5])
    
    plt.ylim(np.min(conf_2p5) - 0.1, np.max(conf_97p5) + 0.1)
    plt.xlim(conf_x[0], conf_x[-1])
    
    cbar = plt.colorbar()
    cbar.set_label('Spectral PEP (Error Probability)')
    cbar.ax.set_yticklabels(pep_col_code.cat.categories.values)
    plt.xlabel("Inferred RT (min)")
    plt.ylabel("Residual RT (min)")

    plt.subplots_adjust(hspace=0.3, wspace=0.4, bottom=0.2)
#    plt.grid()

    fig = plt.gcf()
    fig.set_size_inches(7, 3.5)
    #fname = "./tmp_figs/alignment_" + str(exp) + "_" + exp_names[exp] + ".png"
    #print("Saving figure to", fname, "...")
    #fig.savefig(fname, dpi=160)
    
    #plt.close()
    #fig.clf()
    
    plt.show()

# +
num_experiments = len(dff['exp_id'].unique())
dff['residual'] = dff['Retention time'] - dff['muij']

plots_per_row = 30
if num_experiments < plots_per_row:
    plots_per_row = num_experiments

num_rows = int(np.ceil(num_experiments / plots_per_row))

resi = []

for i in range(0, num_rows):
    ax = plt.subplot2grid((num_rows, 1), (i, 0))

    if (i + 1) * plots_per_row > num_experiments:
        resi = [dff['residual'][(dff['exp_id'] == i) & (~pd.isnull(dff['residual']))] for i in range((i * plots_per_row), num_experiments)]
        ax.boxplot(resi, showfliers=False)
        ax.set_xticklabels(np.arange((i * plots_per_row), num_experiments, 1))
    else:
        resi = [dff['residual'][(dff['exp_id'] == i) & (~pd.isnull(dff['residual']))] for i in range((i * plots_per_row), ((i + 1) * plots_per_row))]
        ax.boxplot(resi, showfliers=False)
        ax.set_xticklabels(np.arange((i * plots_per_row), ((i + 1) * plots_per_row), 1))

    #ax.violinplot(resi, showmedians=True, showextrema=True)
    #ax.boxplot(resi, showfliers=False)
    ax.set_xticks(np.arange(1, plots_per_row + 1, 1))

    ax.set_xlabel('Experiment Number')
    ax.set_ylabel('Residual RT (min)')

#plt.subplots_adjust(hspace=0.6, wspace=0.3)
# plt.tight_layout()

# finalize and save figure
f = plt.gcf()
#f.text(0.5, 0, 'Experiment Number', fontsize=16, ha='center', va='center')
#f.text(0.06, 0.5, 'Residual RT (min)', fontsize=16, ha='center', va='center', rotation='vertical')
f.set_size_inches(12, num_rows * 2)

#fname = os.path.join(figures_path, 'residuals_violin.png')
#logger.info('Saving figure to {} ...'.format(fname))
#f.savefig(fname, dpi=160)
#fig_names.append(fname)

#plt.close()
#f.clf()

# +
# PEP vs. PEP.new scatterplot
inds = (~pd.isnull(dff['pep_new'])) & (dff['pep_new'] > 1e-5) & (dff['PEP'] > 1e-5)

f, ax = plt.subplots()
hst = ax.hist2d(np.log10(dff['PEP'][inds]), np.log10(dff['pep_new'][inds]), bins=(50, 50), cmap=plt.cm.Reds)
ax.plot([-5, 0], [-5, 0], '-r')

ax.grid()

interval = (-5, 0)
ax.set_xlim(interval)
ax.set_ylim(interval)
ax.set_xlabel('Spectral PEP', fontsize=16)
ax.set_ylabel('DART-ID PEP', fontsize=16)
interval = np.arange(-5, 1, 1)
ax.set_xticks(interval)
ax.set_yticks(interval)
ax.set_xticklabels(['$10^{{{}}}$'.format(i) for i in interval], fontsize=12)
ax.set_yticklabels(['$10^{{{}}}$'.format(i) for i in interval], fontsize=12)

f.set_size_inches(7, 5)
# plt.tight_layout()

cbar = plt.colorbar(hst[3], ax=ax)
cbar.set_label('Frequency', fontsize=16, labelpad=20, ha='center', va='top')
cbar.ax.xaxis.set_label_position('top')

# +
# Fold-change increase
num_points=100
x = np.logspace(-5, 0, num=num_points)
y = np.zeros(num_points)
y2 = np.zeros(num_points)
y3 = np.zeros(num_points)
inds = ~pd.isnull(dff['pep_new'])

for i, j in enumerate(x):
    y[i] = np.sum(dff['pep_updated'] < j) / np.sum(dff['PEP'] < j)
    y2[i] = np.sum(dff['PEP'] < j) / dff.shape[0]
    y3[i] = np.sum(dff['pep_updated'] < j) / dff.shape[0]

f, (ax1, ax2) = plt.subplots(2, 1)

ax1.semilogx(x, (y*100)-100, '-b')
# ax1.plot([np.min(x), np.max(x)], [0, 0], '-r', linestyle='dashed', linewidth=2)
ax1.plot([1e-2, 1e-2], [-1000, 1000], '-k', linestyle='dashed')
ax1.grid()
ax1.set_xlim([3e-4, 3e-1])
ax1.set_ylim([-25, np.max(y)*100-50])
ax1.set_xlabel('PEP Threshold', fontsize=16)
ax1.set_ylabel('Increase (%)', fontsize=16)
ax1.set_title('Increase in confident PSMs', fontsize=16)

ax2.semilogx(x, y2, '-b', linewidth=1, label='Spectra PEP')
ax2.semilogx(x, y3, '-g', linewidth=1, label='DART-ID PEP')
#ax2.fill_between(x, 0, y2)
ax2.plot([1e-2, 1e-2], [-1000, 1000], '-k', linestyle='dashed')
ax2.grid()
ax2.set_xlim([3e-4, 3e-1])
ax2.set_ylim([0, 1.05])
ax2.set_xlabel('PEP Threshold', fontsize=16)
ax2.set_ylabel('Fraction', fontsize=16)
ax2.set_title('Fraction of confident PSMs', fontsize=16)
ax2.legend(fontsize=16)

plt.subplots_adjust(hspace=0.6, wspace=0.3)

f.set_size_inches(5, 7)
# plt.tight_layout()
# -



StarcoderdataJupyterScriptsDedupFiltered