-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
e408590
commit dbfd603
Showing
2 changed files
with
49 additions
and
62 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,60 +1,61 @@ | ||
import csv | ||
import pysam | ||
import pandas as pd | ||
import numpy as np | ||
import os | ||
import logging | ||
import gfapy | ||
import matplotlib.pyplot as plt | ||
logging.getLogger("matplotlib.font_manager").disabled = True | ||
import matplotlib as mt | ||
import gfapy | ||
import os | ||
|
||
import pysam | ||
import pandas as pd | ||
import numpy as np | ||
from strainy.params import * | ||
logging.getLogger("matplotlib.font_manager").disabled = True | ||
|
||
|
||
def write_bam(edge, I, AF, cl_file=None,file=None): | ||
infile = pysam.AlignmentFile(StRainyArgs().bam, "rb") | ||
if file==None: | ||
outfile = pysam.AlignmentFile("%s/bam/coloredBAM_unitig_%s.bam" % (StRainyArgs().output_intermediate, edge), "wb", template=infile) | ||
else: | ||
outfile = pysam.AlignmentFile(file,"wb", template=infile) | ||
if cl_file==None: | ||
cl = pd.read_csv("%s/clusters/clusters_%s_%s_%s.csv" % (StRainyArgs().output_intermediate, edge, I, AF),keep_default_na=False) | ||
else: | ||
cl = pd.read_csv(cl_file,keep_default_na=False) | ||
iter = infile.fetch(edge,until_eof=True) | ||
def write_bam(edge, cl, infile,outfile): | ||
"""Creates new bam file based on ifnfile and add YC tag to the alignment based on csv file""" | ||
iterbam = infile.fetch(edge,until_eof=True) | ||
cmap = plt.get_cmap("viridis") | ||
cl.loc[cl["Cluster"] == "NA", "Cluster"] = 0 | ||
#clusters=sorted(set(cl["Cluster"].astype(int))) | ||
clusters = set(cl["Cluster"]) | ||
cmap = cmap(np.linspace(0, 1, len(clusters))) | ||
colors={} | ||
i=0 | ||
colors[0] = "#505050" | ||
|
||
try: | ||
clusters.remove("0") | ||
except: KeyError | ||
except KeyError: | ||
pass | ||
for cluster in clusters: | ||
colors[cluster] = mt.colors.to_hex(cmap[i]) | ||
i = i+1 | ||
cl_dict = dict(zip(cl.ReadName, cl.Cluster)) | ||
|
||
for read in iter: | ||
for read in iterbam: | ||
try: | ||
#clN = int(cl_dict[str(read).split()[0]]) | ||
clN = cl_dict[str(read).split()[0]] | ||
tag = colors[clN] | ||
cl_n = cl_dict[str(read).split()[0]] | ||
tag = colors[cl_n] | ||
read.set_tag("YC", tag, replace=False) | ||
outfile.write(read) | ||
except (KeyError): | ||
except KeyError: | ||
continue | ||
outfile.close() | ||
|
||
|
||
def color(edge,cl_file=None,file=None): | ||
"""Creates colored edge bam based on strainy csv file with clusters IDs by default""" | ||
try: | ||
write_bam(edge, I, StRainyArgs().AF,cl_file,file) | ||
except (FileNotFoundError): | ||
infile = pysam.AlignmentFile(StRainyArgs().bam, "rb") | ||
if file is None: | ||
outfile = pysam.AlignmentFile( | ||
f"{StRainyArgs().output_intermediate}/bam/coloredBAM_unitig_{edge}.bam", | ||
"wb", template=infile) | ||
else: | ||
outfile = pysam.AlignmentFile(file,"wb", template=infile) | ||
if cl_file is None: | ||
cl = pd.read_csv( | ||
f"{StRainyArgs().output_intermediate}/clusters/clusters_{edge}_{I}_{StRainyArgs().AF}.csv", | ||
keep_default_na=False) | ||
else: | ||
cl = pd.read_csv(cl_file,keep_default_na=False) | ||
write_bam(edge,cl,infile,outfile) | ||
except FileNotFoundError: | ||
pass | ||
|