filter for variable sites

filter_for_var.py

#!/usr/bin/env python2
from Bio import AlignIO   
import sys

fformat=sys.argv[1].split('.')[-1]
align = AlignIO.read(open(sys.argv[1]), fformat)        #load file
outname = sys.argv[1].split('.')[:-1]+'.filtered.phylip-relaxed'

#find variable sites 
info=[]
for i in range(len(align[0])):
    site=[a for a in list(align[:, i])]
    if len(list(set(site)))>1:
        info.append(i)
        
#include only these sites       
for i,record in enumerate(align):
    sp_data=[record.seq[j] for j in info]
    align[i].seq="".join(sp_data)

AlignIO.write(align, outname, "phylip-relaxed")