fasta.aligned

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3
  4"""
  5Written by Lucas Sinclair.
  6MIT Licensed.
  7Contact at www.sinclair.bio
  8"""
  9
 10# Built-in modules #
 11import os, multiprocessing, platform
 12from collections import OrderedDict
 13
 14# Internal modules #
 15from fasta import FASTA
 16from plumbing.cache import property_cached
 17from autopaths.file_path import FilePath
 18from autopaths.tmp_path  import new_temp_path, new_temp_dir
 19
 20# Third party modules #
 21import shutil
 22if platform.system() == 'Windows': import pbs3 as sh
 23else: import sh
 24
 25################################################################################
 26class AlignedFASTA(FASTA):
 27    """Also a FASTA file technically, but contains an alignment."""
 28    ext = 'aln'
 29
 30    def __iter__(self): return iter(self.parse())
 31
 32    def parse(self):
 33        self.open()
 34        from Bio import AlignIO
 35        return AlignIO.read(self.handle, self.format)
 36
 37    def add_seq(self, seq):
 38        self.buffer.append(seq)
 39
 40    def flush(self):
 41        from Bio import AlignIO
 42        from Bio.Align import MultipleSeqAlignment
 43        AlignIO.write(MultipleSeqAlignment(self.buffer), self.handle, self.format)
 44
 45    def write(self, reads):
 46        from Bio import AlignIO
 47        if not self.directory.exists: self.directory.create()
 48        self.open('w')
 49        AlignIO.write(reads, self.handle, self.format)
 50        self.close()
 51
 52    @property_cached
 53    def sequences(self):
 54        return OrderedDict(((seq.id, seq) for seq in self))
 55
 56    def gblocks(self,
 57                new_path = None,
 58                seq_type = 'nucl' or 'prot'):
 59        """
 60        Apply the gblocks filtering algorithm to the alignment.
 61        See this page:
 62        http://molevol.cmima.csic.es/castresana/Gblocks/Gblocks_documentation.html
 63        Need to rename all sequences, because it will complain with long names.
 64        """
 65        # Temporary path #
 66        if new_path is None: final = self.__class__(new_temp_path())
 67        else:                final = self.__class__(new_path)
 68        # Mapping every sequence name with a random name #
 69        orig_name_to_temp = {seq.description: 'name' + str(i)
 70                             for i,seq in enumerate(self)}
 71        temp_name_to_orig = {v: k for k, v in orig_name_to_temp.items()}
 72        # Rename every sequence with a random name #
 73        temp_fasta = self.rename_sequences(orig_name_to_temp)
 74        # Options #
 75        if seq_type == 'nucl': t_option = "-t=d"
 76        if seq_type == 'prot': t_option = "-t=p"
 77        # Run it #
 78        result = sh.gblocks91(temp_fasta.path, t_option, '-p=n', "-b4=3",
 79                              "-b3=20", "-b5=a", _ok_code=[0,1])
 80        created_file = temp_fasta.path + '-gb'
 81        assert os.path.exists(created_file)
 82        # Check errors #
 83        if "Execution terminated" in result.stdout:
 84            raise Exception("gblocks crashed again.")
 85        # Back #
 86        temp_fasta.rename_sequences(temp_name_to_orig, final)
 87        # Return #
 88        return final
 89
 90    def build_tree(self, *args, **kwargs):
 91        """
 92        Dispatch a tree build call. Note that you need at least four
 93        taxa to express some evolutionary history on an unrooted tree.
 94        """
 95        # Check length #
 96        assert len(self) > 3
 97        # Default option #
 98        algorithm = kwargs.pop(kwargs, None)
 99        if algorithm is None: algorithm = 'raxml'
100        # Dispatch #
101        if algorithm == 'raxml':
102            return self.build_tree_raxml(*args, **kwargs)
103        if algorithm == 'fasttree':
104            return self.build_tree_fast(*args, **kwargs)
105
106    def build_tree_raxml(self,
107                         new_path    = None,
108                         seq_type    = 'nucl' or 'prot',
109                         num_threads = None,
110                         free_cores  = 2,
111                         keep_dir    = False):
112        """Make a tree with RAxML."""
113        # Check output #
114        if new_path is None: new_path = self.prefix_path + '.tree'
115        # What model to choose #
116        if seq_type == 'nucl': model = "GTRGAMMA"
117        if seq_type == 'prot': model = "PROTGAMMAJTTF"
118        # Threads #
119        if num_threads is None:
120            num_threads = multiprocessing.cpu_count() - free_cores
121        else:
122            num_threads = int(num_threads) - free_cores
123        num_threads = max(1, num_threads)
124        # Run it #
125        temp_dir = new_temp_dir()
126        sh.raxml811('-m', model, "-T", num_threads, '-p', 1, '-s', self.path,
127                    '-n', 'tree', '-w', temp_dir, '-f', 'a', '-x', 1, '-N',
128                    'autoMR')
129        # Move into place #
130        if keep_dir:
131            shutil.rmtree(new_path)
132            shutil.move(temp_dir, new_path)
133        if not keep_dir:
134            shutil.move(temp_dir + 'RAxML_bestTree.tree', new_path)
135        # Return #
136        return FilePath(new_path)
137
138    def build_tree_fast(self, new_path=None, seq_type='nucl' or 'prot'):
139        """Make a tree with FastTree. Names will be truncated however."""
140        # Check output #
141        if new_path is None: new_path = self.prefix_path + '.tree'
142        # Command #
143        command_args = []
144        if seq_type == 'nucl': command_args += ['-nt']
145        command_args += ['-gamma']
146        command_args += ['-out', new_path]
147        command_args += [self.path]
148        # Run it #
149        sh.FastTree(*command_args)
150        # Return #
151        return FilePath(new_path)