Cutting strings into smaller ones based on specific criteria

Question

I've got this largish (for me) script, and I want to see if anybody could tell me if there are any ways to improve it, both in terms of speed, amount of code and the quality of the code. I still consider myself a beginner, and I think that there are parts of my code that aren't maybe all that good - for example, all the variables I have to initiate before the main for-loop in the script, and the for-loop itself. Are there other ways of getting the same result?

This is what the script does:

First, it reads a large .csv-file (around 100,000 rows) with short strings ("peptides", looks like this: AQSIGR) and creates a reference set() from those. Then it reads a smaller file (around 100 rows) with larger strings ("PrESTs", which look like this: TYQIRTTPSATSLPQKTVVMTRSPVTLTSQTTKTDD, shortened) and analyses these in the following ways:

1. At every K or R not followed by P, cut the PrEST to the right of the letter (i.e. giving a first peptide TYQIR from above)
2. Also allow for up to two "missed cleavages", where the first and/or second R or K is ignored (i.e. TYQIR.TTPSATSLPQK and TYQIR.TTPSATSLPQK.TVVMTR, etc. from above)
3. Check every resulting peptide against the peptide reference set created in the beginning: only peptides NOT occurring in the set are of interest (i.e. unique).
4. Store the various peptides and PrEST data in an excel file.

There are various small details in the script as well (like ignoring the first resulting peptide from every PrEST), but the above points is the gist of it:

import pandas as pd
import csv

# ----------------------------------- Input variables -----------------------------------

# Input should be a .csv file with 2 columns (PrEST ID, PrEST Sequence)
data_file = 'Master Thesis 14 PrESTs'
protease_name = 'Trypsin'
#protease_name = 'Lys-C'

# ---------------------------------------------------------------------------------------

data = pd.read_csv(data_file + '.csv', sep=';')
proteasome = 'Non-Unique Reference (' + protease_name + ') (UniProt, canonical).csv'

# Create peptide reference Set
ref_set = set()
with open(proteasome, 'rU') as in_file:
    reader_2 = csv.reader(in_file)
    for row in reader_2:
        (ref_set.add(str(row).replace("'","").replace(",","")
                    .replace("[","").replace("]","").replace(" ","")))
print(str(len(ref_set)) + ' peptides in reference')

def protease(PrEST_seq, a, type):
    if protease_name == 'Trypsin':
        if type == 1:
            if PrEST_seq[a+1:a+2] != 'P' and (PrEST_seq[a:a+1] == 'R' or PrEST_seq[a:a+1] == 'K'):
                return 1
            else:
                return 0
        else:
            if PrEST_seq[a+2:a+3] != 'P' and (PrEST_seq[a+1:a+2] == 'R' or PrEST_seq[a+1:a+2] == 'K'):
                return 1
            else:
                return 0
    if protease_name == 'Lys-C':
        if PrEST_seq[a:a+1] == 'K':
            return 1
        else:
            return 0

# Initiate variables
Peptide_list = [] # List for Peptides (resets for each PrEST)
ID_list = [] # List for PrEST IDs (resets for each PrEST)
Non_Uniques = [] # List for non-unique peptides
Non_Uniques_ID = [] # List for non-unique PrEST IDs
Peptide = '' # Current peptide (no missed cleavages)
Peptide_MC1 = '' # Current peptide with 1 missed cleavage
Peptide_MC2 = '' # Current peptide with 2 missed cleavages

