# ---
# 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
# ---

# ## Statistical analysis of symptom's change

# %config Completer.use_jedi = False
import pymc3 as pm
import numpy as np
import pandas as pd
import statsmodels
import matplotlib.pyplot as plt
import seaborn as sns
import pingouin as pg
import statsmodels.api as sm
import statsmodels.formula.api as smf
from sklearn import preprocessing
import arviz as az

# load file
df = pd.read_csv('KPE_bdi_pcl.csv') 
df.head()

# change names of redcap_enet_name_y
categories_old = ['screening_selfrepo_arm_1', 'visit_1_arm_1','visit_7_week_follo_arm_1',
                  '30_day_follow_up_s_arm_1', '90_day_follow_up_s_arm_1']
categoriesNew=['screen','beforeTrt','afterTrt','FU30','FU90']
cat = pd.Categorical(df.redcap_event_name_y, categories=categories_old)
# set new variable with new categories names
df['time'] = cat.rename_categories(categoriesNew)
df = df.drop('redcap_event_name_y', axis=1) # remove old naming

# remove 90-days (not enough subjects)
df = df[(df.time!='FU90') & (df.time!='screen')]
# generate for something else
df_med = df[(df.time=='beforeTrt')][['scr_id','med_cond', 'scr_dem_gender', 'scr_dem_age']]
df_med.groupby('med_cond').sum()

# ## Demographics

df.groupby(['med_cond','time']).describe()

df_dem = df[['scr_id','med_cond','time','pcl5_total']]
df_dem.groupby(['med_cond','time']).describe()

# +
## simple t test for differences in age etc.
pg.ttest(df_med.scr_dem_age[df_med.med_cond=='Ket'], df_med.scr_dem_age[df_med.med_cond=='Mid'], correction=False)
pg.ttest(df_dem.pcl5_total[(df_dem.time=='beforeTrt') & (df_dem.med_cond=='Ket')],
        df_dem.pcl5_total[(df_dem.time=='beforeTrt') & (df_dem.med_cond=='Mid')], correction=False)
pg.ttest(df_dem.pcl5_total[(df_dem.time=='afterTrt') & (df_dem.med_cond=='Ket')],
        df_dem.pcl5_total[(df_dem.time=='afterTrt') & (df_dem.med_cond=='Mid')], correction=False)

pg.ttest(df_dem.pcl5_total[(df_dem.time=='FU30') & (df_dem.med_cond=='Ket')],
        df_dem.pcl5_total[(df_dem.time=='FU30') & (df_dem.med_cond=='Mid')], correction=False)
# -

## length of diagnosis
lng = pd.read_csv('ptsd_lengthDiagnosis.txt', sep=' ', index_col=False)
lng
lng_all = pd.merge(lng, df_med, left_on='subject', right_on='scr_id')
lng_all[['med_cond','TimeFromEvent']].groupby('med_cond').describe()

pg.ttest(lng_all.TimeFromEvent[lng_all.med_cond=='Ket'],lng_all.TimeFromEvent[lng_all.med_cond=='Mid'], correction=False)

# ## Descriptives

df_clean = df[['scr_id','med_cond','time','pcl5_total']]
# remove screening timepoint - its redundant here
df_clean = df_clean[df_clean.time!='screen'] 
# recategorize time because we removed something
df_clean.time = pd.Categorical(df_clean.time, categories=['beforeTrt','afterTrt','FU30'])
wideDf = df_clean.pivot(index=['scr_id'], columns=['time'], values=['pcl5_total'])
#
wideDf = wideDf.pcl5_total
wideDf = wideDf.merge(df_med, left_on='scr_id', right_on='scr_id', how='outer')
wideDf
sns.pairplot(wideDf, hue="med_cond")

## check changes in symtpoms
wideDf['delta30'] = wideDf.FU30 - wideDf.beforeTrt
print(f'Number of participants improving 5 points or more is {len(wideDf[wideDf.delta30<=-5])}')
print(f'Number of participants improving 10 points or more is {len(wideDf[wideDf.delta30<=-10])}')

wideDf[['beforeTrt','afterTrt','FU30','med_cond']].describe()

# need to scale pcl variable
df_clean['pcl5_total_scaled'] = preprocessing.scale(df_clean.pcl5_total)
df_clean.describe().round(3)

