Export all the light curves from a LEMON database

export-curves.py

#! /usr/bin/env python

""" Take a LEMONdB file and export the light curves of all the astronomical
objects, in all the photometric filters, to an output directory. The curves
are saved as SVG files. """

# Author: Victor Terron (c) 2015
# Email: `echo vt2rron1iaa32s | tr 132 @.e`
# License: GNU GPLv3

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import argparse
import csv
import matplotlib.figure
from matplotlib.backends.backend_gtkagg import FigureCanvasGTKAgg as FigureCanvas
import os
import os.path
import sys

# Fix until LEMON is installed system-wide
LEMON_DIR = os.path.expanduser("~/lemon/")
sys.path.insert(0, LEMON_DIR)

# LEMON modules
import database
import juicer.plot as plot
import methods
from snr import snr_to_error

if __name__ == "__main__":

    # https://stackoverflow.com/a/12151325/184363
    parser = argparse.ArgumentParser(
        description='Export LEMON light curves',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        epilog=("Instead of exporting the light curves of all the objects, "
                "you may use --ra and --dec. If both options are given, only "
                "the curve(s) of the astronomical object closest to these "
                "coordinates will be exported."),
        )

    parser.add_argument('db_path', metavar='LEMON_DB', type=str,
                        help="the LEMON database with the light curves")
    parser.add_argument('output_dir', metavar='OUTPUT_DIR', type=str,
                        help="output directory to which export the curves")

    parser.add_argument('--airmases', dest='show_airmasses', action='store_true',
                        help="also plot airmasses on the y-axis")
    parser.add_argument('--julian', dest='show_julian', action='store_true',
                        help="use Julian dates on the x-axis")
    parser.add_argument('--color', dest='color', type=str, default='blue',
                        help="color to use for the data points")
    parser.add_argument('--csv', dest='csv', action='store_true',
                        help="instead of SVG images, save each light curve as a "
                        "CSV text file, with the Unix date, differential magnitude "
                        "signal-to-noise ratio and error of each measurement.")

    parser.add_argument('--comparison', dest='comparison_stars', action='store_true',
                        help="for each object and filter, save also a text file "
                        "listing the ID, right ascension, declination and weight "
                        "used to generate its articial comparison star")

    parser.add_argument('--ra', dest='ra', type=float, default=None,
                        help="right ascension of the astronomical object")

    parser.add_argument('--dec', dest='dec', type=float, default=None,
                        help="declination of the astronomical object")

    args = parser.parse_args()

    if os.path.exists(args.output_dir):
        msg = "Error: output directory {0} already exists.".format(args.output_dir)
        sys.exit(msg)
    else:
        print("Creating output directory {0}...".format(args.output_dir), end='')
        os.mkdir(args.output_dir)
        print(" OK")

    if bool(args.ra) + bool(args.dec) == 1:
        msg = "Error: --ra and --dec must always be used together."
        sys.exit(msg)

    # Keyword arguments for curve_plot()
    kwargs = dict(airmasses = args.show_airmasses,
                  julian = args.show_julian,
                  delta = 3 * 3600, # three hours
                  color = args.color)

    db = database.LEMONdB(args.db_path)
    print("Stars: {0}".format(len(db)))
    print("Filters: {0}".format(len(db.pfilters)))

    figure = matplotlib.figure.Figure()
    canvas = FigureCanvas(figure)

    if args.ra is not None:
        assert args.dec is not None
        print("Finding closest star to {0} {1}...".format(args.ra, args.dec), end='')
        star_id, distance = db.star_closest_to_world_coords(args.ra, args.dec)
        star_ids = [star_id]
        print(" OK")
        print("Star: ID={0}, distance={1} arcsec".format(star_id, distance))

    else:
        # Export all light curves
        star_ids = db.star_ids

    for star_id in star_ids:
        print("Exporting star ID={0}...".format(star_id), end='')
        for pfilter in db.pfilters:
            curve = db.get_light_curve(star_id, pfilter)
            if curve:

                ra, dec = db.get_star(star_id)[2:4]
                ra_str  = methods.ra_str (ra)
                dec_str = methods.dec_str(dec)

                filename = "{0}_{1}_{2}".format(ra_str, dec_str, pfilter).replace(" ", "_")
                output_path = os.path.join(args.output_dir, filename) # no extension!

                # SVG file
                if not args.csv:
                    figure.clf()
                    plot.curve_plot(figure, curve, **kwargs)
                    figure.suptitle("{0} {1}\nFilter: {2}".format(ra_str, dec_str, pfilter))
                    figure.savefig(output_path + '.svg')

                # CSV file
                else:

                    rows = []
                    for unix_time, magnitude, snr in curve:
                        merr_pos, merr_neg = snr_to_error(snr)
                        r = dict(unix_time = unix_time,
                                 magnitude = magnitude,
                                 snr = snr,
                                 merr_pos = merr_pos,
                                 merr_neg = merr_neg)
                        rows.append(r)
                    rows.sort(key = lambda x: x['unix_time'])

                    with open(output_path + '.csv', 'wt') as fd:
                        fieldnames = ['unix_time', 'magnitude', 'snr', 'merr_pos', 'merr_neg']
                        writer = csv.DictWriter(fd, fieldnames=fieldnames)

                        writer.writeheader()
                        for r in rows:
                            writer.writerow(r)

                if args.comparison_stars:

                    with open(output_path + '-comparison.txt', 'wt') as fd:
                        for cmp_star_id, cmp_star_weight, _ in curve.weights():
                            cmp_ra, cmp_dec = db.get_star(cmp_star_id)[2:4]
                            fmt_args = cmp_star_id, cmp_ra, cmp_dec, cmp_star_weight
                            line = "{0} {1} {2} {3}\n".format(*fmt_args)
                            fd.write(line)

        print(" OK")