<commit_before>#! /usr/bin/env python3

import os
import shutil
import datetime
import sys
import argparse
import bitarray
import itertools

from ete3 import Tree



import logging

DEFAULT_FORMAT = 1


def cigar_from_mask(mask):
	c=[]
	runs=itertools.groupby(mask)
	for run in runs:
		c.append(str(len(list(run[1]))))
		c.append('=' if run[0] else 'X')
	return "".join(c)


class Read:
	def __init__(self, tree, simulate_lca=False, annotate=False,tie_lca=False,dont_translate_blocks=False):
		self.tree=tree
		self.k=tree.k
		self.simulate_lca=simulate_lca
		self.annotate=annotate
		self.tie_lca=tie_lca
		self.dont_translate_blocks=dont_translate_blocks

	def process_krakline(self,krakline,form,crit):
		self.load_krakline(krakline)
		self.find_assignments()
		#print(self.asgs)
		self.filter_assignments()
		self.print_assignments(form,crit)


	def load_krakline(self,krakline):
		parts=krakline.strip().split("\t")
		self.qname,_,qlen,self.krakmers=parts[1:5]
		self.qlen=int(qlen)
		if len(parts)==7:
			self.seq=parts[5]
			self.qual=parts[6]
		else:
			self.seq="*"
			self.qual="*"

		#self.hitmasks=None
		self.asgs={}

		# list of (count,list of nodes)
		self.kmer_blocks=[]

		b_sum=0
		for b in self.krakmers.split(" "):
			(ids,count)=b.split(":")
			count=int(count)
			b_sum+=count
			rnames=ids.split(",")
			if self.simulate_lca:
				rnames=[self.tree.lca(rnames)]
			self.kmer_blocks.append((rnames,count))
		assert self.qlen==b_sum+self.k-1 or self.qlen<self.k, krakline


	def find_assignments(self):
		# hits before top-down propagation
		hitmasks,covmasks=self.tree.masks_from_kmer_blocks(self.kmer_blocks,simulate_lca=self.simulate_lca)

		rnames=hitmasks.keys()

		# hits after top-down propagation
		for rname in rnames:
			self.asgs[rname]={
					'hitmask' : hitmasks[rname].copy(),
					'covmask' : covmasks[rname].copy(),
				}

			node=self.tree.name_dict[rname]
			for p_rname in self.tree.upnodes_dict[rname] & rnames:
				self.asgs[rname]['hitmask']|=hitmasks[p_rname]
				self.asgs[rname]['covmask']|=covmasks[p_rname]


	def filter_assignments(self):
		"""
		Annotate & filter assignment to a node.

		rname=None => unassigned
		"""

		self.max_hit=0
		self.max_cov=0

		self.max_hit_rnames=[]
		self.max_cov_rnames=[]

		for rname in self.asgs:
			asg=self.asgs[rname]

			"""
			1. hit count
			"""
			hit=asg['hitmask'].count()
			asg['h1']=hit

			if hit>self.max_hit:
				self.max_hit=hit
				self.max_hit_rnames=[rname]
			elif hit==self.max_hit:
				self.max_hit_rnames.append(rname)

			"""
			2. coverage
			"""
			cov=asg['covmask'].count()
			asg['c1']=cov

			if cov>self.max_cov:
				self.max_cov=cov
				self.max_cov_rnames=[rname]
			elif cov==self.max_cov:
				self.max_cov_rnames.append(rname)


	def print_assignments(self, form, crit):
		if crit=="h1":
			winners=self.max_cov_rnames
		elif crit=="c1":
			winners=self.max_hit_rnames

		tie_solved=False

		if self.tie_lca and len(winners)>1:
			tie_solved=True
			lca=self.tree.lca(winners)
			winners=[lca]
			# fix what if this node exists!
			asg=self.asgs[lca]={}
			asg['covcigar']="{}I".format(self.qlen)
			asg['hitcigar']="{}I".format(self.qlen)
			if crit=="h1":
				asg['c1']=-1
				asg['h1']=self.max_hit
			elif crit=="c1":
				asg['c1']=self.max_cov
				asg['h1']=-1

		if len(self.max_hit_rnames)>0:
			for rname in winners:
				asg=self.asgs[rname]
				if form=="sam":
					# compute cigar
					if not tie_solved:
						asg['covcigar']=cigar_from_mask(asg['covmask'])<commit_msg>Fix bug c1 <-> h1<commit_after>#! /usr/bin/env python3

import os
import shutil
import datetime
import sys
import argparse
import bitarray
import itertools