PrEST_data = pd.DataFrame()
# ------------------------------------------------ Main PrEST for-loop ------------------------------------------------
for row in data.iterrows():  # For every PrEST (row)
    First = 'Y'
    PrEST_seq = row[1][1]
    Pep_Count = 0
    MC_Pep_Count = 0
    Non_Unique_Count = 0

    # ----------------------------------------- No missed cleavages for-loop ------------------------------------------
    for n in range(len(PrEST_seq)):  # For every AA in every PrEST

        if protease(PrEST_seq, n, 1) == 1:
            if First != 'Y':  # Does not count first peptide + MCs (part of ABP)
                Peptide += PrEST_seq[n:n+1]
                if len(Peptide) >= 6:  # Only appends peptide if longer than 6 AA
                    if Peptide not in ref_set:
                        ID_list.append(row[1][0])
                        Peptide_list.append(Peptide)
                        Pep_Count += 1
                    else:
                        Non_Uniques_ID.append(row[1][0])
                        Non_Uniques.append(Peptide)
                        Non_Unique_Count += 1
                        ID_list.append(row[1][0] + ' (Not unique)')
                        Peptide_list.append(Peptide)

                # ----------------------------------- One missed cleavage while-loop ----------------------------------
                Peptide_MC1 = Peptide
                m = n
                while m+1 <= len(PrEST_seq):
                    m += 1
                    if protease(PrEST_seq, m, 1) == 1:
                        Peptide_MC1 += PrEST_seq[m:m+1]
                        if len(Peptide_MC1) >= 6:
                            if Peptide_MC1 not in ref_set:
                                ID_list.append(row[1][0])
                                Peptide_list.append(Peptide_MC1)
                                MC_Pep_Count += 1
                                break
                            else:
                                Non_Uniques_ID.append(row[1][0])
                                Non_Uniques.append(Peptide_MC1)
                                Non_Unique_Count += 1
                                ID_list.append(row[1][0] + ' (Not unique)')
                                Peptide_list.append(Peptide_MC1)
                    else:
                        Peptide_MC1 += PrEST_seq[m:m+1]

                # ---------------------------------- Two missed cleavages while-loop ----------------------------------
                Peptide_MC2 = Peptide_MC1
                k = m
                while k+1 <= len(PrEST_seq):
                    k += 1
                    if protease(PrEST_seq, k, 1) == 1:
                        Peptide_MC2 += PrEST_seq[k:k+1]
                        if len(Peptide_MC2) >= 6:
                            if Peptide_MC2 not in ref_set:
                                ID_list.append(row[1][0])
                                Peptide_list.append(Peptide_MC2)
                                MC_Pep_Count += 1
                                break
                            else:
                                Non_Uniques_ID.append(row[1][0])
                                Non_Uniques.append(Peptide_MC2)
                                Non_Unique_Count += 1
                                ID_list.append(row[1][0] + ' (Not unique)')
                                Peptide_list.append(Peptide_MC2)
                    else:
                        Peptide_MC2 += PrEST_seq[k:k+1]
                    # -------------------------------------------------------------------------------------------------

                # Resets variables
                Peptide = ''
                Peptide_MC1 = ''
                Peptide_MC2 = ''
            elif First == 'Y':  # Doesn't count first cleavage (contains ABP)
                Peptide = ''
                First = 'N'
        else:  # Non-cleavable AAs - Peptide grows
            Peptide += PrEST_seq[n:n+1]

    # Appends PrEST data
    K = row[1][1].count('K')
    R = row[1][1].count('R')
    PrEST_data = PrEST_data.append(pd.DataFrame(data=[[row[1][0], Pep_Count, MC_Pep_Count, Non_Unique_Count, K, R]],
                                                columns=['PrEST ID', '# Peptides', '# MC Peptides', '# Non-Uniques',
                                                         '# K', 'R']))
# Writes PrEST data to file
PrEST_data = PrEST_data[['PrEST ID', '# Peptides', '# MC Peptides', '# Non-Uniques', '# K', 'R']]
ew = pd.ExcelWriter(data_file + ' Results.xlsx', encoding='iso-8859-1')
PrEST_data.to_excel(ew, sheet_name='PrESTs (' + protease_name + ')', index=False)

# Creates peptide list for Perseus and writes to file
peptides = pd.DataFrame(Peptide_list, columns=['Peptides']).join(pd.DataFrame(ID_list, columns=['PrEST ID']))
peptides['temp'] = peptides['Peptides'].str.len()
peptides = peptides.sort(['PrEST ID', 'temp'], ascending=[True, False]).drop('temp', axis=1)
peptides.to_excel(ew, sheet_name='Peptide list for Perseus', index=False)

