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

#
# # Aggregating vs averaging variant frequencies within groups
# ___

# Load necessary packages and temperature experiment data file 

import pandas as pd
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as sf
import matplotlib.pyplot as plt

from pi_for_temperature_notebook import *

data = pd.read_csv("variant-calls_temperature-expt_2020-04-07_tables.csv", header=0, index_col=0)

# The above errors can be ignored because they arise from 32$ ^{o} $ data, which we will filter out. Also filter out variants that are represented in only one technical replicate, as well as variants that do not have a frequency of 3% or more in at least one technical replicate.

data_no_32 = data.loc[data.temperature != 32]
data_no_32_both_present = data_no_32.dropna(subset=['ADReplicateA','ADReplicateB'])
filtered_data = pd.concat([data_no_32_both_present.loc[data_no_32_both_present.freqPropReplicateA >= 0.03],
                           data_no_32_both_present.loc[data_no_32_both_present.freqPropReplicateB >= 0.03]]).drop_duplicates()


# Previous analyses has shown a tight correlation between technical replicates, as expected by the manner of the replication. Aggregate technical replicate read counts to get a single variant frequency for each sequence variant.

# +
#set new DP and AD columns with aggregate of Replicate RD and AD values
#This aggregates only counts of reads that passed VarScan filters 
filtered_data['DP'] = filtered_data['ADReplicateA'] + filtered_data['ADReplicateB'] +filtered_data['RDReplicateA'] + filtered_data['RDReplicateB']
filtered_data['AD'] = filtered_data['ADReplicateA'] + filtered_data['ADReplicateB']

#calculate frequency of aggregated read counts
filtered_data['freqProp'] = filtered_data['AD'] / filtered_data['DP']

# -

# Add factor data for lineage, species, and segment (for DNA-A and DNA-B segment designations)

# +
#lineage: set to 0 then assign each unique lineage to an unique integer
filtered_data['lineage_factor'] = 0
j = 0
for lineage in filtered_data.lineage.unique():
    filtered_data.loc[filtered_data.lineage == lineage,['lineage_factor']] = j
    j += 1

#species: Set ACMV = 0 and EACMCV = 1
filtered_data['species'] = 0
filtered_data.loc[filtered_data.chrom == 'EACMCV DNA-A',['species']] = 1
filtered_data.loc[filtered_data.chrom == 'EACMCV DNA-B',['species']] = 1

#segment: Set DNA-A = 0 and DNA-B = 1
filtered_data['segment'] = 0
filtered_data.loc[filtered_data.chrom == 'ACMV DNA-B',['segment']] = 1
filtered_data.loc[filtered_data.chrom == 'EACMCV DNA-B',['segment']] = 1

# -

# Load the options for calculation of pi then print them out.

options = get_temp_args()
#set new values for options as necessary
options['coverage'] = 'DP'
options['perSite'] = False

# +
pis = {}

pi_df = get_group_pis(filtered_data,
                     options=options,
                     group_by=['passage','temperature','species','segment','lineage_factor','plantID'])
# -

lm = sf.ols('pi ~ passage + temperature + (C(species)/C(segment)) + C(lineage_factor)',data=pi_df).fit(cov_type='HC1')
print("Least squares summary:")
print(lm.summary2())

# There is strong multicollinearity in the model. Let's inspect a few variables

# +
p = pi_df.boxplot(column='pi',
             by=['passage','temperature'],
             patch_artist=True,
             grid=False,
             return_type='dict')

plt.ylabel('pi',fontsize=16)
plt.xlabel('(passage,temperature)',fontsize=16)
ylim = plt.ylim()

