comparison cpo_galaxy_tree.py @ 6:cabceaa239e4 draft

planemo upload

5:698579246d0d

import time
import urllib.request
import gzip
import collections
import json
import ete3
import numpy

#parses some parameters
parser = optparse.OptionParser("Usage: %prog [options] arg1 arg2 ...")
parser.add_option("-t", "--tree", dest="treePath", type="string", default="./pipelineTest/tree.txt", help="identifier of the isolate")    
parser.add_option("-d", "--distance", dest="distancePath", type="string", default="./pipelineTest/distance.tab", help="absolute file path forward read (R1)")

        with gzip.open(inputpath, 'rb') as f:
            gzContent = f.read()
        with open(outputpath, 'wb') as out:
            out.write(gzContent)
        return True
#endregion

#region functions to parse result files
def ParseWorkflowResults(pathToResult):
    _worflowResult = {}
    r = pandas.read_csv(pathToResult, delimiter='\t', header=None) #read  the kraken2report.tsv
    r = r.replace(numpy.nan, '', regex=True)
    for i in range(len(r.index)):  
        _results = workflowResult()
        if(str(r.iloc[i,0]).lower() == "new"):
            _results.new = True
        else:
            _results.new = False        
        _results.ID = str(r.iloc[i,1])
        _results.ExpectedSpecies = str(r.iloc[i,2])
        _results.MLSTSpecies = str(r.iloc[i,3])
        _results.SequenceType = str(r.iloc[i,4])
        _results.MLSTScheme = (str(r.iloc[i,5]))
        _results.CarbapenemResistanceGenes = (str(r.iloc[i,6]))
        _results.OtherAMRGenes = (str(r.iloc[i,7]))
        _results.TotalPlasmids = int(r.iloc[i,8])
        for j in range(0,_results.TotalPlasmids):
            _plasmid = plasmidObj()
            _plasmid.PlasmidsID =(((str(r.iloc[i,9])).split(";"))[j])
            _plasmid.Num_Contigs = (((str(r.iloc[i,10])).split(";"))[j])
            _plasmid.PlasmidLength	= (((str(r.iloc[i,11])).split(";"))[j])
            _plasmid.PlasmidRepType	= (((str(r.iloc[i,12])).split(";"))[j])
            _plasmid.PlasmidMobility = ((str(r.iloc[i,13])).split(";"))[j]
            _plasmid.NearestReference = ((str(r.iloc[i,14])).split(";"))[j]
            _results.plasmids.append(_plasmid)
        _results.DefinitelyPlasmidContigs = (str(r.iloc[i,15]))
        _results.LikelyPlasmidContigs = (str(r.iloc[i,16]))
        _results.row = "\t".join(str(x) for x in r.ix[i].tolist())
        _worflowResult[_results.ID] = _results
    return _worflowResult

#endregion

