<jupyter_start><jupyter_text>Set up notebook<jupyter_code># Import packages
%matplotlib inline
import os
import re
import string
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import numpy as np
# Define functions (from Kimiko)
def get_ref_and_align(filename):
    na = string.ascii_uppercase+'+ '
    with open(filename) as f:
        seqs = []
        alignments = []
        new_ref = ''
        new_align = ''
        for i, line in enumerate(f):
            if not line:
                pass
            elif line[0] == '>':
                if new_ref:
                    seqs.append(new_ref)
                    new_ref == ''
                if new_align:
                    alignments.append(new_align)
#                     alignments.append(new_align)
                    new_ref = ''
                    new_align = ''
            elif line[0:6] == 'Query ' and not seqs:
                new_ref += line.strip().split()[2]
            elif line.count(' ') >= 12 and all(i in na for i in line.strip()):
                new_align += line.rstrip('\n').lstrip(' ')
        alignments.append(new_align)
        seqs.append(alignments)
        return seqs  
    
def find_motif_site(motif, MaCS):#array of HOG components sequences from reference genome
    # Finds all of the Hog1 concensus sites in the ref
    # str.find() only returns the first instance of the site
    return [m.start() for m in re.finditer(motif,MaCS)]

def number_concensus_sites(seqs):
    sp, tp = find_concensus_site(seqs[0])
    ref_num = len(sp+tp)
    align_num = []
    for seq in seqs[1]:
        sp, tp = find_concensus_site(seq)
        num = len(sp+tp)
        align_num.append(num)
    MaCS_hist = plt.hist(align_num)
#     MaCS_hist[0].set_color('r')
#     plt.hist().get_xticklabels()[0].set_color('red') 
    plt.show()
#     return [ref_num,align_num]

# def plot_concensus_sites(seqs):
#     sp, tp = find_concensus_site(seqs[0])
#     plt.bar(sp,[1,1,1],color='g')
#     plt.show()

def compare_concensus_sites(seqs,filename):#array of HOG components sequences from reference genome
    # Finds all of the Hog1 concensus sites in the ref
    # str.find() only returns the first instance of the site
#     for x in seqs:
#         print(x)
#         _ref = [m.start() for m in re.finditer('SP', x)]
#         print(_ref)
    ref_sp, ref_tp = find_concensus_site(seqs[0])
    if ref_sp or ref_tp:
#     ref['ID'] = 'ref'
        sps = []
        tps = []
        for seq in seqs[1]:
            seq_sp, seq_tp = find_concensus_site(seq)
    #         seq['ID'] = str(i)
            sps += seq_sp
            tps += seq_tp
    #     print(sps)
        plt.clf()
        plt.rcParams['figure.figsize'] = [20.0, 8.0]
        if sps:
            sp_labels, sp_count = zip(*Counter(sps).items())
            plt.bar(sp_labels, sp_count, width=1, color = 'g', label='SP')
        if tps:
            tp_labels, tp_count = zip(*Counter(tps).items())
            plt.bar(tp_labels, tp_count, width=1, color='b', label='TP')
    #     print(sp_count)
    #     indices = np.arange(len(labels))
        plt.axhline(y=len(seqs[1]),color='r', linestyle='--', label='Total number of sequences analyzed')
        plt.legend(bbox_to_anchor=[1, 0.5], loc='center left')
        plt.xlim(0,len(seqs[0]))
        plt.xlabel('Amino Acid Sequence Position')
        plt.title(filename[0].upper() + filename[1:4])
        plt.ylabel('Number of sequences')
        plt.show()
    else:
        print('Bug in code')
#         ref = ref.append(seq, ignore_index=True)
#     _aligns = [[m.start() for m in re.finditer('SP', seqs[1])], [m.start() for m in re.finditer('TP', seqs[1])]]
#     return [m.start() for m in re.finditer('SP',seqs)], [m.start() for m in re.finditer('TP',seqs)] 
#     print(ref)
#     sns.barplot(x='locations', y)
#     return ref, seq

