mothurotuanalysis.py

import os
from subprocess import Popen, PIPE

from jflow.iotypes import OutputFile, OutputFileList, InputFile, InputFileList, Formats
from jflow.abstraction import MultiMap

from weaver.abstraction import Map
from weaver.function import PythonFunction, ShellFunction

from ng6.analysis import Analysis

def get_taxonomy_to_Krona(exec_krona,output_directory,*files):
    import glob,re,os
    tax_stdout = files[0]
    shared_files = files[1]
    taxonomy_files = glob.glob(output_directory+"/*.taxonomy")
    if re.findall('unique_list.\w+.cons',taxonomy_files[0]):
        #case of miseq
        start = "unique_list."
    else:
        #case of 454
        start = "an."
    end = ".cons"
    substrings_labels = map(lambda string : re.search('%s(.*)%s' % (start, end), string).group(1),taxonomy_files)
    result = dict(zip(substrings_labels,taxonomy_files))
    taxonomy_outputs = []
    for label, tax_file in result.items():
        source = tax_file
        destination = os.path.join(output_directory,label+".taxonomy")
        taxonomy_outputs.append(destination)
        os.rename(source, destination)
    
    #get all otus for each samples and each levels
    lines = open(shared_files).readlines()
    firstline = lines.pop(0).split()[3:]
    dico_otus = {}
    for otu in firstline:
        dico_otus[firstline.index(otu)+4] = otu
    dico_samples_level = {}
    for line in lines:
        temp = {}
        line = line.split()
        label = line[0]
        group = line[1]
        nb_otus = line [3:]
        list_otus = []
        for i in range(len(nb_otus)):
            if nb_otus[i] !='0':
                list_otus.append(dico_otus[i+4]+";"+nb_otus[i])
        temp[label] = list_otus
        if dico_samples_level.has_key(group):
            dico_samples_level[group].update(temp)
        else:
            dico_samples_level[group] = temp
    #get all taxonomy for each levels
    dico_level_otu_taxonomy = {}
    for taxanomy_file in taxonomy_outputs:
        label = os.path.splitext(os.path.basename(taxanomy_file))[0]
        lines = open(taxanomy_file).readlines()
        header = lines.pop(0)
        temp = {}
        for line in lines:
            line = line.split()
            num_otu = line[0]
            lineage = line[2]
            temp[num_otu] = lineage
        dico_level_otu_taxonomy[label] = temp
    #write level taxonomy files for each samples
    dico_files = {}
    for sample,levels in dico_samples_level.items():
        for level,otus_desc in levels.items():
            name_file = os.path.join(output_directory,sample+"-"+level+".taxonomy")
            if dico_files.has_key(sample):
                dico_files[sample].append(name_file)
            else:
                dico_files[sample] = []
                dico_files[sample].append(name_file)
            fic_taxonomy = open(name_file,'w')
            fic_taxonomy.write('%s\t%s\t%s\n' % ('OTU','Size','Taxonomy'))
            for item in otus_desc:
                item = item.split(';')
                otu = item[0]
                nb_otu = item[1]
                fic_taxonomy.write('%s\t%s\t%s\n' % (otu,nb_otu,dico_level_otu_taxonomy[level][otu]))
            fic_taxonomy.close()
    for samp,name in dico_files.items():
        cmd = exec_krona+" "+" ".join(name)+" -o "+ os.path.join(output_directory,samp+".krona.html")
        p = os.popen(cmd)
    
