alienzj · September 7, 2017 12:10
diff --git a/README.md b/README.md
diff --git a/classify.py b/classify.py
 #! /usr/bin/env python
 """
 Emulate the kraken k-mer classification step, but on banded MinHash signatures.

 Briefly,

 * load in the k-mer-to-lineage database produced by 'extract.py'.
 * for every hash in the signature, find the lineage for that k-mer
  (here 'lineage' would be the computed last-common-ancestor from the NCBI
  taxonomy, based on GenBank genomes)
 * count & summarize

 Usage:

   kraken/classify.py db.lca.json ecoli_many_sigs/ecoli-1.sig

 where 'db.lca.json' is created by 'extract.py'.

 TODO:

 * add classification of FASTA/FASTQ?
 """

 import argparse
 import collections

 import sourmash_lib, sourmash_lib.signature
 import lca_json

 SCALED=10000                              # should match the LCA compute @CTB

 kraken_rank_code = {
    'genus' : 'G',
    'species': 'S',
    'phylum': 'P',
    'class': 'C',
    'order': 'O',
    'family': 'F',
    'genus': 'G',
    'kingdom': 'K',
    'domain': 'D' }


 def main():
    p = argparse.ArgumentParser()
    p.add_argument('lca_filename')
    p.add_argument('sigfiles', nargs='+')
    p.add_argument('-k', '--ksize', default=31, type=int)
    args = p.parse_args()

    # load lca info
    lca_db = lca_json.LCA_Database(args.lca_filename)
    taxfoo, hashval_to_lca, scaled = lca_db.get_database(args.ksize, SCALED)
    
    # load signatures
    siglist = []
    print('loading signatures from {} signature files'.format(len(args.sigfiles)))
    for sigfile in args.sigfiles:
        sigs = sourmash_lib.signature.load_signatures(sigfile,
                                                      select_ksize=args.ksize)
        sigs = list(sigs)
        siglist.extend(sigs)

    print('loaded {} signatures total at k={}'.format(len(siglist), args.ksize))

    # downsample
    print('downsampling to scaled value: {}'.format(scaled))
    for sig in siglist:
        if sig.minhash.scaled < scaled:
            sig.minhash = sig.minhash.downsample_scaled(scaled)

    # now, extract hash values!
    hashvals = collections.defaultdict(int)
    for sig in siglist:
        for hashval in sig.minhash.get_mins():
            hashvals[hashval] += 1

    found = 0
    total = 0
    by_taxid = collections.defaultdict(int)

    # for every hash, get LCA of labels
    for hashval, count in hashvals.items():
        lca = hashval_to_lca.get(hashval)
        total += count

        if lca is None:
            by_taxid[0] += count
            continue

        by_taxid[lca] += count
        found += count

    print('found LCA classifications for', found, 'of', total, 'hashes')
    not_found = total - found

    # now, propogate counts up the taxonomic tree.
    by_taxid_lca = collections.defaultdict(int)
    for taxid, count in by_taxid.items():
        by_taxid_lca[taxid] += count

        parent = taxfoo.child_to_parent.get(taxid)
        while parent != None and parent != 1:
            by_taxid_lca[parent] += count
            parent = taxfoo.child_to_parent.get(parent)

    total_count = sum(by_taxid.values())

    # sort by lineage length
    x = []
    for taxid, count in by_taxid_lca.items():
        x.append((len(taxfoo.get_lineage(taxid)), taxid, count))

    x.sort()

    # ...aaaaaand output.
    print('{}\t{}\t{}\t{}\t{}\t{}'.format('percent', 'below', 'at node',
                                          'code', 'taxid', 'name'))
    for _, taxid, count_below in x:
        percent = round(100 * count_below / total_count, 2)
        count_at = by_taxid[taxid]

        rank = taxfoo.node_to_info.get(taxid)
        if rank:
            rank = rank[0]
            classify_code = kraken_rank_code.get(rank, '-')
        else:
            classify_code = '-'

        name = taxfoo.taxid_to_names.get(taxid)
        if name:
            name = name[0]
        else:
            name = '-'

        print('{}\t{}\t{}\t{}\t{}\t{}'.format(percent, count_below, count_at,
                                              classify_code, taxid, name))

    if not_found:
        classify_code = 'U'
        percent = round(100 * not_found / total_count, 2)
        count_below = not_found
        count_at = not_found
        taxid = 0
        name = 'not classified'

        print('{}\t{}\t{}\t{}\t{}\t{}'.format(percent, count_below, count_at,
                                              classify_code, taxid, name))


 if __name__ == '__main__':
    main()
diff --git a/create.py b/create.py
 import json

 info = {}
 info['version'] = 1

 db = []
 info['dblist'] = db

 k31 = {}
 k31['ksize'] = 31
 k31['lca_db'] = 'genbank-k31.lca'
 k31['scaled'] = 10000
 k31['nodes'] = 'genbank/nodes.dmp'
 k31['names'] = 'genbank/names.dmp'
 db.append(k31)

 k51 = {}
 k51['ksize'] = 51
 k51['lca_db'] = 'genbank-k51.lca'
 k51['scaled'] = 10000
 k51['nodes'] = 'genbank/nodes.dmp'
 k51['names'] = 'genbank/names.dmp'
 db.append(k51)


 k21 = {}
 k21['ksize'] = 21
 k21['lca_db'] = 'genbank-k21.lca'
 k21['scaled'] = 10000
 k21['nodes'] = 'genbank/nodes.dmp'
 k21['names'] = 'genbank/names.dmp'
 db.append(k21)

 json.dump(info, open('genbank.lca.json', 'wt'))
diff --git a/extract.py b/extract.py
 #! /usr/bin/env python
 """
 Emulate the kraken database-building step, but using banded MinHash signatures
 instead.

 Briefly,

 * load in a list of accessions-to-taxonomic-lineage information (e.g. such
  as the .csv.gz file from https://github.com/dib-lab/2017-ncbi-taxdump);
 * for every hash in each signature, connect that hash to the lineage for that
  signature (signature genbank accession ID -> lineage);
 * after loading in all signatures & hashes, find the last-common-ancestor
  lineage for each hash and associate in a dictionary;
 * save dictionary.

 Usage::

   kraken/extract.py foobar.lca genbank/*.csv.gz nodes.dmp ecoli_many_sigs/ecoli-*.sig --lca-json=db.lca.json

 The E. coli signatures used in the command above can be downloaded like so:

   curl -O -L https://github.com/dib-lab/sourmash/raw/master/data/eschericia-sigs.tar.gz

 'nodes.dmp' comes from 'taxdump.tar.gz' --

    curl -O -L ftp://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz
    tar xzf taxdump.tar.gz nodes.dmp names.dmp
 """

 import sys, os
 import sourmash_lib, sourmash_lib.signature
 import argparse
 from pickle import dump, load
 from collections import defaultdict

 import lca_json
 from ncbi_taxdump_utils import NCBI_TaxonomyFoo


 def traverse_find_sigs(dirnames):
    for dirname in dirnames:
        for root, dirs, files in os.walk(dirname):
            for name in files:
                if name.endswith('.sig') or name.endswith('.sbt'):
                    fullname = os.path.join(root, name)
                    yield fullname


 def main():
    p = argparse.ArgumentParser()
    p.add_argument('lca_output')
    p.add_argument('genbank_csv')
    p.add_argument('nodes_dmp')
    p.add_argument('sigs', nargs='+')
    p.add_argument('-k', '--ksize', default=31, type=int)
    p.add_argument('--scaled', default=10000, type=int)

    p.add_argument('--traverse-directory', action='store_true')

    p.add_argument('-l', '--load-hashvals', action='store_true')

    p.add_argument('--lca-json')
    p.add_argument('--names-dmp', default='')
    args = p.parse_args()

    taxfoo = NCBI_TaxonomyFoo()

    # load the accessions->taxid info
    taxfoo.load_accessions_csv(args.genbank_csv)

    # load the nodes_dmp file to get the tax tree
    print('loading nodes_dmp / taxonomic tree')
    taxfoo.load_nodes_dmp(args.nodes_dmp)

    # track hashval -> set of taxids
    hashval_to_taxids = defaultdict(set)

    # for every minhash in every signature, link it to its NCBI taxonomic ID.
    if args.traverse_directory:
        inp_files = list(traverse_find_sigs(args.sigs))
    else:
        inp_files = list(args.sigs)

    if args.load_hashvals:
        with open(args.lca_output + '.hashvals', 'rb') as hashval_fp:
            print('loading hashvals dict per -l/--load-hashvals...')
            hashval_to_taxids = load(hashval_fp)
            print('loaded {} hashvals'.format(len(hashval_to_taxids)))
    else:

        print('loading signatures & traversing hashes')
        bad_input = 0
        for n, filename in enumerate(inp_files):
            if n % 100 == 0:
                print('... loading file #', n, 'of', len(inp_files), end='\r')

            try:
                sig = sourmash_lib.signature.load_one_signature(filename,
                                                          select_ksize=args.ksize)
            except (FileNotFoundError, ValueError):
                if not args.traverse_directory:
                    raise

                bad_input += 1
                continue

            acc = sig.name().split(' ')[0]   # first part of sequence name
            acc = acc.split('.')[0]          # get acc w/o version

            taxid = taxfoo.get_taxid(acc)
            if taxid == None:
                continue

            sig.minhash = sig.minhash.downsample_scaled(args.scaled)

            mins = sig.minhash.get_mins()

            for m in mins:
                hashval_to_taxids[m].add(taxid)
        print('\n...done')
        if bad_input:
            print('failed to load {} of {} files found'.format(bad_input,
                                                               len(inp_files)))

        with open(args.lca_output + '.hashvals', 'wb') as hashval_fp:
            dump(hashval_to_taxids, hashval_fp)

    ####

    print('traversing tags and finding last-common-ancestor for {} tags'.format(len(hashval_to_taxids)))

    hashval_to_lca = {}
    found_root = 0
    empty_set = 0

    # find the LCA for each hashval and store.
    for n, (hashval, taxid_set) in enumerate(hashval_to_taxids.items()):
        if n % 10000 == 0:
            print('...', n, end='\r')

        # find associated least-common-ancestors.
        lca = taxfoo.find_lca(taxid_set)

        if lca == 1:
            if taxid_set:
                found_root += 1
            else:
                empty_set += 1
            continue

        # save!!
        hashval_to_lca[hashval] = lca
    print('\ndone')

    if found_root:
        print('found root {} times'.format(found_root))
    if empty_set:
        print('found empty set {} times'.format(empty_set))

    print('saving to', args.lca_output)
    with lca_json.xopen(args.lca_output, 'wb') as lca_fp:
        dump(hashval_to_lca, lca_fp)

    # update LCA DB JSON file if provided
    if args.lca_json:
        lca_db = lca_json.LCA_Database()
        if os.path.exists(args.lca_json):
            print('loading LCA JSON file:', args.lca_json)
            lca_db.load(args.lca_json)

        prefix = os.path.dirname(args.lca_json) + '/'

        lca_output = args.lca_output
        if lca_output.startswith(prefix):
            lca_output = lca_output[len(prefix):]

        nodes_dmp = args.nodes_dmp
        if nodes_dmp.startswith(prefix):
            nodes_dmp = nodes_dmp[len(prefix):]

        names_dmp = args.names_dmp
        if names_dmp:
            if names_dmp.startswith(prefix):
                names_dmp = names_dmp[len(prefix):]
        else:
            names_dmp = nodes_dmp.replace('nodes', 'names')

        lca_db.add_db(args.ksize, args.scaled, lca_output, nodes_dmp, names_dmp)
        print('saving LCA JSON file:', args.lca_json)
        lca_db.save(args.lca_json)


 if __name__ == '__main__':
    main()
diff --git a/get_taxid_tara_meren.py b/get_taxid_tara_meren.py
 #! /usr/bin/env python
 import lca_json
 import argparse
 import csv
 import traceback


 want_taxonomy = ['superkingdom', 'phylum', 'order', 'class', 'family', 'genus', 'species']

 def get_taxids_for_name(taxfoo, **ranknames):
    assert len(ranknames) == 1
    srank, sname = list(ranknames.items())[0]

    tid_set = []
    for taxid, (name, _, _) in taxfoo.taxid_to_names.items():
        if name == sname:
            tid_set.append(taxid)

    lineages = [taxfoo.get_lineage_as_dict(tid, want_taxonomy) for tid in tid_set]

    if len(lineages) > 1:
        # if non-unique, collect one that matches the desired rank
        tid_set2 = []
        lin2 = []
        for (tid, lin) in zip(tid_set, lineages):
            if lin.get(srank) == sname:
                tid_set2.append(tid)
                lin2.append(lin)

        lineages = lin2
        tid_set = tid_set2
                
        if len(tid_set) > 2:
            assert 0, (tid_set, lineages)
            return -1

    if not lineages:
        print('not found:', srank, sname)
        return -1
    
    taxid = tid_set[0]
    return taxid


 def main():
    p = argparse.ArgumentParser()
    p.add_argument('lca_db')
    p.add_argument('csv')
    p.add_argument('outcsv')
    args = p.parse_args()

    lca_db = lca_json.LCA_Database(args.lca_db)
    taxfoo = lca_db.get_taxonomy()

    r = csv.reader(open(args.csv, 'rt', encoding='utf8'))
    next(r)

    results = []

    n = 0
    for row in r:
        n += 1
        last_rank = None
        last_name = None

        genome_name, row = row[0], row[1:]

        for rank, name in zip(want_taxonomy, row):
            if name == 'na':
                break
            last_rank, last_name = rank, name
            
        taxid = get_taxids_for_name(taxfoo, **{ last_rank: last_name })
        if taxid == -1:
            taxid = 1

        lineage = taxfoo.get_lineage(taxid, want_taxonomy)
        results.append((genome_name, str(taxid), ";".join(lineage)))
        print(genome_name, taxid, ";".join(lineage))

    with open(args.outcsv, 'wt', encoding='utf8') as outfp:
        w = csv.writer(outfp)
        for name, taxid, lineage in results:
            w.writerow([name, taxid, lineage])

 if __name__ == '__main__':
    main()
diff --git a/get_taxid_tara_tully.py b/get_taxid_tara_tully.py
 #! /usr/bin/env python
 import lca_json
 import argparse
 import csv
 import traceback


 want_taxonomy = ['superkingdom', 'phylum', 'order', 'class', 'family', 'genus', 'species']

 def get_taxids_for_name(taxfoo, **ranknames):
    assert len(ranknames) == 1
    srank, sname = list(ranknames.items())[0]

    tid_set = []
    for taxid, (name, _, _) in taxfoo.taxid_to_names.items():
        if name == sname:
            tid_set.append(taxid)

    lineages = [taxfoo.get_lineage_as_dict(tid, want_taxonomy) for tid in tid_set]

    if len(lineages) > 1:
        # if non-unique, collect one that matches the desired rank
        tid_set2 = []
        lin2 = []
        for (tid, lin) in zip(tid_set, lineages):
            if lin.get(srank) == sname:
                tid_set2.append(tid)
                lin2.append(lin)

        lineages = lin2
        tid_set = tid_set2
                
        if len(tid_set) > 2:
            assert 0, (tid_set, lineages)
            return -1

    if not lineages:
        print('not found:', srank, sname)
        return -1
    
    taxid = tid_set[0]
    return taxid


 def main():
    p = argparse.ArgumentParser()
    p.add_argument('lca_db')
    p.add_argument('csv')
    p.add_argument('outcsv')
    args = p.parse_args()

    lca_db = lca_json.LCA_Database(args.lca_db)
    taxfoo = lca_db.get_taxonomy()

    r = csv.reader(open(args.csv, 'rt', encoding='utf8'))
    next(r)

    results = []

    n = 0
    for row in r:
        n += 1
        last_rank = None
        last_name = None

        genome_name, row = row[0], row[1:]

        for rank, name in zip(want_taxonomy, row):
            if name == 'null' or name.startswith('novel'):
                break
            last_rank, last_name = rank, name
            
        taxid = get_taxids_for_name(taxfoo, **{ last_rank: last_name })
        if taxid == -1:
            taxid = 1

        lineage = taxfoo.get_lineage(taxid, want_taxonomy)
        results.append((genome_name, str(taxid), ";".join(lineage)))
        print(genome_name, taxid, ";".join(lineage))

    with open(args.outcsv, 'wt', encoding='utf8') as outfp:
        w = csv.writer(outfp)
        for name, taxid, lineage in results:
            w.writerow([name, taxid, lineage])

 if __name__ == '__main__':
    main()
