Permalink
Browse files

Make a sandbox for minor dengue specific adjustments

  • Loading branch information...
sidneymbell committed Oct 31, 2017
1 parent 0ce9257 commit 4bb207438035736f2f54b61fb04593c1230829c6
Showing with 146 additions and 0 deletions.
  1. +146 −0 tdb/dengue_download.py
View
@@ -0,0 +1,146 @@
import os, json, datetime, sys
import rethinkdb as r
sys.path.append('') # need to import from base
from base.rethink_io import rethink_io
from vdb.download import download as vdb_download
def get_parser():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-db', '--database', default='tdb', help="database to download from")
parser.add_argument('--rethink_host', default=None, help="rethink host url")
parser.add_argument('--auth_key', default=None, help="auth_key for rethink database")
parser.add_argument('--local', default=False, action="store_true", help ="connect to local instance of rethinkdb database")
parser.add_argument('-v', '--virus', default='dengue', help="virus name")
parser.add_argument('--subtype', help="subtype to be included in download")
parser.add_argument('--ftype', default='tsv', help="output file format, default \"tsv\", options are \"json\" and \"tsv\"")
parser.add_argument('--fstem', default=None, help="default output file name is \"VirusName_Year_Month_Date\"")
parser.add_argument('--path', default='data', help="path to dump output files to")
parser.add_argument('--select', nargs='+', type=str, default=[], help="Select specific fields ie \'--select field1:value1 field2:value1,value2\'")
parser.add_argument('--present', nargs='+', type=str, default=[], help="Select specific fields to be non-null ie \'--present field1 field2\'")
parser.add_argument('--interval', nargs='+', type=str, default=[], help="Select interval of values for fields \'--interval field1:value1,value2 field2:value1,value2\'")
parser.add_argument('--years_back', type=str, default=None, help='number of past years to sample sequences from \'--years_back field:value\'')
parser.add_argument('--relaxed_interval', default=False, action="store_true", help="Relaxed comparison to date interval, 2016-XX-XX in 2016-01-01 - 2016-03-01")
return parser
class download(object):
def __init__(self, database, virus, **kwargs):
'''
parser for virus, output file names, output file format path, interval
'''
self.virus = virus.lower()
self.database = database.lower()
self.measurements = []
def connect_rethink(self, **kwargs):
if self.database not in ['tdb', 'test_tdb', 'cdc_tdb']:
raise Exception("Cant download to this database: " + self.database)
self.rethink_io = rethink_io()
self.rethink_host, self.auth_key = self.rethink_io.assign_rethink(**kwargs)
self.rethink_io.connect_rethink(self.database, self.rethink_host, self.auth_key)
self.rethink_io.check_table_exists(self.database, self.virus)
def count_documents(self):
'''
return integer count of number of documents in table
'''
return r.db(self.database).table(self.virus).count().run()
def download(self, subtype=None, output=True, **kwargs):
'''
download documents from table
'''
import time
start_time = time.time()
self.connect_rethink(**kwargs)
self.vdb_download = vdb_download(database=self.database, virus=self.virus)
select, present, interval, = self.vdb_download.parse_subset_arguments(**kwargs)
if subtype is not None:
select.append(('subtype', [subtype]))
sequence_count = r.table(self.virus).count().run()
print(sequence_count, "measurements in table:", self.virus)
print("Downloading titer measurements from the table: " + self.virus)
measurements = self.rethinkdb_download(self.virus, presents=present, selections=select, intervals=interval, **kwargs)
print("Downloaded " + str(len(measurements)) + " measurements")
for m in measurements:
for k,v in m.items():
try:
m[k] = v.upper()
except:
continue
if output:
self.output(measurements, **kwargs)
print("--- %s minutes to download ---" % ((time.time() - start_time)/60))
def rethinkdb_download(self, table, **kwargs):
'''
Default command merges documents from the sequence table and virus table
Chain rethinkdb filter and has_fields commands to the default command
Return documents from the database that are left after filtering
'''
# take each sequence and merge with corresponding virus document
command = r.table(table)
command = self.vdb_download.add_present_command(command, **kwargs)
command = self.vdb_download.add_selections_command(command, **kwargs)
command = self.vdb_download.add_intervals_command(command, **kwargs)
sequences = list(command.run())
return list(sequences)
def write_json(self, data, fname, indent=1):
'''
writes as list of viruses (dictionaries)
'''
try:
handle = open(fname, 'w')
except:
print("Couldn't open output file")
print(fname)
raise FileNotFoundError
else:
json.dump(data, handle, indent=indent)
handle.close()
print("Wrote to " + fname)
def write_text(self, measurements, fname, text_fields=['virus_strain', 'serum_strain', 'serum_id', 'source', 'titer', 'assay_type']):
try:
handle = open(fname, 'w')
except IOError:
pass
else:
for meas in measurements:
for field in text_fields:
if field in meas and meas[field] is not None:
handle.write(str(meas[field]) + '\t')
else:
handle.write('?' + '\t')
handle.write('\n')
handle.close()
print("Wrote to " + fname)
def output(self, measurements, path, fstem, ftype, **kwargs):
fname = path + '/' + fstem + '_titers' + '.' + ftype
if ftype == 'json':
self.write_json(measurements,fname)
elif ftype == 'tsv':
self.write_text(measurements, fname)
else:
raise Exception("Can't output to that file type, only json or text allowed")
if __name__=="__main__":
parser = get_parser()
args = parser.parse_args()
current_date = str(datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d'))
if args.fstem is None:
if args.subtype is not None:
args.fstem = args.subtype + '_' + current_date
else:
args.fstem = args.virus + '_' + current_date
if not os.path.isdir(args.path):
os.makedirs(args.path)
connTDB = download(**args.__dict__)
connTDB.download(**args.__dict__)

0 comments on commit 4bb2074

Please sign in to comment.