Merge branch 'devel'

VERSION

@@ -1 +1 @@
-master-v0.6.3
+master-v0.6.4

flows/tns.py

@@ -10,10 +10,10 @@
 """
 
 import logging
+from astropy.table import Table
 import astropy.units as u
 import requests
 import json
-from bs4 import BeautifulSoup
 import datetime
 from .config import load_config
 
@@ -135,53 +135,79 @@ def tns_get_obj(name):
 	return None
 
 #--------------------------------------------------------------------------------------------------
-def tns_get_names(months=None, date_begin=None, date_end=None, z_min=None, z_max=None, objtype=None):
-	"""Get SN names from TNS html table"""
+def tns_getnames(months=None, date_begin=None, date_end=None, zmin=None, zmax=None, objtype=[3, 104]):
+	"""
+	Get SN names from TNS.
+
+	Parameters:
+		months (int, optional): Only return objects reported within the last X months.
+		date_begin (date, optional): Discovery date begin.
+		date_end (date, optional): Discovery date end.
+		zmin (float, optional): Minimum redshift.
+		zmax (float, optional): Maximum redshift.
+		objtype (list, optional): Constraint object type.
+			Default is to query for
+			- 3: SN Ia
+			- 104: SN Ia-91T-like
+
+	Returns:
+		list: List of names fulfilling search criteria.
+	"""
 
 	logger = logging.getLogger(__name__)
 
 	# Change formats of input to be ready for query:
 	if isinstance(date_begin, datetime.datetime):
-		date_begin = date_begin.date().isoformat()
-	elif isinstance(date_begin, datetime.date):
-		date_end = date_begin.isoformat()
+		date_begin = date_begin.date()
+	elif isinstance(date_begin, str):
+		date_begin = datetime.datetime.strptime(date_begin, '%Y-%m-%d').date()
+
 	if isinstance(date_end, datetime.datetime):
-		date_end = date_end.date().isoformat()
-	elif isinstance(date_end, datetime.date):
-		date_end = date_end.isoformat()
+		date_end = date_end.date()
+	elif isinstance(date_end, str):
+		date_end = datetime.datetime.strptime(date_end, '%Y-%m-%d').date()
+
 	if isinstance(objtype, (list, tuple)):
 		objtype = ','.join([str(o) for o in objtype])
 
+	# Do some sanity checks:
+	if date_end < date_begin:
+		raise ValueError("Dates are in the wrong order.")
+
+	date_now = datetime.datetime.now(datetime.timezone.utc).date()
+	if months is not None and date_end is not None and date_end < date_now - datetime.timedelta(days=months*30):
+		logger.warning('Months limit restricts days_begin, consider increasing limit_months.')
+
 	# Parameters for query:
 	params = {
-		'discovered_period_value': months,
+		'discovered_period_value': months, # Reported Within The Last
 		'discovered_period_units': 'months',
-		'unclassified_at': 0,
-		'classified_sne': 1,
-		'include_frb': 0,
+		'unclassified_at': 0, # Limit to unclasssified ATs
+		'classified_sne': 1, # Limit to classified SNe
+		'include_frb': 0, # Include FRBs
 		#'name': ,
 		'name_like': 0,
 		'isTNS_AT': 'all',
 		'public': 'all',
 		#'ra':
 		#'decl':
 		#'radius':
-		'coords_unit': 'arcsec',
+		#'coords_unit': 'arcsec',
 		'reporting_groupid[]': 'null',
 		'groupid[]': 'null',
 		'classifier_groupid[]': 'null',
 		'objtype[]': objtype,
 		'at_type[]': 'null',
-		'date_start[date]': date_begin,
-		'date_end[date]': date_end,
+		'date_start[date]': date_begin.isoformat(),
+		'date_end[date]': date_end.isoformat(),
 		#'discovery_mag_min':
 		#'discovery_mag_max':
 		#'internal_name':
 		#'discoverer':
 		#'classifier':
 		#'spectra_count':
-		'redshift_min': z_min,
-		'redshift_max': z_max,
+		'redshift_min': zmin,
+		'redshift_max': zmax,
 		#'hostname':
 		#'ext_catid':
 		#'ra_range_min':
@@ -205,40 +231,48 @@ def tns_get_names(months=None, date_begin=None, date_end=None, z_min=None, z_max
 		#'frb_flux_range_min':
 		#'frb_flux_range_max':
 		'num_page': 500,
-		'display[redshift]': 1,
-		'display[hostname]': 1,
-		'display[host_redshift]': 1,
-		'display[source_group_name]': 1,
-		'display[classifying_source_group_name]': 1,
+		'display[redshift]': 0,
+		'display[hostname]': 0,
+		'display[host_redshift]': 0,
+		'display[source_group_name]': 0,
+		'display[classifying_source_group_name]': 0,
 		'display[discovering_instrument_name]': 0,
 		'display[classifing_instrument_name]': 0,
 		'display[programs_name]': 0,
-		'display[internal_name]': 1,
+		'display[internal_name]': 0,
 		'display[isTNS_AT]': 0,
-		'display[public]': 1,
+		'display[public]': 0,
 		'display[end_pop_period]': 0,
-		'display[spectra_count]': 1,
-		'display[discoverymag]': 1,
-		'display[discmagfilter]': 1,
-		'display[discoverydate]': 1,
-		'display[discoverer]': 1,
+		'display[spectra_count]': 0,
+		'display[discoverymag]': 0,
+		'display[discmagfilter]': 0,
+		'display[discoverydate]': 0,
+		'display[discoverer]': 0,
 		'display[remarks]': 0,
 		'display[sources]': 0,
 		'display[bibcode]': 0,
-		'display[ext_catalogs]': 0
+		'display[ext_catalogs]': 0,
+		'format': 'csv'
 	}
 
+	# Query TNS for names:
 	con = requests.get(url_tns_search, params=params)
 	con.raise_for_status()
 
-	# Parse output using BS4, find html table
-	soup = BeautifulSoup(con.text, 'html.parser')
-	tab = soup.find('table', 'results-table')
-	if tab is None:
-		logger.error('No HTML table obtained from query!')
-		return None
-
-	names = tab.find_all('td', 'cell-name')
-	names_list = [name.text.replace(' ', '') for name in names if name.text.startswith('SN')]
+	# Parse the CSV table:
+	# Ensure that there is a newline in table string.
+	# AstroPy uses this to distinguish file-paths from pure-string inputs:
+	text = str(con.text) + "\n"
+	tab = Table.read(text,
+		format='ascii.csv',
+		guess=False,
+		delimiter=',',
+		quotechar='"',
+		header_start=0,
+		data_start=1)
+
+	# Pull out the names only if they begin with "SN":
+	names_list = [name.replace(' ', '') for name in tab['Name'] if name.startswith('SN')]
+	names_list = sorted(names_list)
 
 	return names_list

requirements.txt

@@ -21,7 +21,6 @@ vtk
 scikit-image
 tqdm
 pytz
-beautifulsoup4
 git+https://github.com/obscode/imagematch.git@photutils#egg=imagematch
 sep
 astroalign > 2.3

run_querytns.py

@@ -13,7 +13,6 @@
 import logging
 from tqdm import tqdm
 import re
-import numpy as np
 from astropy.coordinates import SkyCoord
 from astropy.time import Time
 from datetime import datetime, timedelta, timezone
@@ -60,29 +59,21 @@ def main():
 		parser.error("Error in TNS configuration.")
 		return
 
-	# Extract settings:
-	upload_to_flows = args.autoupload
-	days_before_today = args.days_end  # end = now - dbt
-	days_to_include = args.days_begin  # begin = now - dti
-	months = args.limit_months  # pre-limit TNS search to candidates reported in the last X months
-	if np.floor(days_before_today/30) > months:
-		logger.warning('Months limit restricts days_begin, consider increasing limit_months.')
-
 	# Calculate current date and date range to search:
 	date_now = datetime.now(timezone.utc).date()
-	date_end = date_now - timedelta(days=days_before_today)
-	date_begin = date_now - timedelta(days=days_to_include)
+	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.tns_get_names(
-		months=months,
+	nms = tns.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,
-		z_min=args.zmin,
-		z_max=args.zmax,
-		objtype=args.objtype  # Relevant TNS SN Ia subtypes.
+		zmin=args.zmin,
+		zmax=args.zmax,
+		objtype=args.objtype # Relevant TNS SN Ia subtypes.
 	)
 	logger.debug(nms)
 
@@ -96,38 +87,35 @@ def main():
 	logger.info('Target names obtained: %s', nms)
 
 	# Regular Expression matching any string starting with "ztf"
-	regex = re.compile('^ztf', flags=re.IGNORECASE)
+	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
-	for name in tqdm(nms, **tqdm_settings):
-		sn = re.sub('^SN', '', name)
-		logger.debug('querying TNS for: %s', sn)
-
-		# make GET request to TNS via API
-		reply = tns.tns_get_obj(sn)
-
-		# Parse output
-		if reply:
-			logger.debug('GET query successful')
-
-			# Extract object info
-			name = reply['objname']
-			coord = SkyCoord(ra=reply['radeg'], dec=reply['decdeg'], unit='deg', frame='icrs')
-			redshift = reply['redshift']
-			discovery_date = Time(reply['discoverydate'], format='iso', scale='utc')
-			discovery_mag = reply['discoverymag']
-			host_galaxy = reply['hostname']
-			ztf = list(filter(regex.match, reply['internal_names']))
-			ztf = None if not ztf else ztf[0]
-
-			# Try to upload to FLOWS
-			if upload_to_flows:
-				newtargetid = api.add_target(name, coord,
-					redshift=redshift,
+	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.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]
+
+				# Try to upload to FLOWS
+				newtargetid = api.add_target(reply['objname'], coord,
+					redshift=reply['redshift'],
 					discovery_date=discovery_date,
-					discovery_mag=discovery_mag,
-					host_galaxy=host_galaxy,
+					discovery_mag=reply['discoverymag'],
+					host_galaxy=reply['hostname'],
 					ztf=ztf,
 					status='candidate',
 					project='flows')

tests/test_tns.py

@@ -7,6 +7,7 @@
 """
 
 import pytest
