create_quicklook.py

import os
import sys
from datetime import datetime as dt, timedelta as delta
import logging
import pandas as pd
from netCDF4 import MFDataset, MFTime
import numpy as np
import matplotlib.pyplot as plt
import math

LOG = logging.getLogger(__name__)

# names of the plots used in title (default is `.title()` of plot name)
TITLES = {
    'air_temp': 'Air Temperature',
    'td': 'Air and Dewpoint Temperature',
    'rh': 'Relative Humidity',
    'wind_dir': 'Wind Direction',
    'accum_precip': 'Accumulated Precipitation Since 0Z',
}


class PlotMaker(object):
    """Object for making plots and storing/validating plot metadata"""
    def __init__(self, name, dependencies, title=None, units=None):
        self.name = name
        self.deps = dependencies
        self._full_figure = None
        if title is None:
            title = "{title_prefix}{title_name}{units} {start_time:%Y-%m-%d}"
        self._title = title
        self.units = units

    def missing_deps(self, frame):
        """Get dependency variables missing from the provided frame"""
        for var_name in self.deps:
            if var_name not in frame:
                yield var_name

    def get_title(self, frame, is_subplot):
        if self._title:
            title_prefix = "AO&SS Building Tower " if not is_subplot else ''
            title_name = TITLES.get(self.name, self.name.replace('_', ' ').title())
            unit_str = '({})'.format(self.units) if self.units and is_subplot else ''
            title = self._title.format(title_prefix=title_prefix,
                                       title_name=title_name,
                                       units=unit_str,
                                       start_time=frame.index[0].to_pydatetime())
        else:
            title = ''
        return title

    def get_yticks(self, ymin, ymax, num_plots):
        delta = math.ceil((ymax - ymin) / num_plots)
        new_ticks = np.arange(ymin, (ymin + delta * num_plots), delta)
        if not new_ticks:
            return [ymin, ymin + 0.05, ymin + 0.1]
        return new_ticks

    def get_ylabel(self, is_subplot=False):
        y_label = TITLES.get(self.name, self.name.replace('_', ' ').title())
        if is_subplot:
            return None
        if self.units:
            return "{} ({})".format(y_label, self.units)
        return y_label

    def create_plot(self, frame, fig, is_subplot=None, shared_x=None, title=None):
        """

        :param frame:
        :param fig:
        :param is_subplot: None or (num plots, num columns, num_rows)
        :param shared_x:
        :return:
        """
        if title is None:
            title = self.get_title(frame, is_subplot)

        if is_subplot:
            ax = fig.add_subplot(*is_subplot, sharex=shared_x)
            ax.set_title(title, x=0.5, y=get_subtitle_location(is_subplot[0]), fontsize=8)
        else:
            ax = fig.add_subplot(111, sharex=shared_x)
            fig.suptitle(title, fontsize=13)
            y_label = self.get_ylabel(is_subplot)
            if y_label:
                ax.set_ylabel(y_label)
        plt.sca(ax)

        # get the min for each column then combine them assuming we can
        specific_frame = frame[[x for x in frame.columns if x in self.deps]]
        ymin = np.floor(specific_frame.min().min())
        ymax = np.ceil(specific_frame.max().max())

        ax.plot(specific_frame.index, specific_frame, 'k')
        if ymin == ymax:
            ax.set_ylim(ymin, ymax + 0.1)
        else:
            ax.set_ylim(ymin, ymax)

        if is_subplot:
            new_ticks = self.get_yticks(ymin, ymax, is_subplot[0])
            ax.get_yaxis().get_major_ticks()[-1].set_visible(False)
            ax.set_yticks(new_ticks)
        else:
            ax.set_xlabel('Time (UTC)')

        return ax


