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

# # Homework 9.2: MLE of microtubule catastrophe data (40 pts)
#
# [Dataset download](https://s3.amazonaws.com/bebi103.caltech.edu/data/gardner_time_to_catastrophe_dic_tidy.csv)
#
# <hr>

# Refresh yourself about the microtubule catastrophe data we explored in previous homeworks. We will again work with this data set here.
#
# **a)** In their [paper](http://dx.doi.org/10.1016/j.cell.2011.10.037), <NAME>, and coworkers modeled microtubule catastrophe times as Gamma distributed. Perform a maximum likelihood estimate for the parameters of the Gamma distribution. Because you showed in a previous homework that there is little difference between labeled and unlabeled tubulin, you only need to work this out for the labeled tubulin now and in part (b). Be sure to include confidence intervals or a confidence region for your MLE and discuss the method you used to get the confidence intervals or confidence region.
#

# +
import numpy as np
import pandas as pd
import os
import scipy 
import scipy.optimize
import scipy.stats as st
import tqdm
import math

rg = np.random.default_rng()

data_path = "../data"
fname = os.path.join(data_path, "gardner_time_to_catastrophe_dic_tidy.csv")
df = pd.read_csv(fname)
df
# -

# We made a separate dataframe for the data labeled "True", and obtained the time stamps of those labeled "True".

df_true = df[df["labeled"] == True]
true_time = df_true["time to catastrophe (s)"].values


# Now, we defined the bootstrap to return us alpha and beta.

def bootstrap_param(data, size=1):
    """Parametric bootstrap replicates of parameters of
    Normal distribution."""
    bs_mean = np.empty(size)
    bs_sd = np.empty(size)
    bs_alpha = np.empty(size)
    bs_beta = np.empty(size)

    for i in range(size):
        bs_sample = np.random.choice(data, size=len(data))
        bs_mean[i] = np.mean(bs_sample)
        bs_sd[i] = np.std(bs_sample)
        # Since the mean of the gamma distribution is alpha/beta, and variance is alpha/beta^2, mean/var = beta.
        bs_beta[i] = bs_mean[i]/(bs_sd[i]**2)
        # With the beta value, we solve for alpha by just saying beta * mean = alpha.
        bs_alpha[i] = bs_beta[i] * bs_mean[i]


    return bs_alpha, bs_beta

# We run the bootstrap for our "True" time stamps.

# +
bs_alpha_true, bs_beta_true = bootstrap_param(
    true_time, size=100000
)

bs_alpha_true
# -

# We can obtain the alpha and beta values as an expression of the mean and variance.

# +
mean_true = np.mean(true_time)
var_true = np.var(true_time)

# Since the mean of the gamma distribution is alpha/beta, and variance is alpha/beta^2, mean/var = beta.
beta_true = mean_true / var_true

# With the beta value, we solve for alpha by just saying beta * mean = alpha.
alpha_true = beta_true * mean_true
print(alpha_true, beta_true)
# -

print('α: {:.4f} | 95% conf int α: {}'.format(np.mean(bs_alpha_true),str(np.percentile(bs_alpha_true, [2.5, 97.5]))))
print('β2: {:.4f} | 95% conf int β2: {}'.format(np.mean(bs_beta_true),str(np.percentile(bs_beta_true, [2.5, 97.5]))))


# As the values of the parameters alpha and beta do not overlap, and fall inside the 95% confidence interval for the labeled microtubules, we can reasonably assume that the MLE value we have obtained for the Gamma distribution is reasonable.

# **b)** Obtain a maximum likelihood estimate for the parameters $\beta_1$ and $\beta_2$ from the model you derived in homework 5.2. As a reminder, you derived that the PDF for microtubule catastrophe times is
#
# $$\begin{align}
# f(t;\beta_1, \beta_2) = \frac{\beta_1\beta_2}{\beta_2 - \beta_1}\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right).
# \end{align}$$
#
# Again, include confidence intervals. **Be careful**; this is a *very* tricky calculation. It is possible to analytically compute the MLE. If you choose to do it numerically, you need to think about what happens when $\beta_1 \approx \beta_2$. You may also need to think about how you will handle the [log of sums of exponentials](https://en.wikipedia.org/wiki/LogSumExp).

