bwa.py

#
# Copyright (C) 2012 INRA
# 
# 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/>.
#

import os
from subprocess import Popen, PIPE
from jflow.utils import get_argument_pattern

from ng6.analysis import Analysis
from ng6.utils import Utils

class BWA (Analysis):

    def define_parameters(self, reference_genome, read1, read2=None, group_prefix=None, algorithm="aln",  keep_bam=True):
        """
          @param reference_genome : [str] Path to the reference genome (it must be indexed).
          @param read1 : [list] Paths to reads 1.
          @param read2 : [list] Paths to reads 2.
          @param group_prefix : [list] The component produces one bam by prefix. Each bam is a merge of all files with the correct prefix.
          @param algorithm : [str] Algorithm for the alignment (aln or bwasw or mem).
          @param keep_bam : [bool] True save bam in nG6
        """
        # Parameters
        self.add_parameter_list( "group_prefix", "The component produces one bam by prefix. Each bam is a merge of all files with the correct prefix.", default=group_prefix )
        self.add_parameter( "algorithm", "Algorithm for the alignment : aln or bwasw or mem.", default=algorithm, choices=["aln", "bwasw", "mem"] )
        self.add_parameter("keep_bam", "keep_bam", default=keep_bam, type=bool)
        # Files
        self.add_input_file("reference_genome", "Which reference file should be used", default=reference_genome, required=True)
        self.add_input_file_list( "read1", "Which read1 files should be used.", default=read1, required=True )
        self.add_input_file_list( "read2", "Which read2 files should be used.", default=read2 )
        self.sai1 = None
        self.sai2 = None
        if algorithm == "aln":
            self.add_output_file_list("sai1", "The BWA read1 sai files.", pattern='{basename_woext}.sai', items=self.read1)
            if read2 != None:
                self.add_output_file_list("sai2", "The BWA read2 sai files.", pattern='{basename_woext}.sai', items=self.read2)
        if group_prefix == None:
            self.add_output_file_list("bam_files", "The BWA bam files.", pattern='{basename_woext}.bam', items=self.read1, file_format="bam")
        else:
            self.add_output_file_list("bam_files", "The BWA bam files.", pattern='{basename_woext}.bam', items=group_prefix, file_format="bam")
        
        
        self.add_output_file_list("stderrs", "BWA stderrs files", pattern='{basename_woext}.bwa.stderr', items=self.read1)
        if self.algorithm == 'aln' :
            if self.read2 :
                self.add_output_file_list("stderrs_aln", "BWA stderrs files", pattern='{basename_woext}.bwa_aln.stderr', items=self.read1 + self.read2)
            else :
                self.add_output_file_list("stderrs_aln", "BWA stderrs files", pattern='{basename_woext}.bwa_aln.stderr', items=self.read1)
    
    def get_version(self):
        cmd = [self.get_exec_path("bwa")]
        p = Popen(cmd, stdout=PIPE, stderr=PIPE)
        stdout, stderr = p.communicate()
        return stderr.split()[7]

    def define_analysis(self):
        self.name = "Alignment"
        self.description = "Reads Alignment"
        self.software = "bwa"
        self.options = "Algorithm " + self.algorithm + " | Databank " + os.path.basename(self.reference_genome)
        source_file = self.reference_genome + '_source'    
        if os.path.exists(source_file) :
            source = open(source_file,'r')
            reference_desc = source.readline();
            source.close()
            self.options += " " + reference_desc
    
    def post_process(self):
        if self.keep_bam:
            # Finally create and add the archive to the analysis
            self._save_files(self.bam_files)
            
    def process(self):
        unmerged_bam = self.get_outputs( '{basename_woext}.bam', self.read1 )

        # Algorithm bwasw or mem
        if self.algorithm == "bwasw" or self.algorithm == "mem":
            # Paired-end
            if self.read2:
                self.add_shell_execution(self.get_exec_path("bwa") + " " + self.algorithm + " " + self.reference_genome + 
                                " $1 $2 2>> $4 | " + self.get_exec_path("samtools") + " view -bS - | " + 
                                self.get_exec_path("samtools") + " sort - -o $3 2>> $4;", 
                                cmd_format='{EXE} {IN} {OUT}' , map=True,
                                inputs=[self.read1, self.read2], outputs=[unmerged_bam, self.stderrs], includes=self.reference_genome)
            # Single-end
            else:
                self.add_shell_execution(self.get_exec_path("bwa") + " " + self.algorithm + " " + self.reference_genome + 
                                    " $1 2>> $3 | " + self.get_exec_path("samtools") + " view -bS - | " + 
                                    self.get_exec_path("samtools") + " sort - -o $2 2>> $3 ;", 
                                cmd_format='{EXE} {IN} {OUT}' , map=True,
                                inputs=[self.read1], outputs=[unmerged_bam, self.stderrs], includes=self.reference_genome)

        # Algorithm aln  
        else:
            reads, sais = [], []
            reads.extend(self.read1)
            sais.extend(self.sai1)
            
            # Paired-end
            if self.read2:
                reads.extend(self.read2)
                sais.extend(self.sai2)
                self.add_shell_execution(self.get_exec_path("bwa") + " " + self.algorithm + " " + self.reference_genome + 
                                         " $1 > $2 2>> $3 " , cmd_format='{EXE} {IN} {OUT}', map=True,
                                         inputs=[reads], outputs=[sais, self.stderrs_aln], includes=self.reference_genome)
                self.add_shell_execution(self.get_exec_path("bwa") + " sampe " + self.reference_genome + 
                                         " $1 $2 $3 $4 2>> $6 | " + self.get_exec_path("samtools") + " view -bS - | " + 
                                         self.get_exec_path("samtools") + " sort - -o $5 2>> $6;",
                                         cmd_format='{EXE} {IN} {OUT}', map=True,
                                         inputs=[self.sai1, self.sai2, self.read1, self.read2], outputs=[unmerged_bam, self.stderrs], includes=self.reference_genome)
            # Single-end
            else:
                self.add_shell_execution(self.get_exec_path("bwa") + " " + self.algorithm + " " + self.reference_genome + 
                                         " $1 > $2 2>> $3 " , cmd_format='{EXE} {IN} {OUT}', map=True,
                                         inputs=[reads], outputs=[sais, self.stderrs_aln], includes=self.reference_genome)
                self.add_shell_execution(self.get_exec_path("bwa") + " samse " + self.reference_genome + 
                                         " $1 $2 2>> $4 | " + self.get_exec_path("samtools") + " view -bS - | " + 
                                         self.get_exec_path("samtools") + " sort - -o $3 2>> $4;", 
                                         cmd_format='{EXE} {IN} {OUT}', map=True,  
                                         inputs=[self.sai1, self.read1], outputs=[unmerged_bam, self.stderrs], includes=self.reference_genome)

        if self.group_prefix != None:
            # Create dictionary : key = prefix and value = list of files to merge
            groups_path = Utils.get_filepath_by_prefix( unmerged_bam, self.group_prefix)

            # Create dictionary : key = prefix and value = the output bam
            outputs_path = Utils.get_filepath_by_prefix( self.bam_files, self.group_prefix)

            # Merges bam or rename
            for prefix in self.group_prefix:
                if len(groups_path[prefix]) > 1:
                    [cmd_inputs_pattern, next_arg_number] = get_argument_pattern(groups_path[prefix], 1)
                    self.add_shell_execution(self.get_exec_path("samtools") + ' merge ${' + str(next_arg_number) + '} ' + cmd_inputs_pattern , 
                                             cmd_format='{EXE} {IN} {OUT}', map=False,
                                             inputs=groups_path[prefix], outputs=outputs_path[prefix])
                elif groups_path[prefix] != outputs_path[prefix]:
                    self.add_shell_execution( "ln -fs ${1} ${2}", cmd_format='{EXE} {IN} {OUT}', map=False,
                                              inputs=groups_path[prefix], outputs=outputs_path[prefix])