source: TOOLS/ConsoGENCMIP6/bin/plot_jobs.py @ 2433

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# this must come first
from __future__ import print_function, unicode_literals, division

# standard library imports
from argparse import ArgumentParser
import os
import os.path
import datetime as dt
import numpy as np

# Application library imports
from gencmip6 import *
from gencmip6_path import *


########################################
class BilanDict(dict):
  #---------------------------------------
  def __init__(self):
    self = {}

  #---------------------------------------
  def init_range(self, date_beg, date_end, inc=1):
    """
    """
    delta = date_end - date_beg

    (deb, fin) = (0, delta.days+1)

    dates = (date_beg + dt.timedelta(days=i)
             for i in xrange(deb, fin, inc))

    for date in dates:
      self.add_item(date)

  #---------------------------------------
  def fill_data(self, filein):
    """
    """
    try:
      data = np.genfromtxt(
        filein,
        skip_header=1,
        converters={0: string_to_date,
                    1: string_to_float,
                    2: string_to_percent,
                    3: string_to_percent},
        missing_values="nan",
      )
    except:
      print("Empty file {}".format(filein))
      exit(1)

    for date, conso, real_use, theo_use in data:
      if date in self:
        self.add_item(date, conso, real_use, theo_use)
        self[date].fill()

  #---------------------------------------
  def add_item(self, date, conso=np.nan,
               real_use=np.nan, theo_use=np.nan):
    """
    """
    self[date] = Conso(date, conso, real_use, theo_use)

  #---------------------------------------
  def theo_equation(self):
    """
    """
    (dates, theo_uses) = \
      zip(*((item.date, item.theo_use)
            for item in self.get_items_in_full_range()))

    (idx_min, idx_max) = \
        (np.nanargmin(theo_uses), np.nanargmax(theo_uses))

    x1 = dates[idx_min].timetuple().tm_yday
    x2 = dates[idx_max].timetuple().tm_yday

    y1 = theo_uses[idx_min]
    y2 = theo_uses[idx_max]

    m = np.array([
      [x1, 1.],
      [x2, 1.]
    ], dtype="float")
    n = np.array([
      y1,
      y2
    ], dtype="float")

    try:
      (a, b) = np.linalg.solve(m, n)
    except np.linalg.linalg.LinAlgError:
      (a, b) = (None, None)

    if a and b:
      for date in dates:
        self[date].theo_equ = date.timetuple().tm_yday*a + b

  #---------------------------------------
  def get_items_in_range(self, date_beg, date_end, inc=1):
    """
    """
    items = (item for item in self.itervalues()
                   if item.date >= date_beg and
                      item.date <= date_end)
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]

  #---------------------------------------
  def get_items_in_full_range(self, inc=1):
    """
    """
    items = (item for item in self.itervalues())
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]

  #---------------------------------------
  def get_items(self, inc=1):
    """
    """
    items = (item for item in self.itervalues()
                   if item.isfilled())
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]


class Conso(object):
  #---------------------------------------
  def __init__(self, date, conso=np.nan,
               real_use=np.nan, theo_use=np.nan):
    self.date     = date
    self.conso    = conso
    self.real_use = real_use
    self.theo_use = theo_use
    self.theo_equ = np.nan
    self.filled   = False

  #---------------------------------------
  def __repr__(self):
    return "{:.2f} ({:.2%})".format(self.conso, self.real_use)

  #---------------------------------------
  def isfilled(self):
    return self.filled

  #---------------------------------------
  def fill(self):
    self.filled = True


########################################
def plot_init():
  paper_size  = np.array([29.7, 21.0])
  fig, ax = plt.subplots(figsize=(paper_size/2.54))

  return fig, ax


########################################
def plot_data(ax, xcoord, xlabels, run_jobs, pen_jobs, conso_per_day):
  """
  """
  line_width = 0.

  ax.bar(xcoord, run_jobs, align="center", color="lightgreen",
         linewidth=line_width, label="jobs running")
  ax.bar(xcoord, pen_jobs, bottom=run_jobs, align="center",
         color="indianred", linewidth=line_width, label="jobs pending")

  ax.axhline(y=conso_per_day, color="blue", alpha=0.5,
             label="conso journaliÚre idéale")


########################################
def plot_config(ax, xcoord, xlabels, title):
  """
  """
  # ... Config axes ...
  # -------------------
  # 1) Range
  xmin, xmax = xcoord[0]-1, xcoord[-1]+1
  ax.set_xlim(xmin, xmax)
  # ax.set_ylim(0., ymax)

  # 2) Ticks labels
  # inc_label = 12
  # ax.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
  xcoord_maj = [x for x, l in zip(xcoord, xlabels)
                   if l.hour == 0 or l.hour == 12]
  xlabels_maj = ["{:%d-%m %Hh}".format(l) for l in xlabels
                 if l.hour == 0 or l.hour == 12]
  ax.set_xticks(xcoord, minor=True)
  ax.set_xticks(xcoord_maj, minor=False)
  ax.set_xticklabels(xlabels_maj, rotation="vertical", size="x-small")

  # 3) Define axes title
  ax.set_ylabel("cœurs", fontweight="bold")

  # ... Main title and legend ...
  # -----------------------------
  ax.set_title(title, fontweight="bold", size="large")
  ax.legend(loc="best", fontsize="x-small", frameon=False)