plt.vlines(x=[2.5,4.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

plt.xticks(fontsize=14)
plt.suptitle('')
plt.title('')

for median in p[0]['medians']:
    median.set_color('k')
# -

# Looks like we could omit temperature from the model.

# +
p = pi_df.boxplot(column='pi',
             by=['passage','lineage_factor'],
             patch_artist=True,
             grid=False,
             return_type='dict',
             rot=30)

plt.ylabel('pi',fontsize=16)
plt.xlabel('(passage,lineage)',fontsize=16)
ylim = plt.ylim()

plt.xticks(fontsize=14,ha='right')
plt.suptitle('')
plt.title('')

ylim = plt.ylim()

plt.vlines(x=[5.5,11.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

for median in p[0]['medians']:
    median.set_color('k')
# -

# Looks like we can onit lineage from the model as well.

# +
p = pi_df.boxplot(column='pi',
             by=['species','segment'],
             patch_artist=True,
             grid=False,
             return_type='dict',
             rot=30)

plt.ylabel('pi',fontsize=16)
plt.xlabel('',fontsize=16)
ylim = plt.ylim()

plt.vlines(x=[2.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

plt.xticks([1,2,3,4],
           ['ACMV DNA-A','ACMV DNA-B','EACMCV DNA-A','EACMCV DNA-B'],
           fontsize=14,
           ha='right')
plt.suptitle('')
plt.title('')

for median in p[0]['medians']:
    median.set_color('k')
# -

# Let's remove temperature and lineage from the model

lm = sf.ols('pi ~ passage + (C(species)/C(segment)) ',data=pi_df).fit(cov_type='HC1')
print("Least squares summary:")
print(lm.summary())

# Looks like a moderate improvement gained by reducing variables. Let's aggregate the data to reduce noise.

# +
agg_data = aggregate_groups(filtered_data,['species','segment','pos','alt','passage','ref'])

agg_data.reset_index(inplace=True)
agg_data['pos'] = agg_data.pos.astype(int)   
# -

pi_df = get_group_pis(agg_data,
                     options=options,
                     group_by=['species','segment','passage'])

# +
lm = sf.ols('pi ~ passage + C(species)/C(segment)',data=pi_df).fit(cov_type='HC1')
print("Least squares summary:")
print(lm.summary())

print("\nAnova table:")
table = sm.stats.anova_lm(lm)
print(table)
print("\n")
for ind in table.index:
    
    print("{} explains \t {:.2%} of variance".format(ind,table.loc[ind,'sum_sq'] / table.sum_sq.sum()))
# -

# ### Now let's look at averaging pi within groups

# Let's start with the raw data again, and use the mean between technical replicates as a single frequency per variant. `freqPropMeanNoNA` already exists as the mean of ReplicateA and ReplicateB frequencies, so we can use `filtered_data` as-is. We'll first calculate pi from the average values and inspect the analysis results.

# +
#set the frequency for column name of the mean of the replicates
options['frequency'] = 'freqPropMeanNoNA'

pi_df = get_group_pis(filtered_data,
                     options=options,
                     group_by=['passage','temperature','species','segment','lineage_factor','plantID'])

lm = sf.ols('pi ~ passage + temperature + (C(species)/C(segment)) + C(lineage_factor)',data=pi_df).fit(cov_type='HC1')
print("Least squares summary:")
print(lm.summary())
# -

# Now plot data

# +
p = pi_df.boxplot(column='pi',
             by=['passage','temperature'],
             patch_artist=True,
             grid=False,
             return_type='dict')

plt.ylabel('pi',fontsize=16)
plt.xlabel('(passage,temperature)',fontsize=16)
ylim = plt.ylim()

plt.vlines(x=[2.5,4.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

plt.xticks(fontsize=14)
plt.suptitle('')
plt.title('')

for median in p[0]['medians']:
    median.set_color('k')

# +
p = pi_df.boxplot(column='pi',
             by=['passage','lineage_factor'],
             patch_artist=True,
             grid=False,
             return_type='dict',
             rot=30)

plt.ylabel('pi',fontsize=16)
plt.xlabel('(passage,lineage)',fontsize=16)
ylim = plt.ylim()

plt.xticks(fontsize=14,ha='right')
plt.suptitle('')
plt.title('')

ylim = plt.ylim()

plt.vlines(x=[5.5,11.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

for median in p[0]['medians']:
    median.set_color('k')

# +
p = pi_df.boxplot(column='pi',
             by=['species','segment'],
             patch_artist=True,
             grid=False,
             return_type='dict',
             rot=30)

plt.ylabel('pi',fontsize=16)
plt.xlabel('',fontsize=16)
ylim = plt.ylim()

plt.vlines(x=[2.5],ymin=ylim[0],ymax=ylim[1])
plt.ylim(ylim)

plt.xticks([1,2,3,4],
           ['ACMV DNA-A','ACMV DNA-B','EACMCV DNA-A','EACMCV DNA-B'],
           fontsize=14,
           ha='right')
plt.suptitle('')
plt.title('')

for median in p[0]['medians']:
    median.set_color('k')
# -

# Looks like there is no great difference in averaging or aggregating technical replicate data.

# +
avg_data = average_groups(filtered_data,['species','segment','pos','alt','passage','ref'])

avg_data.reset_index(inplace=True)
avg_data['pos'] = avg_data.pos.astype(int)   


avg_pi_df = get_group_pis(agg_data,
                     options=options,
                     group_by=['species','segment','passage'])

# +
lm_avg = sf.ols('pi ~ passage + C(species)/C(segment)',data=avg_pi_df).fit(cov_type='HC1')
print("Least squares summary:")
print(lm_avg.summary())

print("\nAnova table:")
table_avg = sm.stats.anova_lm(lm_avg)
print(table)

# +
print("Aggregated groups:")
for ind in table.index:
    
    print("{} \t\t {:.2%} of variance".format(ind,table.loc[ind,'sum_sq'] / table.sum_sq.sum()))

print("\nAveraged groups:")
for ind in table.index:
    
    print("{} \t\t {:.2%} of variance".format(ind,table_avg.loc[ind,'sum_sq'] / table_avg.sum_sq.sum()))
# -

# ### Looks there is not a great difference between averaging or aggregating technical replicate data. Also, removing temperature increases model performance. 
# ___
# Not shown here: analyses running `linearmodels.PanelOLS` panel analyses on the temperature experiment fail when including both *temperature* and *passage* due to collinearity.  

StarcoderdataJupyterScriptsDedupFiltered