<jupyter_start><jupyter_text>Mutual information<jupyter_code>def calculate_MI(hit_data):
    MI_by_orf_pair = np.zeros((len(orf_names), len(orf_names)))
    for p in plates:
        for s in strains:
            well_indices = [w for w in range(len(wells)) if wells[w][:2]==p and well2strain[wells[w]]==s]
            tdat = hit_data[:,well_indices]
            jointprobs = dict()
            independence_probs = dict()
            jointprobs['both'] = (tdat @ tdat.T) / len(well_indices)
            pseud = 0.1 / len(well_indices)
            i_probs = np.tile(np.diagonal(jointprobs['both']), (tdat.shape[0], 1)).T
            j_probs = np.tile(np.diagonal(jointprobs['both']), (tdat.shape[0], 1))
            jointprobs['i_only'] = i_probs - jointprobs['both']
            jointprobs['j_only'] = j_probs - jointprobs['both']
            jointprobs['neither'] = 1 - (i_probs + j_probs - jointprobs['both'])
            independence_probs['both'] = (i_probs + pseud) *  (j_probs + pseud)
            independence_probs['i_only'] = (i_probs + pseud) *  (1 - j_probs - pseud)
            independence_probs['j_only'] = (1 - i_probs - pseud) *  (j_probs + pseud)
            independence_probs['neither'] = (1 - i_probs - pseud) *  (1 - j_probs - pseud)
            for outcome in jointprobs:
                MI_by_orf_pair += np.nan_to_num(np.log2(jointprobs[outcome]/independence_probs[outcome])*jointprobs[outcome])
    return MI_by_orf_pair

def draw_data(probs):
    draw = np.random.rand(probs.shape[0], probs.shape[1])
    return np.where(draw<probs, 1, 0)

def simulate_MI(n, probs_to_sim):
    # keep record of total MI and max MI for each run
    MI_rec = []
    for run in range(n):
        sim = draw_data(probs_to_sim)
        MI_mat = calculate_MI(sim)
        MI_mat = MI_mat*np.tril(np.ones_like(MI_mat), k=-1) # set diag and up to zero
        MI_rec.append([np.sum(MI_mat), np.max(MI_mat)])
    return MI_rec, MI_mat, sim


# a dataframe with hits and multiplicity for each ORF in the yeast genome
orf_hit_df = pd.read_csv('../Output/WGS/combined_option/gene_hit_data.tsv', delimiter='\t')
pop_hits = np.sum(np.clip(orf_hit_df[wells], 0, 1), axis=0)
multi_hit_data = orf_hit_df[orf_hit_df['pops_hit']>=6]
orf_names = list(multi_hit_data['ORF'])
mh_dat = np.clip(np.array(multi_hit_data[wells]), 0, 1)
# here we fit probabilities of mutations happening in each background-environment combination using the data
model4_probs = simpler_split_probs(pop_hits, mh_dat, [[w for w in range(len(wells)) if wells[w][:2]==p and well2strain[wells[w]]==s] for p in plates for s in strains])
# then we simulate mutations using this data, and calculate MI between multi-hit ORFs, recording the MI sum and max MI
sim_res, sim_example, sim = simulate_MI(1000, model4_probs)
# we calculate MI between multi-hit orfs in our data
MI_results = calculate_MI(mh_dat)
MI_results2 = MI_results*np.tril(np.ones_like(MI_results), k=-1) # set diag and up to zero
print('Data MI sum:', np.sum(MI_results2), 'Data MI max:', np.max(MI_results2))
f, subs = pl.subplots(1, 2, figsize=(5, 2), dpi=300)
pl.subplots_adjust(wspace=0.5)
subs[0].hist([s[0] for s in sim_res], bins=25, facecolor='#333333')
subs[0].annotate('Data\n$MI_{tot}$', xy=(7333, 0), xytext=(7333, 75), xycoords="data", ha='center', annotation_clip=False, arrowprops=dict(arrowstyle='->', lw=1), fontsize=9)
subs[0].set_xlim([6000, 7800])
subs[1].hist([s[1] for s in sim_res], facecolor='#333333')
subs[1].annotate('Data\n$MI_{max}$', xy=(2.03, 0), xytext=(2.03, 200), xycoords="data", ha='center', annotation_clip=False, arrowprops=dict(arrowstyle='->', lw=1), fontsize=9)
subs[0].set_ylabel('No. simulations')
subs[0].set_xlabel('$MI_{tot}$')
subs[1].set_ylabel('No. simulations')
subs[1].set_xlabel('$MI_{max}$')
sns.despine()<jupyter_output><empty_output><jupyter_text>Let's look at which mutations have the most MI even though it's not significant:<jupyter_code>tmp = []
for i in range(len(orf_names)):
    for j in range(len(orf_names)):
        if i > j:
            tmp.append([o2g.get(orf_names[i], orf_names[i]), o2g.get(orf_names[j], orf_names[j]), MI_results2[i,j]])
