quicklook.py

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

LOG = logging.getLogger(__name__)
FIGURE_TITLE_SIZE = 13
TN_SIZE = (1, 1)

# 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',
}


def remove_yticklines(ax):
    for t in ax.yaxis.get_ticklines():
        t.set_visible(False)


def get_subtitle_location(num_subplots):
    return 1 - 0.055 * num_subplots


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}\n{date_string}"
        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_date_string(self, start_time, end_time):
        delta = (end_time - start_time).total_seconds()
        if delta < timedelta(hours=24).total_seconds():
            return start_time.strftime("%Y-%m-%d")
        else:
            return "{:%Y-%m-%d %H:%M} to {:%Y-%m-%d %H:%M}".format(start_time, end_time)

    def get_title(self, frame, is_subplot, start_time, end_time):
        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 ''
            date_string = self.get_date_string(start_time, end_time)
            title = self._title.format(title_prefix=title_prefix,
                                       title_name=title_name,
                                       units=unit_str,
                                       date_string=date_string)
        else:
            title = ''
        return title

    def get_yticks(self, ymin, ymax, num_plots):
        if ymin == ymax:
            return [ymin, ymin + 0.05, ymin + 0.1]
        delta = math.ceil((ymax - ymin) / num_plots)
        new_ticks = np.arange(ymin, (ymin + delta * num_plots), delta)
        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 _set_ylabel(self, ax, is_subplot=False):
        y_label = self._get_ylabel(is_subplot)
        if y_label and not is_subplot:
            ax.set_ylabel(y_label)

        if is_subplot:
            # put units on the top left of the plot axes
            ax.text(.008, .9, self.units,
                    horizontalalignment='left', va='top',
                    transform=ax.transAxes, size=8)

    def _call_plot(self, frame, ax):
        lines = ax.plot(frame.index, frame, 'k')
        return lines

    def _set_ylim(self, frame, ax):
        ymin = np.floor(frame.min().min())
        ymax = np.ceil(frame.max().max())
        delta = ymax - ymin
        if ymin == ymax:
            ax.set_ylim(ymin - 0.1, ymax + 0.1)
        else:
            ax.set_ylim(ymin - delta * 0.2, ymax + delta * 0.2)
        return ymin, ymax

    def _set_title(self, frame, fig, ax, start_time, end_time, title=None, is_subplot=None):
        if title is None:
            title = self.get_title(frame, is_subplot, start_time, end_time)

        if is_subplot:
            ax.set_title(title, x=0.5, y=get_subtitle_location(is_subplot[0]), fontsize=8)
        else:
            fig.suptitle(title, fontsize=FIGURE_TITLE_SIZE)

    def _get_axes(self, fig, is_subplot, shared_x=None):
        if is_subplot:
            ax = fig.add_subplot(*is_subplot, sharex=shared_x)
        else:
            ax = fig.add_subplot(111, sharex=shared_x)
        plt.sca(ax)
        return ax

    def _set_xlabel(self, ax, is_subplot):
        if not is_subplot:
            ax.set_xlabel('Time (UTC)')

    def _set_yticks(self, ax, ymin, ymax, is_subplot):
        if is_subplot:
            new_ticks = self.get_yticks(ymin, ymax, is_subplot[0])
            # ax.yaxis.get_major_ticks()[-1].set_visible(False)
            ax.set_yticks(new_ticks)
            ax.locator_params(axis='y', nbins=3)
        ax.yaxis.get_major_formatter().set_useOffset(False)

    def _set_xaxis_formatter(self, ax, start_time, end_time, is_subplot):
        if is_subplot:
            return
        ax.set_xlim(start_time, end_time)
        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 - start_time.replace(hour=0, minute=0, second=0, microsecond=0)).total_seconds()
            num_hours = delta_seconds / 3600.
            if interval == md.HOURLY:
                return "{:.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)

    def _convert_fig_to_thumbnail(self, fig, fig_size=TN_SIZE):
        # Resize the figure
        fig.set_size_inches(fig_size)

        # Remove any titles
        fig.suptitle("")

        # Remove colorbar
        if hasattr(fig, "cb"):
            fig.delaxes(fig.cb.ax)
            fig.cb = None

        # Use as much space as possible
        fig.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0)
        if hasattr(fig, "cb"):
            # WARNING: This will not work if there are more than one axes
            for ax in fig.axes:
                ax.set_position([0, 0, 1, 1])

        return fig

    def _convert_ax_to_thumbnail(self, fig):
        # If thumbnails are placed next to each other we don't want ticklines
        # where they touch
        for ax in fig.axes:
            remove_yticklines(ax)

            # Remove any titles
            ax.set_title("")
            # remove y-axis lines
            ax.spines['left'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.spines['top'].set_visible(False)
            ax.spines['bottom'].set_visible(False)

            for t in ax.texts:
                t.set_visible(False)

    def convert_to_thumbnail(self, fig, fig_size=TN_SIZE):
        fig = self._convert_fig_to_thumbnail(fig, fig_size)
        self._convert_ax_to_thumbnail(fig)
        return fig

    def create_plot(self, frame, fig, start_time=None, end_time=None,
                    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:
        """
        specific_frame = frame[[x for x in frame.columns if x in self.deps]]
        if frame.empty or specific_frame.empty or specific_frame.isnull().all().any():
            raise ValueError("No valid data found or missing necessary data to make {}".format(self.name))
        if start_time is None:
            start_time = frame.index[0].to_pydatetime()
        if end_time is None:
            end_time = frame.index[-1].to_pydatetime()

        ax = self._get_axes(fig, is_subplot, shared_x)
        self._set_title(frame, fig, ax,
                        start_time=start_time, end_time=end_time,
                        title=title, is_subplot=is_subplot)

        # make ticks show up on top and bottom inside and out of the axis line
        ax.xaxis.set_tick_params(left=True, right=True, direction='inout')
        lines = self._call_plot(specific_frame, ax)
        ymin, ymax = self._set_ylim(specific_frame, ax)

        self._set_yticks(ax, ymin, ymax, is_subplot)
        self._set_xlabel(ax, is_subplot)
        self._set_ylabel(ax, is_subplot)
        self._set_xaxis_formatter(ax, start_time, end_time, is_subplot)

        return ax


class PrecipPlotMaker(PlotMaker):
    def _set_ylim(self, frame, ax):
        ymin = 0
        ymax = np.ceil(frame.max().max())
        delta = ymax - ymin
        # 0 is the minimum y-min we want
        if ymin == ymax:
            ax.set_ylim(ymin, ymax + 0.2)
        else:
            ax.set_ylim(ymin, ymax + delta * 0.2)
        return ymin, ymax


class TDPlotMaker(PlotMaker):
    def _call_plot(self, frame, ax):
        air_temp = self.deps[0]
        dewpoint = self.deps[1]
        return ax.plot(frame.index, frame[air_temp], 'r', frame.index, frame[dewpoint], 'g')


class WindDirPlotMaker(PlotMaker):
    def _set_ylim(self, frame, ax):
        return 0, 360

    def _set_yticks(self, ax, ymin, ymax, is_subplot):
        ax.yaxis.set_ticks([0, 90, 180, 270])

    def _call_plot(self, frame, ax):
        lines = ax.plot(frame.index, frame, 'k.', markersize=3, linewidth=0)
        return lines


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

    def _convert_ax_to_thumbnail(self, fig):
        if hasattr(fig, '_my_axes'):
            for k, ax in fig._my_axes.items():
                if k not in self.thumbnail_deps:
                    fig.delaxes(ax)
                    continue

        super(MeteorogramPlotMaker, self)._convert_ax_to_thumbnail(fig)

        for idx, ax in enumerate(fig.axes):
            ax.spines['left'].set_visible(False)
            ax.spines['right'].set_visible(False)
            if idx == 0:
                ax.spines['top'].set_visible(False)
            else:
                ax.spines['top'].set_visible(True)

            if idx == len(fig.axes) - 1:
                ax.spines['bottom'].set_visible(False)
            else:
                ax.spines['bottom'].set_visible(True)

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

        if start_time is None:
            start_time = frame.index[0].to_pydatetime()
        if end_time is None:
            end_time = frame.index[-1].to_pydatetime()
        if title is None:
            title = self.get_title(frame, False, start_time, end_time)
        fig.suptitle(title, fontsize=FIGURE_TITLE_SIZE)

        num_plots = len(self.plot_deps)
        shared_x = None
        # some hacky book keeping so we can create a thumbnail
        fig._my_axes = {}
        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)
            fig._my_axes[plot_name] = ax
            if idx == 0:
                shared_x = ax
            if idx != num_plots - 1:
                # Disable the x-axis ticks so we don't interfere with other subplots
                kwargs = {'visible': False}
                for l in ax.get_xticklabels():
                    l.update(kwargs)
            # make the top y-tick label invisible
            # ax.yaxis.get_major_ticks()[-1].label1.update({'visible': False})

        ax.set_xlabel('Time (UTC)')
        # fig.subplots_adjust(hspace=0, bottom=0.125)
        fig.subplots_adjust(hspace=0)
        self._set_xaxis_formatter(ax, start_time, end_time, is_subplot)

        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', 'pressure', 'wind_speed', 'wind_dir', 'accum_precip', 'solar_flux'),
                                        ('td', 'pressure', 'wind_speed', 'wind_dir', 'accum_precip', 'solar_flux'),
                                        ('td', 'wind_speed', 'wind_dir', 'accum_precip')),
    'td': TDPlotMaker('td', ('air_temp', 'dewpoint'), units="°C"),  # air_temp and dewpoint in one plot
    'wind_dir': WindDirPlotMaker('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': PrecipPlotMaker('accum_precip', ('accum_precip',), units='mm'),
    'solar_flux': PlotMaker('solar_flux', ('solar_flux',), units='W/m^2'),
}


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 = datetime(1970, 1, 1) + timedelta(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 + timedelta(seconds=int(s)) for s in offsets]

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


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

    Returns:

    """
    if start_time is None:
        start_time = frame.index[0].to_pydatetime()
    if end_time is None:
        end_time = frame.index[-1].to_pydatetime()

    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()
        try:
            ax = plot_maker.create_plot(plot_frame, fig, start_time=start_time, end_time=end_time)
        except ValueError:
            LOG.error("Could not make '{}'".format(name), exc_info=True)
            continue

        out_fn = output.format(plot_name=name, start_time=start_time, end_time=end_time)
        LOG.info("Saving plot '{}' to filename '{}'".format(name, out_fn))
        fig.savefig(out_fn)

        if thumbnail:
            stem, ext = os.path.splitext(out_fn)
            out_fn = "{}_thumbnail{}".format(stem, ext)
            plot_maker.convert_to_thumbnail(fig)
            LOG.info("Saving thumbnail '{}' to filename '{}'".format(name, out_fn))
            fig.savefig(out_fn)


def _dt_convert(datetime_str):
    """Parse datetime string, return datetime object"""
    try:
        return datetime.strptime(datetime_str, '%Y%m%d')
    except ValueError:
        try:
            return datetime.strptime(datetime_str, '%Y-%m-%d')
        except ValueError:
            return datetime.strptime(datetime_str, '%Y-%m-%dT%H:%M:%S')


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\', \'YYYYMMDD\'")
    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

        create_plot(args.plot_names, frame, args.output, start_time, end_time, args.thumbnail)

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