Using Biopython, I would like to optimize this code

Question

I have the following code:

from Bio import AlignIO
import itertools

out=open("test.csv","a")
align = AlignIO.read("HPV16_CG.aln.fas", "fasta")
n=0
def SNP(line):  
    result=[]
    result.append(str(n+1))
    result.append(line[0])
    result.append(align[y].id.rsplit("|")[3])
    result.append(x)
    return result



while n<len(align[0]):
  line = align[:,n]
  y=0
  for x in line:
    if line[0]!=x:
      print >> out, ','.join(map(str,SNP(line)))
      y=y+1
    else:
      y=y+1
  y=0
  n=n+1
out.close()

f=open("test.csv","rU")

out=open("test_2.csv","a")
lines=f.read().split()

for key, group in itertools.groupby(lines, lambda line: line.partition(',')[0]):
    print >>out, ','.join(group)

out.close()
f.close()

As you can see, I am currently writing two files. I really only need the second file. Does anyone have any suggestions to combine both "subscripts" into one?

The input file "HPV16_CG.aln.fas" looks like this:

>gi|333031|lcl|HPV16REF.1| Alpha-9 - Human Papillomavirus 16, complete genome. 
ACTACAATAATTCATGTATAAAACTAAGGGCGTAACCGAAATCGGTTGAACCGAAACCGG

>gi|333031|gb|K02718.1|PPH16 Human papillomavirus type 16 (HPV16), complete genome
ACTACAATAATTCATGTATAAAACTAAGGGCGTAACCGAAATCGGTTGAACCGAAACCGG

>gi|196170262|gb|FJ006723.1| Human papillomavirus type 16, complete genome
ACTACAATAATTCATGTATAAAACTAAGGGCGTAACCGAAATCGGTTGAACCGAAACCGG

I really appreciate all the help/suggestions to help me improve this!

Answer 1

First of all, read PEP8 (all of it!) and adhere to it. This is the de facto standard of Python code.

In the following, I’ll tacitly assume PEP8 as given.

Secondly, it would help if you had actually written what the expected output looks like. Given your input file, the script produces empty output.

I don’t understand what the while loop is supposed to do. align represents a single alignment record – you are reading a sequence file which contains multiple sequences. Furthermore, what is align[:, n]? Did you maybe intend to access the seq sub-object?

Likewise, the SNP method doesn’t at all make clear what its function is. And it implicitly uses global variables. This is error-prone and unreadable: use parameters instead. The function also accesses some object align[y] which won’t exist when y > 0, because AlignIO.read only reads a single alignment.

Furthermore, if you’re handling CSV files, I’d strongly suggest using the csv module. It’s not really less code in your particular case but it provides a more uniform handling and makes the code easily extensible.

Just to show what this could look like when done cleaner, here’s the production of the first file. I’d love to show the full code but I really don’t understand what the code is supposed to do since it doesn’t work for me, and I’m unfamiliar with the function of align[:,n] (and I think there’s an error in the access to align, anyway.

#!/usr/bin/env python

import csv
import itertools
import sys
from Bio import AlignIO


def snp(line, alignment, n, base, pos):
    """Documentation missing!"""
    id = alignment[pos].id.rsplit('|')[3] # why rsplit?
    return [ str(n + 1), line[0], id, base ]


def main():
    # TODO Use command line argumnents to provide file names

    alignment = Bio.AlignIO.read('HPV16_CG.aln.fas', 'fasta')

    with csv.writer('test.csv') as out:
        for n in range(len(alignment[0])):
            line = align[ : , n]
            for pos, base in line:
                if line[0] != x:
                    out.writerow(snp(line, alignment, n, base, pos))


if __name__ == '__main__':
    sys.exit(main())

In this code I tried to give the variables meaningful names but – again, since I don’t know what this is supposed to do – I had to guess quite a lot based on usage.

And by the way, range(len(x)) is actually kind of an anti-pattern but I can’t for the heck of it remember the “correct” usage (and even the documentation advises this).