annotate data_manager/kraken2_build_database.py @ 13:e9ee4d074d5d draft

planemo upload for repository https://github.com/galaxyproject/tools-iuc/tree/master/data_managers/data_manager_build_kraken2_database/ commit bb6d9b53b197610c5dbbb40c8ab590761755d91a

1 #!/usr/bin/env python
2 
3 import argparse
4 import datetime
5 import errno
6 import json
7 import os
8 import re
9 import shutil
10 import subprocess
11 import sys
12 import tarfile
13 from enum import Enum
14 
15 try:
16     # Python3
17     from urllib.request import urlopen
18     from urllib.error import URLError
19 except ImportError:
20     from urllib2 import urlopen
21     from urllib2 import URLError
22 
23 
24 DATA_TABLE_NAME = "kraken2_databases"
25 
26 
27 class KrakenDatabaseTypes(Enum):
28     standard_local_build = 'standard_local_build'
29     standard_prebuilt = 'standard_prebuilt'
30     minikraken = 'minikraken'
31     special_prebuilt = 'special_prebuilt'
32     special = 'special'
33     custom = 'custom'
34 
35     def __str__(self):
36         return self.value
37 
38 
39 class SpecialDatabaseTypes(Enum):
40     rdp = 'rdp'
41     greengenes = 'greengenes'
42     silva = 'silva'
43 
44     def __str__(self):
45         return self.value
46 
47 
48 class Minikraken2Versions(Enum):
49     v1 = 'v1'
50     v2 = 'v2'
51 
52     def __str__(self):
53         return self.value
54 
55 
56 class StandardPrebuiltSizes(Enum):
57     viral = "viral"
58     minusb = "minusb"
59     standard = "standard"
60     standard_08gb = "standard_08gb"
61     standard_16gb = "standard_16gb"
62     pluspf = "pluspf"
63     pluspf_08gb = "pluspf_08gb"
64     pluspf_16gb = "pluspf_16gb"
65     pluspfp = "pluspfp"
66     pluspfp_08gb = "pluspfp_08gb"
67     pluspfp_16gb = "pluspfp_16gb"
68     eupathdb48 = "eupathdb48"
69 
70     def __str__(self):
71         return self.value
72 
73 
74 def kraken2_build_standard(kraken2_args, target_directory, data_table_name=DATA_TABLE_NAME):
75     now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ")
76 
77     database_value = "_".join([
78         now,
79         "standard",
80         "kmer-len", str(kraken2_args["kmer_len"]),
81         "minimizer-len", str(kraken2_args["minimizer_len"]),
82         "minimizer-spaces", str(kraken2_args["minimizer_spaces"]),
83         "load-factor", str(kraken2_args["load_factor"]),
84     ])
85 
86     database_name = " ".join([
87         "Standard (Local Build)",
88         "(Created:",
89         now + ",",
90         "kmer-len=" + str(kraken2_args["kmer_len"]) + ",",
91         "minimizer-len=" + str(kraken2_args["minimizer_len"]) + ",",
92         "minimizer-spaces=" + str(kraken2_args["minimizer_spaces"]) + ")",
93         "load-factor", str(kraken2_args["load_factor"]),
94     ])
95 
96     database_path = database_value
97 
98     args = [
99         '--threads', str(kraken2_args["threads"]),
00         '--standard',
01         '--kmer-len', str(kraken2_args["kmer_len"]),
02         '--minimizer-len', str(kraken2_args["minimizer_len"]),
03         '--minimizer-spaces', str(kraken2_args["minimizer_spaces"]),
04         '--load-factor', str(kraken2_args["load_factor"]),
05         '--db', database_path
06     ]
07 
08     subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
09 
10     if kraken2_args["clean"]:
11         args = [
12             '--threads', str(kraken2_args["threads"]),
13             '--clean',
14             '--db', database_path
15         ]
16 
17         subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
18 
19     data_table_entry = {
20         'data_tables': {
21             data_table_name: [
22                 {
23                     "value": database_value,
24                     "name": database_name,
25                     "path": database_path,
26                 }
27             ]
28         }
29     }
30 
31     return data_table_entry
32 
33 
34 def kraken2_build_standard_prebuilt(prebuilt_db, prebuilt_date, target_directory, data_table_name=DATA_TABLE_NAME):
35 
36     now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ")
37 
38     prebuild_name = {
39         'viral': "Viral",
40         'minusb': "MinusB (archaea, viral, plasmid, human, UniVec_Core)",
41         'standard': "Standard-Full (archaea, bacteria, viral, plasmid, human,UniVec_Core)",
42         'standard_08gb': "Standard-8 (Standard with DB capped at 8 GB)",
43         'standard_16gb': "Standard-16 (Standard with DB capped at 16 GB)",
44         'pluspf': "PlusPF (Standard plus protozoa and fungi)",
45         'pluspf_08gb': "PlusPF-8 (PlusPF with DB capped at 8 GB)",
46         'pluspf_16gb': "PlusPF-16 (PlusPF with DB capped at 16 GB)",
47         'pluspfp': "PlusPFP (Standard plus protozoa, fungi and plant)",
48         'pluspfp_08gb': "PlusPFP-8 (PlusPFP with DB capped at 8 GB)",
49         'pluspfp_16gb': "PlusPFP-16 (PlusPFP with DB capped at 16 GB)",
50         'eupathdb48': "EuPathDB-46",
51     }
52 
53     database_value = "_".join([
54         now,
55         "standard_prebuilt",
56         prebuilt_db,
57         prebuilt_date
58     ])
59 
60     database_name = " ".join([
61         "Prebuilt Refseq indexes: ",
62         prebuild_name[prebuilt_db],
63         "(Version: ",
64         prebuilt_date,
65         "- Downloaded:",
66         now + ")"
67     ])
68 
69     database_path = database_value
70 
71     # we may need to let the user choose the date when new DBs are posted.
72     date_url_str = prebuilt_date.replace('-', '')
73     # download the pre-built database
74     try:
75         download_url = 'https://genome-idx.s3.amazonaws.com/kraken/k2_%s_%s.tar.gz' % (prebuilt_db, date_url_str)
76         src = urlopen(download_url)
77     except URLError as e:
78         print('url: ' + download_url, file=sys.stderr)
79         print(e, file=sys.stderr)
80         exit(1)
81 
82     with open('tmp_data.tar.gz', 'wb') as dst:
83         shutil.copyfileobj(src, dst)
84     # unpack the downloaded archive to the target directory
85     with tarfile.open('tmp_data.tar.gz', 'r:gz') as fh:
86         for member in fh.getmembers():
87             if member.isreg():
88                 member.name = os.path.basename(member.name)
89                 fh.extract(member, os.path.join(target_directory, database_path))
90 
91     data_table_entry = {
92         'data_tables': {
93             data_table_name: [
94                 {
95                     "value": database_value,
96                     "name": database_name,
97                     "path": database_path,
98                 }
99             ]
00         }
01     }
02 
03     return data_table_entry
04 
05 
06 def kraken2_build_minikraken(minikraken2_version, target_directory, data_table_name=DATA_TABLE_NAME):
07 
08     now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ")
09 
10     database_value = "_".join([
11         now,
12         "minikraken2",
13         minikraken2_version,
14         "8GB",
15     ])
16 
17     database_name = " ".join([
18         "Minikraken2",
19         minikraken2_version,
20         "(Created:",
21         now + ")"
22     ])
23 
24     database_path = database_value
25 
26     # download the minikraken2 data
27     try:
28         download_url = 'https://genome-idx.s3.amazonaws.com/kraken/minikraken2_%s_8GB_201904.tgz' % minikraken2_version
29         src = urlopen(download_url)
30     except URLError as e:
31         print('url: ' + download_url, file=sys.stderr)
32         print(e, file=sys.stderr)
33         exit(1)
34 
35     with open('tmp_data.tar.gz', 'wb') as dst:
36         shutil.copyfileobj(src, dst)
37     # unpack the downloaded archive to the target directory
38     with tarfile.open('tmp_data.tar.gz', 'r:gz') as fh:
39         for member in fh.getmembers():
40             if member.isreg():
41                 member.name = os.path.basename(member.name)
42                 fh.extract(member, os.path.join(target_directory, database_path))
43 
44     data_table_entry = {
45         'data_tables': {
46             data_table_name: [
47                 {
48                     "value": database_value,
49                     "name": database_name,
50                     "path": database_path,
51                 }
52             ]
53         }
54     }
55 
56     return data_table_entry
57 
58 
59 def kraken2_build_special(kraken2_args, target_directory, data_table_name=DATA_TABLE_NAME):
60 
61     now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ")
62 
63     special_database_names = {
64         "rdp": "RDP",
65         "greengenes": "Greengenes",
66         "silva": "Silva",
67     }
68 
69     database_value = "_".join([
70         now,
71         kraken2_args["special_database_type"],
72         "kmer-len", str(kraken2_args["kmer_len"]),
73         "minimizer-len", str(kraken2_args["minimizer_len"]),
74         "minimizer-spaces", str(kraken2_args["minimizer_spaces"]),
75         "load-factor", str(kraken2_args["load_factor"]),
76     ])
77 
78     database_name = " ".join([
79         special_database_names[kraken2_args["special_database_type"]],
80         "(Created:",
81         now + ",",
82         "kmer-len=" + str(kraken2_args["kmer_len"]) + ",",
83         "minimizer-len=" + str(kraken2_args["minimizer_len"]) + ",",
84         "minimizer-spaces=" + str(kraken2_args["minimizer_spaces"]) + ")",
85         "load-factor=" + str(kraken2_args["load_factor"]) + ")",
86     ])
87 
88     database_path = database_value
89 
90     args = [
91         '--threads', str(kraken2_args["threads"]),
92         '--special', kraken2_args["special_database_type"],
93         '--kmer-len', str(kraken2_args["kmer_len"]),
94         '--minimizer-len', str(kraken2_args["minimizer_len"]),
95         '--minimizer-spaces', str(kraken2_args["minimizer_spaces"]),
96         '--load-factor', str(kraken2_args["load_factor"]),
97         '--db', database_path
98     ]
99 
00     subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
01 
02     if kraken2_args["clean"]:
03         args = [
04             '--threads', str(kraken2_args["threads"]),
05             '--clean',
06             '--db', database_path
07         ]
08 
09         subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
10 
11     data_table_entry = {
12         'data_tables': {
13             data_table_name: [
14                 {
15                     "value": database_value,
16                     "name": database_name,
17                     "path": database_path,
18                 }
19             ]
20         }
21     }
22 
23     return data_table_entry
24 
25 
26 def kraken2_build_custom(kraken2_args, custom_database_name, custom_source_info, target_directory, data_table_name=DATA_TABLE_NAME):
27     now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H%M%SZ")
28 
29     database_value = "_".join([
30         now,
31         re.sub(r'[^\w_.-]+', '_', custom_database_name).strip('_'),
32         "kmer-len", str(kraken2_args["kmer_len"]),
33         "minimizer-len", str(kraken2_args["minimizer_len"]),
34         "minimizer-spaces", str(kraken2_args["minimizer_spaces"]),
35         "load-factor", str(kraken2_args["load_factor"]),
36     ])
37 
38     database_name = " ".join([
39         custom_database_name,
40         "(" + custom_source_info + ",",
41         "kmer-len=" + str(kraken2_args["kmer_len"]) + ",",
42         "minimizer-len=" + str(kraken2_args["minimizer_len"]) + ",",
43         "minimizer-spaces=" + str(kraken2_args["minimizer_spaces"]) + ",",
44         "load-factor=" + str(kraken2_args["load_factor"]) + ")",
45     ])
46 
47     database_path = database_value
48 
49     args = [
50         '--threads', str(kraken2_args["threads"]),
51         '--download-taxonomy',
52         '--db', database_path,
53     ]
54 
55     if kraken2_args['skip_maps']:
56         args.append('--skip-maps')
57 
58     subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
59 
60     args = [
61         '--threads', str(kraken2_args["threads"]),
62         '--add-to-library', kraken2_args["custom_fasta"],
63         '--db', database_path,
64     ]
65 
66     subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
67 
68     args = [
69         '--threads', str(kraken2_args["threads"]),
70         '--build',
71         '--kmer-len', str(kraken2_args["kmer_len"]),
72         '--minimizer-len', str(kraken2_args["minimizer_len"]),
73         '--minimizer-spaces', str(kraken2_args["minimizer_spaces"]),
74         '--load-factor', str(kraken2_args["load_factor"]),
75         '--db', database_path,
76     ]
77 
78     subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
79 
80     if kraken2_args["clean"]:
81         args = [
82             '--threads', str(kraken2_args["threads"]),
83             '--clean',
84             '--db', database_path,
85         ]
86 
87         subprocess.check_call(['kraken2-build'] + args, cwd=target_directory)
88 
89     data_table_entry = {
90         'data_tables': {
91             data_table_name: [
92                 {
93                     "value": database_value,
94                     "name": database_name,
95                     "path": database_path,
96                 }
97             ]
98         }
99     }
00 
01     return data_table_entry
02 
03 
04 def main():
05     parser = argparse.ArgumentParser()
06     parser.add_argument('data_manager_json')
07     parser.add_argument('--kmer-len', dest='kmer_len', type=int, default=35, help='kmer length')
08     parser.add_argument('--minimizer-len', dest='minimizer_len', type=int, default=31, help='minimizer length')
09     parser.add_argument('--minimizer-spaces', dest='minimizer_spaces', default=6, help='minimizer spaces')
10     parser.add_argument('--load-factor', dest='load_factor', type=float, default=0.7, help='load factor')
11     parser.add_argument('--threads', dest='threads', default=1, help='threads')
12     parser.add_argument('--database-type', dest='database_type', type=KrakenDatabaseTypes, choices=list(KrakenDatabaseTypes), required=True, help='type of kraken database to build')
13     parser.add_argument('--minikraken2-version', dest='minikraken2_version', type=Minikraken2Versions, choices=list(Minikraken2Versions), help='MiniKraken2 version (only applies to --database-type minikraken)')
14     parser.add_argument('--prebuilt-db', dest='prebuilt_db', type=StandardPrebuiltSizes, choices=list(StandardPrebuiltSizes), help='Prebuilt database to download. Only applies to --database-type standard_prebuilt or special_prebuilt.')
15     parser.add_argument('--prebuilt-date', dest='prebuilt_date', help='Database build date (YYYY-MM-DD). Only applies to --database-type standard_prebuilt.')
16     parser.add_argument('--special-database-type', dest='special_database_type', type=SpecialDatabaseTypes, choices=list(SpecialDatabaseTypes), help='type of special database to build (only applies to --database-type special)')
17     parser.add_argument('--custom-fasta', dest='custom_fasta', help='fasta file for custom database (only applies to --database-type custom)')
18     parser.add_argument('--custom-database-name', dest='custom_database_name', help='Name for custom database (only applies to --database-type custom)')
19     parser.add_argument('--custom-source-info', dest='custom_source_info', help='Description of how this build has been sourced (only applies to --database-type custom)')
20     parser.add_argument('--skip-maps', dest='skip_maps', action='store_true', help='')
21     parser.add_argument('--clean', dest='clean', action='store_true', help='Clean up extra files')
22     args = parser.parse_args()
23 
24     with open(args.data_manager_json) as fh:
25         data_manager_input = json.load(fh)
26 
27     target_directory = data_manager_input['output_data'][0]['extra_files_path']
28 
29     try:
30         os.mkdir(target_directory)
31     except OSError as exc:
32         if exc.errno == errno.EEXIST and os.path.isdir(target_directory):
33             pass
34         else:
35             raise
36 
37     data_manager_output = {}
38 
39     if str(args.database_type) == 'standard_local_build':
40         kraken2_args = {
41             "kmer_len": args.kmer_len,
42             "minimizer_len": args.minimizer_len,
43             "minimizer_spaces": args.minimizer_spaces,
44             "load_factor": args.load_factor,
45             "threads": args.threads,
46             "clean": args.clean,
47         }
48         data_manager_output = kraken2_build_standard(
49             kraken2_args,
50             target_directory,
51         )
52     elif str(args.database_type) in ('standard_prebuilt', 'special_prebuilt'):
53         data_manager_output = kraken2_build_standard_prebuilt(
54             str(args.prebuilt_db),
55             str(args.prebuilt_date),
56             target_directory
57         )
58     elif str(args.database_type) == 'minikraken':
59         data_manager_output = kraken2_build_minikraken(
60             str(args.minikraken2_version),
61             target_directory
62         )
63     elif str(args.database_type) == 'special':
64         kraken2_args = {
65             "special_database_type": str(args.special_database_type),
66             "kmer_len": args.kmer_len,
67             "minimizer_len": args.minimizer_len,
68             "minimizer_spaces": args.minimizer_spaces,
69             "load_factor": args.load_factor,
70             "threads": args.threads,
71             "clean": args.clean,
72         }
73         data_manager_output = kraken2_build_special(
74             kraken2_args,
75             target_directory,
76         )
77     elif str(args.database_type) == 'custom':
78         kraken2_args = {
79             "custom_fasta": args.custom_fasta,
80             "skip_maps": args.skip_maps,
81             "kmer_len": args.kmer_len,
82             "minimizer_len": args.minimizer_len,
83             "minimizer_spaces": args.minimizer_spaces,
84             "load_factor": args.load_factor,
85             "threads": args.threads,
86             "clean": args.clean,
87         }
88         data_manager_output = kraken2_build_custom(
89             kraken2_args,
90             args.custom_database_name,
91             args.custom_source_info,
92             target_directory,
93         )
94     else:
95         sys.exit("Invalid database type")
96 
97     with open(args.data_manager_json, 'w') as fh:
98         json.dump(data_manager_output, fh, sort_keys=True)
99 
00 
01 if __name__ == "__main__":
02     main()