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

import os
import shutil
import datetime
import sys
import argparse

from ete3 import Tree

import logging

DEFAULT_FORMAT = 1

class TreeIndex:

	def __init__(self,tree_newick_fn,format=DEFAULT_FORMAT):
		self.tree_newick_fn=tree_newick_fn
		self.tree=Tree(tree_newick_fn,format=1)

		self.name_dict={}

		for node in self.tree.traverse("postorder"):
			self.name_dict[node.name]=node

		#print (self.name_dict)

	"""
		kmers_assigned_l:
			list of (list_of_nodes, count)

		dict:
			node => hit_vector
	"""

	def dict_from_list(self,kmers_assigned_l,lca=False):
		d={}

		hit_vector=[]

		npos=sum([x[1] for x in kmers_assigned_l])

		pos=0
		for (noden_l, count) in kmers_assigned_l:			
			if noden_l!=["0"]:
				if lca:
					noden_l=[self.lca(noden_l)]

				v=pos*[False] + count*[True] + (npos-pos-count)*[False]

				for noden in noden_l:
					try:
						assert len(d[noden])==len(v)
						d[noden]=d[noden] or v
					except KeyError:
						d[noden]=v

			pos+=count
		return d

	def lca(self,noden_l):
		"""Return LCA for a given list of nodes.
		"""
		assert len(noden_l)>0, noden_l
		if len(noden_l)==1:
			return noden_l[0]
		nodes_l=list(map(lambda x:self.name_dict[x],noden_l))
		lca=nodes_l[0].get_common_ancestor(nodes_l)
		return lca.name

	def name2gi(self,name):
		#print("err",name,file=sys.stderr)
		try:
			gi=self.name_dict[name].gi
		except AttributeError:
			return None
		return gi

	# scores from kmers hits
	def assign(self,kmers_assigned_l,simulate_lca=False):
		all_nodes_hit=set()

		d=self.dict_from_list(kmers_assigned_l,lca=simulate_lca)
		w=d.copy()

		for noden in d:
			if noden=="0":
				continue
			node=self.name_dict[noden]
			while node.up:
				node=node.up
				if node.name in d:
					w[noden]=w[noden] or d[node.name]
		return w

	def print_sam_header(self,file=sys.stdout):
		print("@HD VN:1.5 SO:unsorted",file=file)
		for node in self.tree.traverse("postorder"):
			self.name_dict[node.name]=node

			try:
				ur="\tUR:{}".format(node.fastapath)
			except:
				ur=""

			try:
				sp="\tSP:{}".format(node.sci_name)
			except:
				sp=""

			try:
				as_="\tAS:{}".format(node.gi)
			except:
				as_=""
<commit_msg>spaces => tabs in HD<commit_after>#! /usr/bin/env python3

import os
import shutil
import datetime
import sys
import argparse

from ete3 import Tree

import logging

DEFAULT_FORMAT = 1

class TreeIndex:

	def __init__(self,tree_newick_fn,format=DEFAULT_FORMAT):
		self.tree_newick_fn=tree_newick_fn
		self.tree=Tree(tree_newick_fn,format=1)

		self.name_dict={}

		for node in self.tree.traverse("postorder"):
			self.name_dict[node.name]=node

		#print (self.name_dict)

	"""
		kmers_assigned_l:
			list of (list_of_nodes, count)

		dict:
			node => hit_vector
	"""

	def dict_from_list(self,kmers_assigned_l,lca=False):
		d={}

		hit_vector=[]

		npos=sum([x[1] for x in kmers_assigned_l])

		pos=0
		for (noden_l, count) in kmers_assigned_l:			
			if noden_l!=["0"]:
				if lca:
					noden_l=[self.lca(noden_l)]

				v=pos*[False] + count*[True] + (npos-pos-count)*[False]

				for noden in noden_l:
					try:
						assert len(d[noden])==len(v)
						d[noden]=d[noden] or v
					except KeyError:
						d[noden]=v

			pos+=count
		return d

	def lca(self,noden_l):
		"""Return LCA for a given list of nodes.
		"""
		assert len(noden_l)>0, noden_l
		if len(noden_l)==1:
			return noden_l[0]
		nodes_l=list(map(lambda x:self.name_dict[x],noden_l))
		lca=nodes_l[0].get_common_ancestor(nodes_l)
		return lca.name

	def name2gi(self,name):
		#print("err",name,file=sys.stderr)
		try:
			gi=self.name_dict[name].gi
		except AttributeError:
			return None
		return gi

	# scores from kmers hits
	def assign(self,kmers_assigned_l,simulate_lca=False):
		all_nodes_hit=set()

		d=self.dict_from_list(kmers_assigned_l,lca=simulate_lca)
		w=d.copy()

		for noden in d:
			if noden=="0":
				continue
			node=self.name_dict[noden]
			while node.up:
				node=node.up
				if node.name in d:
					w[noden]=w[noden] or d[node.name]
		return w

	def print_sam_header(self,file=sys.stdout):
		print("@HD\tVN:1.5\tSO:unsorted",file=file)
		for node in self.tree.traverse("postorder"):
			self.name_dict[node.name]=node

			try:
				ur="\tUR:{}".format(node.fastapath)
			except:
				ur=""

			try:
				sp="\tSP:{}".format(node.sci_name)
			except:
				sp=""

			try:
				as_="\tAS:{}".format(node.gi)
			except:
				as_=""

{'fraction_non_alphanumeric': 0.14217578923461277, 'fraction_numerical': 0.006813536225300931, 'mean_word_length': 3.781758957654723, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 8, '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': '1512286', 'n_tokens_mistral': 2161, 'n_tokens_neox': 1834, 'n_words': 333}