+import datetime
 from astropy.coordinates import SkyCoord
 from conftest import capture_cli
 from flows import tns
@@ -35,6 +36,38 @@ def test_tns_get_obj(SETUP_CONFIG):
 	assert res['objname'] == '2019yvr'
 	assert res['name_prefix'] == 'SN'
 
+#--------------------------------------------------------------------------------------------------
+@pytest.mark.parametrize('date_begin,date_end', [
+	('2019-01-01', '2019-02-01'),
+	(datetime.date(2019, 1, 1), datetime.date(2019, 2, 1)),
+	(datetime.datetime(2019, 1, 1, 12, 0), datetime.datetime(2019, 2, 1, 12, 0))
+])
+def test_tns_getnames(SETUP_CONFIG, date_begin, date_end):
+
+	names = tns.tns_getnames(
+		date_begin=date_begin,
+		date_end=date_end,
+		zmin=0,
+		zmax=0.105,
+		objtype=3
+	)
+
+	print(names)
+	assert isinstance(names, list), "Should return a list"
+	for n in names:
+		assert isinstance(n, str), "Each element should be a string"
+		assert n.startswith('SN'), "All names should begin with 'SN'"
+	assert 'SN2019A' in names, "SN2019A should be in the list"
+
+#--------------------------------------------------------------------------------------------------
+def test_tns_getnames_wronginput(SETUP_CONFIG):
+	# Wrong dates should result in ValueError:
+	with pytest.raises(ValueError):
+		tns.tns_getnames(
+			date_begin=datetime.date(2019, 1, 1),
+			date_end=datetime.date(2017, 1, 1)
+		)
+
 #--------------------------------------------------------------------------------------------------
 def test_run_querytns(SETUP_CONFIG):