# ---
# 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 pandas as pd
import patsy as pt
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn import preprocessing
# %matplotlib inline
import re
import pymc3 as pm
import matplotlib.ticker as tk
import re
from sklearn.model_selection import StratifiedKFold
import pickle

# ## Import data

df = pd.read_csv('outputs/ala1_trials_clean.csv')
df = df.rename(columns={'project_name': 'basis', 'cluster__n_clusters': 'n', 'test_mean': 'y'}).\
loc[:, ['basis', 'y', 'n']]

# ## Scale predictors
#

to_scale = ['n']
scaler = preprocessing.MinMaxScaler()
vars_scaled = pd.DataFrame(scaler.fit_transform(df.loc[:, to_scale]), columns=[x+'_s' for x in to_scale])
df = df.join(vars_scaled)
df.T

# ## Create design matrix

y = df.loc[:, 'y']
X = df.loc[:, df.columns.difference(['y'])]
X_c = pt.dmatrix('~ 0 + n_s + C(basis)', data=df, return_type='dataframe')
X_c = X_c.rename(columns=lambda x: re.sub('C|\\(|\\)|\\[|\\]','',x))


# ## Model fitting functions

# +
def gamma(alpha, beta):
    def g(x):
        return pm.Gamma(x, alpha=alpha, beta=beta)
    return g

def hcauchy(beta):
    def g(x):
        return pm.HalfCauchy(x, beta=beta)
    return g


def fit_model_1(y, X, kernel_type='rbf'):
    """
    function to return a pymc3 model
    y : dependent variable
    X : independent variables
    prop_Xu : number of inducing varibles to use
    
    X, y are dataframes. We'll use the column names. 
    """
    with pm.Model() as model:
        # Covert arrays
        X_a = X.values
        y_a = y.values
        X_cols = list(X.columns)
        
        # Globals
        prop_Xu = 0.1 # proportion of observations to use as inducing variables
        l_prior = gamma(1, 0.05)
        eta_prior = hcauchy(2)
        sigma_prior = hcauchy(2)

        # Kernels
        # 3 way interaction
        eta = eta_prior('eta')
        cov = eta**2
        for i in range(X_a.shape[1]):
            var_lab = 'l_'+X_cols[i]
            if kernel_type=='RBF':
                cov = cov*pm.gp.cov.ExpQuad(X_a.shape[1], ls=l_prior(var_lab), active_dims=[i])
            if kernel_type=='Exponential':
                cov = cov*pm.gp.cov.Exponential(X_a.shape[1], ls=l_prior(var_lab), active_dims=[i])
            if kernel_type=='M52':
                cov = cov*pm.gp.cov.Matern52(X_a.shape[1], ls=l_prior(var_lab), active_dims=[i])
            if kernel_type=='M32':
                cov = cov*pm.gp.cov.Matern32(X_a.shape[1], ls=l_prior(var_lab), active_dims=[i])

        # Covariance model
        cov_tot = cov 

        # Model
        gp = pm.gp.MarginalSparse(cov_func=cov_tot, approx="FITC")

        # Noise model
        sigma_n =sigma_prior('sigma_n')

        # Inducing variables
        num_Xu = int(X_a.shape[0]*prop_Xu)
        Xu = pm.gp.util.kmeans_inducing_points(num_Xu, X_a)

        # Marginal likelihood
        y_ = gp.marginal_likelihood('y_', X=X_a, y=y_a,Xu=Xu, noise=sigma_n)
        mp = pm.find_MAP()
        
    return gp, mp, model


# -

# ## Main testing loop

# +
# Inputs
kernels =  ['M32', 'M52', 'RBF', 'Exponential' ]

# Outputs
pred_dfs = []

# iterator
kf = StratifiedKFold(n_splits=10)

for i in range(len(kernels)):
    print(kernels[i])
    for idx, (train_idx, test_idx) in enumerate(kf.split(X.values, X['basis'])):
        print('\tfold: {}'.format(idx))
        # subset dataframes for training and testin
        y_train = y.iloc[train_idx]
        X_train = X_c.iloc[train_idx, :]
        
        y_test = y.iloc[test_idx]
        X_test = X_c.iloc[test_idx, :]  
        
        # Fit gp model
        gp, mp, model = fit_model_1(y=y_train, X=X_train, kernel_type=kernels[i])

        # Get predictions for evalution
        with model:
            # predict latent
            mu, var = gp.predict(X_test.values, point=mp, diag=True,pred_noise=False)
            sd_f = np.sqrt(var)

            # predict target (includes noise)
            _, var = gp.predict(X_test.values, point=mp, diag=True,pred_noise=True)
            sd_y = np.sqrt(var)
            
        res = pd.DataFrame({'f_pred': mu, 'sd_f': sd_f, 'sd_y': sd_y, 'y': y_test.values})
        res.loc[:, 'kernel'] = kernels[i]
        res.loc[:, 'fold_num'] = idx
        
        pred_dfs.append(pd.concat([X_test.reset_index(), res.reset_index()], axis=1))

pred_dfs = pd.concat(pred_dfs)

null_mu = np.mean(y)
null_sd = np.std(y)


# -

# ## Evaluate kernels

# +
def ll(f_pred, sigma_pred, y_true):
    # log predictive density
    tmp = 0.5*np.log(2*np.pi*sigma_pred**2)
    tmp += (f_pred-y_true)**2/(2*sigma_pred**2)
    return tmp


sll = ll(pred_dfs['f_pred'], pred_dfs['sd_y'], pred_dfs['y'])
sll = sll - ll(null_mu, null_sd, pred_dfs['y'])
pred_dfs['msll'] = sll
pred_dfs['smse'] = (pred_dfs['f_pred']-pred_dfs['y'])**2/np.var(y)
pred_dfs.to_pickle('outputs/kernel_cv_fits.p')

msll = pred_dfs.groupby(['kernel'])['msll'].mean()
smse = pred_dfs.groupby(['kernel'])['smse'].mean()

summary = pd.DataFrame(smse).join(other=pd.DataFrame(msll), on=['kernel'], how='left')
summary.to_csv('outputs/kernel_cv_fits_summary.csv')
# -

summary

StarcoderdataJupyterScriptsDedupFiltered