quicklook.py

import logging
import math
import os
import sys
from datetime import datetime, timedelta
from pathlib import Path

import matplotlib as mpl
import matplotlib.dates as md
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from netCDF4 import MFDataset, MFTime

mpl.use("agg")

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 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")
        return f"{start_time:%Y-%m-%d %H:%M} to {end_time:%Y-%m-%d %H:%M}"

    def get_title(self, frame, is_subplot, start_time, end_time):  # noqa: ARG002
        if not self._title:
            return ""
        title_prefix = "AO&SS Building Tower " if not is_subplot else ""
        title_name = TITLES.get(self.name, self.name.replace("_", " ").title())
        unit_str = f"({self.units})" if self.units and is_subplot else ""
        date_string = self.get_date_string(start_time, end_time)
        return self._title.format(
            title_prefix=title_prefix,
            title_name=title_name,
            units=unit_str,
            date_string=date_string,
        )

    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)
        return np.arange(ymin, (ymin + delta * num_plots), delta)

    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 f"{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(0.008, 0.9, self.units, horizontalalignment="left", va="top", transform=ax.transAxes, size=8)

    def _call_plot(self, frame, ax):
        return ax.plot(frame.index, frame, "k")

    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):
        ax = fig.add_subplot(*is_subplot, sharex=shared_x) if is_subplot else 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.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):  # noqa: ARG001
            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.0
            if interval == md.HOURLY:
                return f"{num_hours:.0f}"
            if interval == md.MINUTELY:
                num_minutes = delta_seconds / 60.0
                num_minutes -= int(num_hours) * 60.0
                return f"{int(num_hours):02.0f}:{num_minutes:02.0f}"
            return x.strftime("{%Y-%m-%d}")

        from functools import partial

        xfmt.scaled[1.0 / md.MINUTES_PER_DAY] = plt.FuncFormatter(partial(_fmt, md.MINUTELY))
        xfmt.scaled[1.0 / 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):
        """Create series of plots."""
        specific_frame = frame[[x for x in frame.columns if x in self.deps]]
        if frame.empty or specific_frame.empty or specific_frame.isna().all().any():
            msg = f"No valid data found or missing necessary data to make {self.name}"
            raise ValueError(msg)
        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")
        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):  # noqa: ARG002
        return 0, 360

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

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


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().__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()._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:
            msg = "Meteorogram Plot can not be a subplot or share X-axis"
            raise ValueError(msg)

        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 bookkeeping 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 label in ax.get_xticklabels():
                    label.update(kwargs)
            # make the top y-tick label invisible

        ax.set_xlabel("Time (UTC)")
        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(f"Unknown file variable: {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):
    """Create a series of plots.

    Args:
        plot_names (list of str): Plot names to generate
        frame (pd.DataFrame): DataFrame of data
        output (str): Output pattern for the created files
        start_time (datetime): Start time of the data to use
        end_time (datetime): End time of the data to use
        thumbnail (bool): Additionally generate a thumbnail

    """
    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:
                msg = f"Missing required variable '{var_name}' for plot '{name}'"
                raise ValueError(msg)
            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.isna().any(axis=1)]

        fig = plt.figure()
        try:
            plot_maker.create_plot(plot_frame, fig, start_time=start_time, end_time=end_time)
        except ValueError:
            LOG.error(f"Could not make '{name}'", exc_info=True)
            continue

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

        if thumbnail:
            out_path = Path(out_fn)
            stem = out_path.stem
            ext = out_path.suffix
            out_fn = f"{stem}_thumbnail{ext}"
            plot_maker.convert_to_thumbnail(fig)
            LOG.info(f"Saving thumbnail '{name}' to filename '{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

    from metobscommon.archive import setup_logging

    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(
        "-l",
        "--log-file",
        dest="log_filepath",
        help="Alternate name for log file, default is to not create a file",
    )
    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]
    setup_logging(args.log_filepath, level=levels[min(3, args.verbosity)])

    if not Path(args.output).suffix:
        LOG.warning("File pattern provided does not have a file extension")

    # check the dependencies for the meteorogram
    if args.met_plots:
        if "meteorogram" in args.met_plots:
            msg = "The 'meteorogram' plot can not be a sub-plot of itself."
            raise ValueError(msg)
        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({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, microsecond=999999)
        frame = frame[args.start_time : end_time]

    # allow plotting methods to write inplace on a copy
    frames = [frame.copy()] if not args.daily else (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())