class MothurOTUAnalysis(Analysis):
    """
    """
    
    def define_parameters(self, an_list_files, taxonomy_files, fasta_subsample_files=None, names_files=None, groups_files=None, count_table_files=None, \
                          calc="simpson-npshannon-shannon-jack-ace-chao", tree_calc="thetayc", label="unique-0.03-0.05-0.10", tree_label="unique-0.01", freq=100):
        """
        :param an_list_files: an list files to process
        :type an_list_files: str
        :param taxonomy_files: taxonomy files to process classifyotu
        :type taxonomy_files: str
        :param fasta_files: fasta files to process subsample
        :type fasta_files: str
        :param names_files: names files to process
        :type names_files: string
        :param groups_files: groups files to process
        :type groups_files: str
        :param count_table_files: count table files to process
        :type count_table_files: str
        :param calc : 
        :type calc : str
        :param label : 
        :type label : str
        :param tree_calc : 
        :type calc : str
        :param tree_label : 
        :type label : str
        :param freq : 
        :type freq : int
        """

        #define input files
        self.an_list_files = InputFileList(an_list_files)
        self.taxonomy_files = InputFileList(taxonomy_files, Formats.MOTHUR_TAXONOMY)
        if fasta_subsample_files: self.fasta_subsample_files = InputFileList(fasta_subsample_files, Formats.FASTA)
        else : self.fasta_subsample_files = None
        if names_files : self.names_files = InputFileList(names_files, Formats.MOTHUR_NAMES)
        else : self.names_files = None
        if groups_files : self.groups_files = InputFileList(groups_files, Formats.MOTHUR_GROUPS)
        else : self.groups_files = None
        if count_table_files : self.count_table_files = InputFileList(count_table_files, Formats.MOTHUR_COUNT_TABLE)
        else : self.count_table_files = None
        
        #define parameters
        self.calc = calc
        self.tree_calc = tree_calc
        self.tree_label = tree_label
        self.freq = freq
        labels = label.split('-')
        user_labels = [lines.split('\t')[0] for lines in open(self.an_list_files[0]).readlines()]
        good_labels = list(set(user_labels).intersection(set(labels)))
        self.label = good_labels
        #define makeshared output files
        self.shared_files = OutputFileList(self.get_outputs('{basename_woext}.shared', self.an_list_files))
        self.makeshared_stdout = OutputFileList(self.get_outputs('{basename_woext}.makeshared.stdout', self.an_list_files))
        #define summarysingle output files
        self.summary_single_files = OutputFileList(self.get_outputs('{basename_woext}.groups.summary', self.shared_files))
        self.summarysingle_stdout = OutputFileList(self.get_outputs('{basename_woext}.summarysingle.stdout', self.shared_files))
        #define rarefactionsingle output files
        self.rarefaction_single_files = OutputFileList(self.get_outputs('{basename_woext}.groups.rarefaction', self.shared_files))
        self.rarefactionsingle_stdout = OutputFileList(self.get_outputs('{basename_woext}.rarefactionsingle.stdout', self.shared_files))
        #define treeshared output files