def Main():
    metadata = ParseWorkflowResults(metadataPath)
    distance = read(distancePath)
    treeFile = "".join(read(treePath))

    distanceDict = {}
    for i in range(len(distance)):
        temp = distance[i].split("\t")
        distanceDict[temp[0]] = temp[1:]
    #region step5: tree construction
    t = ete3.Tree(treeFile)
    t.set_outgroup(t&"Reference")

    ts = ete3.TreeStyle()
    ts.show_leaf_name = True
    ts.show_branch_length = True
    ts.scale = 2000 #pixel per branch length unit
    ts.branch_vertical_margin = 15 #pixel between branches
    style2 = ete3.NodeStyle()
    style2["fgcolor"] = "#000000"
    style2["shape"] = "circle"
    style2["vt_line_color"] = "#0000aa"
    style2["hz_line_color"] = "#0000aa"
    style2["vt_line_width"] = 2
    style2["hz_line_width"] = 2
    style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted
    style2["hz_line_type"] = 0
    for n in t.traverse():
        n.set_style(style2)
    '''
    #region create detailed tree
    
    plasmidCount = 0
    for n in t.traverse():
        if (n.is_leaf() and not n.name == "Reference"):
            mData = metadata[n.name.replace(".fa","")]
            face = ete3.faces.TextFace(mData.MLSTSpecies,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_left = 10
            face.margin_right = 10
            n.add_face(face, 0, "aligned")
            face = ete3.faces.TextFace(mData.SequenceType,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 10
            n.add_face(face, 1, "aligned")
            face = ete3.faces.TextFace(mData.CarbapenemResistanceGenes,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 10
            n.add_face(face, 2, "aligned")
            index = 3
            if (mData.TotalPlasmids > plasmidCount):
                plasmidCount = mData.TotalPlasmids
            for i in range(0, mData.TotalPlasmids):
                face = ete3.faces.TextFace(mData.plasmids[i].PlasmidRepType,fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 10
                n.add_face(face, index, "aligned")
                index+=1
                face = ete3.faces.TextFace(mData.plasmids[i].PlasmidMobility,fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 10
                n.add_face(face, index, "aligned")
                index+=1

    face = ete3.faces.TextFace("Species",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    face.margin_left = 10
    (t&"Reference").add_face(face, 0, "aligned")
    face = ete3.faces.TextFace("Sequence Type",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    (t&"Reference").add_face(face, 1, "aligned")
    face = ete3.faces.TextFace("Carbapenamases",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    (t&"Reference").add_face(face, 2, "aligned")
    index = 3
    for i in range(0, plasmidCount):
        face = ete3.faces.TextFace("plasmid " + str(i) + " replicons",fsize=10,tight_text=True)
        face.border.margin = 5
        face.margin_right = 10
        (t&"Reference").add_face(face, index, "aligned")
        index+=1
        face = ete3.faces.TextFace("plasmid " + str(i) + " mobility",fsize=10,tight_text=True)
        face.border.margin = 5
        face.margin_right = 10
        (t&"Reference").add_face(face, index, "aligned")
        index+=1

    t.render("./pipelineTest/tree.png", w=5000,units="mm", tree_style=ts)
    
    #endregion
    '''
    #region create box tree
     #region step5: tree construction
    treeFile = "".join(read("./pipelineTest/tree.txt"))
    t = ete3.Tree(treeFile)
    t.set_outgroup(t&"Reference")

    ts = ete3.TreeStyle()
    ts.show_leaf_name = True
    ts.show_branch_length = True
    ts.scale = 2000 #pixel per branch length unit
    ts.branch_vertical_margin = 15 #pixel between branches
    style2 = ete3.NodeStyle()
    style2["fgcolor"] = "#000000"
    style2["shape"] = "circle"
    style2["vt_line_color"] = "#0000aa"
    style2["hz_line_color"] = "#0000aa"
    style2["vt_line_width"] = 2
    style2["hz_line_width"] = 2
    style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted
    style2["hz_line_type"] = 0
    for n in t.traverse():
        n.set_style(style2)
    plasmidIncs = {}
    for key in metadata:
        for plasmid in metadata[key].plasmids:
            for inc in plasmid.PlasmidRepType.split(","):
                if (inc.lower().find("inc") > -1):

                        plasmidIncs[inc] = [metadata[key].ID]
                    else:
                        if metadata[key].ID not in plasmidIncs[inc]:
                            plasmidIncs[inc].append(metadata[key].ID)
    #plasmidIncs = sorted(plasmidIncs)
    for n in t.traverse():
        if (n.is_leaf() and n.name == "Reference"):
            face = ete3.faces.TextFace("New?",fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            (t&"Reference").add_face(face, 0, "aligned")
            for i in range(len(plasmidIncs)): #this loop adds the columns (aka the incs) to the reference node
                face = ete3.faces.TextFace(list(plasmidIncs.keys())[i],fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 5
                face.margin_left = 5
                (t&"Reference").add_face(face, i + 1, "aligned")
            face = ete3.faces.TextFace("MLSTScheme",fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            (t&"Reference").add_face(face, len(plasmidIncs) + 0 + 1, "aligned")
            face = ete3.faces.TextFace("Sequence Type",fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            (t&"Reference").add_face(face, len(plasmidIncs) + 1 + 1, "aligned")
            face = ete3.faces.TextFace("Carbapenamases",fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            (t&"Reference").add_face(face, len(plasmidIncs) + 2 + 1, "aligned")
            for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds the columns (aka the incs) to the reference node
                face = ete3.faces.TextFace(distanceDict[list(distanceDict.keys())[0]][i],fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 5
                face.margin_left = 5
                (t&"Reference").add_face(face, len(plasmidIncs) + 2 + i + 1 + 1, "aligned")
        elif (n.is_leaf() and not n.name == "Reference"):
            if (metadata[n.name.replace(".fa","")].new == True):
                face = ete3.faces.RectFace(30,30,"green","green") # TextFace("Y",fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 5
                face.margin_left = 5
                face.vt_align = 1
                face.ht_align = 1
                n.add_face(face, 0, "aligned")
            for incs in plasmidIncs: #this loop adds presence/absence to the sample nodes
                if (n.name.replace(".fa","") in plasmidIncs[incs]):
                    face = ete3.faces.RectFace(30,30,"black","black") # TextFace("Y",fsize=10,tight_text=True)
                    face.border.margin = 5
                    face.margin_right = 5
                    face.margin_left = 5
                    face.vt_align = 1
                    face.ht_align = 1
                    n.add_face(face, list(plasmidIncs.keys()).index(incs) + 1, "aligned")
            mData = metadata[n.name.replace(".fa","")]
            face = ete3.faces.TextFace(mData.MLSTSpecies,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            n.add_face(face, len(plasmidIncs) + 0 + 1, "aligned")
            face = ete3.faces.TextFace(mData.SequenceType,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 5
            face.margin_left = 5
            n.add_face(face, len(plasmidIncs) + 1 + 1, "aligned")
            face = ete3.faces.TextFace(mData.CarbapenemResistanceGenes,fsize=10,tight_text=True)
            face.margin_right = 5
            face.margin_left = 5
            n.add_face(face, len(plasmidIncs) + 2 + 1, "aligned")
            for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds distance matrix
                face = ete3.faces.TextFace(list(distanceDict[n.name])[i],fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 5
                face.margin_left = 5
                n.add_face(face, len(plasmidIncs) + 2 + i + 1 + 1, "aligned")

    t.render("./tree.png", w=5000,units="mm", tree_style=ts)

    #endregion
#endregion

6:cabceaa239e4

import time
import urllib.request
import gzip
import collections
import json
import numpy
import ete3 as e


#parses some parameters
parser = optparse.OptionParser("Usage: %prog [options] arg1 arg2 ...")
parser.add_option("-t", "--tree", dest="treePath", type="string", default="./pipelineTest/tree.txt", help="identifier of the isolate")    
parser.add_option("-d", "--distance", dest="distancePath", type="string", default="./pipelineTest/distance.tab", help="absolute file path forward read (R1)")

        with gzip.open(inputpath, 'rb') as f:
            gzContent = f.read()
        with open(outputpath, 'wb') as out:
            out.write(gzContent)
        return True
def addFace(name):
    #if its the reference branch, populate the faces with column headers
    face = e.faces.TextFace(name,fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 5
    face.margin_left = 5
    return face
#endregion

#region functions to parse result files
def ParseWorkflowResults(pathToResult):
    _worflowResult = {}
    r = pandas.read_csv(pathToResult, delimiter='\t', header=0) 
    r = r.replace(numpy.nan, '', regex=True)
    for i in range(len(r.index)):  
        _results = workflowResult()
        if(str(r.loc[r.index[i], 'new']).lower() == "new"):
            _results.new = True
        else:
            _results.new = False        
        _results.ID = str(r.loc[r.index[i], 'ID'])
        _results.ExpectedSpecies = str(r.loc[r.index[i], 'Expected Species'])
        _results.MLSTSpecies = str(r.loc[r.index[i], 'MLST Species'])
        _results.SequenceType = str(r.loc[r.index[i], 'Sequence Type'])
        _results.MLSTScheme = (str(r.loc[r.index[i], 'MLST Scheme']))
        _results.CarbapenemResistanceGenes = (str(r.loc[r.index[i], 'Carbapenem Resistance Genes']))
        _results.OtherAMRGenes = (str(r.loc[r.index[i], 'Other AMR Genes']))
        _results.TotalPlasmids = int(r.loc[r.index[i], 'Total Plasmids'])
        for j in range(0,_results.TotalPlasmids):
            _plasmid = plasmidObj()
            _plasmid.PlasmidsID =(((str(r.loc[r.index[i], 'Plasmids ID'])).split(";"))[j])
            _plasmid.Num_Contigs = (((str(r.loc[r.index[i], 'Num_Contigs'])).split(";"))[j])
            _plasmid.PlasmidLength	= (((str(r.loc[r.index[i], 'Plasmid Length'])).split(";"))[j])
            _plasmid.PlasmidRepType	= (((str(r.loc[r.index[i], 'Plasmid RepType'])).split(";"))[j])
            _plasmid.PlasmidMobility = ((str(r.loc[r.index[i], 'Plasmid Mobility'])).split(";"))[j]
            _plasmid.NearestReference = ((str(r.loc[r.index[i], 'Nearest Reference'])).split(";"))[j]
            _results.plasmids.append(_plasmid)
        _results.DefinitelyPlasmidContigs = (str(r.loc[r.index[i], 'Definitely Plasmid Contigs']))
        _results.LikelyPlasmidContigs = (str(r.loc[r.index[i], 'Likely Plasmid Contigs']))
        _results.row = "\t".join(str(x) for x in r.ix[i].tolist())
        _worflowResult[_results.ID] = _results
    return _worflowResult

#endregion

def Main():
    metadata = ParseWorkflowResults(metadataPath)
    distance = read(distancePath)
    treeFile = "".join(read(treePath))

    distanceDict = {} #store the distance matrix as rowname:list<string>
    for i in range(len(distance)):
        temp = distance[i].split("\t")
        distanceDict[temp[0]] = temp[1:]
    #region step5: tree construction

    '''
    #region create detailed tree
    
    plasmidCount = 0
    for n in t.traverse():
        if (n.is_leaf() and not n.name == "Reference"):
            mData = metadata[n.name.replace(".fa","")]
            face = faces.TextFace(mData.MLSTSpecies,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_left = 10
            face.margin_right = 10
            n.add_face(face, 0, "aligned")
            face = faces.TextFace(mData.SequenceType,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 10
            n.add_face(face, 1, "aligned")
            face = faces.TextFace(mData.CarbapenemResistanceGenes,fsize=10,tight_text=True)
            face.border.margin = 5
            face.margin_right = 10
            n.add_face(face, 2, "aligned")
            index = 3
            if (mData.TotalPlasmids > plasmidCount):
                plasmidCount = mData.TotalPlasmids
            for i in range(0, mData.TotalPlasmids):
                face = faces.TextFace(mData.plasmids[i].PlasmidRepType,fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 10
                n.add_face(face, index, "aligned")
                index+=1
                face = faces.TextFace(mData.plasmids[i].PlasmidMobility,fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 10
                n.add_face(face, index, "aligned")
                index+=1

    face = faces.TextFace("Species",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    face.margin_left = 10
    (t&"Reference").add_face(face, 0, "aligned")
    face = faces.TextFace("Sequence Type",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    (t&"Reference").add_face(face, 1, "aligned")
    face = faces.TextFace("Carbapenamases",fsize=10,tight_text=True)
    face.border.margin = 5
    face.margin_right = 10
    (t&"Reference").add_face(face, 2, "aligned")
    index = 3
    for i in range(0, plasmidCount):
        face = faces.TextFace("plasmid " + str(i) + " replicons",fsize=10,tight_text=True)
        face.border.margin = 5
        face.margin_right = 10
        (t&"Reference").add_face(face, index, "aligned")
        index+=1
        face = faces.TextFace("plasmid " + str(i) + " mobility",fsize=10,tight_text=True)
        face.border.margin = 5
        face.margin_right = 10
        (t&"Reference").add_face(face, index, "aligned")
        index+=1

    t.render("./pipelineTest/tree.png", w=5000,units="mm", tree_style=ts)
    
    #endregion
    '''
    #region create box tree
    #region step5: tree construction
    treeFile = "".join(read(treePath))
    t = e.Tree(treeFile)
    t.set_outgroup(t&"Reference")

    #set the tree style
    ts = e.TreeStyle()
    ts.show_leaf_name = False
    ts.show_branch_length = True
    ts.scale = 2000 #pixel per branch length unit
    ts.branch_vertical_margin = 15 #pixel between branches
    style2 = e.NodeStyle()
    style2["fgcolor"] = "#000000"
    style2["shape"] = "circle"
    style2["vt_line_color"] = "#0000aa"
    style2["hz_line_color"] = "#0000aa"
    style2["vt_line_width"] = 2
    style2["hz_line_width"] = 2
    style2["vt_line_type"] = 0 # 0 solid, 1 dashed, 2 dotted
    style2["hz_line_type"] = 0
    for n in t.traverse():
        n.set_style(style2)

    #find the plasmid origins
    plasmidIncs = {}
    for key in metadata:
        for plasmid in metadata[key].plasmids:
            for inc in plasmid.PlasmidRepType.split(","):
                if (inc.lower().find("inc") > -1):

                        plasmidIncs[inc] = [metadata[key].ID]
                    else:
                        if metadata[key].ID not in plasmidIncs[inc]:
                            plasmidIncs[inc].append(metadata[key].ID)
    #plasmidIncs = sorted(plasmidIncs)
    for n in t.traverse(): #loop through the nodes of a tree
        if (n.is_leaf() and n.name == "Reference"):
            #if its the reference branch, populate the faces with column headers
            index = 0
            (t&"Reference").add_face(addFace("SampleID"), index, "aligned")
            index = index + 1
            (t&"Reference").add_face(addFace("New?"), index, "aligned")
            index = index + 1
            for i in range(len(plasmidIncs)): #this loop adds the columns (aka the incs) to the reference node
                (t&"Reference").add_face(addFace(list(plasmidIncs.keys())[i]), i + index, "aligned")
            index = index + len(plasmidIncs)
            (t&"Reference").add_face(addFace("MLSTScheme"), index, "aligned")
            index = index + 1
            (t&"Reference").add_face(addFace("Sequence Type"), index, "aligned")
            index = index + 1 
            (t&"Reference").add_face(addFace("Carbapenamases"), index, "aligned")
            index = index + 1
            for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds the distance matrix
                (t&"Reference").add_face(addFace(distanceDict[list(distanceDict.keys())[0]][i]), index + i, "aligned")
            index = index + len(distanceDict[list(distanceDict.keys())[0]])
        elif (n.is_leaf() and not n.name == "Reference"): 
            #not reference branches, populate with metadata
            index = 0
            mData = metadata[n.name.replace(".fa","")]
            n.add_face(addFace(mData.ID), index, "aligned")
            index = index + 1
            if (metadata[n.name.replace(".fa","")].new == True): #new column
                face = e.RectFace(30,30,"green","green") # TextFace("Y",fsize=10,tight_text=True)
                face.border.margin = 5
                face.margin_right = 5
                face.margin_left = 5
                face.vt_align = 1
                face.ht_align = 1
                n.add_face(face, index, "aligned")
            index = index + 1
            for incs in plasmidIncs: #this loop adds presence/absence to the sample nodes
                if (n.name.replace(".fa","") in plasmidIncs[incs]):
                    face = e.RectFace(30,30,"black","black") # TextFace("Y",fsize=10,tight_text=True)
                    face.border.margin = 5
                    face.margin_right = 5
                    face.margin_left = 5
                    face.vt_align = 1
                    face.ht_align = 1
                    n.add_face(face, list(plasmidIncs.keys()).index(incs) + index, "aligned")
            index = index + len(plasmidIncs)
            n.add_face(addFace(mData.MLSTSpecies), index, "aligned")
            index = index + 1
            n.add_face(addFace(mData.SequenceType), index, "aligned")
            index = index + 1
            n.add_face(addFace(mData.CarbapenemResistanceGenes), index, "aligned")
            index = index + 1
            for i in range(len(distanceDict[list(distanceDict.keys())[0]])): #this loop adds distance matrix
                n.add_face(addFace(list(distanceDict[n.name])[i]), index + i, "aligned")
                
    t.render("./tree.png", w=5000,units="mm", tree_style=ts) #save it as a png. or an phyloxml

    #endregion
#endregion