td = pd.DataFrame(tmp, columns=['ORF1', 'ORF2', 'MI'])
td.sort_values(by='MI', ascending=False).iloc[:10]
def multinomial_sim(weights, max_hits, reps=10000):
    """
    weights are the weighted number of mutations in each populations
    max hits is the maximum number of hits we want to consider
    """
    sm = np.zeros((max_hits+1, len(weights)+1))
    for i in range(2, max_hits+1):
        tmp = []
        for rep in range(reps):
            tmp_pops_hit = len([j for j in np.random.multinomial(i, pvals=weights) if j != 0])
            sm[i][tmp_pops_hit] += 1
    
    smn = []
    for row in sm:
        smn.append(row/np.nansum(row))
    smn = np.nan_to_num(smn, nan=0)
    
    return sm, smn

def get_gene_hit_dispersion_matrix(genes_hit_mat):
    """
    input: matrix with genes as rows, wells as columns, w the # times each gene is hit in each pop
    output: matrix with no. of mutations as rows and no. unique pops hit as columns, and the same matrix where rows sum to 1
    """
    no_hits = np.sum(genes_hit_mat, axis=1)
    pops_hit = np.sum(np.clip(genes_hit_mat, 0, 1), axis=1)
    bm = np.zeros((np.max(no_hits)+1, genes_hit_mat.shape[1]+1))
    for i in range(len(no_hits)):
        bm[no_hits[i]][pops_hit[i]] += 1

    bmn = []
    for row in bm:
        bmn.append(row/np.nansum(row))
    bmn = np.nan_to_num(bmn, nan=0)
    
    return bm, bmn
    

f, subps = pl.subplots(3, 3, figsize=(12, 9), dpi=100)
cbar_ax = f.add_axes([.99, .3, .02, .4])
cbar_ax.set_title('Excess\nProbability', fontsize=14)
for p in range(3):
    plate = plates[p]
    for s in range(3):
        sub = subps[s][p]
        strain = strains[s]
        use_wells = [w for w in wells if w[:2]==plate and well2strain[w]==strain]
        mutation_mat = np.array(orf_hit_df[use_wells])
        pop_weights = np.sum(mutation_mat, axis=0) / np.sum(mutation_mat)
        max_hits = np.max(np.sum(mutation_mat, axis=1))
        sim_counts, sim_res = multinomial_sim(pop_weights, max_hits)
        bm, bmn = get_gene_hit_dispersion_matrix(mutation_mat)
        
        print(max_hits, bm[2:,:].shape, sim_res.shape)
        mymap = sns.diverging_palette(240, 10, as_cmap=True)
        use = bmn-sim_res
        mask=np.where(bm==0, np.ones_like(use), np.zeros_like(use))
        x = np.min([len(pop_weights)+1, max_hits+1])
        y = max_hits+1
        if p == 2 and s == 0:
            sns.heatmap(use[2:y,1:x], mask=mask[2:y,1:x], cmap=mymap, annot=bm[2:y,1:x], annot_kws={'size': 7}, fmt='g', ax=sub, center=0, vmax=1, vmin=-1, cbar_ax=cbar_ax)
        else:
            sns.heatmap(use[2:y,1:x], mask=mask[2:y,1:x], cmap=mymap, annot=bm[2:y,1:x], annot_kws={'size': 7}, fmt='g', ax=sub, center=0, vmax=1, vmin=-1, cbar=False)
        sub.set_xticks([i+0.5 for i in range(x-1)])
        sub.set_yticks([i+0.5 for i in range(y-2)])
        sub.set_xticklabels(list(range(1,x)))
        sub.set_yticklabels(labels=list(range(2,y)), rotation='horizontal')
        sub.set_ylim([y-2, -0.5])
        sub.set_xlim([-0.5, x-1])

        if s == 0:
            sub.set_title(plate2env[plate], y = 1.1, fontsize=14)
        elif s == 2:
            sub.set_xlabel('No. populations\nwith mutation')
        if p == 0:
            sub.annotate(strains_for_print[strain], xy=(-0.5, 0.5), xycoords="axes fraction", ha='right', va='center', fontsize=14)
            sub.set_ylabel('No. mutations\nin gene')
sns.despine()<jupyter_output>10 (9, 13) (11, 13)
3 (2, 7) (4, 7)
10 (9, 13) (11, 13)
6 (5, 13) (7, 13)
3 (2, 7) (4, 7)
10 (9, 13) (11, 13)
8 (7, 13) (9, 13)
5 (4, 7) (6, 7)
10 (9, 13) (11, 13)
{'fraction_non_alphanumeric': 0.16159129195283142, 'fraction_numerical': 0.036283529869120125, 'mean_word_length': 3.489819662594532, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 12, '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}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '118089', 'license': 'MIT', 'n_tokens_mistral': 3240, 'n_tokens_neox': 3054, 'n_words': 667, 'path': '.ipynb_checkpoints/Parallelism_and_MI_testing-checkpoint.ipynb', 'repo_name': 'ChiaWei-Liao/VLTE_PIPELINES', 'size': '7717'}