plot_ts

#!/usr/bin/env python

import  netCDF4 as nc
import matplotlib.pyplot as mpl
import sys
from argparse import ArgumentParser
# from scipy import stats
import flo_utils as fu
import seaborn as sns
import numpy as np
from smoother import smooth
from datetime import datetime # , timedelta
import re
import json

js_config = json.load(open(fu.home+"/MPI/offsets.json"))
sns.set()
#plt.xkcd()
#mpl.rcParams['figure.figsize']=[3,2]
def get_array(filename, varname, skip = 1):
    fu.debug_cerr("reading %s from %s "%(varname,filename))
    infile = nc.Dataset(filename, "r")
    fu.debug_cerr(infile.variables[varname].shape)
    var = np.squeeze(infile.variables[varname][:])
    attrs = {x : infile.variables[varname].getncattr(x) for x in infile.variables[varname].ncattrs()}
    attrs ["file_name"] = filename
    attrs ["var_name"] = varname
    return var[::skip],attrs


def light(color):
    return [.75*x + .25 for x in color]
def dark(color):
    return [.75 * x  for x in color]


def ensmean (xvecs, yvecs, time_scaling):
  etmin=max([x[0] for x in xvecs])
  etmax=min([x[-1] for x in xvecs])
  dest = np.arange(etmin, etmax+1, time_scaling)
  print("INTERPOLATION TARGET ",(etmin, etmax+1, time_scaling))
  target = np.zeros(len(dest))
  for (n, x )  in enumerate(xvecs) :
    print(x.shape)
    yin = yvecs[n][x <= etmax ]
    xin = x[x <= etmax ]
    yin = yin[xin >= etmin ]
    xin = xin[xin >= etmin ]
    interpolated = np.interp(dest, xin, yin)
    target = target + interpolated
  target = target / len (xvecs)
  return dest, target

def get_color(name, num):
  print((name, num))
  cmap=sns.color_palette()
  linecol = js_config.get("colors",{}).get(name, num)
  print (linecol)
  print((type (linecol)))
  if type (linecol) is int or type(linecol) is np.int64:
    print((linecol, len(cmap)))
    linecol =cmap[linecol]
  return linecol

def plot_vars(xvecs, yvecs, xattrs, yattrs, plot_opts = {}):
    fig = mpl.figure(figsize=(16/2.54,9/2.54))
#    fig = mpl.figure()
    mpl.hold("on")
    def cleanup_name(filename):
        return re.sub("_"," ", re.sub(".nc$", "", filename))
    if plot_opts.get("names", False):
      labels = plot_opts.get("names")
    else:
      labels = [y.get("label", cleanup_name(y["file_name"])) for y in yattrs]
    ax=mpl.subplot(1,1,1)
    cmap=sns.husl_palette(max(len(xvecs),6),s=.6) # sns.color_palette("Set2", 8)
    sns.set_palette(cmap)
    print((len(cmap)))
    namedict = {}
    patterndict = {}
    dashdict = js_config.get("dashes",{})
#    colordict={"I01":(cmap[1]), "H87":(cmap[3]), "I02":(cmap[4]), "I09":(cmap[5])}
    mpl.subplots_adjust(bottom=0.2, left=.15)

    yscale=plot_opts.get('yscale', 1.)
    xscale=plot_opts.get('xscale', 1.)
    addto_x=plot_opts.get('addto_x', 0.)
    addto_y=plot_opts.get('addto_y', 0.)
    plots = [ mpl.plot(xvec * xscale + addto_x, yvec * yscale+addto_y, patterndict.get(lb[:5],'-'),  label=namedict.get(lb,lb), color=get_color(lb,i%len(cmap))) for (xvec,yvec,lb,i) in zip(xvecs,yvecs,labels,np.arange(len(xvecs))) ]
    dashes = [dashdict.get(x, False) for x in labels]
    for (p, dash) in zip(plots, dashes):
        if(dash):
            p[0].set_dashes(dash)