#         list_tree_files = [os.path.splitext(self.shared_files[0])[0]+"."+tree_calc+"."+label+".tre" for label in self.label]
#         self.tree_files = OutputFileList(self.get_outputs('{basename}', list_tree_files))
#         self.treeshared_stdout = OutputFileList(self.get_outputs('{basename_woext}.treeshared.stdout', self.shared_files))
        #define subsample output files
        if fasta_subsample_files : 
            self.subsample_fasta_files = OutputFileList(self.get_outputs('{basename_woext}.subsample.fasta', self.fasta_subsample_files), Formats.FASTA)
        if an_list_files : 
            self.subsample_list_files = OutputFileList(self.get_outputs('{basename_woext}.subsample.list', self.an_list_files))
        if count_table_files : 
            self.subsample_count_files = OutputFileList(self.get_outputs('{basename_woext}.subsample.count_table', self.count_table_files), Formats.MOTHUR_COUNT_TABLE)
        self.subsample_stdout = OutputFileList(self.get_outputs('{basename_woext}.subsample.stdout', self.count_table_files))
        #define classifyotu output files
        self.cons_taxonomy_files = OutputFileList([os.path.join(self.output_directory,l+".taxonomy") for l in self.label],Formats.MOTHUR_TAXONOMY)
        self.classifyotu_stdout = OutputFileList(self.get_outputs('{basename_woext}.classifyotu.stdout', self.an_list_files))


    def define_analysis(self):
        self.name = "OTUAnalysis"
        self.description = "Organizational Taxon Unit analysis."
        self.software = "mothur"
        label = "-".join(self.label)
        if self.groups_files:
            self.options += '#make.shared(list=FILE.an.list,group=FILE.groups,label='+label+');'
        if self.count_table_files :
            self.options += '#make.shared(list=FILE.an.list,count=FILE.count_table_files,label='+label+');'
        self.options += '#summary.single(shared=FILE.shared,calc='+self.calc+',label='+label+');'
        self.options += '#rarefaction.single(shared=FILE.shared,label='+label+',freq='+str(self.freq)+');'
        self.options += '#tree.shared(shared=FILE.shared,calc='+self.tree_calc+',label='+self.tree_label+');'
        if self.fasta_subsample_files and self.count_table_files:
             self.options += '#sub.sample(fasta=FILE.fasta,count=FILE.count_table);'
        elif self.an_list_files and self.count_table_files:
            self.options += '#sub.sample(list=FILE.list,count=FILE.count_table);'
        elif self.fasta_subsample_files and not self.count_table_files:
            self.options += ' #sub.sample(fasta=FILE.fasta);'
        elif self.an_list_files and not self.count_table_files:
             self.options += '#sub.sample(list=FILE.list);'
        if self.names_files and not self.groups_files:
            self.options += '#classify.otu(list=FILE.list,taxonomy=FILE.taxonomy,name=FILE.names,label='+str(label)+');'
        elif self.names_files and self.groups_files: 
            self.options += '#classify.otu(list=FILE.list,taxonomy=FILE.taxonomy,name=FILE.names,group=$4,label='+str(label)+');'
        elif self.count_table_files:
            self.options += '#classify.otu(list=FILE.list,taxonomy=FILE.taxonomy,count=FILE.count_table,label='+str(label)+');'
        else:
            self.options += '#classify.otu(list=FILE.list,taxonomy=FILE.taxonomy,label='+str(label)+');'
        self.options += 'ImportMothurTaxonomy.pl '+" ".join([os.path.basename(tax) for tax in self.cons_taxonomy_files])+' -o FILE.html;'
        
    def get_version(self):
        cmd = [self.get_exec_path("mothur"), "-version"]
        p = Popen(cmd, stdout=PIPE, stderr=PIPE)
        stdout, stderr = p.communicate()
        return stdout.split()[1].split('=')[1]

    def get_rarefaction_values(self,rar_file, labels):

        """
        Parse the classify log file
        @param rar_file    : the rar file path
        @param labels : the labels to display
        @return            : the rarefaction values
        """
        lines = open(rar_file).readlines()
        firstline = lines.pop(0).split()[1:]
        samples = set([i.split('-')[1] for i in firstline])
        dico_labels = {}
        dico_rarefaction = {}
        for sample in samples:
            indx = {}
            for val in labels:
                try:
                    indx[val] = firstline.index(val+"-"+sample)+1
                except:
                    pass
            dico_labels[sample] = indx
        for sample,labels in dico_labels.items():
            for label,col in labels.items():
                x = []
                y = []
                for line in lines:
                    line = line.split("\t")[:-1]
                    x.append(line[0])
                    if line[int(col)] == "NA" or line[int(col)] == "":
                        y.append(y[-1])
                    else:
                        y.append(line[col])
                x = ",".join(x)
                y = ",".join(y)
                if dico_rarefaction.has_key(sample):
                    dico_rarefaction[sample].append([x,y,label])
                else:
                    dico_rarefaction[sample] = []
                    dico_rarefaction[sample].append([x,y,label])
        return dico_rarefaction

    def _get_otu_values(self, shared_file):
        """ Parse mothur shared file
        @param shared_file: the shared file
        @return: the json value
        """
        otus_seq = {}
        nb_seq = {}
        header = []
        labels = []
        for line in open(shared_file, "r").readlines():
            line = line.strip()
            parts = line.split()
            if line.startswith("label"):
                header = parts[3:]
            else:
                sample = parts[1].replace('.','_')
                values = ""
                nbs = 0
                for i, head in enumerate(header[:int(parts[2])]):
                    if parts[i+3] != "0":
                        values += head + "=" + parts[i+3] + ";"
                        nbs += int(parts[i+3])
                values = values[:-1]
                if otus_seq.has_key(sample):
                    otus_seq[sample][parts[0]] = values
                else:
                    otus_seq[sample] = {parts[0]:values}
                if not nb_seq.has_key(sample):
                    nb_seq[sample] = nbs
                if parts[0] not in labels:
                    labels.append(parts[0])
        return [otus_seq, nb_seq, labels]
        
    def post_process(self):
        result_files = []
        rarefaction_files = []
        summary_files = []
        shared_file = None
        krona_files = []
        for file in os.listdir(self.output_directory):
            file = file.strip()
            if file.endswith(".rarefaction"):
                rarefaction_files.append(os.path.join(self.output_directory, file))
            elif file.endswith(".summary"):
                summary_files.append(os.path.join(self.output_directory, file))
            elif file.endswith(".shared"):
                shared_file = os.path.join(self.output_directory, file)
            elif file.endswith(".html"):
                krona_files.append(os.path.join(self.output_directory, file))
                
        # First add OTU's values
        [otus, nb_seq, labels] = self._get_otu_values(shared_file)
        for sample in otus.keys():
            for level in otus[sample].keys():
                self._add_result_element(sample, "otus", otus[sample][level], level)
                self._add_result_element(sample, "sobs", len(otus[sample][level].split(";")), level)

        # nb seq
        for sample in nb_seq.keys():
            self._add_result_element(sample, "nb_seq", nb_seq[sample])
        
        #Add Krona files
        for krona_fic in krona_files:
            sample_name = os.path.basename(krona_fic).split('.')[0]
            self._add_result_element(sample_name, "krona", self._save_file(krona_fic))
            if os.path.isdir(krona_fic+".files"):
                self._save_directory(krona_fic+".files")

        # First add rarefaction infos
        for rarefaction_file in rarefaction_files:
            rarefaction_values = self.get_rarefaction_values(rarefaction_file, labels)
            for sample_name, vals in rarefaction_values.items():
                sample_name = sample_name.replace('.','_')
                for val in vals:
                    self._add_result_element(sample_name, "rarefaction_x", val[0], val[2])
                    self._add_result_element(sample_name, "rarefaction_y", val[1], val[2])

        # Then the OTU's table
        for summary_file in summary_files:
            header = []
            estimators = []
            for line in open(summary_file, "r").readlines():
                line = line.strip()
                parts = line.split()
                if line.startswith("label"):
                    header = line.split()
                    for head in header:
                        if head != "label" and head != "group" and not head.endswith("_hci") and not head.endswith("_lci"):
                            estimators.append(head)
                elif parts[0] in labels:
                    if header[1] == "group":
                        sample_name = parts[1].replace('.','_')
                    else:
                        sample_name = ".".join(os.path.basename(summary_file).split(".")[:-8])
                            
                    #self._add_result_element(sample_name, "sobs", parts[2], parts[0])
                        
                    for estimator in estimators:
                        es_value = parts[header.index(estimator)]
                        try:
                            eshci = parts[header.index(estimator+"_hci")]
                            eslci = parts[header.index(estimator+"_lci")]
                            es_value += ";" + eslci + ";" + eshci
                        except:
                            pass
                        self._add_result_element(sample_name.replace('.','_'), estimator, es_value, parts[0])
  
        result_files.extend(rarefaction_files)
        result_files.extend(summary_files)
        # Finaly create and add the archive to the analyse
        #self._create_and_archive(result_files, archive_name)
    
    def process(self):
        label = "-".join(self.label)
        #Divertsity
        if self.groups_files :
            makeshared = ShellFunction(self.get_exec_path("mothur") + ' "#make.shared(list=$1,group=$2,label='+label+',outputdir='+self.output_directory+'/)" > $3',\
                                       cmd_format='{EXE} {IN} {OUT}')
            makeshared = MultiMap(makeshared, inputs=[self.an_list_files,self.groups_files], outputs=[self.makeshared_stdout,self.shared_files])
        if self.count_table_files :
            makeshared = ShellFunction(self.get_exec_path("mothur") + ' "#make.shared(list=$1,count=$2,label='+label+',outputdir='+self.output_directory+'/)" > $3',\
                                       cmd_format='{EXE} {IN} {OUT}')
            makeshared = MultiMap(makeshared, inputs=[self.an_list_files,self.count_table_files], outputs=[self.makeshared_stdout,self.shared_files])
        # Alpha diversity analysis  
        summarysingle = ShellFunction(self.get_exec_path("mothur") + ' "#summary.single(shared=$1,calc='+self.calc+',label='+label+',outputdir='+\
                                   self.output_directory+'/)" > $2',cmd_format='{EXE} {IN} {OUT}')
        summarysingle = MultiMap(summarysingle, inputs=[self.shared_files], outputs=[self.summarysingle_stdout,self.summary_single_files])
        rarefactionsingle = ShellFunction(self.get_exec_path("mothur") + ' "#rarefaction.single(shared=$1,label='+label+',freq='+str(self.freq)+',outputdir='+\
                                          self.output_directory+'/)" > $2',cmd_format='{EXE} {IN} {OUT}')
        rarefactionsingle = MultiMap(rarefactionsingle, inputs=[self.shared_files], outputs=[self.rarefactionsingle_stdout,self.rarefaction_single_files])
        # Beta diversity analysis
        #TODO : create an interface in NG6
