taxaforge 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- taxaforge-0.2.0/LICENSE.txt +21 -0
- taxaforge-0.2.0/PKG-INFO +93 -0
- taxaforge-0.2.0/README.md +71 -0
- taxaforge-0.2.0/pyproject.toml +94 -0
- taxaforge-0.2.0/taxaforge.py +529 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) [2024] [Pandikunta Anand Reddy]
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
taxaforge-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: taxaforge
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: CLI tool to build taxonomic classifier databases (Kraken2, ganon2, ...) with a single command
|
|
5
|
+
Maintainer-email: ChillarAnand <anand@avilpage.com>
|
|
6
|
+
Requires-Python: >=3.8
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Classifier: Development Status :: 4 - Beta
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Classifier: Programming Language :: Python
|
|
13
|
+
Classifier: Typing :: Typed
|
|
14
|
+
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
|
|
15
|
+
Classifier: Topic :: Utilities
|
|
16
|
+
License-File: LICENSE.txt
|
|
17
|
+
Requires-Dist: ncbi-genome-download
|
|
18
|
+
Requires-Dist: click
|
|
19
|
+
Project-URL: Documentation, https://avilpage.com/taxaforge.html
|
|
20
|
+
Project-URL: Source, https://github.com/AvilPage/TaxaForge
|
|
21
|
+
|
|
22
|
+
Installation
|
|
23
|
+
============
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
pip install taxaforge
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
Usage
|
|
30
|
+
=====
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
taxaforge --help
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
To create standard Kraken2 database
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
taxaforge build --db-type standard
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
Before creating a standard database, you can try a smaller database like fungi.
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
taxaforge build --db-type fungi
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
To build a ganon2 database instead, pass `--tool ganon2`
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
taxaforge build --tool ganon2 --db-type fungi
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
To use locally downloaded files, run the following command
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
taxaforge build --db-name k2_test --genomes-dir /path/to/genomes --taxonomy-dir /path/to/taxonomy
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
To limit the number of genomes in the database, use the `--limit` option
|
|
61
|
+
|
|
62
|
+
```bash
|
|
63
|
+
taxaforge build --db-name k2_test_100 --genomes-dir /path/to/genomes --limit 1000
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
Config
|
|
67
|
+
======
|
|
68
|
+
|
|
69
|
+
Read/write config, stored in the OS default config location (`~/.config/taxaforge/config.ini` on Linux, `~/Library/Application Support/taxaforge/config.ini` on macOS).
|
|
70
|
+
|
|
71
|
+
```bash
|
|
72
|
+
taxaforge config set threads 8
|
|
73
|
+
taxaforge config get threads
|
|
74
|
+
taxaforge config
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
Why TaxaForge?
|
|
79
|
+
==============
|
|
80
|
+
|
|
81
|
+
TaxaForge aims to provide a simple and easy to use tool to build wide variety of taxonomic classifier databases with a single command.
|
|
82
|
+
|
|
83
|
+
Why not kraken2-build/ganon directly?
|
|
84
|
+
|
|
85
|
+
kraken2-build and ganon each build databases for their own tool only. TaxaForge wraps the shared download/taxonomy pipeline once and dispatches to the right tool via `--tool`, so adding support for more classifiers is a matter of plugging in a new build step.
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
Documentation
|
|
89
|
+
=============
|
|
90
|
+
|
|
91
|
+
- [Kraken2 Database Builder](https://avilpage.com/kdb.html)
|
|
92
|
+
- [Mastering Kraken2](https://avilpage.com/tags/kraken2.html)
|
|
93
|
+
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
Installation
|
|
2
|
+
============
|
|
3
|
+
|
|
4
|
+
```bash
|
|
5
|
+
pip install taxaforge
|
|
6
|
+
```
|
|
7
|
+
|
|
8
|
+
Usage
|
|
9
|
+
=====
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
taxaforge --help
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
To create standard Kraken2 database
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
taxaforge build --db-type standard
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
Before creating a standard database, you can try a smaller database like fungi.
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
taxaforge build --db-type fungi
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
To build a ganon2 database instead, pass `--tool ganon2`
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
taxaforge build --tool ganon2 --db-type fungi
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
To use locally downloaded files, run the following command
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
taxaforge build --db-name k2_test --genomes-dir /path/to/genomes --taxonomy-dir /path/to/taxonomy
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
To limit the number of genomes in the database, use the `--limit` option
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
taxaforge build --db-name k2_test_100 --genomes-dir /path/to/genomes --limit 1000
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
Config
|
|
46
|
+
======
|
|
47
|
+
|
|
48
|
+
Read/write config, stored in the OS default config location (`~/.config/taxaforge/config.ini` on Linux, `~/Library/Application Support/taxaforge/config.ini` on macOS).
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
taxaforge config set threads 8
|
|
52
|
+
taxaforge config get threads
|
|
53
|
+
taxaforge config
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
Why TaxaForge?
|
|
58
|
+
==============
|
|
59
|
+
|
|
60
|
+
TaxaForge aims to provide a simple and easy to use tool to build wide variety of taxonomic classifier databases with a single command.
|
|
61
|
+
|
|
62
|
+
Why not kraken2-build/ganon directly?
|
|
63
|
+
|
|
64
|
+
kraken2-build and ganon each build databases for their own tool only. TaxaForge wraps the shared download/taxonomy pipeline once and dispatches to the right tool via `--tool`, so adding support for more classifiers is a matter of plugging in a new build step.
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
Documentation
|
|
68
|
+
=============
|
|
69
|
+
|
|
70
|
+
- [Kraken2 Database Builder](https://avilpage.com/kdb.html)
|
|
71
|
+
- [Mastering Kraken2](https://avilpage.com/tags/kraken2.html)
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "taxaforge"
|
|
3
|
+
version = "0.2.0"
|
|
4
|
+
description = "CLI tool to build taxonomic classifier databases (Kraken2, ganon2, ...) with a single command"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = {file = "LICENSE.txt"}
|
|
7
|
+
maintainers = [{name = "ChillarAnand", email = "anand@avilpage.com"}]
|
|
8
|
+
classifiers = [
|
|
9
|
+
"Development Status :: 4 - Beta",
|
|
10
|
+
"Intended Audience :: Developers",
|
|
11
|
+
"License :: OSI Approved :: MIT License",
|
|
12
|
+
"Operating System :: OS Independent",
|
|
13
|
+
"Programming Language :: Python",
|
|
14
|
+
"Typing :: Typed",
|
|
15
|
+
"Topic :: Scientific/Engineering :: Bio-Informatics",
|
|
16
|
+
"Topic :: Utilities",
|
|
17
|
+
]
|
|
18
|
+
requires-python = ">=3.8"
|
|
19
|
+
dependencies = [
|
|
20
|
+
"ncbi-genome-download",
|
|
21
|
+
"click",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Documentation = "https://avilpage.com/taxaforge.html"
|
|
26
|
+
Source = "https://github.com/AvilPage/TaxaForge"
|
|
27
|
+
|
|
28
|
+
[build-system]
|
|
29
|
+
requires = ["flit_core<4"]
|
|
30
|
+
build-backend = "flit_core.buildapi"
|
|
31
|
+
|
|
32
|
+
[tool.flit.module]
|
|
33
|
+
name = "taxaforge"
|
|
34
|
+
|
|
35
|
+
[tool.flit.sdist]
|
|
36
|
+
include = [
|
|
37
|
+
"requirements/",
|
|
38
|
+
"tests/",
|
|
39
|
+
]
|
|
40
|
+
exclude = [
|
|
41
|
+
"docs/_build/",
|
|
42
|
+
]
|
|
43
|
+
|
|
44
|
+
[tool.pytest.ini_options]
|
|
45
|
+
testpaths = ["tests"]
|
|
46
|
+
filterwarnings = [
|
|
47
|
+
"error",
|
|
48
|
+
]
|
|
49
|
+
|
|
50
|
+
[tool.coverage.run]
|
|
51
|
+
branch = true
|
|
52
|
+
source = ["click", "tests"]
|
|
53
|
+
|
|
54
|
+
[tool.coverage.paths]
|
|
55
|
+
source = ["src", "*/site-packages"]
|
|
56
|
+
|
|
57
|
+
[tool.mypy]
|
|
58
|
+
python_version = "3.8"
|
|
59
|
+
files = ["src/click", "tests/typing"]
|
|
60
|
+
show_error_codes = true
|
|
61
|
+
pretty = true
|
|
62
|
+
strict = true
|
|
63
|
+
|
|
64
|
+
ignore_missing_imports = true
|
|
65
|
+
|
|
66
|
+
[tool.pyright]
|
|
67
|
+
pythonVersion = "3.8"
|
|
68
|
+
include = ["src/click", "tests/typing"]
|
|
69
|
+
typeCheckingMode = "basic"
|
|
70
|
+
|
|
71
|
+
[tool.ruff]
|
|
72
|
+
extend-exclude = ["examples/"]
|
|
73
|
+
src = ["src"]
|
|
74
|
+
fix = true
|
|
75
|
+
show-fixes = true
|
|
76
|
+
output-format = "full"
|
|
77
|
+
|
|
78
|
+
[tool.ruff.lint]
|
|
79
|
+
select = [
|
|
80
|
+
"B", # flake8-bugbear
|
|
81
|
+
"E", # pycodestyle error
|
|
82
|
+
"F", # pyflakes
|
|
83
|
+
"I", # isort
|
|
84
|
+
"UP", # pyupgrade
|
|
85
|
+
"W", # pycodestyle warning
|
|
86
|
+
]
|
|
87
|
+
ignore-init-module-imports = true
|
|
88
|
+
|
|
89
|
+
[tool.ruff.lint.isort]
|
|
90
|
+
force-single-line = true
|
|
91
|
+
order-by-type = false
|
|
92
|
+
|
|
93
|
+
[project.scripts]
|
|
94
|
+
taxaforge = "taxaforge:cli"
|
|
@@ -0,0 +1,529 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
import concurrent.futures
|
|
3
|
+
import configparser
|
|
4
|
+
import datetime
|
|
5
|
+
import hashlib
|
|
6
|
+
import logging
|
|
7
|
+
import multiprocessing
|
|
8
|
+
import os
|
|
9
|
+
import shlex
|
|
10
|
+
import shutil
|
|
11
|
+
import subprocess
|
|
12
|
+
import sys
|
|
13
|
+
import urllib.request
|
|
14
|
+
from pathlib import Path
|
|
15
|
+
|
|
16
|
+
import click
|
|
17
|
+
import ncbi_genome_download
|
|
18
|
+
from tqdm import tqdm
|
|
19
|
+
|
|
20
|
+
logger = logging.getLogger(__name__)
|
|
21
|
+
logger.setLevel(logging.INFO)
|
|
22
|
+
logger.addHandler(logging.StreamHandler())
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
NCBI_SERVER = "https://ftp.ncbi.nlm.nih.gov"
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
DB_TYPE_CONFIG = {
|
|
29
|
+
'standard': ("archaea", "bacteria", "viral", "plasmid", "human", "UniVec_Core")
|
|
30
|
+
}
|
|
31
|
+
REQUIRED_BINS = {
|
|
32
|
+
'kraken2': "kraken2-build",
|
|
33
|
+
'ganon2': "ganon",
|
|
34
|
+
'ganon': "ganon",
|
|
35
|
+
}
|
|
36
|
+
hashes = set()
|
|
37
|
+
md5_file = None
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def hash_file(filename, buf_size=8192):
|
|
41
|
+
md5 = hashlib.md5()
|
|
42
|
+
with open(filename, "rb") as in_file:
|
|
43
|
+
while True:
|
|
44
|
+
data = in_file.read(buf_size)
|
|
45
|
+
if not data:
|
|
46
|
+
break
|
|
47
|
+
md5.update(data)
|
|
48
|
+
digest = md5.hexdigest()
|
|
49
|
+
return digest
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def run_basic_checks(tool, use_k2=False):
|
|
53
|
+
if not shutil.which("ncbi-genome-download"):
|
|
54
|
+
logger.error("ncbi-genome-download not found in PATH. Exiting.")
|
|
55
|
+
sys.exit(1)
|
|
56
|
+
|
|
57
|
+
if tool not in REQUIRED_BINS:
|
|
58
|
+
logger.error(f"Unknown tool: {tool}. Supported tools: {', '.join(REQUIRED_BINS)}")
|
|
59
|
+
sys.exit(1)
|
|
60
|
+
|
|
61
|
+
binary = "k2" if (tool == 'kraken2' and use_k2) else REQUIRED_BINS[tool]
|
|
62
|
+
if not shutil.which(binary):
|
|
63
|
+
logger.error(f"{binary} not found in PATH. Exiting.")
|
|
64
|
+
sys.exit(1)
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def create_cache_dir():
|
|
68
|
+
# Unix ~/.cache/taxaforge
|
|
69
|
+
# macOS ~/Library/Caches/taxaforge
|
|
70
|
+
if sys.platform == "darwin":
|
|
71
|
+
cache_dir = Path.home() / "Library" / "Caches" / "taxaforge"
|
|
72
|
+
if sys.platform == "linux":
|
|
73
|
+
cache_dir = Path.home() / ".cache" / "taxaforge"
|
|
74
|
+
|
|
75
|
+
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
76
|
+
return cache_dir
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def create_config_dir():
|
|
80
|
+
# Unix ~/.config/taxaforge
|
|
81
|
+
# macOS ~/Library/Application Support/taxaforge
|
|
82
|
+
if sys.platform == "darwin":
|
|
83
|
+
config_dir = Path.home() / "Library" / "Application Support" / "taxaforge"
|
|
84
|
+
if sys.platform == "linux":
|
|
85
|
+
config_dir = Path.home() / ".config" / "taxaforge"
|
|
86
|
+
|
|
87
|
+
config_dir.mkdir(parents=True, exist_ok=True)
|
|
88
|
+
return config_dir
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
CONFIG_SECTION = "taxaforge"
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def get_config_path():
|
|
95
|
+
return create_config_dir() / "config.ini"
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def load_config():
|
|
99
|
+
parser = configparser.ConfigParser()
|
|
100
|
+
parser.read(get_config_path())
|
|
101
|
+
if not parser.has_section(CONFIG_SECTION):
|
|
102
|
+
parser.add_section(CONFIG_SECTION)
|
|
103
|
+
return parser
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def save_config(parser):
|
|
107
|
+
with open(get_config_path(), "w") as out_file:
|
|
108
|
+
parser.write(out_file)
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def download_file(url, position):
|
|
112
|
+
filename = url.rsplit("/", 1)[-1]
|
|
113
|
+
existing = os.path.getsize(filename) if os.path.exists(filename) else 0
|
|
114
|
+
|
|
115
|
+
request = urllib.request.Request(url)
|
|
116
|
+
if existing:
|
|
117
|
+
request.add_header("Range", f"bytes={existing}-")
|
|
118
|
+
|
|
119
|
+
with urllib.request.urlopen(request) as response:
|
|
120
|
+
resumed = response.status == 206
|
|
121
|
+
total = int(response.headers.get("Content-Length", 0)) + (existing if resumed else 0)
|
|
122
|
+
|
|
123
|
+
with open(filename, "ab" if resumed else "wb") as out_file, tqdm(
|
|
124
|
+
total=total, initial=existing if resumed else 0, unit="B", unit_scale=True,
|
|
125
|
+
desc=filename, position=position, leave=True
|
|
126
|
+
) as bar:
|
|
127
|
+
while True:
|
|
128
|
+
chunk = response.read(1024 * 1024)
|
|
129
|
+
if not chunk:
|
|
130
|
+
break
|
|
131
|
+
out_file.write(chunk)
|
|
132
|
+
bar.update(len(chunk))
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def download_files(urls, max_workers=4):
|
|
136
|
+
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
137
|
+
futures = [
|
|
138
|
+
executor.submit(download_file, url, index % max_workers)
|
|
139
|
+
for index, url in enumerate(urls)
|
|
140
|
+
]
|
|
141
|
+
for future in concurrent.futures.as_completed(futures):
|
|
142
|
+
future.result()
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def download_taxanomy(cache_dir, skip_maps=None, protein=None):
|
|
146
|
+
taxonomy_path = os.path.join(cache_dir, "taxonomy")
|
|
147
|
+
os.makedirs(taxonomy_path, exist_ok=True)
|
|
148
|
+
os.chdir(taxonomy_path)
|
|
149
|
+
|
|
150
|
+
if not skip_maps:
|
|
151
|
+
if not protein:
|
|
152
|
+
# Define URLs for nucleotide accession to taxon map
|
|
153
|
+
urls = [
|
|
154
|
+
f"{NCBI_SERVER}/pub/taxonomy/accession2taxid/nucl_gb.accession2taxid.gz",
|
|
155
|
+
f"{NCBI_SERVER}/pub/taxonomy/accession2taxid/nucl_wgs.accession2taxid.gz"
|
|
156
|
+
]
|
|
157
|
+
else:
|
|
158
|
+
# Define URL for protein accession to taxon map
|
|
159
|
+
urls = ["ftp://ftp.ncbi.nlm.nih.gov/pub/taxonomy/accession2taxid/prot.accession2taxid.gz"]
|
|
160
|
+
else:
|
|
161
|
+
logger.info("Skipping maps download")
|
|
162
|
+
|
|
163
|
+
# Download taxonomy tree data
|
|
164
|
+
urls.append(f"{NCBI_SERVER}/pub/taxonomy/taxdump.tar.gz")
|
|
165
|
+
|
|
166
|
+
logger.info(f"Downloading {len(urls)} taxonomy files")
|
|
167
|
+
download_files(urls)
|
|
168
|
+
|
|
169
|
+
logger.info("Extracting taxdump.tar.gz")
|
|
170
|
+
cmd = f"tar -k -xvf taxdump.tar.gz"
|
|
171
|
+
run_cmd(cmd)
|
|
172
|
+
|
|
173
|
+
logger.info("Decompressing taxonomy data")
|
|
174
|
+
cmd = f"find {cache_dir}/taxonomy -name '*.gz' | xargs -n 1 gunzip -k"
|
|
175
|
+
run_cmd(cmd)
|
|
176
|
+
|
|
177
|
+
logger.info("Finished downloading taxonomy data")
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
def run_cmd(cmd, return_output=False, no_output=False):
|
|
181
|
+
if not no_output:
|
|
182
|
+
logger.info(f"Running command: {cmd}")
|
|
183
|
+
|
|
184
|
+
if return_output:
|
|
185
|
+
return subprocess.check_output(cmd, shell=True).decode("utf-8").strip().split("\n")
|
|
186
|
+
|
|
187
|
+
try:
|
|
188
|
+
if no_output:
|
|
189
|
+
subprocess.run(cmd, shell=True, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
190
|
+
else:
|
|
191
|
+
subprocess.run(cmd, shell=True, check=True)
|
|
192
|
+
except subprocess.CalledProcessError:
|
|
193
|
+
pass
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
def download_genomes(cache_dir, cwd, db_type, db_name, threads, force=False):
|
|
197
|
+
organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
|
|
198
|
+
if force:
|
|
199
|
+
shutil.rmtree(cwd / db_name, ignore_errors=True)
|
|
200
|
+
|
|
201
|
+
os.makedirs(cwd / db_name, exist_ok=True)
|
|
202
|
+
|
|
203
|
+
for organism in organisms:
|
|
204
|
+
logger.info(f"Downloading genomes for {organism}")
|
|
205
|
+
os.chdir(cache_dir)
|
|
206
|
+
ncbi_genome_download.download(
|
|
207
|
+
section='refseq', groups=organism, file_formats='fasta',
|
|
208
|
+
progress_bar=True, parallel=threads,
|
|
209
|
+
assembly_levels=['complete'],
|
|
210
|
+
output=cache_dir
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
cmd = f"find {cache_dir}/refseq/{organism} -name '*.gz' | xargs -n 1 -P {threads} gunzip -k"
|
|
214
|
+
run_cmd(cmd)
|
|
215
|
+
logger.info(f"Finished downloading {organism} genomes")
|
|
216
|
+
|
|
217
|
+
os.chdir(cwd)
|
|
218
|
+
logger.info("Finished downloading all genomes")
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
def build_db(
|
|
222
|
+
cache_dir, cwd, db_type, db_name, threads, kmer_len, min_len,
|
|
223
|
+
fast_build, rebuild, load_factor, use_k2
|
|
224
|
+
):
|
|
225
|
+
run_cmd(f"cd {cwd}")
|
|
226
|
+
|
|
227
|
+
if not os.path.exists(f"{db_name}/taxonomy"):
|
|
228
|
+
cmd = f"ln -s {cache_dir}/taxonomy {db_name}/"
|
|
229
|
+
run_cmd(cmd)
|
|
230
|
+
|
|
231
|
+
if rebuild:
|
|
232
|
+
cmd = f"rm -rf {db_name}/*.k2d"
|
|
233
|
+
run_cmd(cmd)
|
|
234
|
+
|
|
235
|
+
# TODO: Fix issue with macos threads
|
|
236
|
+
if sys.platform == "darwin":
|
|
237
|
+
threads = 1
|
|
238
|
+
|
|
239
|
+
if use_k2:
|
|
240
|
+
cmd = f"k2 build"
|
|
241
|
+
else:
|
|
242
|
+
cmd = f"kraken2-build --build"
|
|
243
|
+
|
|
244
|
+
cmd += f" --db {db_name} --threads {threads} --kmer-len {kmer_len} --minimizer-len {min_len} --load-factor {load_factor}"
|
|
245
|
+
if fast_build:
|
|
246
|
+
cmd += " --fast-build"
|
|
247
|
+
|
|
248
|
+
run_cmd(cmd)
|
|
249
|
+
|
|
250
|
+
cmd = f"du -sh {db_name}/*.k2d"
|
|
251
|
+
run_cmd(cmd)
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
def build_ganon2(cache_dir, cwd, genomes_dir, db_type, db_name, threads, kmer_len, min_len, level, rebuild):
|
|
255
|
+
os.chdir(cwd)
|
|
256
|
+
|
|
257
|
+
if genomes_dir:
|
|
258
|
+
input_dirs = [str(genomes_dir)]
|
|
259
|
+
else:
|
|
260
|
+
organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
|
|
261
|
+
input_dirs = [f"{cache_dir}/refseq/{organism}" for organism in organisms]
|
|
262
|
+
|
|
263
|
+
if rebuild:
|
|
264
|
+
cmd = f"rm -f {db_name}.*"
|
|
265
|
+
run_cmd(cmd)
|
|
266
|
+
|
|
267
|
+
cmd = (
|
|
268
|
+
f"ganon build-custom --input {' '.join(input_dirs)} --input-recursive "
|
|
269
|
+
f"--taxonomy-files {cache_dir}/taxonomy/nodes.dmp {cache_dir}/taxonomy/names.dmp "
|
|
270
|
+
f"--db-prefix {db_name} --threads {threads} --kmer-size {kmer_len} "
|
|
271
|
+
f"--window-size {min_len} --level {level}"
|
|
272
|
+
)
|
|
273
|
+
run_cmd(cmd)
|
|
274
|
+
|
|
275
|
+
cmd = f"du -sh {db_name}.*"
|
|
276
|
+
run_cmd(cmd)
|
|
277
|
+
|
|
278
|
+
|
|
279
|
+
def get_files(genomes_dir, cache_dir, db_type, db_name, threads):
|
|
280
|
+
if genomes_dir:
|
|
281
|
+
logger.info(f"Adding {genomes_dir} genomes to library")
|
|
282
|
+
|
|
283
|
+
cmd = f"find {genomes_dir} -name '*.gz' | xargs -n 1 -P {threads} gunzip -k"
|
|
284
|
+
run_cmd(cmd)
|
|
285
|
+
|
|
286
|
+
cmd = f"find {genomes_dir} -name '*.gbff'"
|
|
287
|
+
files = run_cmd(cmd, return_output=True)
|
|
288
|
+
for file in files:
|
|
289
|
+
if os.path.exists(f"{file}.fna"):
|
|
290
|
+
continue
|
|
291
|
+
cmd = f"any2fasta -u {file} > {file}.fna"
|
|
292
|
+
run_cmd(cmd)
|
|
293
|
+
|
|
294
|
+
cmd = f"find {genomes_dir} -type f -name '*.fna'"
|
|
295
|
+
files = run_cmd(cmd, return_output=True)
|
|
296
|
+
logger.info(f"Found {len(files)} genomes to add to {db_name} library")
|
|
297
|
+
else:
|
|
298
|
+
organisms = DB_TYPE_CONFIG.get(db_type, [db_type])
|
|
299
|
+
files = []
|
|
300
|
+
for organism in organisms:
|
|
301
|
+
cmd = f"find {cache_dir}/refseq/{organism} -name '*.fna'"
|
|
302
|
+
org_files = run_cmd(cmd, return_output=True)
|
|
303
|
+
logger.info(f"Found {len(org_files)} genomes for {organism}")
|
|
304
|
+
files.extend(org_files)
|
|
305
|
+
|
|
306
|
+
return files
|
|
307
|
+
|
|
308
|
+
|
|
309
|
+
def save_md5_file(*args, **kwargs):
|
|
310
|
+
global md5_file
|
|
311
|
+
with open(md5_file, "w") as out_file:
|
|
312
|
+
for line in hashes:
|
|
313
|
+
out_file.write(line + "\n")
|
|
314
|
+
logger.info(f"Saved {len(hashes)} md5 hashes")
|
|
315
|
+
|
|
316
|
+
|
|
317
|
+
def add_to_library(
|
|
318
|
+
cache_dir, cwd, genomes_dir, db_type, db_name,
|
|
319
|
+
limit, batch_size, threads, use_k2
|
|
320
|
+
):
|
|
321
|
+
os.chdir(cwd)
|
|
322
|
+
os.makedirs(cwd / db_name / "library", exist_ok=True)
|
|
323
|
+
|
|
324
|
+
files = get_files(genomes_dir, cache_dir, db_type, db_name, threads)
|
|
325
|
+
if limit:
|
|
326
|
+
logger.info(f"Limiting number of genomes to {limit}")
|
|
327
|
+
files = files[:limit]
|
|
328
|
+
|
|
329
|
+
step = batch_size
|
|
330
|
+
dynamic_step = len(files) // 10
|
|
331
|
+
step = min(step, dynamic_step)
|
|
332
|
+
if step == 0:
|
|
333
|
+
step = 1
|
|
334
|
+
|
|
335
|
+
logger.info(f"Using step size of {step}")
|
|
336
|
+
|
|
337
|
+
file_count = len(files)
|
|
338
|
+
start = datetime.datetime.now()
|
|
339
|
+
|
|
340
|
+
if use_k2:
|
|
341
|
+
for index, file in enumerate(files, start=1):
|
|
342
|
+
if index % step == 0:
|
|
343
|
+
duration = datetime.datetime.now() - start
|
|
344
|
+
average_speed = duration / step
|
|
345
|
+
eta = (file_count - index) * average_speed
|
|
346
|
+
logger.info(f"{datetime.datetime.now()}: Added {index} genomes in {duration}. ETA: {eta}")
|
|
347
|
+
start = datetime.datetime.now()
|
|
348
|
+
|
|
349
|
+
cmd = f"k2 add-to-library --db {db_name} --files {file}"
|
|
350
|
+
run_cmd(cmd, no_output=True)
|
|
351
|
+
|
|
352
|
+
logger.info(f"Added downloaded genomes to library")
|
|
353
|
+
end = datetime.datetime.now()
|
|
354
|
+
print(f"Time taken: {end - start}")
|
|
355
|
+
return
|
|
356
|
+
|
|
357
|
+
global hashes
|
|
358
|
+
global md5_file
|
|
359
|
+
md5_file = cwd / db_name / "library" / "added.md5"
|
|
360
|
+
|
|
361
|
+
if os.path.exists(md5_file):
|
|
362
|
+
with open(md5_file, "r") as in_file:
|
|
363
|
+
hashes = {line.strip() for line in in_file}
|
|
364
|
+
|
|
365
|
+
logger.info(f"Found {len(hashes)} md5 hashes in {md5_file}")
|
|
366
|
+
|
|
367
|
+
for index, file in enumerate(files, start=1):
|
|
368
|
+
if index % step == 0:
|
|
369
|
+
duration = datetime.datetime.now() - start
|
|
370
|
+
average_speed = duration / step
|
|
371
|
+
eta = (file_count - index) * average_speed
|
|
372
|
+
logger.info(f"{datetime.datetime.now()}: Added {index} genomes in {duration}. ETA: {eta}")
|
|
373
|
+
start = datetime.datetime.now()
|
|
374
|
+
|
|
375
|
+
if not os.path.exists(f"{file}.md5"):
|
|
376
|
+
md5sum = hash_file(file)
|
|
377
|
+
with open(f"{file}.md5", "w") as fh:
|
|
378
|
+
fh.write(md5sum)
|
|
379
|
+
else:
|
|
380
|
+
with open(f"{file}.md5", "r") as in_file:
|
|
381
|
+
md5sum = in_file.read()
|
|
382
|
+
|
|
383
|
+
if md5sum in hashes:
|
|
384
|
+
continue
|
|
385
|
+
|
|
386
|
+
cmd = f"kraken2-build --db {db_name} --add-to-library {file} --threads {threads}"
|
|
387
|
+
run_cmd(cmd, no_output=True)
|
|
388
|
+
|
|
389
|
+
with open(md5_file, "a") as out_file:
|
|
390
|
+
out_file.write(md5sum + "\n")
|
|
391
|
+
|
|
392
|
+
hashes.add(md5sum)
|
|
393
|
+
|
|
394
|
+
end = datetime.datetime.now()
|
|
395
|
+
print(f"Time taken: {end - start}")
|
|
396
|
+
|
|
397
|
+
logger.info(f"Added downloaded genomes to library")
|
|
398
|
+
|
|
399
|
+
|
|
400
|
+
@click.group(no_args_is_help=True, epilog=f"Config file: {get_config_path()}")
|
|
401
|
+
def cli():
|
|
402
|
+
pass
|
|
403
|
+
|
|
404
|
+
|
|
405
|
+
@cli.command(no_args_is_help=True, context_settings={"ignore_unknown_options": True})
|
|
406
|
+
@click.option('--tool', default='kraken2', type=click.Choice(list(REQUIRED_BINS)), help='Classifier to build the database for')
|
|
407
|
+
@click.option('--db-type', default=None, help='database type to build')
|
|
408
|
+
@click.option('--db-name', default=None, help='database name to build')
|
|
409
|
+
@click.option('--genomes-dir', default=None, help='Directory containing genomes')
|
|
410
|
+
@click.option('--cache-dir', default=create_cache_dir(), help='Cache directory')
|
|
411
|
+
@click.option('--threads', default=multiprocessing.cpu_count(), help='Number of threads to use', type=int)
|
|
412
|
+
@click.option('--load-factor', default=0.7, help='Proportion of the hash table to be populated. Used only for kraken2')
|
|
413
|
+
@click.option('--kmer-len', default=35, help='Kmer length in bp/aa. Used only in build task', type=int)
|
|
414
|
+
@click.option('--min-len', default=31, help='Minimizer/window length in bp/aa. Used only in build task', type=int)
|
|
415
|
+
@click.option('--level', default='leaves', type=click.Choice(['leaves', 'species', 'genus', 'assembly']), help='Taxonomic level to group sequences by. Used only for ganon2')
|
|
416
|
+
@click.option('--limit', default=None, help='Limit number of genomes to use', type=int)
|
|
417
|
+
@click.option('--batch-size', default=1000, help='Number of genomes to add to library at a time. Used only for kraken2', type=int)
|
|
418
|
+
@click.option('--force', is_flag=True, help='Force download and build')
|
|
419
|
+
@click.option('--rebuild', is_flag=True, help='Clean existing build files and re-build')
|
|
420
|
+
@click.option('--fast-build', is_flag=True, help='Non deterministic but faster build. Used only for kraken2')
|
|
421
|
+
@click.option('--use-k2', is_flag=True, help='Use k2 CLI instead of kraken2-build. Used only for kraken2')
|
|
422
|
+
@click.argument('ganon_args', nargs=-1, type=click.UNPROCESSED)
|
|
423
|
+
@click.pass_context
|
|
424
|
+
def build(
|
|
425
|
+
context,
|
|
426
|
+
tool: str, db_type: str, db_name, cache_dir, genomes_dir,
|
|
427
|
+
threads, load_factor, kmer_len: int, min_len, level: str, limit: int, batch_size: int,
|
|
428
|
+
force: bool, rebuild, fast_build: bool, use_k2: bool, ganon_args
|
|
429
|
+
):
|
|
430
|
+
if tool in ('ganon', 'ganon2') and ganon_args:
|
|
431
|
+
logger.info(f"Passing through unrecognized options to native ganon build: {' '.join(ganon_args)}")
|
|
432
|
+
run_basic_checks(tool, use_k2)
|
|
433
|
+
cmd = "ganon build " + " ".join(shlex.quote(arg) for arg in ganon_args)
|
|
434
|
+
run_cmd(cmd)
|
|
435
|
+
return
|
|
436
|
+
|
|
437
|
+
logger.info(f"Building {tool} database of type {db_type}")
|
|
438
|
+
run_basic_checks(tool, use_k2)
|
|
439
|
+
cwd = Path(os.getcwd())
|
|
440
|
+
|
|
441
|
+
if cache_dir == '.':
|
|
442
|
+
cache_dir = cwd
|
|
443
|
+
|
|
444
|
+
if not db_name:
|
|
445
|
+
db_name = f"{tool}_{context.params['db_type']}"
|
|
446
|
+
|
|
447
|
+
if force:
|
|
448
|
+
if tool == 'kraken2':
|
|
449
|
+
run_cmd(f"rm -rf {db_name}")
|
|
450
|
+
run_cmd(f"mkdir -p {db_name}")
|
|
451
|
+
else:
|
|
452
|
+
run_cmd(f"rm -f {db_name}.*")
|
|
453
|
+
|
|
454
|
+
logger.info(f"Using cache directory {cache_dir}")
|
|
455
|
+
|
|
456
|
+
download_taxanomy(cache_dir)
|
|
457
|
+
|
|
458
|
+
if not genomes_dir:
|
|
459
|
+
download_genomes(cache_dir, cwd, db_type, db_name, threads, force)
|
|
460
|
+
|
|
461
|
+
if tool == 'kraken2':
|
|
462
|
+
add_to_library(
|
|
463
|
+
cache_dir, cwd, genomes_dir, db_type, db_name,
|
|
464
|
+
limit, batch_size, threads, use_k2
|
|
465
|
+
)
|
|
466
|
+
build_db(
|
|
467
|
+
cache_dir, cwd, db_type, db_name, threads, kmer_len, min_len,
|
|
468
|
+
fast_build, rebuild, load_factor, use_k2
|
|
469
|
+
)
|
|
470
|
+
elif tool == 'ganon2':
|
|
471
|
+
build_ganon2(
|
|
472
|
+
cache_dir, cwd, genomes_dir, db_type, db_name, threads,
|
|
473
|
+
kmer_len, min_len, level, rebuild
|
|
474
|
+
)
|
|
475
|
+
|
|
476
|
+
|
|
477
|
+
@cli.group(name='config', invoke_without_command=True)
|
|
478
|
+
@click.pass_context
|
|
479
|
+
def config(context):
|
|
480
|
+
if context.invoked_subcommand is not None:
|
|
481
|
+
return
|
|
482
|
+
|
|
483
|
+
parser = load_config()
|
|
484
|
+
items = parser.items(CONFIG_SECTION)
|
|
485
|
+
if not items:
|
|
486
|
+
logger.info(f"No config set. Config file: {get_config_path()}")
|
|
487
|
+
return
|
|
488
|
+
|
|
489
|
+
for key, value in items:
|
|
490
|
+
print(f"{key} = {value}")
|
|
491
|
+
|
|
492
|
+
|
|
493
|
+
@config.command(name='get')
|
|
494
|
+
@click.argument('key')
|
|
495
|
+
def config_get(key):
|
|
496
|
+
parser = load_config()
|
|
497
|
+
if not parser.has_option(CONFIG_SECTION, key):
|
|
498
|
+
logger.error(f"{key} not set")
|
|
499
|
+
sys.exit(1)
|
|
500
|
+
print(parser.get(CONFIG_SECTION, key))
|
|
501
|
+
|
|
502
|
+
|
|
503
|
+
@config.command(name='set')
|
|
504
|
+
@click.argument('key')
|
|
505
|
+
@click.argument('value')
|
|
506
|
+
def config_set(key, value):
|
|
507
|
+
parser = load_config()
|
|
508
|
+
parser.set(CONFIG_SECTION, key, value)
|
|
509
|
+
save_config(parser)
|
|
510
|
+
logger.info(f"Set {key} = {value}")
|
|
511
|
+
|
|
512
|
+
|
|
513
|
+
DOCTOR_BINS = ["ncbi-genome-download", "kraken2-build", "k2", "ganon", "any2fasta", "wget", "tar", "gunzip"]
|
|
514
|
+
|
|
515
|
+
|
|
516
|
+
@cli.command(name='doctor')
|
|
517
|
+
def doctor():
|
|
518
|
+
for binary in DOCTOR_BINS:
|
|
519
|
+
path = shutil.which(binary)
|
|
520
|
+
status = path if path else "MISSING"
|
|
521
|
+
print(f"{binary:<20} {status}")
|
|
522
|
+
|
|
523
|
+
print()
|
|
524
|
+
print(f"Cache dir: {create_cache_dir()}")
|
|
525
|
+
print(f"Config file: {get_config_path()}")
|
|
526
|
+
|
|
527
|
+
|
|
528
|
+
if __name__ == '__main__':
|
|
529
|
+
cli()
|