-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathrun_querytns.py
130 lines (110 loc) · 5.66 KB
/
run_querytns.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Query TNS for new targets and upload to candidate marshal.
https://wis-tns.weizmann.ac.il/
TNS bot apikey must exist in config
@TODO: Move to flows API
.. codeauthor:: Emir Karamehmetoglu <[email protected]>
.. codeauthor:: Rasmus Handberg <[email protected]>
"""
import argparse
import logging
from tqdm import tqdm
import re
from astropy.coordinates import SkyCoord
from astropy.time import Time
from datetime import datetime, timedelta, timezone
from tendrils import api
from tendrils.utils import load_tns_config, tns_getnames, TNSConfigError, tns_get_obj
# --------------------------------------------------------------------------------------------------
def main():
# Parse command line arguments:
parser = argparse.ArgumentParser(description='Query TNS and upload to Flows candidates.')
parser.add_argument('-d', '--debug', action='store_true', help='Print debug messages.')
parser.add_argument('-q', '--quiet', action='store_true', help='Only report warnings and errors.')
parser.add_argument('--zmax', type=float, default=0.105, help='Maximum redshift.')
parser.add_argument('--zmin', type=float, default=0.000000001, help='Minimum redshift.')
parser.add_argument('-b', '--days_begin', type=int, default=30, help='Discovery day at least X days before today.')
parser.add_argument('-e', '--days_end', type=int, default=3, help='Discovery day at most X days before today.')
parser.add_argument('-o', '--objtype', type=str, default=[3, 104],
help='TNS objtype int given as comma separed string with no spaces')
parser.add_argument('-m', '--limit_months', type=int, default=2, help='Integer number of months to limit TNS search (for speed). \
Should be greater than days_begin.')
parser.add_argument('--autoupload', action='store_true',
help="Automatically upload targets to Flows website. Only do this, if you know what you are doing!")
args = parser.parse_args()
# Set logging level:
logging_level = logging.INFO
if args.quiet:
logging_level = logging.WARNING
elif args.debug:
logging_level = logging.DEBUG
# Setup logging:
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
console = logging.StreamHandler()
console.setFormatter(formatter)
logger = logging.getLogger(__name__)
if not logger.hasHandlers():
logger.addHandler(console)
logger.setLevel(logging_level)
tqdm_settings = {'disable': None if logger.isEnabledFor(logging.INFO) else True}
# Try to load TNS config - only used for early stopping
try:
load_tns_config()
except TNSConfigError:
parser.error("Error in TNS configuration.")
return
# Calculate current date and date range to search:
date_now = datetime.now(timezone.utc).date()
date_end = date_now - timedelta(days=args.days_end)
date_begin = date_now - timedelta(days=args.days_begin)
logger.info('Date begin = %s, date_end = %s', date_begin, date_end)
# Query TNS for SN names
logger.info('Querying TNS for all targets, this may take awhile')
nms = tns_getnames(months=args.limit_months, # pre-limit TNS search to candidates reported in the last X months
date_begin=date_begin, date_end=date_end, zmin=args.zmin, zmax=args.zmax,
objtype=args.objtype# Relevant TNS SN Ia subtypes.
)
logger.debug(nms)
if not nms:
logger.info("No targets were found.")
return
# Remove already existing names using flows api
included_names = ['SN' + target['target_name'] for target in api.get_targets()]
nms = list(set(nms) - set(included_names))
logger.info('Target names obtained: %s', nms)
# Regular Expression matching any string starting with "ztf"
regex_ztf = re.compile('^ztf', flags=re.IGNORECASE)
regex_sn = re.compile(r'^sn\s*', flags=re.IGNORECASE)
# Query TNS for object info using API, then upload to FLOWS using API.
num_uploaded = 0
if args.autoupload:
for name in tqdm(nms, **tqdm_settings):
sn = regex_sn.sub('', name)
logger.debug('querying TNS for: %s', sn)
# make GET request to TNS via API
reply = tns_get_obj(sn)
# Parse output
if reply:
logger.debug('GET query successful')
# Extract object info
coord = SkyCoord(ra=reply['radeg'], dec=reply['decdeg'], unit='deg', frame='icrs')
discovery_date = Time(reply['discoverydate'], format='iso', scale='utc')
ztf = list(filter(regex_ztf.match, reply['internal_names']))
ztf = None if not ztf else ztf[0]
if 'object_type' in reply and 'name' in reply['object_type']:
sntype = regex_sn.sub('', reply['object_type']['name'])
else:
sntype = None
# Try to upload to FLOWS
newtargetid = api.add_target(reply['objname'], coord, redshift=reply['redshift'],
discovery_date=discovery_date, discovery_mag=reply['discoverymag'],
host_galaxy=reply['hostname'], ztf=ztf, sntype=sntype, status='candidate',
project='flows')
logger.debug('Uploaded to FLOWS with targetid=%d', newtargetid)
num_uploaded += 1
logger.info("%d targets uploaded to Flows.", num_uploaded)
# --------------------------------------------------------------------------------------------------
if __name__ == '__main__':
main()