submit_tm.py

# A script to submit DECam pointings to Treasure Map
# Utilize the des20a environment to run
# Authors: R. Morgan and M. Gill

import datetime
import getpass
import glob
import json
import logging
from optparse import OptionParser
import os
import requests
import sys
import yaml
sys.path.append('TreasureMapPy/treasuremap')

from astropy.coordinates import SkyCoord
import pandas as pd
from treasuremap import Pointings

def get_creators():
    """
    Load in a dictionary of authors and affiliations.
    
    :return: creators: a list of author and affiliation dictionaries
    """
    with open("authors.yaml", 'r') as creator_file:
        creators = yaml.safe_load(creator_file)
        
    # Put submitter's name first in the list
    try:
        name = glob.glob('api_tokens/' + USERNAME + '/*.name')[0].split('/')[-1].split('.name')[0].replace('_', ' ')
    except IndexError:
        name = "Mandeep Gill"
    authors = [{"name": k, "affiliation": v} for k, v in creators['AUTHORS'].items()]
    name_idx = authors.index({"name": name, "affiliation": creators["AUTHORS"][name]})
    value = authors.pop(name_idx)
    
    return [value] + authors
    
    

# Get username of user
USERNAME = getpass.getuser()

# Get the time for stamping the log files
time = datetime.datetime.now()

# Set up logging
log_dir = "logs/{}/".format(time.strftime("%y-%m-%d_%H-%M-%S"))
rc = os.system("mkdir " + log_dir)
if rc != 0:
    print("Unable to make log directory " + log_dir)
    print("Check your user's write permissions or make the directory manually")
    sys.exit()

logging.basicConfig(filename=log_dir + "submit_{}.log".format(time.strftime("%y%m%d_%H%M%S")),
                    filemode="a+",
                    format="|%(levelname)s\t| %(asctime)s -- %(message)s",
                    datefmt="20%y-%m-%d %I:%M:%S %p",
                    level=logging.DEBUG)
logging.info("[" + USERNAME + "] " + "submit_tm.py started")
logging.debug("[" + USERNAME + "] " + "program command: " + ' '.join(sys.argv))

# Handle command line arguments
parser = OptionParser(__doc__)
parser.add_option('--infile', default=None, help="Name of file with pointing info")
parser.add_option('--preview', action='store_true', help="Prompt before submitting")
parser.add_option('--test', action='store_true', help="Run without submitting")
parser.add_option('--graceid', default=None, help="Name of the GraceDB event")
options, args = parser.parse_args(sys.argv[1:])

if not options.infile:
    print("Use '--infile' to specify the file with pointing info")
    logging.critical("[" + USERNAME + "] " + "Missing --infile argument")
    logging.info("[" + USERNAME + "] " + "Program terminating")
    logging.shutdown()
    sys.exit()

if not options.graceid:
    print("Use '--graceid' to specify the GraceDB event name")
    logging.critical("[" + USERNAME + "] " + "Missing --graceid argument")
    logging.info("[" + USERNAME + "] " + "Program terminating")
    logging.shutdown()
    sys.exit()

logging.debug("[" + USERNAME + "] " + "--infile set to {}".format(options.infile))
logging.debug("[" + USERNAME + "] " + "--graceid set to {}".format(options.graceid))
logging.debug("[" + USERNAME + "] " + "--preview set to {}".format(options.preview))
logging.debug("[" + USERNAME + "] " + "--test set to {}".format(options.test))

# Copy the infile to the log directory
logging.info("[" + USERNAME + "] " + "Copying infile to log directory")
copy_rc = os.system("cp {} ".format(options.infile) + log_dir)
if copy_rc == 0:
    logging.debug("[" + USERNAME + "] " + "Copy was successful")
else:
    logging.warning("[" + USERNAME + "] " + "Copy FAILED")

if options.test:
    options.preview = True
    logging.info("[" + USERNAME + "] " + "Setting --preview to True for testing")

# API token - Get your own by making a TreasureMap account
try:
    API_TOKEN = glob.glob('api_tokens/{}/*.api_token'.format(USERNAME))[0].split('/')[-1].split('.')[0]
    logging.debug("[" + USERNAME + "] " + "API_TOKEN set to {}".format(API_TOKEN))
except IndexError:
    logging.error("[" + USERNAME + "] " + "Unable to find user-specific API_TOKEN")
    logging.info("[" + USERNAME + "] " + "M. Gill's token will be used as a default")
    try:
        API_TOKEN = glob.glob('api_tokens/mssgill/*.api_token')[0].split('/')[-1].split('.')[0]
        logging.debug("[" + USERNAME + "] " + "API_TOKEN set to {}".format(API_TOKEN))
    except IndexError:
        logging.error("[" + USERNAME + "] " + "Unable to find M. Gill's default API_TOKEN")
        logging.critical("[" + USERNAME + "] " + "Program needs a valid API_token and will terminate")
        logging.shutdown()
        sys.exit()

# Read the pointing information into a Pandas DataFrame
try:
    pointings_df = pd.read_csv(options.infile)
    logging.info("[" + USERNAME + "] " + "Successfully read {}".format(options.infile))
except FileNotFoundError:
    logging.error("[" + USERNAME + "] " + "Unable to find specified infile")
    logging.critical("[" + USERNAME + "] " + "Progams needs a valid pointing file")
    logging.info("[" + USERNAME + "] " + "Program terminating")
    logging.shutdown()
    sys.exit()

