modified export-curves.py from vterron's code.

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 ConfigParser
import datetime
import functools
import operator
import os.path
import random
import re
import sys
import time
import warnings
import astropy.time
import re
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,
                                 utctimer = methods.utctime(unix_time, suffix = False),
                                 jd    = astropy.time.Time(unix_time, format='unix').jd,
                                 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', 'utctimer', 'jd', '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")