class MeteorogramPlotMaker(PlotMaker):
    def __init__(self, name, dependencies, plot_deps=None, title=None):
        self.plot_deps = plot_deps
        super(MeteorogramPlotMaker, self).__init__(name, dependencies, title=title)

    def create_plot(self, frame, fig, is_subplot=False, shared_x=None):
        if is_subplot or shared_x:
            raise ValueError("Meteorogram Plot can not be a subplot or share X-axis")

        title = self.get_title(frame, False)
        fig.suptitle(title, fontsize=13)

        num_plots = len(self.plot_deps)
        shared_x = None
        for idx, plot_name in enumerate(self.plot_deps):
            plot_maker = PLOT_TYPES.get(plot_name, PlotMaker(plot_name, (plot_name,)))
            title_name = TITLES.get(plot_name, plot_name.replace('_', ' ').title())
            ax = plot_maker.create_plot(frame, fig,
                                        is_subplot=(num_plots, 1, idx + 1),
                                        shared_x=shared_x,
                                        title=title_name)
            if idx == 0:
                shared_x = ax
            if idx != num_plots - 1:
                # Disable the x-axis ticks so we don't interfere with other subplots
                ax.get_xaxis().get_major_ticks()[-1].set_visible(False)
                ax.get_xaxis().get_major_ticks()[0].set_visible(False)
        ax.set_xlabel('Time (UTC)')

        return ax

# map plot name -> variable dependencies
# if not listed then plot name is assumed to be the same as the variable needed
PLOT_TYPES = {
    'meteorogram': MeteorogramPlotMaker('meteorogram',
                                        ('air_temp', 'dewpoint', 'rh', 'wind_speed', 'wind_dir', 'accum_precip'),
                                        ('td', 'rh', 'wind_speed', 'wind_dir', 'accum_precip')),
    'td': PlotMaker('td', ('air_temp', 'dewpoint'), units="°C"),  # air_temp and dewpoint in one plot
    'wind_dir': PlotMaker('wind_dir', ('wind_dir',), units='°'),  # special tick labels
    'rh': PlotMaker('rh', ('rh',), units='%'),
    'air_temp': PlotMaker('air_temp', ('air_temp',), units='°C'),
    'pressure': PlotMaker('pressure', ('pressure',), units='hpa'),
    'dewpoint': PlotMaker('dewpoint', ('air_temp',), units='°C'),
    'wind_speed': PlotMaker('wind_speed', ('wind_speed',), units='m/s'),
    'accum_precip': PlotMaker('accum_precip', ('accum_precip',), units='mm'),
    'solar_flux': PlotMaker('solar_flux', ('solar_flux',), units='W/m^2'),
}


def convert_to_thumbnail(fig):
    # TODO
    pass


def get_data(input_files, columns):
    data_dict = {}
    files = MFDataset(input_files)

    # get the data from the files
    for name in columns:
        if name not in files.variables:
            LOG.warning("Unknown file variable: {}".format(name))
            continue
        data_dict[name] = files.variables[name][:]
        data_dict['qc_' + name] = files.variables['qc_' + name][:]

    # convert base_time epoch format into date_time object
    base_time = files.variables['base_time'][:]
    base_time_obj = dt(1970, 1, 1) + delta(seconds=int(base_time))
    
    # convert per-file offsets to offsets based on the first file's base_time
    offsets = MFTime(files.variables['time_offset'])[:]
    # for each offset, convert that into a datetime object
    data_dict['stamps'] = [base_time_obj + delta(seconds=int(s)) for s in offsets]

    return pd.DataFrame(data_dict).set_index(['stamps'])


# The purpose of this method is to determines all the 12:00:00 days
# within a start and end date
# @param cur_dt - start time
# @param dates - end time
# @return dates in lists

def get_dates_in_range(cur_dt, end):
    curr = cur_dt
    while curr <= end:
        yield curr
        curr += delta(days=1)

# The purpose of this method is to get the min and max of 
# an array of stamps and determines all the 12:00:00 days
# in that array
# @param dewpoint stamps - time stamps for valid dewpoint data. Only specified if
# air temp and dew point need to be ploted in the same plot
# @param dates - time stamps for data
# @return datesInRange - all days with 12:00:00
def find_half_days(dewpoint_stamps, dates):
    date_max = np.amax(dates)
    date_min = np.amin(dates)

    if(dewpoint_stamps):
        dew_max = np.amax(dewpoint_stamps)
        dew_min = np.amin(dewpoint_stamps)

        date_max = max(date_max, dew_max)
        date_min = min(date_min, dew_min)    

    cur_dt = date_min.replace(hour=12, minute=0, second=0)
    
    if cur_dt < date_min:
        cur_dt += delta(days=1)

    datesInRange = list(get_dates_in_range(cur_dt, date_max))    

    return datesInRange           

