Commit d5d610e4 authored by rchikhi's avatar rchikhi
Browse files

added barcode graph from ema script

parent f75e2095
#arguments: samfile produced by e.g. EMA with MI information, mapped to contigs (not unitigs)
# returns a barcode graph
# additionnally takes into account MI information, but only on molecules fully aligned inside a single contig
import sys
import itertools
import gzip
import pysam
import operator
from collections import Counter, defaultdict
from networkx import nx
g = nx.Graph()
from statistics import mean
# parameters
min_barcode_reads = 20 # drop barcodes having less than X reads
min_overlap_length = 1000 # minimum overlap length between molecules
contained_threshold = 3000 # minimum distance to contig extremity to consider a molecule to be fully within a contig (=resolved)
# -----------
molecule_span = dict()
def compute_molecule_id(barcode,mi):
return barcode+"-MI"+str(mi)
def update_molecule_span(molecule, read_contig, read_start, read_end):
global molecule_span
if molecule in molecule_span:
previous_span = molecule_span[molecule]
pcontig, pstart, pend = previous_span
if pcontig == "multicontig": return
if pcontig != read_contig:
molecule_span[molecule] = ("multicontig",0,0)
pstart = min(pstart, read_start)
pend = max(pend, read_end)
pcontig, pstart, pend = read_contig, read_start, read_end
molecule_span[molecule] = (pcontig, pstart, pend)
# init some global variables
read_barcode = dict()
read_mi = dict()
barcode_read_count = Counter()
nb_reads_inserted, nb_bx_key_errors, nb_mi_key_errors = 0,0,0
nb_bad_mapq, nb_discarded, nb_mapped = 0,0,0
# open sam file
sam_file = sys.argv[1]
if sam_file.endswith('.sam'):
samfile = pysam.AlignmentFile(sam_file, "r")
samfile = pysam.AlignmentFile(sam_file, "rb")
# parse the sam header to get contig lengths
read_barcode = dict()
barcode_read_count = Counter()
molecule_read_count = Counter()
contig_molecules = dict()
for record in samfile.header['SQ']:
contig = record['SN']
length = int(record['LN'])
contig_dict[contig] = []
contig_lengths[contig] = length
contig_molecules[contig] = set()
# parse the reads
for read in samfile.fetch():
if read.is_unmapped or read.is_secondary:
nb_mapped += 1
mapq = read.mapping_quality
if mapq < 50: #not 60, but 50
nb_discarded += 1
nb_bad_mapq += 1
read_id = read.query_name
read_contig = read.reference_name
read_start = read.reference_start
read_end = read.reference_end
barcode_read_count[barcode] += 1
except KeyError:
nb_bx_key_errors += 1
if read.has_tag('MI'):
except KeyError:
nb_mi_key_errors += 1
if read_id in read_barcode and read_id in read_mi: # make sure both barcode and MI are reported for this read
# compute molecule span
molecule = compute_molecule_id(read_barcode[read_id],read_mi[read_id])
update_molecule_span(molecule, read_contig, read_start, read_end)
molecule_read_count[molecule] += 1
print("parsed %d alignments, %d discarded reads incl %d bad mapq" % (nb_mapped,nb_discarded, nb_bad_mapq))
print(nb_reads_inserted,"reads taken into account, ","%d/%d" %(nb_bx_key_errors,nb_mi_key_errors),"BX/MI skipped")
print("number of barcoded reads indexed:",len(read_barcode))
print("number of barcodes :",len(barcode_read_count))
# we say that a molecule is _resolved_ if it's far away from a contig extremity
# those correspond to molecules that are fully spanned in contigs
# all other molecules are unresolved and we'll just assign them together to their barcode
# i.e. all unresolved molecules are the same molecule (=the barcode)
resolved_molecules = set()
for molecule in molecule_span:
mcontig, mstart, mend = molecule_span[molecule]
contig_len = contig_lengths[mcontig]
if mstart > contained_threshold and mend < contig_len-contained_threshold:
print(len(resolved_molecules),"/",len(molecule_span), "contig-isolated molecules")
def get_resolved_barcode(mol):
if mol in resolved_molecules:
return mol
#def compute_molecule_id(barcode,mi):
# return barcode+"-MI"+str(mi)
return mol.split("-MI")[0]
# construct graph nodes as sufficiently seen barcodes
nb_reads_per_barcode = []
for barcode in barcode_read_count:
count = barcode_read_count[barcode]
if count > min_barcode_reads:
nb_reads_per_barcode += [count]
print("mean number of reads per barcode",sum(nb_reads_per_barcode)/len(nb_reads_per_barcode))
link_count = Counter()
link_contigs = dict()
debug_contig_mols = open("debug_contig_mols.txt","w")
t = int(len(contig_dict)/10)
# iterate over all contigs one by one
# to construct barcode graph edges
# this code only has one constant: min_overlap_length
for contig_number, contig in enumerate(contig_dict):
barcodes = []
barcodes_positions = defaultdict(list)
active_mols = []
last_endpoint = 0
molecule_overlaps = set()
# record all endpoints of molecules
molecule_endpoints = []
for molecule in contig_molecules[contig]:
pcontig, pstart, pend = molecule_span[molecule]
molecule_endpoints += [(pstart, START, molecule)]
molecule_endpoints += [(pend, END, molecule)]
molecule_endpoints = sorted(molecule_endpoints)
# sliding window over all endpoints of molecules
for endpoint_pos, endpoint_type, molecule in molecule_endpoints:
print("molecule",molecule,"endpoint:",endpoint_pos,"start" if endpoint_type == START else "end","active mols",len(active_mols))
#find all existing overlaps
molecule_overlaps.update(set([(x,y) for (ast,ae,x) in active_mols for (bst,be,y) in active_mols if x != y and ast < bst and ae-bst >= min_length]))
# determine which molecules are now out
new_active_mols = []
for active_mol in active_mols:
astart, aend, amol = active_mol
# keep only mols that finish after the current position
if aend > endpoint_pos:
new_active_mols += [active_mol]
#print("endpoint_type/pos","END" if endpoint_type==1 else "START",endpoint_pos,"molecule",molecule,"astart/aend/amol",astart,aend,amol)
# add maybe new active mol
if endpoint_type == START:
pcontig, pstart, pend = molecule_span[molecule]
new_active_mols += [(pstart, pend, molecule)]
active_mols = new_active_mols
# some debug
debug_advanced = False
if debug_advanced:
str_mols = []
for molecule in molecule_span:
(pcontig,start,end) = molecule_span[molecule]
str_mols += [(start,end) for molecule in contig_molecules[contig]]
str_mols = str(sorted(str_mols))
debug_contig_mols.write("contig number %d length %d, %d molecules: %s\n" % (contig_number,contig_lengths[contig],len(molecule_endpoints)/2,str_mols))
debug_contig_mols.write("contig number %d length %d, %d molecules\n" % (contig_number,contig_lengths[contig],len(molecule_endpoints)/2))
# now iterate over all found overlaps between molecules
for mol1, mol2 in set(molecule_overlaps):
# since we're constructing a barcode graph: get barcode of unresolved molecules
a = get_resolved_barcode(mol1)
b = get_resolved_barcode(mol2)
a,b = min([(a,b),(b,a)])
if label not in link_count:
link_count[label] = 0
link_contigs[label] = ""
link_count[label] += molecule_read_count[mol1] + molecule_read_count[mol2]
link_contigs[label] += " " + contig
g.add_edge(a,b, weight=link_count[label], contigs=link_contigs[label])
if contig_number % t == 0:
sys.stderr.write('..%d%%..' % ((contig_number/t)*10))
print(len(g.nodes()),"graph nodes",len(g.edges()),"edges")
# remove isolated nodes
for node in g.nodes():
if == 0:
list_isolated += [node]
print(len(list_isolated),"isolated nodes removed")
nx.write_graphml(g, sys.argv[1]+".graphml")
nx.write_gpickle(g, sys.argv[1]+".gpickle")
# write graph in paul's format
debug_graph = open("debug_graph.tsv","w")
for edge in g.edges():
a,b = edge
a,b = min([(a,b),(b,a)])
debug_graph.write("%s %s %d\n" % (a,b,link_count[label]))
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment