Commit 2a4d8168 authored by Sean Solari's avatar Sean Solari
Browse files

Bug Fixes

parent 9bc1f0c5
......@@ -22,3 +22,9 @@ results/
# Distribution files.
*/expam.egg*
./expam/
# Compiled C files.
src/expam/ext/kmers/_build/src
src/expam/ext/map/_build/src
src/expam/ext/sets/_build/src
......@@ -33,7 +33,7 @@ from setuptools.extension import Extension
from Cython.Build import cythonize
import numpy as np
EXPAM_VERSION = (1, 0, 4)
EXPAM_VERSION = (1, 0, 5)
SOURCE = os.path.dirname(os.path.abspath(__file__))
......
#!/usr/bin/env python3
from argparse import ArgumentParser
from multiprocessing import Process
import os
import re
from typing import List
from typing import Iterable, List, Mapping, Set
from ete3 import NCBITaxa
from expam.classify import ResultsPathConfig
from expam.classify.config import load_results_config
from expam.classify.taxonomy import TaxonomyNCBI
from expam.classify.taxonomy import EntrezRequestor, TaxonomyNCBI
from expam.database import FileLocationConfig
from expam.database.config import make_database_config
from expam.utils import ls
def find_compatible_names(db_path: str, taxids_to_replace: List[str]):
def load_expam_taxonomy_data(db_path: str):
config: FileLocationConfig = make_database_config(db_path)
tax_obj = TaxonomyNCBI(config)
......@@ -31,6 +32,10 @@ def find_compatible_names(db_path: str, taxids_to_replace: List[str]):
taxid_to_name[taxid] = name
name_to_taxid[name] = taxid
return taxid_to_lineage, taxid_to_name, name_to_taxid
def find_compatible_names(taxid_to_lineage: Mapping[str, Iterable[str]], taxid_to_name: Mapping[str, str], name_to_taxid: Mapping[str, str], taxids_to_replace: Iterable[str]):
ete3_tax_obj = NCBITaxa()
compatible = dict()
......@@ -127,15 +132,10 @@ def distribute(iterable, n):
return lists
def make_results_compatible(db_path: str, results_config: ResultsPathConfig, bad_taxids: List[str], n: int = 5):
taxid_map = find_compatible_names(db_path, bad_taxids)
bad_taxids_set = set(bad_taxids)
print(taxid_map)
def make_summary_results_compatible(bad_taxids: Set[str], output_files: List[str], n: int = 5):
# Replace names in tax summary files.
summary_files = distribute(ls(results_config.tax, ext=".csv"), n)
procs = [CleanSummary(summary_files[i], bad_taxids_set) for i in range(n)]
summary_files = distribute(output_files, n)
procs = [CleanSummary(summary_files[i], bad_taxids) for i in range(n)]
for proc in procs:
proc.start()
......@@ -143,8 +143,10 @@ def make_results_compatible(db_path: str, results_config: ResultsPathConfig, bad
for proc in procs:
proc.join()
def make_raw_results_compatible(taxid_map: Mapping[str, str], output_files: List[str], n: int = 5):
# Replace names in tax raw output files.
raw_files = distribute(ls(results_config.tax_raw, ext=".csv"), n)
raw_files = distribute(output_files, n)
procs = [CleanRaw(raw_files[i], taxid_map) for i in range(n)]
for proc in procs:
......@@ -154,12 +156,82 @@ def make_results_compatible(db_path: str, results_config: ResultsPathConfig, bad
proc.join()
def main(db_path: str, results_path: str, bad_taxids_file: str, n: int):
with open(bad_taxids_file, 'r') as f:
bad_taxids = [re.findall(r"\d+", line)[0] for line in f]
def load_taxids(taxids_file: str):
with open(taxids_file, 'r') as f:
taxids = [re.findall(r"\d+", line)[0] for line in f]
return taxids
def expam_main(db_path: str, results_path: str, bad_taxids_file: str, n: int):
bad_taxids = load_taxids(bad_taxids_file)
results_config: ResultsPathConfig = load_results_config(results_path, create=False)
make_results_compatible(db_path, results_config, bad_taxids, n)
# Map summary files.
summary_files: List[str] = ls(results_config.tax, ext=".csv")
make_summary_results_compatible(set(bad_taxids), summary_files, n=n)
# Map raw output files.
raw_files: List[str] = ls(results_config.tax_raw, ext=".csv")
taxid_to_lineage, taxid_to_name, name_to_taxid = load_expam_taxonomy_data(db_path)
taxid_map = find_compatible_names(taxid_to_lineage, taxid_to_name, name_to_taxid, bad_taxids)
make_raw_results_compatible(taxid_map, raw_files, n=n)
def load_taxid_cache(cache_path: str) -> Mapping[str, str]:
with open(cache_path, 'r') as f:
cache = dict(line.strip().split('\t') for line in f)
return cache
def cache_taxids(cache_path: str, cache: Mapping[str, str]):
with open(cache_path, 'w') as f:
f.write('\n'.join('\t'.join(v) for v in cache.items()))
def retrieve_taxid_information(taxids: Set[str]):
# Get lineages for each input taxid.
requestor: EntrezRequestor = EntrezRequestor()
taxid_lineage, name_taxid_rank = requestor.request_labels("taxonomy", "xml", list(taxids))
taxid_lineage = {taxid: lineage.split(',') for taxid, lineage in taxid_lineage}
taxid_to_name = {}
name_to_taxid = {}
for name, data in name_taxid_rank.items():
taxid, rank = data.split(",")
taxid_to_name[taxid] = name
name_to_taxid[name] = taxid
return taxid_lineage, taxid_to_name, name_to_taxid
def kraken_main(results_path: str, bad_taxids_file: str, n: int):
# Load previously mapped taxids.
_kraken_taxid_cache = os.path.join(os.path.dirname(bad_taxids_file), 'taxid_cache.txt')
if os.path.exists(_kraken_taxid_cache):
taxid_map = load_taxid_cache(_kraken_taxid_cache)
else:
taxid_map = {}
# Get any new taxids.
bad_taxids = load_taxids(bad_taxids_file)
new_bad_taxids = {tid for tid in bad_taxids if tid not in taxid_map}
if new_bad_taxids:
taxid_lineage, taxid_to_name, name_to_taxid = retrieve_taxid_information(new_bad_taxids)
new_taxid_map = find_compatible_names(taxid_lineage, taxid_to_name, name_to_taxid, new_bad_taxids)
taxid_map.update(new_taxid_map)
cache_taxids(_kraken_taxid_cache, taxid_map)
# Fix taxids in kraken output.
kraken_raw_files = ls(results_path, ext='fq')
assert all(f.endswith("__output.fq") for f in kraken_raw_files)
# make_raw_results_compatible(taxid_map, kraken_raw_files, n=n)
else:
print("No new taxids to fix. Exiting...")
if __name__ == "__main__":
......@@ -168,8 +240,13 @@ if __name__ == "__main__":
parser.add_argument('-o', '--out', dest="results_path", help="Path to results.")
parser.add_argument('-T', dest="bad_taxids_file", help="Path to file containing taxids to replace.")
parser.add_argument('-N', default=5, dest="n_procs", help="Number of worker processes to use.")
parser.add_argument('--kraken', default=False, dest='kraken', action="store_true", help="Process Kraken(2) files.")
args = parser.parse_args()
main(args.db_path, args.results_path, args.bad_taxids_file, int(args.n_procs))
if not args.kraken:
expam_main(args.db_path, args.results_path, args.bad_taxids_file, int(args.n_procs))
else:
kraken_main(args.results_path, args.bad_taxids_file, int(args.n_procs))
\ No newline at end of file
......@@ -4,13 +4,14 @@ import os
import numpy as np # Pairwise distances.
from expam.sequences import get_sequence
from expam.utils import die
try:
from sourmash import signature as sm_signature
import sourmash
except ModuleNotFoundError:
print("Could not import sourmash! Install sourmash to use this feature.")
die("Could not import sourmash! Install sourmash to use this feature.")
def make_signature(genome_url, k, s):
......
......@@ -39,14 +39,11 @@ def ls(path, ext=None):
return files
def isfiletype(path, ext, check_compressed=False):
if ext[0] != '.':
ext = '.' + ext
def isfiletype(path: str, ext: str, check_compressed=False):
if check_compressed:
path = removecompressedextension(path)
if path[-len(ext):] == ext:
if path.endswith(ext):
return True
else:
return False
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment