diff --git a/demo/demo_ts.py b/demo/demo_ts.py new file mode 100755 index 0000000..07f79c8 --- /dev/null +++ b/demo/demo_ts.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python +import sys + +plotterdir = './repo/plotter' +sys.path.insert(0, plotterdir) + +from plotter.reader import reader, get_format +import plotter.plotter_solo as plotter_solo + +import matplotlib.colors as colors + +import numpy as np +from pathlib import Path + +fnames = [ + '/scratch1/01923/ling/hysplit/hysplit.v5.0.0_CentOS/scripts/outconc.S2.const.hrrr.2m75kghr.txt', + '/scratch1/00576/yosuke/projects/astra/calpuff/work_yk/toy_mmif/calpost/tseries/tseries_ch4_1min_conc_toy_min_sys1_onesrc_byweek_20190925_20190927.dat', +] + +fmts = [get_format(fn[0]) if isinstance(fn, list) else get_format(fn) for fn in fnames] + +# i have to provide coords for toy region because inconsitency in hrrr +# projection and calpuff projecton +xc = np.repeat(np.arange(34) * 0.1 - 464.35, 47) +yc = np.tile(np.arange(47) * 0.1 - 906.65, 34) + +# above should work for hysplit, but i did differently (linear location of x +# and y, instead of all points' x,y) +xh = np.arange(34) * 0.1 - 464.35 +yh = np.arange(47) * 0.1 - 906.65 + + +xs = [] +ys = [] +dats = [] +for fn, fmt in zip(fnames, fmts): + print(fmt) + x = {'calpost': xc, 'hysplit': xh}[fmt] + y = {'calpost': yc, 'hysplit': yh}[fmt] + dat = reader(fn, x=x, y=y) + xs.append(dat['x']) + ys.append(dat['y']) + dats.append(dat) + +# find common time period +tstamps = [_['ts'] for _ in dats] + +if tstamps[0][0] < tstamps[1][0]: + s0 = np.where(tstamps[1][0] == tstamps[0]) + s1 = 0 + print('s0', s0) + assert len(s0) == 1 + s0 = int(s0[0]) +else: + s1 = np.where(tstamps[0][0] == tstamps[1]) + s0 = 0 + print('s1', s1) + assert len(s1) == 1 + s1 = int(s1[0]) + +if tstamps[0][-1] < tstamps[1][-1]: + e1 = np.where(tstamps[0][-1] == tstamps[1]) + e0 = None + print('e1', e1) + assert len(e1) == 1 + e1 = int(e1[0]) + 1 +else: + e0 = np.where(tstamps[1][-1] == tstamps[0]) + e1 = None + print('e0', e0) + assert len(e0) == 1 + e0 = int(e0[0]) + 1 + +assert np.all(tstamps[0][s0:e0] == tstamps[1][s1:e1]) + +# extract necessary data +arrays = [dat['v'][s:e] for dat, s, e in zip(dats, (s0, s1), (e0, e1))] +tss = [ts[s:e] for ts, s, e in zip(tstamps, (s0, s1), (e0, e1))] +ts = tstamps[0][s0:e0] + +# conversion factors +convfacs = [{'calpost': 1. / 16.043 * 0.024465403697038 * 1e9, + 'hysplit': 1., }[_] for _ in fmts] + +arrays = [arr*cf for arr, cf in zip(arrays, convfacs)] + +# array has nan, so mask them +arrays = [np.ma.masked_invalid(arr) for arr in arrays] + +dct_arrays = {k: v for k, v in zip(fmts, arrays)} + +# # calpost knows location but hysplit needt to be told +# xs = [{'calpost': None, 'hysplit': get_receptor_coords.df.x}[_] for _ in +# fmts] +# ys = [{'calpost': None, 'hysplit': get_receptor_coords.df.y}[_] for _ in +# fmts] +# +# dats = [reader(fn, x=x, y=y,) +# for fn, x, y in zip(fnames, xs, ys)] + +if True: + # solo, ts + plotter_options = {'tseries': True} + oname = 'ts_test.mp4' + p = plotter_solo.Plotter(array=dct_arrays, tstamps=ts, + plotter_options=plotter_options) + p.savefig(Path(oname).with_suffix('.png')) + p.savemp4(oname, wdir=None) + +import plotter.plotter_multi as plotter_multi +from plotter.plotter_util import LambertConformalTCEQ +from plotter.plotter_background import BackgroundManager +import cartopy.io.img_tiles as cimgt + +# Mrinali/Gary's surfer color scale +cmap = colors.ListedColormap([ + '#D6FAFE', '#02FEFF', '#C4FFC4', '#01FE02', + '#FFEE02', '#FAB979', '#EF6601', '#FC0100', ]) +cmap.set_under('#FFFFFF') +cmap.set_over('#000000') +# Define a normalization from values -> colors +bndry = [1, 10, 50, 100, 200, 500, 1000, 2000] +norm = colors.BoundaryNorm(bndry, len(bndry)) +contour_options = { + 'levels': bndry, + 'cmap': cmap, + 'norm': norm, + 'alpha': .5, + 'extend': 'max', + } +colorbar_options = { + 'label': r'$CH_4$ (ppbV)', + } + +if True: + tile_plotter_options = { + 'background_manager': BackgroundManager( + add_image_options=[cimgt.GoogleTiles(style='satellite'), 13], + ), + 'contour_options': contour_options, + 'colorbar_options': None, + 'footnote': '', + # 'footnote_options': {'text':''}, + } + figure_options = { + 'colorbar_options': { + 'label': r'$CH_4$ (ppbV)', + }, + 'footnote_options': {'text': "{tstamp}", 'y': .05}, # 'fontsize': 'small'}, + 'figsize': (10, 4), + } + + listof_plotter_options = [ + tile_plotter_options.copy(), + {'tseries': True}, + tile_plotter_options.copy(), + ] + listof_plotter_options[0].update({ + 'title': 'Hysplit', + }) + listof_plotter_options[2].update({ + 'title': 'Calpuff', + }) + + oname = 'ts_test_trio.mp4' + + assert np.all(xs[0] == xs[1]) + assert np.all(ys[0] == ys[1]) + + p = plotter_multi.Plotter(arrays=[dct_arrays['hysplit'], dct_arrays, + dct_arrays['calpost']], tstamps=ts, + x=xs[0]*1000, y=ys[0]*1000, + projection=LambertConformalTCEQ(), + plotter_options=listof_plotter_options, + figure_options=figure_options) + + p.savefig(Path(oname).with_suffix('.png'), tidx=10) + p.savemp4(oname, wdir=None) diff --git a/plotter/plotter_multi.py b/plotter/plotter_multi.py index 88e8a46..0d4a130 100644 --- a/plotter/plotter_multi.py +++ b/plotter/plotter_multi.py @@ -1,11 +1,13 @@ try: from . import plotter_core as pc from . import plotter_vprof as pv + from . import plotter_tseries as pt from . import plotter_util as pu from . import plotter_footnote as pf except ImportError: import plotter_core as pc import plotter_vprof as pv + import plotter_ts as pt import plotter_util as pu import plotter_footnote as pf @@ -108,8 +110,17 @@ def __init__(self, arrays, tstamps, projection=None, extent=None, # create plots if z is None: - self.plotters = [pc.PlotterCore(arr, tstamps, projection=projection, extent=extent, - x=x, y=y, plotter_options=po) for arr, po in zip(arrays, plotter_options)] + self.plotters = [] + for arr,po in zip(arrays, plotter_options): + if 'tseries' in po: + self.plotters.append( + pt.PlotterTseries(arr, tstamps, plotter_options=po) + ) + else: + self.plotters.append( + pc.PlotterCore(arr, tstamps, projection=projection, extent=extent, + x=x, y=y, plotter_options=po) + ) else: self.plotters = [] for arr, po in zip(arrays, plotter_options): diff --git a/plotter/plotter_solo.py b/plotter/plotter_solo.py index f4a58bf..e98d611 100644 --- a/plotter/plotter_solo.py +++ b/plotter/plotter_solo.py @@ -1,9 +1,11 @@ try: from . import plotter_core as pc from . import plotter_vprof as pv + from . import plotter_tseries as pt except ImportError: import plotter_core as pc import plotter_vprof as pv + import plotter_ts as pt import matplotlib.pyplot as plt import matplotlib as mpl @@ -31,7 +33,11 @@ def __init__(self, array, tstamps, projection=None, extent=None, x=None, self.tstamps = tstamps if z is None: - self.plotter = pc.PlotterCore(array, tstamps, projection=projection, + if 'tseries' in plotter_options: + self.plotter = pt.PlotterTseries(array, tstamps, + plotter_options=plotter_options) + else: + self.plotter = pc.PlotterCore(array, tstamps, projection=projection, extent=extent, x=x, y=y, plotter_options=plotter_options) else: if 'kdx' in plotter_options: @@ -61,6 +67,7 @@ def savefig(self, oname, tidx=None, footnote=None, *args, **kwargs): :param dict kwargs: extra arguments passed to plt.savefig() """ self.plotter.update(tidx, footnote) + plt.figure(self.plotter.fig.number) plt.savefig(oname, *args, **kwargs) def __call__(self, oname, *args, **kwargs): diff --git a/plotter/plotter_tseries.py b/plotter/plotter_tseries.py new file mode 100644 index 0000000..010a9b2 --- /dev/null +++ b/plotter/plotter_tseries.py @@ -0,0 +1,199 @@ + +import plotnine as p9 +from plotnine import ( + ggplot, aes, geom_point, labs, + geom_line, + geom_vline, + scale_x_continuous, + scale_color_brewer, + scale_alpha_manual, + scale_y_log10, + scale_size_manual, +) +p9.options.dpi = 600 + +import matplotlib.pylab as plt +import matplotlib.image as img + +import pandas as pd +import numpy as np +import tempfile +import warnings + + +class PlotterTseries: + def __init__(self, array, tstamps, plotter_options=None): + """ + + :rtype: PlotterTseries + :param dict array: dict of 3-d array of data values + :param np.ndarray tstamps: 1-d array of datetime, dimensions(t) + :param dict plotter_options: all the arguments passed to plotter + """ + + if plotter_options is None: + plotter_options = {} + + # i have to know the axes being used, even user wants default + # so i grab default axes here and hold onto it + # TODO this seems to creats open figure and unless i close this + # somehow it hangs there wasting memory. what should I do? + # shouldnt this be get instad of setdefault? + # self.fig = plotter_options.setdefault('fig', plt.figure()) + self.fig = plotter_options.get('fig', plt.figure()) + pos = plotter_options.get('pos', None) + + # cat all the data into one dataframe + self.df = self.arr2df(array, tstamps) + self.tstamps = tstamps + + # create plots + if pos: + self.ax = self.fig.add_subplot(*pos) + else: + self.ax = self.fig.add_subplot() + + # get quantiles + # TODO allow user to do whatever somehow + qpoints = [.9, .99, 1] + qlabels = ["{0:.0%}".format(_) for _ in qpoints] + nq = len(qpoints) + self.df2 = self.df.groupby(['tag', 't']).quantile(q=qpoints).unstack() + + # massage the dataframe + self.df2.columns = qlabels + self.df2 = self.df2.reset_index().melt(id_vars=['tag', 't'], + var_name='q', + value_name='v') + self.df2['q'] = pd.Categorical(self.df2['q'], categories=qlabels) + self.df2['g'] = self.df2['tag'] + self.df2['q'].astype(str) + self.nq = nq + +# this makes this plotter object unpickable +# # base tseries plot +# self.gg = ( +# ggplot(data=self.df2, mapping=aes('t', 'v' )) + +# geom_line(aes(color='tag',alpha='q', size='q', group='g')) + +# # TODO nice time axis needed +# # scale_x_continuous(breaks=np.arange(0, hrmax, 24), +# # minor_breaks = np.arange(0, hrmax, 6), +# # ) + +# scale_color_brewer(name='Model', type='qual', palette='Set1') + +# scale_alpha_manual(values=np.linspace(1, .2, num=nq)) + +# scale_size_manual(values=np.geomspace(2, .5, num=nq)) + +# scale_y_log10(limits=[1.0, self.df2['v'].max()]) + +# labs( +# # title=title, +# # x='hour', +# y='conc (ppb)', +# alpha='Quantile', +# size='Quantile', +# ) +# ) + + self.hasdata = False + + def update(self, tidx=None, footnote=None, title=None): + """ + Update plot to data at tidx + + :param int tidx: time index + :param str footnote: footnote overwrite + :param str title: title overwrite + """ + if tidx is None: + tidx = 0 + t = self.tstamps[tidx] + # print(self.df2.loc[self.df2['t'] == t, :]) + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + gg = ( + ggplot(data=self.df2, mapping=aes('t', 'v')) + + geom_line(aes(color='tag', alpha='q', size='q', group='g')) + + # # TODO nice time axis needed + # scale_x_continuous(breaks=np.arange(0, hrmax, 24), + # minor_breaks = np.arange(0, hrmax, 6), + # ) + + scale_color_brewer(name='Model', type='qual', palette='Set1') + + scale_alpha_manual(values=np.linspace(1, .2, num=self.nq)) + + scale_size_manual(values=np.geomspace(2, .5, num=self.nq)) + + scale_y_log10(limits=[1.0, self.df2['v'].max()]) + + labs( + # title=title, + # x='hour', + y='conc (ppb)', + alpha='Quantile', + size='Quantile', + ) + + geom_vline(aes(xintercept=t)) + + geom_point(data=self.df2.loc[self.df2['t'] == t, :], + mapping=aes('t', 'v', color='tag'), + inherit_aes=False) + ) + + arr = self.gg2arr(gg) +# print(arr) +# print(arr.shape) +# print(np.quantile(arr[:,:,0], q=np.linspace(0,1))) +# print(np.quantile(arr[:,:,1], q=np.linspace(0,1))) +# print(np.quantile(arr[:,:,2], q=np.linspace(0,1))) +# print(np.quantile(arr[:,:,3], q=np.linspace(0,1))) + + if self.hasdata: + # print('b') + self.im.set_data(arr) + else: + # print('a') + self.im = self.ax.imshow(arr) + self.ax.set_axis_off() + # print(self.im) + self.hasdata = True + + @staticmethod + def arr2df(arrays, ts): + """ + take arr, ts from reader, return dataframe for ggplot + + :rtype: pd.DataFrame + :param dict arrays: + :param np.ndarray ts: 1-d array of datetime, dimensions(t) + """ + + # flatten spatial dimensions + vs = {k: arr.reshape(arr.shape[0], -1) for k, arr in arrays.items()} + + # drop nans + vs = {k: v[:, ~np.isnan(v[0, :])] for k, v in vs.items()} + + n = list(vs.values())[0].shape[-1] + dfs = [] + for k, v in vs.items(): + vv = v.reshape(-1) + df = pd.DataFrame( + dict( + tag=k, + t=np.repeat(ts, n), + v=vv, + )) + dfs.append(df) + + df = pd.concat(dfs, axis=0) + return df + + @staticmethod + def gg2arr(gg): + """ + ggplot into image, and then read + + :param p9.ggplot, gg: + """ + # fig = plt.figure() # figsize=figsize) + plt.autoscale(tight=True) + # gg.save('gg.png')#, height=height, width=width, verbose=False) + # arr = img.imread('gg.png') + with tempfile.NamedTemporaryFile() as f: + gg.save(f.name, format='png') + arr = img.imread(f.name, format='png') + return arr + + diff --git a/plotter/plotter_util.py b/plotter/plotter_util.py index 1dd6902..070f0f7 100644 --- a/plotter/plotter_util.py +++ b/plotter/plotter_util.py @@ -168,10 +168,12 @@ def __init__(self, p, png_fmt, is_multi): if is_multi: for plotter in self.p.plotters: if hasattr(plotter, 'background_manager'): - plotter.background_manager.purge_bgfile_hook() + if plotter.background_manager is not None: + plotter.background_manager.purge_bgfile_hook() else: if hasattr(self.p.plotter, 'background_manager'): - self.p.plotter.background_manager.purge_bgfile_hook() + if self.p.plotter.background_manager is not None: + self.p.plotter.background_manager.purge_bgfile_hook() self.png_fmt = png_fmt