#    for nnn,xxx in enumerate(((plots[0]),plots[2], plots[3])):
#      xxx[0].set_dashes([15,nnn*2+4])
    if (plot_opts.get("ensmean", False)):
      x, y  = ensmean(xvecs, yvecs, plot_opts.get("multime", 1.))
      plots.append ( mpl.plot(x, y*yscale,  '-',  label="Mean", color='k' , linewidth=3))
    if "no_legend" in plot_opts:
        if plot_opts["no_legend"] == False :
            legend = mpl.legend(frameon=1)
            frame = legend.get_frame()
            frame.set_facecolor('white')
            frame.set_edgecolor('white')

    else:
        legend = mpl.legend(frameon=1)
        frame = legend.get_frame()
        frame.set_facecolor('white')
        frame.set_edgecolor('white')

    #    mpl.axvline(0,zorder=-999, color="gray")
    if "xlabel" in list(plot_opts.keys()):
        mpl.xlabel(plot_opts["xlabel"])
    if "ylabel" in list(plot_opts.keys()):
        mpl.ylabel(plot_opts["ylabel"])
    if "title" in list(plot_opts.keys()):
        mpl.title(plot_opts["title"])
    if "xlim" in list(plot_opts.keys()):
        mpl.xlim( plot_opts["xlim"] )
    if "ylim" in list(plot_opts.keys()):
        mpl.ylim(plot_opts["ylim"])
    if "tick_interval" in list(plot_opts.keys()):
        ti=plot_opts["tick_interval"]
        mpl.xticks=np.arange(ti[0], ti[1],ti[2])
    if plot_opts.get("despine", False):
      sns.despine()
    if "output" in list(plot_opts.keys()):
        mpl.savefig(plot_opts["output"])
    return fig

def filter(vecs, plot_opts={}):
    if "smooth_window_length" in list(plot_opts.keys()):
        npoints = plot_opts["smooth_window_length"]
    else:
        npoints = 0
    if "smooth_filter" in list(plot_opts.keys()):
        window = plot_opts["smooth_filter"]
    else:
        window = "flat"
    if npoints:
        vecs = [ smooth(x,npoints, window) for x in vecs]
    return vecs

def show():
    mpl.show()

def sort_labels(xattr, yattr):
    xlabel = ""
    ylabel = ""
    title = ""
    if "long_name" in list(xattr.keys()):
        xlabel = xattr["long_name"]
        if "units" in list(xattr.keys()):
            xlabel = xlabel + " in %s"%xattr["units"]
    if "long_name" in list(yattr.keys()):
        ylabel = yattr["long_name"]
        if "units" in list(yattr.keys()):
            ylabel = ylabel + " in %s"%yattr["units"]
    return {"xlabel"  : xlabel , "ylabel" : ylabel , "title" : title}



def parse_args():
  parser = ArgumentParser()
  parser.description = "Scatterplot two variables from a set of files"
  parser.add_argument("FILES", nargs='*')
  parser.add_argument("-v", "--verbose",
                    help='''Be verbose''', action="store_true")
  parser.add_argument("-x", "--xvar",
                    help='''xvar''', default="time")
  parser.add_argument("-y", "--yvar",
                    help='''yvar''', default="ivol")
  parser.add_argument("-o", "--output",
                    help='''output file to save image to''', default=None)
  parser.add_argument("-X", "--x_file",
                    help='''file for x-var''', default=None)
  parser.add_argument("-Y", "--y_file",
                    help='''file for y-var''', default=None)
  parser.add_argument("--skip",
                    help='''skipping when reading''', default=1, type = int)

  parser.add_argument("-t", "--title",
                    help='''plot title''', default = None)
  parser.add_argument("--xlabel",
                    help='''plot x-Axis label''', default = None)
  parser.add_argument("--ylabel",
                    help='''plot y-Axis label''', default = None)
  parser.add_argument("--mask_thk",
                    help='''mask with thickness''', action = "store_true")
  parser.add_argument("--xlim",
                    help='''x limits''', type = float , nargs = 2)
  parser.add_argument("--ylim",
                    help='''y limits''', type = float , nargs = 2)
  parser.add_argument("-s", "--smooth_window_length",
                    help='''smooth function with window length N''', type = int)
  parser.add_argument("--smooth_filter",
                    help='''smooth filter function''', default = None)
  parser.add_argument("--no_legend",
                    help='''suppress legend''', action = "store_true")
  parser.add_argument("--white_grid",
                    help='''white grid''', action = "store_true")
  parser.add_argument("--white",
                    help='''white backgrund''', action = "store_true")
  parser.add_argument("--ticks",
                    help='''ticks and white backgrund''', action = "store_true")
  parser.add_argument("--despine",
                    help='''call despine''', action = "store_true")
  parser.add_argument("--tick_interval",
                    help='''draw vertical lines at years specified by start range step ''', nargs=3, type=float )
  parser.add_argument("--yscale",
                    help='''scale y axis by factor ''', type=float , default=1.)
  parser.add_argument("--xscale",
                    help='''scale x axis by factor ''', type=float , default=1.)
  parser.add_argument("--addto_x",
                    help='''add to x values ''', type=float , default=0.)
  parser.add_argument("--addto_y",
                    help='''add to y values ''', type=float , default=0.)
  parser.add_argument("--multime",
                    help='''multiply raw x axis by factor ''', type=float , default=1.)
  parser.add_argument("--font_scale",
                    help='''scale fontsize by factor ''', type=float , default=1.)
  parser.add_argument("--names",
                    help='''Comma separated list of experiment names''')
  parser.add_argument("--ensmean", help='''Add ensembel mean plot''', action = "store_true")
  options = parser.parse_args()
  options_dict = vars(options)
  if options_dict.get("names", False):
    print(options_dict["names"])
    options_dict["names"]= options_dict["names"].split(",")
  return options_dict

