<commit_before>from itertools import izip
from matplotlib import cm
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style('darkgrid')

colors = sns.color_palette(['#a6cee3','#1f78b4','#b2df8a','#33a02c','#fb9a99',
                         '#e31a1c','#fdbf6f','#ff7f00','#cab2d6'])

fs = 12
plt.locator_params(nbins=4)
fmts = ['.pdf','.svg','.png']
figheight = 4
######################################
#### make list of mutation effects across different periods
######################################
def mutation_list(flu = 'H3N2'):
    from glob import glob
    import pandas as pd
    flist = glob('../auspice/'+flu+'/*to*/HI_mutation_effects.tsv')
    mutation_effects = {}
    mutation_counts = {}
    for fname in flist:
        interval = fname.split('/')[-2]
        years = map(int, interval.split('to'))
        if years[1]-years[0] not in [10,11]:
            continue
        tmp_effects, tmp_counts = {},{}
        with open(fname) as infile:
            for line in infile:
                mut, val, count = line.strip().split('\t')
                tmp_effects[mut] = float(val)
                tmp_counts[mut] =  int(count)
        mutation_effects[interval]=tmp_effects
        mutation_counts[interval]=tmp_counts
    M = pd.DataFrame(mutation_effects)
    C = pd.DataFrame(mutation_counts)
    mutation_statistics = []
    for r,c in izip(M.iterrows(), C.iterrows()):
        vals = np.array([x for x in r[1] if not np.isnan(x)])
        mutation_statistics.append((r, c, vals.max(), vals.mean(), vals.std()))
    mutation_statistics.sort(key=lambda x:x[2], reverse=True)
    N = pd.DataFrame.from_items([(x[0][0],x[0][1]) for x in mutation_statistics]).transpose()
    D = pd.DataFrame.from_items([(x[1][0],x[1][1]) for x in mutation_statistics]).transpose()

    sep = '\t'
    le = '\n'
#    sep = ' & '
#    le = '\\\\ \n'
    ndigits = 2
    with open('HI_mutation_effects_all.tsv', 'w') as ofile:
        ofile.write("mutation\t"+sep.join(N.columns)+le)
        for r,c in izip(N.iterrows(), D.iterrows()):
            tmp = sep.join([r[0]]+map(str,[(round(x,ndigits),y) for x,y in izip(r[1], c[1])])).replace('nan','---')
            print tmp
            ofile.write(tmp+le)
    return N, D

######################################
#### make a figure that plots the cumulative antigenic change vs time
######################################
def cumulative_antigenic(myflu, n=10):
    from random import sample
    leaf_sample = sample([leaf for leaf in myflu.tree.leaf_iter()
                          if leaf.num_date>2014], n)<commit_msg>made mutation table function more flexible<commit_after>from itertools import izip
from matplotlib import cm
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style('darkgrid')

colors = sns.color_palette(['#a6cee3','#1f78b4','#b2df8a','#33a02c','#fb9a99',
                         '#e31a1c','#fdbf6f','#ff7f00','#cab2d6'])

fs = 12
plt.locator_params(nbins=4)
fmts = ['.pdf','.svg','.png']
figheight = 4
######################################
#### make list of mutation effects across different periods
######################################
def mutation_list(flu = 'H3N2', format='tsv', nconstraints=False):
    from glob import glob
    import pandas as pd
    flist = glob('../auspice/'+flu+'/*to*/HI_mutation_effects.tsv')
    mutation_effects = {}
    mutation_counts = {}
    for fname in flist:
        interval = fname.split('/')[-2]
        years = map(int, interval.split('to'))
        if years[1]-years[0] not in [10,11]: # skip all non-10year interval
            continue
        tmp_effects, tmp_counts = {},{}
        with open(fname) as infile:
            for line in infile:
                mut, val, count = line.strip().split('\t')
                tmp_effects[mut] = float(val)
                tmp_counts[mut] =  int(count)
        mutation_effects[interval]=tmp_effects
        mutation_counts[interval]=tmp_counts
    M = pd.DataFrame(mutation_effects)
    C = pd.DataFrame(mutation_counts)
    mutation_statistics = []
    for r,c in izip(M.iterrows(), C.iterrows()):
        vals = np.array([x for x in r[1] if not np.isnan(x)])
        mutation_statistics.append((r, c, vals.max(), vals.mean(), vals.std()))
    mutation_statistics.sort(key=lambda x:x[2], reverse=True)
    N = pd.DataFrame.from_items([(x[0][0],x[0][1]) for x in mutation_statistics]).transpose()
    D = pd.DataFrame.from_items([(x[1][0],x[1][1]) for x in mutation_statistics]).transpose()

    if format=='tex':
        sep = ' & '
        le = '\\\\ \n'
    else:
        sep = '\t'
        le = '\n'
    ndigits = 2
    with open('HI_mutation_effects_all.tsv', 'w') as ofile:
        ofile.write("mutation\t"+sep.join(N.columns)+le)
        for r,c in izip(N.iterrows(), D.iterrows()):
            if nconstraints:
                tmp = sep.join([r[0]]+map(str,[(round(x,ndigits),y) for x,y in izip(r[1], c[1])])).replace('nan','---')
            else:
                tmp = sep.join([r[0]]+map(str,[round(x,ndigits) for x,y in izip(r[1], c[1])])).replace('nan','---')
            print tmp
            ofile.write(tmp+le)
    return N, D

######################################
#### make a figure that plots the cumulative antigenic change vs time
######################################
def cumulative_antigenic(myflu, n=10):
    from random import sample
    leaf_sample = sample([leaf for leaf in myflu.tree.leaf_iter()
                          if leaf.num_date>2014], n)
{'fraction_non_alphanumeric': 0.20759861843301217, 'fraction_numerical': 0.021632430467187783, 'mean_word_length': 3.0426157237325495, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3689468', 'n_tokens_mistral': 1960, 'n_tokens_neox': 1858, 'n_words': 399}