# List for non-unique peptides
NU_peptides = pd.DataFrame(Non_Uniques, columns=['Peptides']).join(pd.DataFrame(Non_Uniques_ID, columns=['PrEST ID']))
NU_peptides['temp'] = NU_peptides['Peptides'].str.len()
NU_peptides = NU_peptides.sort(['PrEST ID', 'temp'], ascending=[True, False]).drop('temp', axis=1)
NU_peptides.to_excel(ew, sheet_name='Non-unique peptides', index=False)

ew.save()

Here are some input rows from the peptide reference .csv-file (100,000-ish rows):

Reference peptides
IEPVFHVMGFSVTGRVLNGPEGDGVPEAVVTLNNQIK
LDAKKRR
EILRNPMEAMYPHIFYFHFK
KGPPPPPPKK
MVVEVDSMPAASSVK
ERNRNK
QDWTILHSYVHLNADELEALQMCTGYVAGFVDLEVSNRPDLYDVFVNLAESEITIAPLAK
EQFGQDGPISVHCSAGVGR
ADIGVAMGIAGSDVSK

And here is some rows from the input file to be analysed (the Xs are not really there, but since this is confidential gene sequences I changed them here):

PrEST ID  Sequence
HPRR06    QGAWKTISNGFGFKDAVFDGSSCISPTIVQQFGYQRRASDXXXXXXXXXSNTIRVFLPNKQRTVVNVRNXXXXXXXXXKALKVRGLQPECCAVFRLLHEHKGKKARLDWNTXXXXXXXXXLQVDFLDHVPLTTHNFAR
HPRR07    ERTFHVETPEEREEWTTAIQTVADGXXXXXXXXXPSDNSGAEEMEVSLAKPKHRVTMNEFEYLKLLGKGTFGKVILVKEKATGRYYAMKILKKEVIVA
HPRR08    MDPMTVGRIEGDCESLNFSEVSSSSKDVENGGKDKPPQPGAKTSSRNDYIHSGLYSSFTLNSLNSSNVKLFKLIKTENPAEKLAEKKSPQEPTPSVIKFVTTPSKKPPVEPVAATISIGPS
HPRR09    VPVSNQSSLQFSNPSGSLVTPSLVTSSXXXXXXXXXRNSVSPGLPQRPASAGAMLGGDLNSANGACPSPVGNGYVSARASPGLLPVANGNSLNKVIPAKSPPPPTHSXXXXXXXXXDLRVITSQAG
HPRR10    DLQWLVQPALVSSVAPSQTRAPHPFGVPAPSAGAYSRXXXXXXXXXRAQSIGRRGKVEQLSPEEEEKRRIRRERNKMAAAKCRNRRRELTDTLQAETDQLEDEKSALQTEIANLXXXXXXXXXEFILAAHRPACKIPDDL
HPRR52    STIAESEDSQESVDSVTDSQKRREILSRRPSYRKILNDLSSXXXXXXXXXVPTPIYQTSSGQYIAITQGGAIQLANNGTDGVQG
HPRR14    MPKKKPTPIQLNPAPDGSAVNGTSSAXXXXXXXXXELELDEQQRKRLEAFLTQKQ
HPRR35    MEDSHKSTTSETAXXXXXXXXHIAQQVSSLSESEESQDSSDSIGSSQKAHGILARRPSY
HPRR02    EGFLCVFAINNTKXXXXXXXXYREQIKRVKDSEDV

The main output (first sheet of excel file, PrEST_data DataFrame) looks like this:

PrEST ID  # Peptides  # MC Peptides   # Non-Uniques   # K # R
HPRR52    10  31  2   10  9
HPRR10    5   26  5   12  5
HPRR09    7   25  0   13  2
HPRR08    5   12  0   3   6
HPRR06    7   45  4   10  15
HPRR07    1   13  6   3   6
HPRR14    2   10  3   7   2
HPRR02    2   3   0   2   2
HPRR35    1   2   1   3   2