diff --git a/lca_json.py b/lca_json.py
 import json
 import gzip
 import os
 from pickle import load

 from ncbi_taxdump_utils import NCBI_TaxonomyFoo


 class LCA_Database(object):
    def __init__(self, filename=None):
        if filename:
            self.load(filename)
        else:
            self.lca = {}
            self.lca['version'] = 1
            self.lca['dblist'] = []

    def load(self, filename):
        with xopen(filename, 'rt') as json_fp:
            info = json.load(json_fp)
            assert info['version'] == 1

        info['basepath'] = os.path.dirname(filename)

        self.lca = info

    def save(self, filename):
        if 'basepath' in self.lca:
            del self.lca['basepath']

        with open(filename, 'wt') as json_fp:
            json.dump(self.lca, json_fp)

    def add_db(self, ksize, scaled, lca_path, nodes_path, names_path):
        db_info = {}
        db_info['ksize'] = int(ksize)
        db_info['scaled'] = int(scaled)
        db_info['lca_db'] = lca_path
        db_info['nodes'] = nodes_path
        db_info['names'] = names_path

        self.lca['dblist'].append(db_info)

    def get_database(self, ksize, scaled):
        lca_info = self.lca

        assert lca_info['version'] == 1
        basepath = lca_info['basepath']

        matching_ksizes = []
        for db in lca_info['dblist']:
            if db['ksize'] == ksize:
                matching_ksizes.append(db)

        # ignore scaled matching for now, take first one
        assert len(matching_ksizes) == 1
        entry = matching_ksizes[0]

        taxfoo = self.get_taxonomy(entry)
        
        lca_file = os.path.join(basepath, entry['lca_db'])
        print('loading k-mer DB from:', lca_file)
        with xopen(lca_file, 'rb') as hashval_fp:
            hashval_to_lca = load(hashval_fp)

        return taxfoo, hashval_to_lca, entry['scaled']

    def get_taxonomy(self, entry=None):
        if entry == None:
            entry = self.lca['dblist'][0]

        basepath = self.lca['basepath']
        taxfoo = NCBI_TaxonomyFoo()

        # load the nodes_dmp file to get the tax tree
        nodes_file = os.path.join(basepath, entry['nodes'])
        print('loading taxonomic nodes from:', nodes_file)
        taxfoo.load_nodes_dmp(nodes_file)

        names_file = os.path.join(basepath, entry['names'])
        print('loading taxonomic names from:', names_file)
        taxfoo.load_names_dmp(names_file)

        return taxfoo


 ### utility functions


 def xopen(filename, mode):
    if filename.endswith('.gz'):
        return gzip.open(filename, mode)
    return open(filename, mode)
diff --git a/ncbi_taxdump_utils.py b/ncbi_taxdump_utils.py
 """
 Utilities to deal with NCBI taxonomic foo.
 """

 import gzip
 import csv
 import os
 from pickle import dump, load


 want_taxonomy = ['superkingdom', 'phylum', 'order', 'class', 'family', 'genus', 'species']


 class NCBI_TaxonomyFoo(object):
    def __init__(self):
        self.child_to_parent = None
        self.node_to_info = None
        self.taxid_to_names = None
        self.accessions = None

    def load_nodes_dmp(self, filename, do_save_cache=True):
        cache_file = filename + '.cache'
        if os.path.exists(cache_file):
            with xopen(cache_file, 'rb') as cache_fp:
                self.child_to_parent, self.node_to_info = load(cache_fp)
        else:
            self.child_to_parent, self.node_to_info = parse_nodes(filename)
            if do_save_cache:
                self.save_nodes_cache(cache_file)

    def save_nodes_cache(self, cache_file):
        with xopen(cache_file, 'wb') as cache_fp:
            dump((self.child_to_parent, self.node_to_info), cache_fp)

    def load_names_dmp(self, filename, do_save_cache=True):
        cache_file = filename + '.cache'
        if os.path.exists(cache_file):
            with xopen(cache_file, 'rb') as cache_fp:
                self.taxid_to_names = load(cache_fp)
        else:
            self.taxid_to_names = parse_names(filename)
            self.save_names_cache(cache_file)

    def save_names_cache(self, cache_file):
        with xopen(cache_file, 'wb') as cache_fp:
            dump(self.taxid_to_names, cache_fp)

    def load_accessions_csv(self, filename):
        self.accessions = load_genbank_accessions_csv(filename)

    # get taxid
    def get_taxid(self, acc):
        # @CTB hack hack split off NZ from accession
        if acc.startswith('NZ_'):
            acc = acc[3:]

        info = self.accessions.get(acc)
        if not info:
            return None

        taxid = info['taxid']
        taxid = int(taxid)
        return taxid

    # code to find the last common ancestor from a set of taxids
    def find_lca(self, taxid_set):
        # empty? exit.
        if not taxid_set:
            return 1

        taxid_set = set(taxid_set)  #  copy

        if len(taxid_set) == 1:
            return taxid_set.pop()

        # get the first full path
        taxid = taxid_set.pop()
        path = []
        while taxid != 1:
            path.insert(0, taxid)
            taxid = self.child_to_parent.get(taxid, 1)   # @CTB reexamine

        # find the first shared taxid in each follow-on path
        while taxid_set:
            taxid = taxid_set.pop()

            path2 = []
            while taxid != 1:
                if taxid in path:
                    path2.insert(0, taxid)
                taxid = self.child_to_parent.get(taxid, 1)

            path = path2

        if path:
            return path[-1]
        return 1

    def get_lineage(self, taxid, want_taxonomy=None):
        """
        Extract the text taxonomic lineage in order (kingdom on down)
        """
        taxid = int(taxid)

        lineage = []
        while 1:
            if taxid not in self.node_to_info:
                print('cannot find taxid {}; quitting.'.format(taxid))
                break
            rank = self.node_to_info[taxid][0]
            name = self.taxid_to_names[taxid][0]
            if not want_taxonomy or rank in want_taxonomy:
                lineage.insert(0, name)
            taxid = self.child_to_parent[taxid]
            if taxid == 1:
                break

        return lineage


    def get_lineage_as_dict(self, taxid, want_taxonomy=None):
        """
        Extract the text taxonomic lineage in order (kingdom on down)
        """
        taxid = int(taxid)

        lineage = {}
        while 1:
            if taxid not in self.node_to_info:
                print('cannot find taxid {}; quitting.'.format(taxid))
                break
            rank = self.node_to_info[taxid][0]
            name = self.taxid_to_names[taxid][0]
            if not want_taxonomy or rank in want_taxonomy:
                lineage[rank] = name
            taxid = self.child_to_parent[taxid]
            if taxid == 1:
                break

        return lineage


 ### internal utility functions
 def xopen(filename, mode):
    if filename.endswith('.gz'):
        return gzip.open(filename, mode)
    return open(filename, mode)

 def parse_nodes(filename):
    "Parse the NCBI nodes_dmp file."
    child_to_parent = dict()
    node_to_info = dict()

    with xopen(filename, 'rt') as fp:
        for n, line in enumerate(fp):
            x = line.split('\t|\t')

            node_id, parent_node_id, rank, embl, div_id, div_flag, gencode, mgc_inherit, mgc_flag, mgc_id, hidden_flag, subtree_flag, comments = x

            node_id = int(node_id)
            parent_node_id = int(parent_node_id)

            child_to_parent[node_id] = parent_node_id
            node_to_info[node_id] = rank, embl, div_id, div_flag, comments

    return child_to_parent, node_to_info


 def parse_names(filename):
    """
    Parse an NCBI names.dmp file.
    """
    taxid_to_names = dict()
    with xopen(filename, 'rt') as fp:
        for n, line in enumerate(fp):
            line = line.rstrip('\t|\n')
            x = line.split('\t|\t')

            taxid, name, uniqname, name_class = x
            taxid = int(taxid)

            if name_class == 'scientific name':
                taxid_to_names[taxid] = (name, uniqname, name_class)

    return taxid_to_names


 def load_genbank_accessions_csv(filename):
    """
    Load a file containing genbank accession -> taxid + lineage string.

    See https://github.com/dib-lab/2017-ncbi-taxdump for scripts to create
    this file from a large collection of sequences.
    """
    # load the genbank CSV (accession -> lineage)
    print('loading genbank accession -> lineage info')
    with xopen(filename, 'rt') as fp:
        accessions = {}
        for row in csv.DictReader(fp, fieldnames=['acc', 'taxid', 'lineage']):
            acc = row['acc']
            accessions[acc] = row

    return accessions
