trim_msa_based_on_mask.py

#!/usr/bin/env python

###############################################################################
#                                                                             #
#    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 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
#                                                                             #
###############################################################################

from __future__ import print_function

__prog_name__ = 'trim_msa_based_on_mask.py'
__prog_desc__ = 'Trim an MSA based on a prexisting mask. We assumed the mask has the same length as the MSA'

__author__ = 'Pierre Chaumeil'
__copyright__ = 'Copyright 2017'
__credits__ = ['Pierre Chaumeil']
__license__ = 'GPL3'
__version__ = '0.0.1'
__maintainer__ = 'Pierre Chaumeil'
__email__ = 'p.chaumeil@uq.edu.au'
__status__ = 'Development'

import sys
import argparse
from biolib.seq_io import read_fasta


class MSATrimmer(object):
    def __init__(self):
        """Initialization."""
        pass

    def run(self, msa, mask, outf):

        outfwriter = open(outf, 'w')
        dict_genomes = read_fasta(msa, False)
        with open(mask, 'r') as f:
            maskstr = f.readline()
        print(maskstr)
        print(len(maskstr))

        from future.utils import iteritems

        for k, v in dict_genomes.iteritems():
            aligned_seq = ''.join([v[i] for i in range(0, len(maskstr)) if maskstr[i] == '1'])
            fasta_outstr = ">%s\n%s\n" % (k, aligned_seq)
            outfwriter.write(fasta_outstr)
        outfwriter.close()


if __name__ == '__main__':
    print(__prog_name__ + ' v' + __version__ + ': ' + __prog_desc__)
    print('  by ' + __author__ + ' (' + __email__ + ')' + '\n')

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--msa', help='Untrimmed multi sequence alignment.')
    parser.add_argument('--mask', help='Mask file generated by gtdb.')
    parser.add_argument('--output_file', help='Output file.')

    args = parser.parse_args()

    try:
        msatrimmer = MSATrimmer()
        msatrimmer.run(args.msa, args.mask, args.output_file)
    except SystemExit:
        print("\nControlled exit resulting from an unrecoverable error or warning.")
    except:
        print("\nUnexpected error:", sys.exc_info()[0])
        raise