03_full_length_different_probabilities_reduced_AA.py

import os
import shutil
import pandas as pd
import numpy as np
from multiprocessing import Pool, cpu_count
from functools import partial
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_context("paper", font_scale=2.0)
sns.set_style('whitegrid')

from functions import *

from pandarallel import pandarallel

pandarallel.initialize(progress_bar=False)

with SequenceFile("data/uniprot-9606.fasta", digital=True, alphabet=alphabet) as seq_file:
    sequences = list(seq_file)


uniprot = fasta_reader('data/uniprot-9606.fasta')


# final
aa_knowns = ['L', 'LS', 'C', 'K', 'Y', 'CK', 'KY', 'YC', 'LSE', 'CKY', \
'W', 'WM', 'WMC', 'LSEA', 'WMCH', 'WMCHY', 'LSEAG', ]

prob_range = [0.9, 0.8, 0.6, 0.4, 0.2, 0.1, 0.09, 0.08, 0.07, 0.06, 0.05]

for aa_known in aa_knowns:
    for prob in prob_range:
        
        # output file name
        fname = f'full_length_{aa_known}_max_prob_{prob}_del.pkl.gz'
        if not os.path.isfile('results/scan_results/'+fname):
        
            temp_df = uniprot.copy()

            temp_df['temp'] = temp_df.values.tolist()
            temp_df['posteriors'] =  temp_df.temp.parallel_apply(\
                                 lambda x: generate_reads(seq=x[1], aa_known=aa_known, max_prob=prob))


            temp_df['temp1'] = temp_df[['Accession', 'posteriors']].values.tolist()
            temp_df['hmm'] = temp_df.temp1.parallel_apply(\
                     lambda x: hmm_build(x[1], x[0], f'full_length_{x[0]}_{aa_known}_max_prob_{prob}'))


            hmms = np.array_split(temp_df.hmm, 10)
            hmms = [i.tolist() for i in hmms]

            results = []
            for i, v in enumerate(hmms):
                results.append(score(v, sequences, background))
            #             print(f'\t\t\t\tdone: {i}', end='\r')
            fname = f'full_length_{aa_known}_max_prob_{prob}_del.pkl.gz'

            res_df = pd.concat(results)
            res_df.to_pickle('results/scan_results/' + fname)
            del temp_df, res_df, hmms
        shutil.rmtree('temp/', ignore_errors=True)
        try:
            os.makedirs('temp/')
        except FileExistsError:
            pass