# The purpose of this method is to create a thumbnail plot
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param dewpoint_stamps - timestamps for valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param dewpoint_data - valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
#@param output - output filename pattern
def thumbnail_plot(dates, data, ymin, ymax,
    dewpoint_stamps, dewpoint_data, output, wind_dir):

    if(dewpoint_stamps):
        plt.plot(dates, data, 'r', dewpoint_stamps, dewpoint_data, 'b')
        plt.axis([min(dates), max(dates), ymin, ymax])
        #plt.plot(dewpoint_stamps, dewpoint_data, 'b')

    elif(wind_dir):
        plt.plot(dates, data, 'ko', markersize=0.1)
        plt.axis([min(dates), max(dates), 0, 360])

    else:
        plt.plot(dates, data, 'k')
        plt.axis([min(dates), max(dates), ymin, ymax])

    #scale
    fig = plt.gcf()
    dpi = fig.get_dpi()
    fig.set_size_inches(80/float(dpi), 80/float(dpi))

    #create axes
    axes = plt.gca()
    axes.axes.get_xaxis().set_ticklabels([])
    axes.axes.get_xaxis().set_tick_params(length=7)
    axes.axes.get_yaxis().set_ticklabels([])
    axes.axes.get_yaxis().set_visible(False)
    axes.spines['right'].set_visible(False)
    axes.spines['bottom'].set_visible(False)

    #got the basic shape of the plot
    #still need to draw a line in the middle
    datesInRange = find_half_days(dewpoint_stamps, dates)
    
    for date in datesInRange:
        plt.axvline(x=date, color='k')       

# The purpose of this method is to create a thumbnail
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param create_air_dew_plot - boolean specifying if air temp and dewpoint
# @param output - output filename pattern 
# should be in same plot
def create_thumbnail(plot_names, frame, ymin, ymax, create_air_dew_plot, output):
    #see if we're going to need subplots
    #subplots needed when len is greater than 2 for one variable
    #or greater than 4 for two variables

    need_subplots = len(list(frame.columns.values)) > 2 and not create_air_dew_plot
    need_subplots = need_subplots or len(list(frame.columns.values)) > 4 and create_air_dew_plot

    if need_subplots:
        plots_created = 1

        if create_air_dew_plot:
            numberPlots = (len(list(frame.columns.values)) - 2) / 2
            plotNumber = numberPlots * 100 + 10

        else:
            numberPlots = len(list(frame.columns.values)) / 2
            plotNumber = numberPlots * 100 + 10    

    stamps = frame.index
    #get dewpoint data
    if create_air_dew_plot:
        all_data = frame['dewpoint']
        qc_data = frame['qc_dewpoint']

        good_list = get_good_data(qc_data, stamps, all_data)
        dewpoint_data = good_list[0]
        dewpoint_stamps = good_list[1]

    for name in plot_names:
        if name == 'dewpoint' and create_air_dew_plot:
            continue

        if name not in list(frame.columns.values):
            continue        

        elif need_subplots:
            plt.subplot(plotNumber + plots_created)
            plots_created += 1


        all_data = frame[name]
        qc_data = frame['qc_' + name]

        good_list = get_good_data(qc_data, stamps, all_data)
        good_data = good_list[0]
        good_stamps= good_list[1]

        #if we don't need two lines in same plot, plot
        if not (create_air_dew_plot and name == 'air_temp'):
            if name != 'wind_dir':    
                thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                    None, None, output, False)

            if name == 'wind_dir':
               thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                    None, None, output, True)      

        else:        
            thumbnail_plot(good_stamps, good_data, ymin[name], ymax[name], 
                dewpoint_stamps, dewpoint_data, output, False)

    plt.savefig(output.format(plot_name='test', start_time=good_stamps[0]))