Answer 1

I have a limited understanding of the different parts of your code so I'll just comment on some details that can be greatly improved.

---

General:

A few things are not quite pythonic. You can make your code go through some automatic checks with pep8online (PEP8 is the Style Guide for Python Code) and pylint.

---

About protease :

def protease(PrEST_seq, a, type):
    if protease_name == 'Trypsin':
        if type == 1:
            if PrEST_seq[a+1:a+2] != 'P' and (PrEST_seq[a:a+1] == 'R' or PrEST_seq[a:a+1] == 'K'):
                return 1
            else:
                return 0
        else:
            if PrEST_seq[a+2:a+3] != 'P' and (PrEST_seq[a+1:a+2] == 'R' or PrEST_seq[a+1:a+2] == 'K'):
                return 1
            else:
                return 0
    if protease_name == 'Lys-C':
        if PrEST_seq[a:a+1] == 'K':
            return 1
        else:
            return 0

• from the implementation and the way you use it, it seems like this function always returns 0 or 1 and should probably return a boolean. Make yourself a favor and actually use booleans. You can now replace all instances of if protease(foo, bar, foobar) == 1: with if protease(foo, bar, foobar):. Also, in the function itself if condition: return 1 else return 0 becomes return condition.

• You have 2 pretty similar conditions. By adding 1 to a if type != 1, you could make the 2 conditions actually identical.

• You do not need to get PrEST_seq[a+1:a+2] twice, just check if its values is in some list.

• For a string s, s[i:i+1] and s[i] seems to be the same thing when i is smaller than the length of s. Prefer the shorter if you can. You can also consider try-catch.

• It seems like you do not always return a value which is a bit odd.

Taking all these comments into account leads to the following code :

def protease(PrEST_seq, a, type):
    if protease_name == 'Trypsin':
        if type != 1:
            a+=1
        return PrEST_seq[a+1] != 'P' and (PrEST_seq[a] in ['R','K']):
    assert(protease_name == 'Lys-C')
    return PrEST_seq[a] == 'K'

---

About First:

You are using First as a character which is either Y or not. Here again a boolean would be more than enough.

• First = 'Y' becomes first = True
• if First != 'Y': becomes if not first:
• elif First == 'Y': is not required and you can do Peptide = '' and first = False no matter what the value of first was in the first place.

---

About loops

Your while loops can easily be transformed into a for loop using range.

For instance,

            m = n
            while m+1 <= len(PrEST_seq):
                m += 1
                stuff_about(m)

could be written :

            for m in range(n+1, len(PrEST_seq)+1):
                stuff_about(m)

(but my guess is that what you really want to do is for m in range(n+1, len(PrEST_seq)+1): and this could be an explanation for your 'out-of-range' exceptions)

---

About duplicated code :

You sometimes have duplicated logic, for instance on both branches of a test. If you keep things simple, it will be easier to understand and to maintain.

Also, you should not store the number of element in a list in a variable if you can just as easily check the length of the list.

For instance :

            if len(Peptide) >= 6:  # Only appends peptide if longer than 6 AA
                if Peptide not in ref_set:
                    ID_list.append(row[1][0])
                    Peptide_list.append(Peptide)
                    Pep_Count += 1
                else:
                    Non_Uniques_ID.append(row[1][0])
                    Non_Uniques.append(Peptide)
                    Non_Unique_Count += 1
                    ID_list.append(row[1][0] + ' (Not unique)')
                    Peptide_list.append(Peptide)

could become :

            if len(Peptide) >= 6:  # Only appends peptide if longer than 6 AA
                Peptide_list.append(Peptide)

                if Peptide not in ref_set:
                    ID_list.append(row[1][0])
                    Pep_Count += 1
                else:
                    Non_Uniques_ID.append(row[1][0])
                    Non_Uniques.append(Peptide)
                    ID_list.append(row[1][0] + ' (Not unique)')