o-trim-uninformative-columns-from-alignment

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (C) 2010 - 2012, A. Murat Eren
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# Please read the COPYING file.

import sys

import Oligotyping.lib.fastalib as u
import Oligotyping.utils.utils as utils

progress = utils.Progress()
run = utils.Run()

# get read length
fasta = u.SequenceSource(sys.argv[1], lazy_init = False)
next(fasta)
len_fasta_entry = len(fasta.seq)

# reset fasta.
fasta.reset()

# make sure all reads have equal length
while next(fasta):
    if len(fasta.seq) != len_fasta_entry:
        sys.stderr.write('All reads must have equal number of characters, but it is not the case. Sorry and bye.\n')
        sys.exit()
fasta.reset()

# -*-

nucleotide_positions = set(range(0, len_fasta_entry))
invalid_columns = set(range(0, len_fasta_entry))

progress.new('Step 1')
while next(fasta):
    if fasta.pos % 100 == 1:
        progress.update('%.2d%% -- pos: %d' % (fasta.pos * 100 / fasta.total_seq, fasta.pos))

    for i in nucleotide_positions:
        if fasta.seq[i] != '-' and i in invalid_columns:
            invalid_columns.remove(i)

progress.end()
fasta.reset()

# -*-

columns_to_keep = [x for x in range(0, len_fasta_entry) if x not in invalid_columns]

f = open(sys.argv[1] + '-TRIMMED', 'w')

progress.new('Step 2')
while next(fasta):
    if fasta.pos % 100 == 1:
        progress.update('%.2d%% -- pos: %d' % (fasta.pos * 100 / fasta.total_seq, fasta.pos))

    new_seq = ''.join(fasta.seq[i] for i in columns_to_keep)

    f.write('>' + fasta.id + '\n')
    f.write(new_seq + '\n')

f.close()