def plot_seq_MCS(seqs,filename):
    seqs = [seqs[0]] + seqs[1]
    grid =[]
    for x in seqs:
        temp = []
        for i in range(len(x)-1):
        #     print(x[i:i+2])
            if x[i:i+2] == 'SP':
                temp.append(2)
            elif x[i:i+2] == 'TP':
                temp.append(1)
            else:
                temp.append(0)
        grid.append(temp)
    colors = ['#DCDCDC','#ee9537','#3db7d6']
    grid_df = pd.DataFrame(grid)
    
    plt.clf()
    plt.rcParams['figure.figsize'] = [30.0, 20.0]
    sns.heatmap(grid_df, cmap=colors, cbar=False)
    plt.title(filename[0].upper() + filename[1:4]) #add percent conserved
    plt.xlabel('Amino Acid Sequence Position')
    plt.ylabel('Sequence number') # change to strain name!
    plt.show()
    
def get_MaCS_df(seq):
    refs_sp = [m.start() for m in re.finditer('SP', seq)]
    refs_tp = [m.start() for m in re.finditer('TP', seq)]
    _refs_locs = refs_sp + refs_tp
    _ = {'locations': _refs_locs, 'type': ['SP' for x in range(len(refs_sp))] + ['TP' for x in range(len(refs_tp))]}
    df = pd.DataFrame.from_dict(_)
    return df<jupyter_output><empty_output><jupyter_text>Make pandas data frame<jupyter_code>#### USER DEFINED THINGS ####
# Directory with protein sequences
dir_to_check = '/Users/AmyAllen/Documents/A-pipeline-for-fast-conserved-protein-motif-identification/aligned_cp_for_counting/trimmed_data/'
# Motif sequence 
motif = 'TP'
#### END USER DEFINED THINGS #### 

# Set all empty lists to be filled in for loops
protein_list = []
short_protein_list = []
num = []
motif_loc = []
all_locations = []
motif_list = []
n_loc = []
conserved = []

# Get all files from direcotry 
files = os.listdir(dir_to_check)

# For loop to loop through all files in directory 
for filename in files:
    seq = get_ref_and_align(dir_to_check+filename) # Import protein sequence from filename
    protein = filename.split('.')[0].split('_')[0] # Get name of protein (string before the _ or .)
    
    # For loop to loop through each instance of the protein 
    for i in range(len(seq[1])):
        motif_loc+=(find_motif_site(motif, seq[1][i])) # Find motif sites in protein  
        
    locations = list(dict.fromkeys(motif_loc)) # Store motif sites 
    n_loc += [len(locations)] # make list of number of locations in filename 
    all_locations+=locations # make list of all locations 
    
    protein_num = []
    # For loop to loop through all locations 
    for loc in locations:
        num.append(motif_loc.count(loc)) # Count number of instances (all proteins)
        protein_num.append(motif_loc.count(loc)) # Count number of instances (protein associated with file name)
        
    conserved.append(np.mean(protein_num)/len(seq[1])*100) # Calculate how conserved the sites in the protein are between instances 
    protein_list+=[protein] * len(locations) # List of proteins for detailed
    short_protein_list+=[protein] # List of proteins for short
    motif_list+=[motif] * len(locations) # list of motifs for deailed 

# Make dictionaries of key information for the detailed and summary data tables
detail_data = {'Protein': protein_list, 'Motif': motif_list, 'Location' : all_locations, 'Frequency' : num}
summary_data = {'Protein' : short_protein_list, 'Number of locations' : n_loc, 'Conservedness' : conserved}

# Make data frames 
detailed_data_frame = pd.DataFrame.from_dict(detail_data)
summary_data_frame = pd.DataFrame.from_dict(summary_data)
# Check detailed data frame
detailed_data_frame.head()
# Check summary data frame 
summary_data_frame.head()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12979035919348378, 'fraction_numerical': 0.009213513152623849, 'mean_word_length': 2.7866531850353895, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, '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': '159907', 'license': 'MIT', 'n_tokens_mistral': 2591, 'n_tokens_neox': 2497, 'n_words': 630, 'path': 'src/Tabulate.ipynb', 'repo_name': 'NCBI-Codeathons/Using-single-cell-SRA-sequence-to-verify-improve-submitted-metadata-available-in-SRA', 'size': '7489'}