def full_plot_stamps(stamps):
    # assume the timestamps are monotonic
    first_stamp = stamps[0].replace(hour=0, minute=0, second=0)
    return [((s - first_stamp).total_seconds()/3600) for s in stamps]

# The purpose of this method is to change how the y ticks are shown
# @param min - minimum of the y range right now
# @param max - maximum of the y range right now
# @return new tick labels

def get_new_labels(mini, maxi):
    delta = math.ceil((maxi - mini) / 6)
    return np.arange(mini, (mini + delta*6), delta)
        
# The purpose of this method is to create a full_size plot
# @param frame - all data I need in a data frame
# @param ymin - lower limit of the y axis if specified. set only when --daily is given
# @param ymax - upper limit of yaxis if specified. Set only when --daily is given
# @param dewpoint_stamps - timestamps for valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param dewpoint_data - valid dewpoint data. Only specified if 
# air temp and dew point need to be plotted in the same plot
# @param output - output filename pattern
# @param wind_direction - says different yaxis labels
# @param accum_precip - different yaxis handling
# @param see if we need subplots
def full_plot(fig, axes, dates, data, ymin, ymax,
    dewpoint_stamps, dewpoint_data, output, wind_dir,
    accum_precip, need_subplots):

    dates = full_plot_stamps(dates)

    if(dewpoint_stamps):
        dewpoint_stamps = full_plot_stamps(dewpoint_stamps)

        plt.plot(dates, data, 'r', dewpoint_stamps, dewpoint_data, 'b')
        plt.axis([min(dates), max(dates), ymin, ymax])
        

        #labels = [math.ceil(float(item.get_text())) for item in ticks]
        #plt.plot(dewpoint_stamps, dewpoint_data, 'b')


    elif(wind_dir):
        plt.plot(dates, data, 'ko', markersize=1)
        plt.axis([min(dates), max(dates), 0, 360])

        axes = plt.gca()
        axes.get_yaxis().set_ticks([0,90,180,270])


    else:
        # import ipdb; ipdb.set_trace()
        plt.plot(dates, data, 'k')
        # plt.axis([min(dates), max(dates), ymin, ymax])

    if not need_subplots:
        return

    #if not wind_direction, reset yaxis ticks
    if not wind_dir:
        if not accum_precip:
            mini, maxi = axes.get_ylim()
            new_labels = get_new_labels(mini, maxi)

        else:
            mini, maxi = axes.get_ylim()
            delta = ((maxi - mini)/6)
            new_labels = np.arange(mini, (mini + delta*6), delta)
            
            if len(new_labels) == 0:
                new_labels = [0, 0.05, 0.1]   

        plt.yticks(new_labels)
        axes.get_yaxis().get_major_ticks()[-1].set_visible(False)

    #hid last tick because for subplots above the last one, it's peeping out...
    # will show for the last subplot or when subplots are not needed
    # in other words, this code should only affect plots above the last plot
    # only when sublots are needed
    axes.get_xaxis().get_major_ticks()[-1].set_visible(False)
    axes.get_xaxis().get_major_ticks()[0].set_visible(False)

# The purpose of this method is to get the subtitle's y coordinate
# based upon how many subplots there are
# @param numSubPlots - number of subplots
# @return y-coordinate of subtitle
def get_subtitle_location(numSubPlots):
    return 1 - 0.05*numSubPlots