#         treeshared = ShellFunction(self.get_exec_path("mothur") + ' "#tree.shared(shared=$1,calc='+self.tree_calc+',label='+self.tree_label+',outputdir='+self.output_directory+\
#                                    '/)" > $2',cmd_format='{EXE} {IN} {OUT}')
#         treeshared(inputs=self.shared_files[0], outputs=[self.treeshared_stdout,self.tree_files])
        #Classify OTUs
        if self.fasta_subsample_files and self.count_table_files:
            subsample = ShellFunction(self.get_exec_path("mothur") + ' "#sub.sample(fasta=$1,count=$2,outputdir='+self.output_directory+'/)" > $3',cmd_format='{EXE} {IN} {OUT}')
            subsample = MultiMap(subsample, inputs=[self.fasta_files,self.count_table_files], outputs=[self.subsample_stdout,self.subsample_fasta_files,self.subsample_count_files])
        elif self.an_list_files and self.count_table_files:
            subsample = ShellFunction(self.get_exec_path("mothur") + ' "#sub.sample(list=$1,count=$2,outputdir='+self.output_directory+'/)" > $3',cmd_format='{EXE} {IN} {OUT}')
            subsample = MultiMap(subsample, inputs=[self.an_list_files,self.count_table_files], outputs=[self.subsample_stdout,self.subsample_list_files,self.subsample_count_files])
        elif self.fasta_subsample_files and not self.count_table_files:
            subsample = ShellFunction(self.get_exec_path("mothur") + ' "#sub.sample(fasta=$1,outputdir='+self.output_directory+'/)" > $2',cmd_format='{EXE} {IN} {OUT}')
            subsample = MultiMap(subsample, inputs=[self.fasta_files], outputs=[self.subsample_stdout,self.subsample_fasta_files])
        elif self.an_list_files and not self.count_table_files:
            subsample = ShellFunction(self.get_exec_path("mothur") + ' "#sub.sample(list=$1,outputdir='+self.output_directory+'/)" > $2',cmd_format='{EXE} {IN} {OUT}')
            subsample = MultiMap(subsample, inputs=[self.an_list_files], outputs=[self.subsample_stdout,self.subsample_list_files])
        if self.names_files and not self.groups_files:
            classifyotu = ShellFunction(self.get_exec_path("mothur") + ' "#classify.otu(list=$1,taxonomy=$2,name=$3,outputdir='+self.output_directory+'/,\
            label='+str(label)+')" > $4',cmd_format='{EXE} {IN} {OUT}')
            classifyotu = MultiMap(classifyotu, inputs=[self.an_list_files,self.taxonomy_files,self.names_files,], outputs=[self.classifyotu_stdout])
        elif self.names_files and self.groups_files:  
            classifyotu = ShellFunction(self.get_exec_path("mothur") + ' "#classify.otu(list=$1,taxonomy=$2,name=$3,group=$4,outputdir='+self.output_directory+'/,\
            label='+str(label)+')" > $5',cmd_format='{EXE} {IN} {OUT}')
            classifyotu = MultiMap(classifyotu, inputs=[self.an_list_files,self.taxonomy_files,self.names_files,self.groups_files], outputs=[self.classifyotu_stdout])
        elif self.count_table_files:
            classifyotu = ShellFunction(self.get_exec_path("mothur") + ' "#classify.otu(list=$1,taxonomy=$2,count=$3,outputdir='+self.output_directory+'/,\
            label='+str(label)+')" > $4',cmd_format='{EXE} {IN} {OUT}')
            classifyotu = MultiMap(classifyotu, inputs=[self.an_list_files,self.taxonomy_files,self.count_table_files,], outputs=[self.classifyotu_stdout])
        else:
            classifyotu = ShellFunction(self.get_exec_path("mothur") + ' "#classify.otu(list=$1,taxonomy=$2,outputdir='+self.output_directory+'/,label='+str(label)+')"\
             > $3',cmd_format='{EXE} {IN} {OUT}')
            classifyotu = MultiMap(classifyotu, inputs=[self.an_list_files,self.taxonomy_files], outputs=[self.classifyotu_stdout])
        krona = PythonFunction(get_taxonomy_to_Krona, cmd_format="{EXE} {ARG} {IN} {OUT}")
        krona(arguments=[self.get_exec_path("ImportMothurTaxonomy.pl"),self.output_directory],inputs=[self.classifyotu_stdout,self.shared_files],\
              outputs=[self.cons_taxonomy_files])