def convert_year(year, cal):
    if cal:
        return     nc.date2num(datetime (int(year),6,30),units="days since 0001-01-01 00:00:00", calendar=cal)
    else:
        return year


def specials(figure, xattrs, plot_opts, ax):
    attr = xattrs[0]
    years = list(range(1980,2101,20))
    if "tick_interval" in list(plot_opts.keys()):
        ti=plot_opts["tick_interval"]
        years=np.arange(ti[0], ti[1],ti[2])
    print(years)
    mpl.xticks(years )


def process_xvecs(xvecs, xattrs, plot_opts):
  if plot_opts["xvar"] == "time" or plot_opts["xvar"] == "t":
    attr = xattrs[0]
    if "calendar" in list(attr.keys()) and attr["calendar"] and not "years" in xattrs[0]["units"]:
      print((xattrs[0]["units"]))
      xvecs = [ nc.date2num(nc.num2date(vec , units=attr["units"] , calendar=attr["calendar"]), units="days since 0000-01-01 00:00:00", calendar=attr["calendar"] ) / 360. for (vec, attr) in zip (xvecs, xattrs) ]
      xvecs = [ x * plot_opts["multime"] for x in xvecs ]
    if plot_opts.get ('offsets', False):
        offsets=plot_opts.get('offsets')
    else:
      if plot_opts.get("names", False):
        names = plot_opts.get("names")
        offsets_dict = js_config.get("offsets", {})
        print(offsets_dict)
        offsets = [ offsets_dict.get(x, 0) for x in names]
        print(offsets)
        time_scaling_dict = js_config.get("time_scaling", {})
        time_scaling = [ time_scaling_dict.get(x, 1) for x in names]
        for x in range (len(xvecs)):
          # if offsets[x]:
          #   data_labels[x] = data_labels [x] + " - %i years"%(round( offsets[x]))
          # if time_scaling[x] != 1 :
          #     data_labels[x] = "(%s) * %d"%(data_labels[x],time_scaling[x])
          xvecs[x] = (xvecs[x]-offsets[x]) * time_scaling[x]
  return xvecs


def main(argv):
  options = parse_args()
  if options.get("verbose", False):
    fu.debug = True
    fu.debug_cerr(dir(options))
  xvecs = [ ]
  yvecs = [ ]
  xattrs = [ ]
  yattrs = [ ]
  if options.get("x_file", False) and not options.get("y_file"):
      cerr("Need y-file to match x-file")
      sys.exit(1)
  if options.get("x_file", False) and options.get("y_file", False):
    (xvec, xattr) = get_array(options.get("x_file", False), options.get("xvar"), skip = options.get("skip"))
    (yvec, yattr) = get_array(options.get("y_file", False), options.get("yvar"), skip = options.get("skip"))

    options["FILES"] = []
    xvecs.append(xvec)
    yvecs.append(yvec)
  if options.get("y_file", False) and not options.get("x_file", False):
    (yvec, yattr) = get_array(options.get("y_file"), options.get("yvar"), skip = options.get("skip"))
    (xvec, xattr) = (np.arange(len(yvec)), {})
    options["FILES"] = []
    xvecs.append(xvec)
    yvecs.append(yvec)
  for filename in options.get("FILES", False):
    (xvec, xattr) = get_array(filename, options.get("xvar"), skip = options.get("skip"))
    (yvec, yattr) = get_array(filename, options.get("yvar"), skip = options.get("skip"))
    fu.debug_cerr (xvec.size)
    xvecs.append(xvec)
    yvecs.append(yvec)
    xattrs.append(xattr)
    yattrs.append(yattr)

  plot_opts = sort_labels(xattr, yattr)
  for x in list(options.keys()):
    if not options[x] is  None :
      plot_opts[x] = options[x]

  xvecs = process_xvecs(xvecs, xattrs,  plot_opts)

  if options.get("smooth_window_length", False):
      xvecs=list(filter(xvecs, plot_opts))
      yvecs=list(filter(yvecs, plot_opts))

  if options.get("white_grid", False):
    sns.set(style="whitegrid") # , context="poster
  if options.get("white", False):
    sns.set(style="white") # , context="poster
  if options.get("ticks", False):
    sns.set(style="ticks") # , context="poster
  sns.set(style=sns.axes_style(), font_scale=options.get("font_scale", 1.))

  fig = plot_vars(xvecs, yvecs, xattrs, yattrs, plot_opts)
  if not options.get("output", False):
    mpl.show()


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