<commit_before>#!/usr/bin/env python
"""
This library was adapted from:
        https://github.com/luo-chengwei/utilitomics

This library works with the taxonomy db download from NCBI FTP:

ftp://ftp.ncbi.nih.gov:/pub/taxonomy/taxdump.tar.gz

basically you can init an taxonomy tree obj by:

t_tree = ncbi.NCBITree(path)   # tree being the tax dump db unzip path
then you can basically get the taxonomy ranks using:

path = t_tree.get_name_path_with_taxon_id(taxon_id)
"""

import os
import networkx as nx
import csv
from itertools import chain

from shogun.utilities.pickleable import Pickleable
from shogun.utilities.downloadable import download
from shogun.downloaders.download_ncbi_taxonomy import NCBITaxdmp


class NCBITree(Pickleable):
    def __init__(self, _downloader=NCBITaxdmp()):
        # Private variables (should be set in settings)
        self._downloader = _downloader
        super().__init__()

    @download
    def _parse(self):
        # Initialize variables
        self.tree = nx.DiGraph()
        self.name2taxon_id = {}
        self.taxon_id2name = {}
        ncbi_taxdmp_dir = self._downloader.path

        with open(os.path.join(ncbi_taxdmp_dir, 'names.dmp'), 'r') as handle:
            csv_handle = csv.reader(handle, delimiter="\t")
            for cols in csv_handle:
                taxon_id = int(cols[0])
                name = cols[2]
                self.name2taxon_id[name] = taxon_id
                if cols[-2] == 'scientific name':
                    self.taxon_id2name[taxon_id] = name

        # construct node tree
        edges = []
        nodes = {}
        with open(os.path.join(ncbi_taxdmp_dir, 'nodes.dmp'), 'r') as handle:
            csv_handle = csv.reader(handle, delimiter="\t")
            for cols in csv_handle:
                parent_node = int(cols[2])
                child_node = int(cols[0])
                rank = cols[4]
                nodes[child_node] = rank
                if child_node != parent_node:
                    edges.append((child_node, parent_node))

        self.tree.add_edges_from(edges)
        nx.set_node_attributes(self.tree, 'rank', nodes)

    def get_taxon_id_lineage_with_taxon_id(self, taxon_id):
        try:
            path = [taxon_id]
            current_node = taxon_id
            return chain(path, self.tree.successors_iter(current_node))
        except nx.exception.NetworkXError:
            return []

    def get_taxon_id_lineage_with_name(self, name):
        if name not in self.name2taxon_id:
            return []
        return self.get_taxon_id_lineage_with_taxon_id(self.name2taxon_id[name])

    def get_name_lineage_with_taxon_id(self, taxon_id):
        tid_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in tid_lineage:
            rank = self.tree.node[x]['rank']
            try:
                name = self.taxon_id2name[x]
            except KeyError:
                name = ''
            name_lineage.append((name, rank))
        return name_lineage

    def get_name_lineage_with_name(self, name):
        if name not in self.name2taxon_id:
            return []
        path = self.get_name_lineage_with_taxon_id(self.name2taxon_id[name])
        return path

    def get_rank_with_taxon_id(self, taxon_id, rank):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = self.get_name_lineage_with_taxon_id(taxon_id)
        for tid, (name, x) in zip(taxon_id_lineage, name_lineage):
            if x == rank:
                return tid, name
        return None, None

    def get_name_at_rank(self, name, rank):
        if name not in self.name2taxon_id:
            return None, None
        return self.get_rank_with_taxon_id(self.name2taxon_id[name], rank)

    def get_lineage(self, taxon_id, ranks={'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'}):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            try:
                name = self.taxon_id2name[x]
            except KeyError:
                name = ''
            if rank in ranks:
                name_lineage.append((name, x, rank))
        return name_lineage

    def get_lineage_depth(self, taxon_id, depth, ranks=['superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        ranks = set(ranks[depth:])
        lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            if rank in ranks:
                lineage.append(x)
        return lineage

    def mp_lineage(self, taxon_id, ranks={'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'}
                   , nodes=('k__', 'p__', 'c__', 'o__', 'f__', 'g__', 's__')):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            try:
                name = self.taxon_id2name[x]
            except KeyError:
                name = ''
            if rank in ranks:
                name_lineage.append(name)
        return mp_format(reversed(name_lineage), nodes=nodes)


def mp_format(taxa_array, nodes=('k__', 'p__', 'c__', 'o__', 'f__', 'g__', 's__')):
    return ';'.join([i + j.replace(' ', '_') for i, j in zip(nodes, taxa_array)])


def main():
    ncbi_tree = NCBITree()
    print(ncbi_tree.get_taxon_id_lineage_with_taxon_id(640701851))

if __name__ == '__main__':
    main()<commit_msg>Fixed an outstanding bug on the MP_Format<commit_after>#!/usr/bin/env python
"""
This library was adapted from:
        https://github.com/luo-chengwei/utilitomics

This library works with the taxonomy db download from NCBI FTP:

ftp://ftp.ncbi.nih.gov:/pub/taxonomy/taxdump.tar.gz

basically you can init an taxonomy tree obj by:

t_tree = ncbi.NCBITree(path)   # tree being the tax dump db unzip path
then you can basically get the taxonomy ranks using:

path = t_tree.get_name_path_with_taxon_id(taxon_id)
"""

import os
import networkx as nx
import csv
from functools import lru_cache
from collections import defaultdict

from shogun.utilities.pickleable import Pickleable
from shogun.utilities.downloadable import download
from shogun.downloaders.download_ncbi_taxonomy import NCBITaxdmp


class NCBITree(Pickleable):
    def __init__(self, mp_ranks=None, _downloader=NCBITaxdmp()):
        # Private variables (should be set in settings)
        self._downloader = _downloader
        if mp_ranks is None:
            self.mp_ranks = dict(zip(('superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'),
                     ('k__', 'p__', 'c__', 'o__', 'f__', 'g__', 's__')))
        else:
            self.mp_ranks = mp_ranks
        super().__init__()

    @download
    def _parse(self):
        # Initialize variables
        self.tree = nx.DiGraph()
        self.name2taxon_id = defaultdict(int)
        self.taxon_id2name = defaultdict(str)
        ncbi_taxdmp_dir = self._downloader.path

        with open(os.path.join(ncbi_taxdmp_dir, 'names.dmp'), 'r') as handle:
            csv_handle = csv.reader(handle, delimiter="\t")
            for cols in csv_handle:
                taxon_id = int(cols[0])
                name = cols[2]
                self.name2taxon_id[name] = taxon_id
                if cols[-2] == 'scientific name':
                    self.taxon_id2name[taxon_id] = name

        # construct node tree
        edges = []
        nodes = {}
        with open(os.path.join(ncbi_taxdmp_dir, 'nodes.dmp'), 'r') as handle:
            csv_handle = csv.reader(handle, delimiter="\t")
            for cols in csv_handle:
                parent_node = int(cols[2])
                child_node = int(cols[0])
                rank = cols[4]
                nodes[child_node] = rank
                if child_node != parent_node:
                    edges.append((child_node, parent_node))

        self.tree.add_edges_from(edges)
        nx.set_node_attributes(self.tree, 'rank', nodes)

    def get_taxon_id_lineage_with_taxon_id(self, taxon_id):
        if taxon_id in self.tree:
            for i in nx.dfs_preorder_nodes(self.tree, taxon_id):
                yield i

    def get_taxon_id_lineage_with_name(self, name):
        if name not in self.name2taxon_id:
            return []
        return self.get_taxon_id_lineage_with_taxon_id(self.name2taxon_id[name])

    def get_name_lineage_with_taxon_id(self, taxon_id):
        tid_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in tid_lineage:
            rank = self.tree.node[x]['rank']
            name = self.taxon_id2name[x]
            name_lineage.append((name, rank))
        return name_lineage

    def get_name_lineage_with_name(self, name):
        if name not in self.name2taxon_id:
            return []
        path = self.get_name_lineage_with_taxon_id(self.name2taxon_id[name])
        return path

    def get_rank_with_taxon_id(self, taxon_id, rank):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = self.get_name_lineage_with_taxon_id(taxon_id)
        for tid, (name, x) in zip(taxon_id_lineage, name_lineage):
            if x == rank:
                return tid, name
        return None, None

    def get_name_at_rank(self, name, rank):
        if name not in self.name2taxon_id:
            return None, None
        return self.get_rank_with_taxon_id(self.name2taxon_id[name], rank)

    def get_lineage(self, taxon_id, ranks={'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'}):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            name = self.taxon_id2name[x]
            if rank in ranks:
                name_lineage.append((name, x, rank))
        return name_lineage

    def get_lineage_depth(self, taxon_id, depth, ranks=('superkingdom', 'phylum', 'class', 'order', 'family', 'genus',
                                                        'species')):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        ranks = set(ranks[depth:])
        lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            if rank in ranks:
                lineage.append(x)
        return lineage

    # Note that this will create a global cache for all instances of NCBITree.
    # Will be fine unless if you want to compare trees.
    @lru_cache(maxsize=128)
    def mp_lineage(self, taxon_id):
        taxon_id_lineage = self.get_taxon_id_lineage_with_taxon_id(taxon_id)
        name_lineage = []
        for x in taxon_id_lineage:
            rank = self.tree.node[x]['rank']
            name = self.taxon_id2name[x]
            if rank in self.mp_ranks:
                name_lineage.append(self.mp_ranks[rank] + name.replace(' ', '_'))
        return ';'.join(reversed(name_lineage))


def main():
    ncbi_tree = NCBITree()
    print(ncbi_tree.get_taxon_id_lineage_with_taxon_id(640701851))

if __name__ == '__main__':
    main()
{'fraction_non_alphanumeric': 0.10313980003508157, 'fraction_numerical': 0.00517453078407297, 'mean_word_length': 2.4253529588465006, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4544589', 'n_tokens_mistral': 3951, 'n_tokens_neox': 3791, 'n_words': 792}