Add thumbnail conversion to quicklooks

aosstower/tower_quicklooks/create_quicklook.py

@@ -11,6 +11,7 @@ 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 = {
@@ -22,6 +23,11 @@ TITLES = {
 }
 
 
+def remove_yticklines(ax):
+    for t in ax.yaxis.get_ticklines():
+        t.set_visible(False)
+
+
 class PlotMaker(object):
     """Object for making plots and storing/validating plot metadata"""
     def __init__(self, name, dependencies, title=None, units=None):
@@ -81,6 +87,7 @@ class PlotMaker(object):
             ax.set_ylabel(y_label)
 
     def _call_plot(self, frame, ax):
+        print(self.name, frame.columns)
         lines = ax.plot(frame.index, frame, 'k')
         return lines
 
@@ -94,11 +101,7 @@ class PlotMaker(object):
             ax.set_ylim(ymin - delta * 0.2, ymax + delta * 0.2)
         return ymin, ymax
 
-    def _set_title(self, frame, fig, ax, start_time=None, end_time=None, title=None, is_subplot=None):
-        if start_time is None:
-            start_time = frame.index[0].to_pydatetime()
-        if end_time is None:
-            end_time = frame.index[-1].to_pydatetime()
+    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)
 
@@ -124,6 +127,73 @@ class PlotMaker(object):
             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)
+
+    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):
@@ -135,6 +205,12 @@ class PlotMaker(object):
         :param shared_x:
         :return:
         """
+        if start_time is None:
+            start_time = frame.index[0].to_pydatetime()
+        if end_time is None:
+            end_time = frame.index[-1].to_pydatetime()
+
+        # TODO: If frame is empty, make an empty plot
         ax = self._get_axes(fig, is_subplot, shared_x)
         self._set_title(frame, fig, ax,
                         start_time=start_time, end_time=end_time,
@@ -150,6 +226,7 @@ class PlotMaker(object):
         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
 
@@ -187,10 +264,34 @@ class WindDirPlotMaker(PlotMaker):
 
 
 class MeteorogramPlotMaker(PlotMaker):
-    def __init__(self, name, dependencies, plot_deps=None, title=None):
+    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:
@@ -206,6 +307,8 @@ class MeteorogramPlotMaker(PlotMaker):
 
         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())
@@ -213,6 +316,7 @@ class MeteorogramPlotMaker(PlotMaker):
                                         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:
@@ -224,7 +328,9 @@ class MeteorogramPlotMaker(PlotMaker):
             # 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, bottom=0.125)
+        fig.subplots_adjust(hspace=0)
+        self._set_xaxis_formatter(ax, start_time, end_time, is_subplot)
 
         return ax
 
@@ -232,8 +338,9 @@ class MeteorogramPlotMaker(PlotMaker):
 # 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')),
+                                        ('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='%'),
@@ -246,11 +353,6 @@ PLOT_TYPES = {
 }
 
 
-def convert_to_thumbnail(fig):
-    # TODO
-    pass
-
-
 def get_data(input_files, columns):
     data_dict = {}
     files = MFDataset(input_files)
@@ -522,7 +624,8 @@ def get_subtitle_location(numSubPlots):
     return 1 - 0.055*numSubPlots
 
 
-def create_full_plot(plot_names, frame, output, start_time=None, end_time=None):
+def create_full_plot(plot_names, frame, output,
+                     start_time=None, end_time=None, thumbnail=False):
     """
     
     Args:
@@ -536,25 +639,6 @@ def create_full_plot(plot_names, frame, output, start_time=None, end_time=None):
     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
-
     if start_time is None:
         start_time = frame.index[0].to_pydatetime()
     if end_time is None:
@@ -579,31 +663,17 @@ def create_full_plot(plot_names, frame, output, start_time=None, end_time=None):
 
         fig = plt.figure()
         ax = plot_maker.create_plot(plot_frame, fig, start_time=start_time, end_time=end_time)
-        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 - 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 "{:.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=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)
+            fig.savefig(out_fn)
+
         # all_data = frame[name]
         # qc_data = frame['qc_' + name]
         #
@@ -694,6 +764,7 @@ def main():
 
     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))
+    print(plot_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:
@@ -722,10 +793,7 @@ def main():
             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)
+        create_full_plot(args.plot_names, frame, args.output, start_time, end_time, args.thumbnail)
 
 if __name__ == "__main__":
     sys.exit(main())