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

import os
import metag
import datetime
import sys
import argparse

from tree_formatter import *

import logging

logger = logging.getLogger()
handler = logging.StreamHandler()
formatter = logging.Formatter(
        '%(asctime)s %(levelname)-8s %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)

#t=read_newick("id_tree_bin.newick",format=10)
#t.render("test.pdf")

#for leaf in t.iter_leaves():
#	#print (leaf.name, leaf.named_lineage)
#	if hasattr(leaf,"fastapath"):
#		print(leaf.fastapath)
#	print(leaf.features)
#	print(leaf.infasta_seqnum)
#	print()

#def NodeIndex:
#	def __init__(self,node,output_fasta_fn):
#		self.node=node
#		self.output_fasta_fn=output_fasta_fn
#		self.kmers_set=set()
#		self.fasta_fo=None
#
#	def add_kmer(self,kmer):
#		self.kmers_set.add(kmers_set)
#
#	def add_kmers(self,kmers_set):
#		self.kmers_set|=kmers_set
#
#	def __del__(self):
#		if len(self.kmers)>0:
#			metag.set_to_fasta(
#					fasta_fn=self.output_fasta_fn,
#					set_of_kmers=self.kmers_set,
#					assemble=True
#				)


def size_in_mb(file_fn):
	return os.path.getsize(file_fn)/(1024**2)

class TreeIndex:

	def __init__(self,tree_newick_fn,format=10,directory="./"):
		self.tree=read_newick(tree_newick_fn,format=format)
		self.newick_directory=os.path.dirname(tree_newick_fn)
		self.directory=directory
		os.makedirs(directory,exist_ok=True)


	#def _rightmost_descendant(self,node):
	#	while not node.is_leaf():
	#		node=node.get_children()[-1]
	#	return node
	#
	#def _rightmost_leaf_of_subtree(self,leaf):
	#	node=leaf
	#	while node.get_ancestors() is not None:
	#		#print(node.name)
	#		ancestors=node.get_ancestors()
	#		if len(ancestors)==0:
	#			break
	#		father=ancestors[0]
	#		fathers_children=father.get_children()
	#		assert(len(fathers_children)<=2)
	#		if len(fathers_children)!=1:
	#			if fathers_children[1]==node:
	#				break
	#		node=father
	#
	#	return _rightmost_descendant(node)
	#
	#def _save_kmers_to_dict(kmers_set,rightmost_leaf):
	#	for kmer in kmers_set:

	@staticmethod
	def _node_debug(node):
		if hasattr(node,"common_name") and node.common_name!="":
			return "{}_{}".format(node.name,node.common_name)
		elif hasattr(node,"sci_name") and node.sci_name!="":
			return "{}_{}".format(node.name,node.sci_name)
		else:
			return "{}".format(node.name)

	def create_fasta(self,node,kmers_set):
		fasta_fn=os.path.join(self.directory,"{}.fa".format(node.name))
		logger.info('Creating FASTA "{}" (assembling {} kmers)'.format(fasta_fn,kmers_set//2))
		logger.debug('... from k-mers "{}"'.format(", ".format(kmers_set)))
		metag.set_to_fasta(
				fasta_fn=fasta_fn,
				set_of_kmers=kmers_set,
				assemble=True,
				contig_prefix="node{}".format(node.name),
			)
		logger.info('... FASTA created "{}"'.format(fasta_fn))

	def get_shared_kmers(self,node,k):
		logger.info('BEGIN get shared k-mers for node "{}"'.format(self._node_debug(node)))
		if node.is_leaf():
			kmers_set=set()
			if hasattr(node,"fastapath"):
				fastas_fn=node.fastapath.split("@")
				#print( datetime.datetime.now().time() )
				for fasta_fn in fastas_fn:
					fasta_fn=os.path.join(self.newick_directory,fasta_fn)
					if os.path.isfile(fasta_fn):
						logger.info('Loading FASTA "{}" (size: {} MB)'.format(fasta_fn,size_in_mb(fasta_fn)))
						kmers_set|=metag.set_from_fasta(fasta_fn,k)
						logger.info('... FASTA loaded "{}"'.format(fasta_fn))
					else:
						logger.warning('FASTA "{}" does not exist'.format(fasta_fn))
			logger.info('END get shared k-mers for node "{}" ({} kmers)'.format(self._node_debug(node),len(kmers_set)//2))
			logger.debug('... kmers (from fasta files): "{}"'.format(", ".join(kmers_set)))
			return kmers_set
		else:
			children=node.get_children()
			list_of_full_sets=[self.get_shared_kmers(child,k=k) for child in children]
			intersection=set.intersection(*list_of_full_sets)
			list_of_reduced_sets=[full_set-intersection for full_set in list_of_full_sets]
			for (i,reduced_set) in enumerate(list_of_reduced_sets):
				if len(reduced_set)>0:
					self.create_fasta(children[i],reduced_set)
			logger.info('END get shared k-mers for node "{}" ({} kmers)'.format(self._node_debug(node),len(intersection)//2))
			logger.debug('... kmers (from children): "{}"'.format(", ".join(intersection)))
			return intersection

	def build_index(self,k):
		logger.info('Going to build index for k={}'.format(k))
		self.get_shared_kmers(self.tree.get_tree_root(),k=k)


if __name__ == "__main__":

	parser = argparse.ArgumentParser(description='Build index.')
	parser.add_argument(
			'-n','--newick-tree',
			type=str,
			metavar='str',
			dest='newick_fn',
			required=True,
			help='Taxonomic tree (Newick).',
		)
	parser.add_argument(
			'-k',
			type=int,
			metavar='int',
			dest='k',
			required=True,
			help='K-mer length k.',
		)
	parser.add_argument(
			'-o','--output-dir',
			type=str,
			metavar='str',
			dest='output_dir_fn',
			required=True,
			help='Output directory.',
		)

	args = parser.parse_args()

	k=args.k
	assert k>0
	newick_fn=args.newick_fn
	output_dir_fn=args.output_dir_fn

	logger.info("Starting index construction")
	logger.info("       newick : {}".format(newick_fn))
	logger.info("   output dir : {}".format(output_dir_fn))
	logger.info("            k : {}".format(k))

	ti=TreeIndex(
			tree_newick_fn=newick_fn,
			directory=output_dir_fn
		)
	ti.build_index(k=k)
<commit_msg>Minifix.<commit_after>#! /usr/bin/env python3

import os
import metag
import datetime
import sys
import argparse

from tree_formatter import *

import logging

logger = logging.getLogger()
handler = logging.StreamHandler()
formatter = logging.Formatter(
        '%(asctime)s %(levelname)-8s %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)

#t=read_newick("id_tree_bin.newick",format=10)
#t.render("test.pdf")

#for leaf in t.iter_leaves():
#	#print (leaf.name, leaf.named_lineage)
#	if hasattr(leaf,"fastapath"):
#		print(leaf.fastapath)
#	print(leaf.features)
#	print(leaf.infasta_seqnum)
#	print()

#def NodeIndex:
#	def __init__(self,node,output_fasta_fn):
#		self.node=node
#		self.output_fasta_fn=output_fasta_fn
#		self.kmers_set=set()
#		self.fasta_fo=None
#
#	def add_kmer(self,kmer):
#		self.kmers_set.add(kmers_set)
#
#	def add_kmers(self,kmers_set):
#		self.kmers_set|=kmers_set
#
#	def __del__(self):
#		if len(self.kmers)>0:
#			metag.set_to_fasta(
#					fasta_fn=self.output_fasta_fn,
#					set_of_kmers=self.kmers_set,
#					assemble=True
#				)


def size_in_mb(file_fn):
	return os.path.getsize(file_fn)/(1024**2)

class TreeIndex:

	def __init__(self,tree_newick_fn,format=10,directory="./"):
		self.tree=read_newick(tree_newick_fn,format=format)
		self.newick_directory=os.path.dirname(tree_newick_fn)
		self.directory=directory
		os.makedirs(directory,exist_ok=True)


	#def _rightmost_descendant(self,node):
	#	while not node.is_leaf():
	#		node=node.get_children()[-1]
	#	return node
	#
	#def _rightmost_leaf_of_subtree(self,leaf):
	#	node=leaf
	#	while node.get_ancestors() is not None:
	#		#print(node.name)
	#		ancestors=node.get_ancestors()
	#		if len(ancestors)==0:
	#			break
	#		father=ancestors[0]
	#		fathers_children=father.get_children()
	#		assert(len(fathers_children)<=2)
	#		if len(fathers_children)!=1:
	#			if fathers_children[1]==node:
	#				break
	#		node=father
	#
	#	return _rightmost_descendant(node)
	#
	#def _save_kmers_to_dict(kmers_set,rightmost_leaf):
	#	for kmer in kmers_set:

	@staticmethod
	def _node_debug(node):
		if hasattr(node,"common_name") and node.common_name!="":
			return "{}_{}".format(node.name,node.common_name)
		elif hasattr(node,"sci_name") and node.sci_name!="":
			return "{}_{}".format(node.name,node.sci_name)
		else:
			return "{}".format(node.name)

	def create_fasta(self,node,kmers_set):
		fasta_fn=os.path.join(self.directory,"{}.fa".format(node.name))
		logger.info('Creating FASTA "{}" (assembling {} kmers)'.format(fasta_fn,len(kmers_set)//2))
		logger.debug('... from k-mers "{}"'.format(", ".format(kmers_set)))
		metag.set_to_fasta(
				fasta_fn=fasta_fn,
				set_of_kmers=kmers_set,
				assemble=True,
				contig_prefix="node{}".format(node.name),
			)
		logger.info('... FASTA created "{}"'.format(fasta_fn))

	def get_shared_kmers(self,node,k):
		logger.info('BEGIN get shared k-mers for node "{}"'.format(self._node_debug(node)))
		if node.is_leaf():
			kmers_set=set()
			if hasattr(node,"fastapath"):
				fastas_fn=node.fastapath.split("@")
				#print( datetime.datetime.now().time() )
				for fasta_fn in fastas_fn:
					fasta_fn=os.path.join(self.newick_directory,fasta_fn)
					if os.path.isfile(fasta_fn):
						logger.info('Loading FASTA "{}" (size: {} MB)'.format(fasta_fn,size_in_mb(fasta_fn)))
						kmers_set|=metag.set_from_fasta(fasta_fn,k)
						logger.info('... FASTA loaded "{}"'.format(fasta_fn))
					else:
						logger.warning('FASTA "{}" does not exist'.format(fasta_fn))
			logger.info('END get shared k-mers for node "{}" ({} kmers)'.format(self._node_debug(node),len(kmers_set)//2))
			logger.debug('... kmers (from fasta files): "{}"'.format(", ".join(kmers_set)))
			return kmers_set
		else:
			children=node.get_children()
			list_of_full_sets=[self.get_shared_kmers(child,k=k) for child in children]
			intersection=set.intersection(*list_of_full_sets)
			list_of_reduced_sets=[full_set-intersection for full_set in list_of_full_sets]
			for (i,reduced_set) in enumerate(list_of_reduced_sets):
				if len(reduced_set)>0:
					self.create_fasta(children[i],reduced_set)
			logger.info('END get shared k-mers for node "{}" ({} kmers)'.format(self._node_debug(node),len(intersection)//2))
			logger.debug('... kmers (from children): "{}"'.format(", ".join(intersection)))
			return intersection

	def build_index(self,k):
		logger.info('Going to build index for k={}'.format(k))
		self.get_shared_kmers(self.tree.get_tree_root(),k=k)


if __name__ == "__main__":

	parser = argparse.ArgumentParser(description='Build index.')
	parser.add_argument(
			'-n','--newick-tree',
			type=str,
			metavar='str',
			dest='newick_fn',
			required=True,
			help='Taxonomic tree (Newick).',
		)
	parser.add_argument(
			'-k',
			type=int,
			metavar='int',
			dest='k',
			required=True,
			help='K-mer length k.',
		)
	parser.add_argument(
			'-o','--output-dir',
			type=str,
			metavar='str',
			dest='output_dir_fn',
			required=True,
			help='Output directory.',
		)

	args = parser.parse_args()

	k=args.k
	assert k>0
	newick_fn=args.newick_fn
	output_dir_fn=args.output_dir_fn

	logger.info("Starting index construction")
	logger.info("       newick : {}".format(newick_fn))
	logger.info("   output dir : {}".format(output_dir_fn))
	logger.info("            k : {}".format(k))

	ti=TreeIndex(
			tree_newick_fn=newick_fn,
			directory=output_dir_fn
		)
	ti.build_index(k=k)

{'fraction_non_alphanumeric': 0.16415302556556924, 'fraction_numerical': 0.004230274048188339, 'mean_word_length': 4.891115926327194, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 9, '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': '5773759', 'n_tokens_mistral': 4762, 'n_tokens_neox': 4313, 'n_words': 620}