# $$\begin{align}\tag{1}
# f(t;\beta_1, \beta_2) = \frac{\beta_1 \beta_2}{\beta_2 - \beta_1}\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)
# \end{align}$$
#
# $$\begin{align}\tag{2}
# L(t;\beta_1, \beta_2)= \ln{f(t;\beta_1, \beta_2)} = \ln{\left(\beta_1\beta_2\right)} - \ln{\left(\beta_2-\beta_1\right)}+\ln{\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)}
# \end{align}$$
#
# Our third term is the log of sum of exponentials, and we can change this term to:
#
# $$\begin{align}\tag{3}
# y = \ln{\left(\mathrm{e}^{-\beta_1 t} - \mathrm{e}^{-\beta_2 t}\right)}\\
# \Leftrightarrow \mathrm{e}^{y} = \sum_{i=1}^{n}\mathrm{e}^{x_{i}}\\
# \Leftrightarrow \mathrm{e}^{\beta_1 t}\mathrm{e}^{y} = \mathrm{e}^{-\beta_1 t}\sum_{i=1}^{n}\mathrm{e}^{x_{i}}\\
# \Leftrightarrow \mathrm{e}^{y + \beta_1 t} = \sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t} \\
# \Leftrightarrow y + \beta_1 t = \ln{\sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t}}\\
# \Leftrightarrow y = -\beta_1 t + \ln{\sum_{i=1}^{n}\mathrm{e}^{x_{i} + \beta_1 t}}\\
# \Leftrightarrow y = -\beta_1 t + \ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}}
# \end{align}$$
#
# As such, our initial log equation will become as follows:
# $$\begin{align}\tag{4}
# \ln{f(t;\beta_1, \beta_2)} = \ln{\left(\beta_1\beta_2\right)} - \ln{\left(\beta_2-\beta_1\right)}+\ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}}\\
# = \ln{\beta_1} + \ln{\beta_2} - \ln{(\beta_2-\beta_1)} -\beta_1 t - ln{\mathrm{e}^{\left(\beta_1-\beta_2\right)t}} \\
# = \ln{\beta_1} + \ln{\beta_2} - \ln{(\beta_2-\beta_1)} -\beta_1 t - \left(\beta_1-\beta_2\right)t
# \end{align}$$
#
# Taking the partial derivatives of each terms gives:
# $$\begin{align}\tag{5}
# \frac{\partial L}{\partial \beta_1} = \frac{1}{\beta_1} + 0 - \left(\frac{1}{\beta_2-\beta_1}\left(0-1\right)\right) - t - t + 0 \\
# = \frac{1}{\beta_1} + \frac{1}{\beta_2-\beta_1} - 2t
# \end{align}$$
# $$\begin{align}\tag{6}
# \frac{\partial L}{\partial \beta_2} = 0 + \frac{1}{\beta_2}-\left(\frac{1}{\beta_2-\beta_1}\left(1-0\right)\right)-0-0+t\\
# = \frac{1}{\beta_2} - \frac{1}{\beta_2-\beta_1} + t
# \end{align}$$
# $$\begin{align}\tag{7}
# \frac{\partial L}{\partial t} = 0 + 0 - 0 -\beta_1 -\beta_1 + \beta_2\\
# = \beta_2 - 2\beta_1\\
# \end{align}$$
#
# We can evaluate the turning point by setting $\frac{\partial L}{\partial t} = 0$, which gives the expression: 
# $$\begin{align}\tag{8}
# \beta_1 = \frac{1}{2}\beta_2 \Leftrightarrow 2\beta_1 = \beta_2
# \end{align}$$
#
# Whereas we can rearrange our initial equation to express $t$ in terms of $\beta_1$ and $\beta_2$: 
# $$\begin{align}\tag{9}
# t = \frac{\ln{\beta_1} + \ln{\beta_2} - \ln{\left(\beta_2-\beta_1\right)}}{2\beta_1-\beta_2}
# \end{align}$$
#
# As such, the turning point are points where-in: 
# $$\begin{align}\tag{10}
# \beta_1 = \frac{1}{2} \wedge \beta_2 = 1
# \end{align}$$
# $$\begin{align}\tag{11}
# 2\beta_1 - \beta_2 \neq 0 \wedge t = \frac{\ln{\beta_1} + \ln{\beta_2} - \ln{\left(\beta_2-\beta_1\right)}}{2\beta_1-\beta_2}
# \end{align}$$
#

# Using a numerical approach, we defined the log likelihood function using the gamma distribution, as well as the corresponding MLE.

# +
def log_like_iid_gamma(params, n):
    
    beta1, beta2 = params
    mlesum = 0

    for t in n:
        mlesum += np.log(beta1)+ np.log(beta2) - np.log((beta2-beta1)) + np.log(np.exp(-beta1*t) - np.exp(-beta2*t))
        if (math.isnan(mlesum)):
            return -np.inf
    return mlesum

def mle_iid_custom(n):
    """Perform maximum likelihood estimates for parameters for i.i.d."""
#     with warnings.catch_warnings():
#         warnings.simplefilter("ignore")

    res = scipy.optimize.minimize(
        fun=lambda params, n: -log_like_iid_gamma(params, n),
        x0=np.array([0.0035, 0.0045]),
        args=(n,),
        method='Powell',
    )
    if res.success:
        return res.x
    else:
        raise RuntimeError('Convergence failed with message', res.message)


# -

# Now we run our MLE to obtain our beta values.

mle_labeled = mle_iid_custom(df_true["time to catastrophe (s)"].values)

print('β1: ' + str(mle_labeled[0]))
print('β2: ' + str(mle_labeled[1]))


# We can define a function to draw our bootstrap replicates, as such:

# +
def draw_bs_sample(data):
    """Draw a bootstrap sample from a 1D data set."""
    return rg.choice(data, size=len(data))


def draw_bs_reps_mle(mle_fun, data, args=(), size=1, progress_bar=False):

    if progress_bar:
        iterator = tqdm.tqdm(range(size))
    else:
        iterator = range(size)

    return np.array([mle_fun(draw_bs_sample(data), *args) for _ in iterator])


# -

# We run our bootstrap to obtain the confidence intervals for our beta values.

bs_reps = draw_bs_reps_mle(
    mle_iid_custom, df_true["time to catastrophe (s)"].values, size=100,
)

# +
conf_int = np.percentile(bs_reps, [2.5, 97.5], axis=0)
conf_int = np.transpose(conf_int)

print('β1: {:.4f} | 95% conf int β1: {}'.format(np.mean(bs_reps[:,0]),str(conf_int[0])))
print('β2: {:.4f} | 95% conf int β2: {}'.format(np.mean(bs_reps[:,1]),str(conf_int[1])))
# -

# Comparing the parameter estimates from both the analytical and numerical approach gives very different results; in our analytical approach we were not able to find a single point for the MLE, whereas using bootstrapping methods we can obtain precise beta values with a small confidence interval range.

# ## Attributes: 
#
# Done altogether as a group

StarcoderdataJupyterScriptsDedupFiltered