diff --git a/tara-meren-SuppTable3.csv b/tara-meren-SuppTable3.csv
diff --git a/tara-tully-Table4.csv b/tara-tully-Table4.csv
diff --git a/tara_meren_taxids.csv b/tara_meren_taxids.csv
diff --git a/tara_tully_taxids.csv b/tara_tully_taxids.csv
	#! /usr/bin/env python
	"""
	Emulate the kraken k-mer classification step, but on banded MinHash signatures.

	Briefly,

	* load in the k-mer-to-lineage database produced by 'extract.py'.
	* for every hash in the signature, find the lineage for that k-mer
	(here 'lineage' would be the computed last-common-ancestor from the NCBI
	taxonomy, based on GenBank genomes)
	* count & summarize

	Usage:

	kraken/classify.py db.lca.json ecoli_many_sigs/ecoli-1.sig

	where 'db.lca.json' is created by 'extract.py'.

	TODO:

	* add classification of FASTA/FASTQ?
	"""

	import argparse
	import collections

	import sourmash_lib, sourmash_lib.signature
	import lca_json

	SCALED=10000 # should match the LCA compute @CTB

	kraken_rank_code = {
	'genus' : 'G',
	'species': 'S',
	'phylum': 'P',
	'class': 'C',
	'order': 'O',
	'family': 'F',
	'genus': 'G',
	'kingdom': 'K',
	'domain': 'D' }


	def main():
	p = argparse.ArgumentParser()
	p.add_argument('lca_filename')
	p.add_argument('sigfiles', nargs='+')
	p.add_argument('-k', '--ksize', default=31, type=int)
	args = p.parse_args()

	# load lca info
	lca_db = lca_json.LCA_Database(args.lca_filename)
	taxfoo, hashval_to_lca, scaled = lca_db.get_database(args.ksize, SCALED)

	# load signatures
	siglist = []
	print('loading signatures from {} signature files'.format(len(args.sigfiles)))
	for sigfile in args.sigfiles:
	sigs = sourmash_lib.signature.load_signatures(sigfile,
	select_ksize=args.ksize)
	sigs = list(sigs)
	siglist.extend(sigs)

	print('loaded {} signatures total at k={}'.format(len(siglist), args.ksize))

	# downsample
	print('downsampling to scaled value: {}'.format(scaled))
	for sig in siglist:
	if sig.minhash.scaled < scaled:
	sig.minhash = sig.minhash.downsample_scaled(scaled)

	# now, extract hash values!
	hashvals = collections.defaultdict(int)
	for sig in siglist:
	for hashval in sig.minhash.get_mins():
	hashvals[hashval] += 1

	found = 0
	total = 0
	by_taxid = collections.defaultdict(int)

	# for every hash, get LCA of labels
	for hashval, count in hashvals.items():
	lca = hashval_to_lca.get(hashval)
	total += count

	if lca is None:
	by_taxid[0] += count
	continue

	by_taxid[lca] += count
	found += count

	print('found LCA classifications for', found, 'of', total, 'hashes')
	not_found = total - found

	# now, propogate counts up the taxonomic tree.
	by_taxid_lca = collections.defaultdict(int)
	for taxid, count in by_taxid.items():
	by_taxid_lca[taxid] += count

	parent = taxfoo.child_to_parent.get(taxid)
	while parent != None and parent != 1:
	by_taxid_lca[parent] += count
	parent = taxfoo.child_to_parent.get(parent)

	total_count = sum(by_taxid.values())

	# sort by lineage length
	x = []
	for taxid, count in by_taxid_lca.items():
	x.append((len(taxfoo.get_lineage(taxid)), taxid, count))

	x.sort()

	# ...aaaaaand output.
	print('{}\t{}\t{}\t{}\t{}\t{}'.format('percent', 'below', 'at node',
	'code', 'taxid', 'name'))
	for _, taxid, count_below in x:
	percent = round(100 * count_below / total_count, 2)
	count_at = by_taxid[taxid]

	rank = taxfoo.node_to_info.get(taxid)
	if rank:
	rank = rank[0]
	classify_code = kraken_rank_code.get(rank, '-')
	else:
	classify_code = '-'

	name = taxfoo.taxid_to_names.get(taxid)
	if name:
	name = name[0]
	else:
	name = '-'

	print('{}\t{}\t{}\t{}\t{}\t{}'.format(percent, count_below, count_at,
	classify_code, taxid, name))

	if not_found:
	classify_code = 'U'
	percent = round(100 * not_found / total_count, 2)
	count_below = not_found
	count_at = not_found
	taxid = 0
	name = 'not classified'

	print('{}\t{}\t{}\t{}\t{}\t{}'.format(percent, count_below, count_at,
	classify_code, taxid, name))


	if __name__ == '__main__':
	main()
	import json

	info = {}
	info['version'] = 1

	db = []
	info['dblist'] = db

	k31 = {}
	k31['ksize'] = 31
	k31['lca_db'] = 'genbank-k31.lca'
	k31['scaled'] = 10000
	k31['nodes'] = 'genbank/nodes.dmp'
	k31['names'] = 'genbank/names.dmp'
	db.append(k31)

	k51 = {}
	k51['ksize'] = 51
	k51['lca_db'] = 'genbank-k51.lca'
	k51['scaled'] = 10000
	k51['nodes'] = 'genbank/nodes.dmp'
	k51['names'] = 'genbank/names.dmp'
	db.append(k51)


	k21 = {}
	k21['ksize'] = 21
	k21['lca_db'] = 'genbank-k21.lca'
	k21['scaled'] = 10000
	k21['nodes'] = 'genbank/nodes.dmp'
	k21['names'] = 'genbank/names.dmp'
	db.append(k21)

	json.dump(info, open('genbank.lca.json', 'wt'))
	#! /usr/bin/env python
	"""
	Emulate the kraken database-building step, but using banded MinHash signatures
	instead.

	Briefly,

	* load in a list of accessions-to-taxonomic-lineage information (e.g. such
	as the .csv.gz file from https://github.com/dib-lab/2017-ncbi-taxdump);
	* for every hash in each signature, connect that hash to the lineage for that
	signature (signature genbank accession ID -> lineage);
	* after loading in all signatures & hashes, find the last-common-ancestor
	lineage for each hash and associate in a dictionary;
	* save dictionary.

	Usage::

	kraken/extract.py foobar.lca genbank/.csv.gz nodes.dmp ecoli_many_sigs/ecoli-.sig --lca-json=db.lca.json

	The E. coli signatures used in the command above can be downloaded like so:

	curl -O -L https://github.com/dib-lab/sourmash/raw/master/data/eschericia-sigs.tar.gz

	'nodes.dmp' comes from 'taxdump.tar.gz' --

	curl -O -L ftp://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz
	tar xzf taxdump.tar.gz nodes.dmp names.dmp
	"""

	import sys, os
	import sourmash_lib, sourmash_lib.signature
	import argparse
	from pickle import dump, load
	from collections import defaultdict

	import lca_json
	from ncbi_taxdump_utils import NCBI_TaxonomyFoo


	def traverse_find_sigs(dirnames):
	for dirname in dirnames:
	for root, dirs, files in os.walk(dirname):
	for name in files:
	if name.endswith('.sig') or name.endswith('.sbt'):
	fullname = os.path.join(root, name)
	yield fullname


	def main():
	p = argparse.ArgumentParser()
	p.add_argument('lca_output')
	p.add_argument('genbank_csv')
	p.add_argument('nodes_dmp')
	p.add_argument('sigs', nargs='+')
	p.add_argument('-k', '--ksize', default=31, type=int)
	p.add_argument('--scaled', default=10000, type=int)

	p.add_argument('--traverse-directory', action='store_true')

	p.add_argument('-l', '--load-hashvals', action='store_true')

	p.add_argument('--lca-json')
	p.add_argument('--names-dmp', default='')
	args = p.parse_args()

	taxfoo = NCBI_TaxonomyFoo()

	# load the accessions->taxid info
	taxfoo.load_accessions_csv(args.genbank_csv)

	# load the nodes_dmp file to get the tax tree
	print('loading nodes_dmp / taxonomic tree')
	taxfoo.load_nodes_dmp(args.nodes_dmp)

	# track hashval -> set of taxids
	hashval_to_taxids = defaultdict(set)

	# for every minhash in every signature, link it to its NCBI taxonomic ID.
	if args.traverse_directory:
	inp_files = list(traverse_find_sigs(args.sigs))
	else:
	inp_files = list(args.sigs)

	if args.load_hashvals:
	with open(args.lca_output + '.hashvals', 'rb') as hashval_fp:
	print('loading hashvals dict per -l/--load-hashvals...')
	hashval_to_taxids = load(hashval_fp)
	print('loaded {} hashvals'.format(len(hashval_to_taxids)))
	else:

	print('loading signatures & traversing hashes')
	bad_input = 0
	for n, filename in enumerate(inp_files):
	if n % 100 == 0:
	print('... loading file #', n, 'of', len(inp_files), end='\r')

	try:
	sig = sourmash_lib.signature.load_one_signature(filename,
	select_ksize=args.ksize)
	except (FileNotFoundError, ValueError):
	if not args.traverse_directory:
	raise

	bad_input += 1
	continue

	acc = sig.name().split(' ')[0] # first part of sequence name
	acc = acc.split('.')[0] # get acc w/o version

	taxid = taxfoo.get_taxid(acc)
	if taxid == None:
	continue

	sig.minhash = sig.minhash.downsample_scaled(args.scaled)

	mins = sig.minhash.get_mins()

	for m in mins:
	hashval_to_taxids[m].add(taxid)
	print('\n...done')
	if bad_input:
	print('failed to load {} of {} files found'.format(bad_input,
	len(inp_files)))

	with open(args.lca_output + '.hashvals', 'wb') as hashval_fp:
	dump(hashval_to_taxids, hashval_fp)

	####

	print('traversing tags and finding last-common-ancestor for {} tags'.format(len(hashval_to_taxids)))

	hashval_to_lca = {}
	found_root = 0
	empty_set = 0

	# find the LCA for each hashval and store.
	for n, (hashval, taxid_set) in enumerate(hashval_to_taxids.items()):
	if n % 10000 == 0:
	print('...', n, end='\r')

	# find associated least-common-ancestors.
	lca = taxfoo.find_lca(taxid_set)

	if lca == 1:
	if taxid_set:
	found_root += 1
	else:
	empty_set += 1
	continue

	# save!!
	hashval_to_lca[hashval] = lca
	print('\ndone')

	if found_root:
	print('found root {} times'.format(found_root))
	if empty_set:
	print('found empty set {} times'.format(empty_set))

	print('saving to', args.lca_output)
	with lca_json.xopen(args.lca_output, 'wb') as lca_fp:
	dump(hashval_to_lca, lca_fp)

	# update LCA DB JSON file if provided
	if args.lca_json:
	lca_db = lca_json.LCA_Database()
	if os.path.exists(args.lca_json):
	print('loading LCA JSON file:', args.lca_json)
	lca_db.load(args.lca_json)

	prefix = os.path.dirname(args.lca_json) + '/'

	lca_output = args.lca_output
	if lca_output.startswith(prefix):
	lca_output = lca_output[len(prefix):]

	nodes_dmp = args.nodes_dmp
	if nodes_dmp.startswith(prefix):
	nodes_dmp = nodes_dmp[len(prefix):]

	names_dmp = args.names_dmp
	if names_dmp:
	if names_dmp.startswith(prefix):
	names_dmp = names_dmp[len(prefix):]
	else:
	names_dmp = nodes_dmp.replace('nodes', 'names')

	lca_db.add_db(args.ksize, args.scaled, lca_output, nodes_dmp, names_dmp)
	print('saving LCA JSON file:', args.lca_json)
	lca_db.save(args.lca_json)


	if __name__ == '__main__':
	main()
	import json
	import gzip
	import os
	from pickle import load

	from ncbi_taxdump_utils import NCBI_TaxonomyFoo


	class LCA_Database(object):
	def __init__(self, filename=None):
	if filename:
	self.load(filename)
	else:
	self.lca = {}
	self.lca['version'] = 1
	self.lca['dblist'] = []

	def load(self, filename):
	with xopen(filename, 'rt') as json_fp:
	info = json.load(json_fp)
	assert info['version'] == 1

	info['basepath'] = os.path.dirname(filename)

	self.lca = info

	def save(self, filename):
	if 'basepath' in self.lca:
	del self.lca['basepath']

	with open(filename, 'wt') as json_fp:
	json.dump(self.lca, json_fp)

	def add_db(self, ksize, scaled, lca_path, nodes_path, names_path):
	db_info = {}
	db_info['ksize'] = int(ksize)
	db_info['scaled'] = int(scaled)
	db_info['lca_db'] = lca_path
	db_info['nodes'] = nodes_path
	db_info['names'] = names_path

	self.lca['dblist'].append(db_info)

	def get_database(self, ksize, scaled):
	lca_info = self.lca

	assert lca_info['version'] == 1
	basepath = lca_info['basepath']

	matching_ksizes = []
	for db in lca_info['dblist']:
	if db['ksize'] == ksize:
	matching_ksizes.append(db)

	# ignore scaled matching for now, take first one
	assert len(matching_ksizes) == 1
	entry = matching_ksizes[0]

	taxfoo = self.get_taxonomy(entry)

	lca_file = os.path.join(basepath, entry['lca_db'])
	print('loading k-mer DB from:', lca_file)
	with xopen(lca_file, 'rb') as hashval_fp:
	hashval_to_lca = load(hashval_fp)

	return taxfoo, hashval_to_lca, entry['scaled']

	def get_taxonomy(self, entry=None):
	if entry == None:
	entry = self.lca['dblist'][0]

	basepath = self.lca['basepath']
	taxfoo = NCBI_TaxonomyFoo()

	# load the nodes_dmp file to get the tax tree
	nodes_file = os.path.join(basepath, entry['nodes'])
	print('loading taxonomic nodes from:', nodes_file)
	taxfoo.load_nodes_dmp(nodes_file)

	names_file = os.path.join(basepath, entry['names'])
	print('loading taxonomic names from:', names_file)
	taxfoo.load_names_dmp(names_file)

	return taxfoo


	### utility functions


	def xopen(filename, mode):
	if filename.endswith('.gz'):
	return gzip.open(filename, mode)
	return open(filename, mode)
	"""
	Utilities to deal with NCBI taxonomic foo.
	"""

	import gzip
	import csv
	import os
	from pickle import dump, load


	want_taxonomy = ['superkingdom', 'phylum', 'order', 'class', 'family', 'genus', 'species']


	class NCBI_TaxonomyFoo(object):
	def __init__(self):
	self.child_to_parent = None
	self.node_to_info = None
	self.taxid_to_names = None
	self.accessions = None

	def load_nodes_dmp(self, filename, do_save_cache=True):
	cache_file = filename + '.cache'
	if os.path.exists(cache_file):
	with xopen(cache_file, 'rb') as cache_fp:
	self.child_to_parent, self.node_to_info = load(cache_fp)
	else:
	self.child_to_parent, self.node_to_info = parse_nodes(filename)
	if do_save_cache:
	self.save_nodes_cache(cache_file)

	def save_nodes_cache(self, cache_file):
	with xopen(cache_file, 'wb') as cache_fp:
	dump((self.child_to_parent, self.node_to_info), cache_fp)

	def load_names_dmp(self, filename, do_save_cache=True):
	cache_file = filename + '.cache'
	if os.path.exists(cache_file):
	with xopen(cache_file, 'rb') as cache_fp:
	self.taxid_to_names = load(cache_fp)
	else:
	self.taxid_to_names = parse_names(filename)
	self.save_names_cache(cache_file)

	def save_names_cache(self, cache_file):
	with xopen(cache_file, 'wb') as cache_fp:
	dump(self.taxid_to_names, cache_fp)

	def load_accessions_csv(self, filename):
	self.accessions = load_genbank_accessions_csv(filename)

	# get taxid
	def get_taxid(self, acc):
	# @CTB hack hack split off NZ from accession
	if acc.startswith('NZ_'):
	acc = acc[3:]

	info = self.accessions.get(acc)
	if not info:
	return None

	taxid = info['taxid']
	taxid = int(taxid)
	return taxid

	# code to find the last common ancestor from a set of taxids
	def find_lca(self, taxid_set):
	# empty? exit.
	if not taxid_set:
	return 1

	taxid_set = set(taxid_set) # copy

	if len(taxid_set) == 1:
	return taxid_set.pop()

	# get the first full path
	taxid = taxid_set.pop()
	path = []
	while taxid != 1:
	path.insert(0, taxid)
	taxid = self.child_to_parent.get(taxid, 1) # @CTB reexamine

	# find the first shared taxid in each follow-on path
	while taxid_set:
	taxid = taxid_set.pop()

	path2 = []
	while taxid != 1:
	if taxid in path:
	path2.insert(0, taxid)
	taxid = self.child_to_parent.get(taxid, 1)

	path = path2

	if path:
	return path[-1]
	return 1

	def get_lineage(self, taxid, want_taxonomy=None):
	"""
	Extract the text taxonomic lineage in order (kingdom on down)
	"""
	taxid = int(taxid)

	lineage = []
	while 1:
	if taxid not in self.node_to_info:
	print('cannot find taxid {}; quitting.'.format(taxid))
	break
	rank = self.node_to_info[taxid][0]
	name = self.taxid_to_names[taxid][0]
	if not want_taxonomy or rank in want_taxonomy:
	lineage.insert(0, name)
	taxid = self.child_to_parent[taxid]
	if taxid == 1:
	break

	return lineage


	def get_lineage_as_dict(self, taxid, want_taxonomy=None):
	"""
	Extract the text taxonomic lineage in order (kingdom on down)
	"""
	taxid = int(taxid)

	lineage = {}
	while 1:
	if taxid not in self.node_to_info:
	print('cannot find taxid {}; quitting.'.format(taxid))
	break
	rank = self.node_to_info[taxid][0]
	name = self.taxid_to_names[taxid][0]
	if not want_taxonomy or rank in want_taxonomy:
	lineage[rank] = name
	taxid = self.child_to_parent[taxid]
	if taxid == 1:
	break

	return lineage


	### internal utility functions
	def xopen(filename, mode):
	if filename.endswith('.gz'):
	return gzip.open(filename, mode)
	return open(filename, mode)

	def parse_nodes(filename):
	"Parse the NCBI nodes_dmp file."
	child_to_parent = dict()
	node_to_info = dict()

	with xopen(filename, 'rt') as fp:
	for n, line in enumerate(fp):
	x = line.split('\t\|\t')

	node_id, parent_node_id, rank, embl, div_id, div_flag, gencode, mgc_inherit, mgc_flag, mgc_id, hidden_flag, subtree_flag, comments = x

	node_id = int(node_id)
	parent_node_id = int(parent_node_id)

	child_to_parent[node_id] = parent_node_id
	node_to_info[node_id] = rank, embl, div_id, div_flag, comments

	return child_to_parent, node_to_info


	def parse_names(filename):
	"""
	Parse an NCBI names.dmp file.
	"""
	taxid_to_names = dict()
	with xopen(filename, 'rt') as fp:
	for n, line in enumerate(fp):
	line = line.rstrip('\t\|\n')
	x = line.split('\t\|\t')

	taxid, name, uniqname, name_class = x
	taxid = int(taxid)

	if name_class == 'scientific name':
	taxid_to_names[taxid] = (name, uniqname, name_class)

	return taxid_to_names


	def load_genbank_accessions_csv(filename):
	"""
	Load a file containing genbank accession -> taxid + lineage string.

	See https://github.com/dib-lab/2017-ncbi-taxdump for scripts to create
	this file from a large collection of sequences.
	"""
	# load the genbank CSV (accession -> lineage)
	print('loading genbank accession -> lineage info')
	with xopen(filename, 'rt') as fp:
	accessions = {}
	for row in csv.DictReader(fp, fieldnames=['acc', 'taxid', 'lineage']):
	acc = row['acc']
	accessions[acc] = row

	return accessions