def create_full_plot(plot_names, frame, output, start_time=None, end_time=None):
    """
    
    Args:
        plot_names: 
        frame: 
        output: 
        start_time: 
        end_time: 
        daily: Whether or not this plot should represent one day of data

    Returns:

    """
    #see if we're going to need subplots
    #subplots needed when len is greater than 2 for one variable
    #or greater than 4 for two variables

    # need_subplots = len(list(frame.columns.values)) > 2 and not create_air_dew_plot
    # need_subplots = need_subplots or len(list(frame.columns.values)) > 4 and create_air_dew_plot

    #get need supplot vars
    # if need_subplots:
    #     plots_created = 1
    #
    #     if create_air_dew_plot:
    #         numberPlots = (len(list(frame.columns.values)) - 2) / 2
    #         plotNumber = numberPlots * 100 + 10
    #
    #     else:
    #         numberPlots = len(list(frame.columns.values)) / 2
    #         plotNumber = numberPlots * 100 + 10

    for name in plot_names:
        plot_maker = PLOT_TYPES.get(name, PlotMaker(name, (name,)))
        var_names = []
        for var_name in plot_maker.deps:
            if var_name not in frame:
                raise ValueError("Missing required variable '{}' for plot '{}'".format(var_name, name))
            var_names.append(var_name)
            # write NaNs where QC values are not 0
            qc_name = 'qc_' + var_name
            if qc_name in frame:
                frame[var_name].mask(frame[qc_name] != 0)
                var_names.append(qc_name)

        # create a frame that doesn't include any of the bad values
        plot_frame = frame[var_names]
        plot_frame = plot_frame[~plot_frame.isnull().any(axis=1)]

        fig = plt.figure()
        ax = plot_maker.create_plot(plot_frame, fig)
        ax.set_xlim(start_time, end_time)
        import matplotlib.dates as md
        xloc = md.AutoDateLocator(minticks=5, maxticks=8, interval_multiples=True)
        xfmt = md.AutoDateFormatter(xloc)
        def _fmt(interval, x, pos=None):
            x_num = md.num2date(x).replace(tzinfo=None)
            delta_seconds = (x_num - plot_frame.index[0].replace(hour=0, minute=0, second=0, microsecond=0)).total_seconds()
            num_hours = delta_seconds / 3600.
            if interval == md.HOURLY:
                return "{:02.0f}".format(num_hours)
            elif interval == md.MINUTELY:
                num_minutes = delta_seconds / 60.
                num_minutes -= int(num_hours) * 60.
                return "{:02.0f}:{:02.0f}".format(int(num_hours), num_minutes)
            else:
                return x.strftime("{%Y-%m-%d}")
        from functools import partial
        xfmt.scaled[1. / md.MINUTES_PER_DAY] = plt.FuncFormatter(partial(_fmt, md.MINUTELY))
        xfmt.scaled[1. / md.HOURS_PER_DAY] = plt.FuncFormatter(partial(_fmt, md.HOURLY))
        ax.xaxis.set_major_locator(xloc)
        ax.xaxis.set_major_formatter(xfmt)

        out_fn = output.format(plot_name=name, start_time=plot_frame.index[0].to_pydatetime())
        LOG.info("Saving plot '{}' to filename '{}'".format(name, out_fn))
        fig.savefig(out_fn)

        # all_data = frame[name]
        # qc_data = frame['qc_' + name]
        #
        # good_list = get_good_data(qc_data, stamps, all_data)
        # good_data = good_list[0]
        # good_stamps = good_list[1]
        #
        # dateString = good_stamps[0].strftime('%m/%d/%Y')
        #
        # # create plots
        # if need_subplots:
        #     if plots_created == 1:
        #         plt.figure().suptitle('AO&SS Building Tower Meteorogram ' + dateString, fontsize=13)
        #         ax1 = plt.subplot(plotNumber + plots_created)
        #         ax1.set_title(TITLES[name], x=0.5, y=get_subtitle_location(numberPlots), fontsize=8)
        #     else:
        #         curr_plot = plt.subplot(plotNumber + plots_created, sharex=ax1)
        #         curr_plot.set_title(TITLES[name], x=0.5, y=get_subtitle_location(numberPlots), fontsize=8)
        #
        #     plots_created += 1
        #
        # else:
        #     plt.figure().suptitle('AO&SS Building Tower ' + IND_TITLES[name] + ' ' + dateString, fontsize=13)
        #     plt.ylabel(TITLES[name])
        #
        # #if we don't need two lines in same plot, plot
        # if not (create_air_dew_plot and name == 'air_temp'):
        #     full_plot(good_stamps, good_data, ymin[name], ymax[name],
        #                 None, None, output, name == 'wind_dir',
        #                 name == 'accum_precip', need_subplots)
        #
        # else:
        #     full_plot(good_stamps, good_data, ymin[name], ymax[name],
        #         dewpoint_stamps, dewpoint_data, output, False, False,
        #         need_subplots)

    # if need_subplots:
    #     plt.subplots_adjust(hspace=0, bottom=0.125)
    #
    # plt.xlabel('Time (UTC)')