########################################
def get_arguments():
  parser = ArgumentParser()
  parser.add_argument("-v", "--verbose", action="store_true",
                      help="verbose mode")
  parser.add_argument("-s", "--show", action="store_true",
                      help="interactive mode")
  parser.add_argument("-d", "--dods", action="store_true",
                      help="copy output on dods")

  return parser.parse_args()


########################################
if __name__ == '__main__':

  # .. Initialization ..
  # ====================
  # ... Command line arguments ...
  # ------------------------------
  args = get_arguments()
  if args.verbose:
    print(args)

  # ... Turn interactive mode off ...
  # ---------------------------------
  if not args.show:
    import matplotlib
    matplotlib.use('Agg')

  import matplotlib.pyplot as plt
  # from matplotlib.backends.backend_pdf import PdfPages

  if not args.show:
    plt.ioff()

  # ... Files and directories ...
  # -----------------------------
  (file_param, file_utheo) = \
      get_input_files(DIR["SAVEDATA"], [OUT["PARAM"], OUT["UTHEO"]])

  img_name = "jobs"
  today = os.path.basename(file_param).strip(OUT["PARAM"])

  if args.verbose:
    print(file_param)
    print(file_utheo)
    # print(file_data)
    print(today)
    print(img_name)

  file_list = glob.glob(os.path.join(DIR["SAVEDATA"],
                                     "OUT_JOBS_PENDING_*"))

  fg_first = True

  for filein in sorted(file_list):
    try:
      data = np.genfromtxt(
        filein,
        skip_header=1,
        converters={
          0: string_to_datetime,
          1: int,
          2: int,
        },
        missing_values="nan",
      )
    except Exception as rc:
      print("Problem with file {} :\n{}".format(filein, rc))
      exit(1)

    if fg_first:
      fg_first = False
      full_data = data
    else:
      full_data = np.append(full_data, values=data)

  # .. Get project info ..
  # ======================
  gencmip6 = Project()
  gencmip6.fill_data(file_param)
  gencmip6.get_date_init(file_utheo)

  # .. Fill in conso data ..
  # ========================
  # # ... Initialization ...
  # # ----------------------
  # bilan = BilanDict()
  # bilan.init_range(gencmip6.date_init, gencmip6.deadline)
  # # ... Extract data from file ...
  # # ------------------------------
  # bilan.fill_data(file_data)
  # # ... Compute theoratical use from known data  ...
  # # ------------------------------------------------
  # bilan.theo_equation()

  # # .. Extract data depending on C.L. arguments ..
  # # ==============================================
  # if args.full:
  #   selected_items = bilan.get_items_in_full_range(args.inc)
  # elif args.range:
  #   selected_items = bilan.get_items_in_range(
  #     args.range[0], args.range[1], args.inc
  #   )
  # else:
  #   selected_items = bilan.get_items(args.inc)

  # .. Compute data to be plotted ..
  # ================================
  # nb_items = len(selected_items)
  nb_items = full_data.size

  xcoord   = np.linspace(1, nb_items, num=nb_items)
  if full_data.size > 1:
    xlabels  = np.array([date for date, run, pen in full_data])
    run_jobs = np.array([run for date, run, pen in full_data],
                          dtype=int)
    pen_jobs = np.array([pen for date, run, pen in full_data],
                          dtype=int)
  else:
    xlabels , run_jobs, pen_jobs = full_data.tolist()
    xlabels = list(xlabels)
    run_jobs = list(run_jobs )
    pen_jobs = list(pen_jobs )

  # cumul     = np.array([item.conso for item in selected_items],
  #                       dtype=float)
  # consos    = []
  # consos.append(cumul[0])
  # consos[1:nb_items] = cumul[1:nb_items] - cumul[0:nb_items-1]
  # consos    = np.array(consos, dtype=float)

  conso_per_day = gencmip6.alloc / (gencmip6.days * 24.)

  # theo_uses = np.array([100.*item.theo_use for item in selected_items],
  #                      dtype=float)
  # real_uses = np.array([100.*item.real_use for item in selected_items],
  #                      dtype=float)
  # theo_equs = np.array([100.*item.theo_equ for item in selected_items],
  #                      dtype=float)

  # .. Plot stuff ..
  # ================
  # ... Initialize figure ...
  # -------------------------
  (fig, ax) = plot_init()

  # ... Plot data ...
  # -----------------
  plot_data(ax, xcoord, xlabels, run_jobs, pen_jobs, conso_per_day)

  # # ... Tweak figure ...
  # # --------------------
  # if args.max:
  #   ymax = gencmip6.alloc
  # else:
  #   ymax = np.nanmax(consos) + np.nanmax(consos)*.1

  title = "Consommation {}\n({:%d/%m/%Y} - {:%d/%m/%Y})".format(
    gencmip6.project.upper(),
    gencmip6.date_init,
    gencmip6.deadline
  )

  plot_config(ax, xcoord, xlabels, title)

  # ... Save figure ...
  # -------------------
  img_in  = os.path.join(DIR["PLOT"], "{}.pdf".format(img_name))
  img_out = os.path.join(DIR["SAVEPLOT"],
                         "{}_{}.pdf".format(img_name, today))

  plot_save(img_in, img_out, title)

  # ... Publish figure on dods ...
  # ------------------------------
  if args.dods:
    dods_cp(img_in)

  if args.show:
    plt.show()

  exit(0)