from ete3 import Tree



import logging

DEFAULT_FORMAT = 1


def cigar_from_mask(mask):
	c=[]
	runs=itertools.groupby(mask)
	for run in runs:
		c.append(str(len(list(run[1]))))
		c.append('=' if run[0] else 'X')
	return "".join(c)


class Read:
	def __init__(self, tree, simulate_lca=False, annotate=False,tie_lca=False,dont_translate_blocks=False):
		self.tree=tree
		self.k=tree.k
		self.simulate_lca=simulate_lca
		self.annotate=annotate
		self.tie_lca=tie_lca
		self.dont_translate_blocks=dont_translate_blocks

	def process_krakline(self,krakline,form,crit):
		self.load_krakline(krakline)
		self.find_assignments()
		#print(self.asgs)
		self.filter_assignments()
		self.print_assignments(form,crit)


	def load_krakline(self,krakline):
		parts=krakline.strip().split("\t")
		self.qname,_,qlen,self.krakmers=parts[1:5]
		self.qlen=int(qlen)
		if len(parts)==7:
			self.seq=parts[5]
			self.qual=parts[6]
		else:
			self.seq="*"
			self.qual="*"

		#self.hitmasks=None
		self.asgs={}

		# list of (count,list of nodes)
		self.kmer_blocks=[]

		b_sum=0
		for b in self.krakmers.split(" "):
			(ids,count)=b.split(":")
			count=int(count)
			b_sum+=count
			rnames=ids.split(",")
			if self.simulate_lca:
				rnames=[self.tree.lca(rnames)]
			self.kmer_blocks.append((rnames,count))
		assert self.qlen==b_sum+self.k-1 or self.qlen<self.k, krakline


	def find_assignments(self):
		# hits before top-down propagation
		hitmasks,covmasks=self.tree.masks_from_kmer_blocks(self.kmer_blocks,simulate_lca=self.simulate_lca)

		rnames=hitmasks.keys()

		# hits after top-down propagation
		for rname in rnames:
			self.asgs[rname]={
					'hitmask' : hitmasks[rname].copy(),
					'covmask' : covmasks[rname].copy(),
				}

			node=self.tree.name_dict[rname]
			for p_rname in self.tree.upnodes_dict[rname] & rnames:
				self.asgs[rname]['hitmask']|=hitmasks[p_rname]
				self.asgs[rname]['covmask']|=covmasks[p_rname]


	def filter_assignments(self):
		"""
		Annotate & filter assignment to a node.

		rname=None => unassigned
		"""

		self.max_hit=0
		self.max_cov=0

		self.max_hit_rnames=[]
		self.max_cov_rnames=[]

		for rname in self.asgs:
			asg=self.asgs[rname]

			"""
			1. hit count
			"""
			hit=asg['hitmask'].count()
			asg['h1']=hit

			if hit>self.max_hit:
				self.max_hit=hit
				self.max_hit_rnames=[rname]
			elif hit==self.max_hit:
				self.max_hit_rnames.append(rname)

			"""
			2. coverage
			"""
			cov=asg['covmask'].count()
			asg['c1']=cov

			if cov>self.max_cov:
				self.max_cov=cov
				self.max_cov_rnames=[rname]
			elif cov==self.max_cov:
				self.max_cov_rnames.append(rname)


	def print_assignments(self, form, crit):
		if crit=="c1":
			winners=self.max_cov_rnames
		elif crit=="h1":
			winners=self.max_hit_rnames

		tie_solved=False

		if self.tie_lca and len(winners)>1:
			tie_solved=True
			lca=self.tree.lca(winners)
			winners=[lca]
			# fix what if this node exists!
			asg=self.asgs[lca]={}
			asg['covcigar']="{}I".format(self.qlen)
			asg['hitcigar']="{}I".format(self.qlen)
			if crit=="h1":
				asg['c1']=-1
				asg['h1']=self.max_hit
			elif crit=="c1":
				asg['c1']=self.max_cov
				asg['h1']=-1

		if len(self.max_hit_rnames)>0:
			for rname in winners:
				asg=self.asgs[rname]
				if form=="sam":
					# compute cigar
					if not tie_solved:
						asg['covcigar']=cigar_from_mask(asg['covmask'])
{'fraction_non_alphanumeric': 0.15262707024557395, 'fraction_numerical': 0.00885208452312964, 'mean_word_length': 4.105685131195335, 'pattern_counts': {'":': 12, '<': 6, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7043148', 'n_tokens_mistral': 3404, 'n_tokens_neox': 2929, 'n_words': 442}