## Start with simple GLM models (using statsmodel)
glm = smf.mixedlm('pcl5_total_scaled ~ C(time) * C(med_cond)', data=df_clean, groups = df_clean.scr_id, 
                 missing='drop')
model1 = glm.fit()
print(model1.summary())

# so in simple mixed model we see effect for time after treatment and in 30 and 90 days follow-up
# use pinguoin for ANOVA of the mixed model
pg.mixed_anova(data=df_clean, dv='pcl5_total', within='time', between='med_cond', subject='scr_id').round(3)


# ## Go Bayes
# - Organize indexes and define model
# - Try with/without hyperpriors

# set groups as 0 and 1
group = {'Ket': 1,'Mid': 0} 
df_clean['groupIdx'] =[group[item] for item in df_clean.med_cond] 
# set all the rest
df_clean['scr_id'] = df_clean.scr_id.astype('category')
sub_idx = df_clean.scr_id.cat.codes.values # get county category index values (613 x 1) vec with nums [0-7]
print(sub_idx)
n_subs = len(np.unique(df_clean.scr_id))
print(n_subs)
# set time as index also
#df_clean.groupIdx = pd.Categorical(df_clean.groupIdx)
df_clean.time = pd.Categorical(df_clean.time, 
                    categories=['beforeTrt','afterTrt','FU30'])
timeIDX = df_clean.time.cat.codes.values
print(timeIDX)

# lets define our own model using pymc3
with pm.Model() as model1:
  
    
    a = pm.Normal('a', mu=0, sd=1, shape=n_subs)
       
    b = pm.Normal('b', mu=0, sd=1, shape=3) # time
    
    bMed = pm.Normal('bMed', mu=0, sd=.5) # two groups
  
    # Model error
    eps = pm.HalfCauchy('eps', 3)

    y_hat = a[sub_idx] + b[timeIDX] + bMed*df_clean.groupIdx 

    # Data likelihood
    likelihood = pm.Normal('likelihood', mu=y_hat,
                           sigma=eps, observed=df_clean.pcl5_total_scaled)
    prior_checks = pm.sample_prior_predictive(samples=100, random_seed=123)
    trace = pm.sample(draws=2000, tune=2000, target_accept=0.9, return_inferencedata=False)

sns.displot(prior_checks['bMed'])

az.summary(trace)

az.plot_trace(trace)

with model1:
    ppc = pm.sample_posterior_predictive(
        trace, random_seed=124
    )

# calculate difference between before/after
afterDiff = trace["b"][:,1] - trace["b"][:,0]
sns.distplot(afterDiff)
# lets see what are the chances of no difference or more than zero
sum(afterDiff[1000:]>=0) / len(afterDiff[1000:])  # no chance so significant difference

# lets test for 30 days F/U
Diff30 = trace["b"][:, 2] - trace["b"][:, 0]
sns.distplot(Diff30)
# lets see what are the chances of no difference or more than zero
sum(Diff30>=0) / len(Diff30[:]) 

sns.distplot(trace["b"][:-1000,0], color = 'blue', label='before')
sns.distplot(trace["b"][:-1000,1], color = 'orange', label = 'after')
sns.distplot(trace["b"][:-1000,2], color = 'green', label = '30Days')
plt.legend()

# test the same on group
#groupDiff = trace['bMed'][:,1] - trace['bMed'][:,0]
sns.distplot(trace["bMed"])

# ## Plotting

# +
## Plot symtpoms
sns.set_style("whitegrid")
ax = sns.catplot('time','pcl5_total', data=df_clean, kind='point', ci=68, n_boot=5000,
           set_ylabel=("Total PCL5 score"))

ax.set_ylabels("Total PCL5 score")
ax.set_xlabels("Time of Assesment")
#ax.savefig("SymtpomsChange.png", dpi=300)
# -

## Plot symtpoms
sns.set_style("whitegrid")
ax = sns.catplot('time','pcl5_total',hue = 'med_cond', data=df_clean, kind='point', ci=68, n_boot=5000,
           set_ylabel=("Total PCL5 score"))
ax.set_ylabels("Total PCL5 score")
ax.set_xlabels("Time of Assesment")
plt.hlines(y=33, xmin=-0.2, xmax=2, linestyles='--')

StarcoderdataJupyterScriptsDedupFiltered