# Instantiate the treasuremap.Pointing class
#  -- looks like treasuremap requires a different object per band
pointings = {}
logging.info("[" + USERNAME + "] " + "Initializing treasuremap.Pointings")
for flt in set(pointings_df['band'].values):
    pointings[flt] = Pointings(status="completed",
                               graceid=options.graceid,
                               instrumentid=38, # 38 == DECam
                               band=flt,
                               api_token=API_TOKEN)
logging.debug("[" + USERNAME + "] " + "Made pointings for " + ','.join(list(pointings.keys())) + " bands")

logging.info("[" + USERNAME + "] " + "Starting processing of DECam pointings")
# Iterate through our pointings and add them
for index, row in pointings_df.iterrows():
    # Format RA and Dec as proper data type
    coord = SkyCoord(row['ra'], row['dec'], frame="icrs",  unit="deg")

    # Add the pointing based on the filter used
    pointings[row['band']].add_pointing(ra=coord.ra.deg, 
                                        dec=coord.dec.deg,
                                        time=row['time'],
                                        depth=row['depth'],
                                        depth_unit=row['depth_unit'])
    logging.debug("[" + USERNAME + "] " + "Added pointing for index {}".format(index))
logging.info("[" + USERNAME + "] " + "Finished making pointings")


# Build all jsons
logging.info("[" + USERNAME + "] " + "Starting generation of json data")
for pointing in pointings.values():
    pointing.build_json()
logging.info("[" + USERNAME + "] " + "Finished building json data")

# If the preview argument was used, wait for approval before submitting
if options.preview:
    logging.info("[" + USERNAME + "] " + "Preview argument passed, requesting user approval")
    for flt in pointings.keys():
        logging.debug("[" + USERNAME + "] " + "Showing user {} band pointings".format(flt))
        # Display pointings
        print('\n\n\n\tShowing {}-band pointings\n\n\n'.format(flt))
        print(json.dumps(pointings[flt].json_data, indent=4))

        # Ask if the pointings look okay
        user_input = input("Does the JSON data look okay? (y/n) ").strip().lower()
        while user_input not in ['y', 'n']:
            print("Please enter 'y' or 'n'")
            user_input = input("Does the JSON data look okay? (y/n) ").strip().lower()
        logging.debug("[" + USERNAME + "] " + "user_input set to {}".format(user_input))

        if user_input == 'n':
            print("Please make your corrections and restart")
            logging.info("[" + USERNAME + "] " + "User has identified error in pointings")
            logging.info("[" + USERNAME + "] " + "Program will terminate to allow for correction")
            logging.shutdown()
            sys.exit()

if not options.test:
    # Submit jsons
    logging.info("[" + USERNAME + "] " + "Beginning submission process")
    requests_dict = {}
    for flt in pointings.keys():
        try:
            request = pointings[flt].submit()
            logging.info("[" + USERNAME + "] " + "Submitted {} band pointing".format(flt))
            requests_dict[flt] = request
        except Exception:
            logging.info("[" + USERNAME + "] " + "There was a prolem with the submisison.")
            logging.exception("[" + USERNAME + "] " + "The traceback for the submission is below")
            logging.warning("[" + USERNAME + "] " + "The {} band pointings may not have submitted properly".format(flt))
    logging.info("[" + USERNAME + "] " + "Finished submisison")
    
    # Request a DOI and save TreasureMap response
    logging.info("[" + USERNAME + "] " + "Requesting a DOI for submitted pointings")
    creators = get_creators()
    json_data = {"api_token": API_TOKEN, "graceid": options.graceid, "creators": creators}
    r = requests.post(url="http://treasuremap.space/api/v0/request_doi", json=json_data)
    logging.info("[" + USERNAME + "] " + "DOI response saved to doi.json")
    response_file = open(log_dir + "doi.json", 'w+')
    response_file.write(json.dumps(r.text, indent=4))

    # Save pointings
    logging.info("[" + USERNAME + "] " + "Saving submission pointings")
    pointing_filename = log_dir + "pointings.json"
    pointing_file = open(pointing_filename, 'w+')
    logging.debug("[" + USERNAME + "] " + "pointing file set to {}".format(pointing_filename))

    for flt in pointings.keys():
        pointing_file.write(json.dumps(pointings[flt].pointings, indent=4))
        pointing_file.write('\n\n')

        logging.info("[" + USERNAME + "] " + "Wrote pointings for {} band".format(flt))

    # Save requests
    logging.info("[" + USERNAME + "] " + "Saving submission requests")
    request_filename = log_dir + 'requests.json'
    request_file = open(request_filename, 'w+')
    logging.debug("[" + USERNAME + "] " + "request file set to {}".format(request_filename))
    
    for flt in requests_dict.keys():
        request_file.write(json.dumps(requests_dict[flt], indent=4))
        request_file.write('\n\n')
        
        logging.info("[" + USERNAME + "] " + "Wrote requests for {} band".format(flt))

    # Close open file streams
    request_file.close()
    pointing_file.close()
    response_file.close()

else:
    logging.info("[" + USERNAME + "] " + "Skipping submission process due to --test argument")

# Clean up directory
if copy_rc == 0:
    os.system('rm {}'.format(options.infile))
    logging.info("[" + USERNAME + "] " + "removed pointings file, since it was copied to log dir")
    
# Conclude the program
logging.info("[" + USERNAME + "] " + "Progam finished")
logging.shutdown()