def _dt_convert(datetime_str):
    try:
        return dt.strptime(datetime_str, '%Y-%m-%dT%H:%M:%S')
    except ValueError:
        return dt.strptime(datetime_str, '%Y-%m-%d')


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Use data from level_b1 netCDF files to create netCDF files")
    parser.add_argument('-v', '--verbose', action='count',
                       default=int(os.environ.get("VERBOSITY", 2)), 
                       dest='verbosity',
                       help=('each occurence increases verbosity 1 level through'
                             + ' ERROR-WARNING-INFO-DEBUG (default INFO)'))
    parser.add_argument('-s', '--start-time', type=_dt_convert,
    help="Start time of plot. If only -s is given, a plot of " +
        "only that day is created. Formats allowed: \'YYYY-MM-DDTHH:MM:SS\', \'YYYY-MM-DD\'")
    parser.add_argument('-e', '--end-time', type=_dt_convert,
    help="End time of plot. If only -e is given, a plot of only that day is " +
          "created. Formats allowed: \'YYYY-MM-DDTHH:MM:SS\', \'YYYY-MM-DD\'")
    parser.add_argument('--met-plots', nargs='+',
                        help="Override plots to use in the combined meteorogram plot")
    parser.add_argument("input_files", nargs="+", help="aoss_tower_level_b1 files")
    parser.add_argument('-o', '--output', default="{plot_name}_{start_time:%Y%m%d_%H%M%S}.png", help="filename pattern")
    parser.add_argument('-t', '--thumbnail', action='store_true', help="if specified, script creates a thumbnail")
    parser.add_argument('-p', '--plot-names', nargs="+",
                        required=True,
                        help="the variable names or plot types to create")
    parser.add_argument('-d', '--daily', action='store_true',
        help="creates a plot for every day. Usually used to create plots " +
        "that will line up for aoss tower quicklooks page")
    args = parser.parse_args()

    levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
    level = levels[min(3, args.verbosity)]
    logging.basicConfig(level=level)

    if not os.path.splitext(args.output)[-1]:
        LOG.warning("File pattern provided does not have a file extension")

    # check the dependencies for the meteorogram
    if args.met_plots:
        assert 'meteorogram' not in args.met_plots
        PLOT_TYPES['meteorogram'].deps = args.met_plots

    plot_deps = [PLOT_TYPES[k].deps if k in PLOT_TYPES else (k,) for k in args.plot_names]
    plot_deps = list(set(d for deps in plot_deps for d in deps))
    frame = get_data(args.input_files, plot_deps)
    bad_plot_names = set(args.plot_names) - (set(frame.columns) | set(PLOT_TYPES.keys()))
    if bad_plot_names:
        raise ValueError("Unknown plot name(s): {}".format(", ".join(bad_plot_names)))

    #frame only contains data from start-end times
    if args.start_time and args.end_time:
        frame = frame[args.start_time: args.end_time]
    elif args.start_time:
        #frame only contains data from start-end of that day
        end_time = args.start_time.replace(hour=23, minute=59, second=59)
        frame = frame[args.start_time: end_time]

    if not args.daily:
        # allow plotting methods to write inplace on a copy
        frames = [frame.copy()]
    else:
        frames = (group[1] for group in frame.groupby(frame.index.day))

    for frame in frames:
        if args.daily:
            # modify start and end time to the current day
            start_time = frame.index[0].to_pydatetime().replace(hour=0, minute=0, second=0, microsecond=0)
            end_time = frame.index[0].to_pydatetime().replace(hour=23, minute=59, second=59, microsecond=999999)
        else:
            start_time = args.start_time
            end_time = args.end_time

        if args.thumbnail:
            create_thumbnail(args.plot_names, frame, args.output, start_time, end_time)
        else:
            create_full_plot(args.plot_names, frame, args.output, start_time, end_time)

if __name__ == "__main__":
    sys.exit(main())