import datetime

import warnings

import numpy as np

import matplotlib

from matplotlib import colors

from matplotlib.patches import Patch

from matplotlib.colors import LogNorm

from matplotlib import rc_context

from hapiplot.plot.datetick import datetick

def heatmap(x, y, z, **kwargs):

"""Plot a heatmap using pcolormesh and do typical configuration.

plt, fig, canvas, ax, im, cb = heatmap(x, y, z, **kwargs)

x and y are bin centers, edges, or ranges; x is associated with the

columns of z and y is associated with the rows of z.

If len(x) == z.shape[1], x is assumed to correspond to the center of the

rectangle and the labels are placed at the center of the bin. The bin

widths are x[i+1]-x[i]. If the spacing of x values is not uniform,

the bin width of x[-1] is assumed to be equal to that of x[-2].

The same logic is applied to y if len(y) == z.shape[0].

If len(x) == z.shape[1] + 1, x is assumed to correspond to edges of the

rectangle and the labels are placed at x[i] and the bin widths are x[i+1]-x[i].

The same logic is applied to y if len(y) == z.shape[0] + 1.

If x.shape = (z.shape[1], 2), each row gives the range of the rectangle

associated with the columns of z and a "No data" color is used for

ranges that do not have data.

The same logic is applied to y if y.shape = (z.shape[0], 2).

kwargs:

Axes

----

* xlabel

* ylabel

* title

* logx

* logy

Colorbar

------

* zlabel

* ztitle

* logz

* clim - Ignore if ...

* cmap

* cmap.name (ignored if cmap given)

* cmap.numcolors (ignored if cmap given)

* cmap.clim

Logarithm

---------

* logx - All x values must be either >= 0 or <= 0.

* logy - All y values must be either >= 0 or <= 0.

* logz - All z values must be either >= 0 or <= 0. The zero values rectangles are colored using logz0color.

* logz0.legend

* log.auto

Tile edges, gaps, and NaN

-------------------------

* edgecolor - Edge color of tile border

* logz0.color

* logz0.alpha

* logz0.hatch

* logz0.legend - Show legend entry for logz0

* gap.color

* gap.alpha

* gap.hatch

* gap.hatch.color

* gap.legend - Show legend entry for gaps (True by default and if gaps)

* nan.color

* nan.alpha

* nan.hatch

* nan.hatch.color

* nan.legend - Show legend entry for nans (True by default and if NaNs)

"""

###########################################################################

def warning(message):

print("\x1b[31mheatmap() warning: \x1b[0m" + message)

def calcEdges(y, coord):

"""Caclulate bin ranges given bin centers."""

# TODO: Set bin width to 1 if y is not datetime.

# TODO: If datetime, guess cadence by inspecting and warn?

if len(y) == 1:

# Put tick at center of bin; make bin have width = 1.

y = np.array([y[0]-0.5,y[0]+0.5])

else:

# y are bin centers

dy = np.diff(y)

dyu = np.unique(dy)

if len(dyu) > 1:

if coord == 'y':

warning('Only bin centers given for y and bin separation distance is not constant. ' + \

'Bin width assumed based on separation distance and data pickers will not work properly.')

else:

warning('Only bin centers given for x and bin separation distance is not constant. ' + \

'Bin width assumed based on separation distance and data pickers will not work properly.')

y = np.append(y, y[-1] + dy[-1])

else:

y = np.append(y, y[-1] + dy[0])

y = y-dyu[0]/2

return y

def calcGaps(y, z, coord):

if coord == 'x':

z = np.transpose(z)

havegaps = False

Igaps = []

ynew = []

for k in range(0, len(y)-1):

ynew.append(y[k][0])

if y[k][1] != y[k+1][0]:

havegaps = True

break

if len(y) == 1:

ynew.append(y[0][0])

ynew.append(y[0][1])

y = ynew

elif not havegaps:

# Set top-most two bin edges.

ynew.append(y[k+1][0])

ynew.append(y[k+1][1])

y = ynew

else:

# Insert NaN rows where gaps

znew = z[0,:]

nanrow = np.nan*z[0,:]

ynew = []

ynew.append(y[0][0])

if y[0][1] != y[1][0]: # Gap between bins

ynew.append(y[0][1])

znew = np.vstack((znew, nanrow))

Igaps.append(1)

for k in range(1, len(y)-1):

ynew.append(y[k][0])

znew = np.vstack((znew, z[k,:]))

if y[k][1] != y[k+1][0]: # Gap between bins

ynew.append(y[k][1])

#print('inserting nan row or column')

znew = np.vstack((znew, nanrow))

Igaps.append(znew.shape[0]-1)

# Set top-most two bin edges

ynew.append(y[k+1][0])

ynew.append(y[k+1][1])

znew = np.vstack((znew, z[k+1,:]))

z = znew

Igaps = np.asarray(Igaps, dtype=np.int32)

y = np.array(ynew)

if coord == 'x':

z = np.transpose(z)

return y, z, Igaps

def adjustCenters(y, yc):

# Adjusts centers and their labels if nonuniform diff(y)

dy = np.diff(y)

dyu = np.unique(dy)

if len(dyu) > 1:

# If centers not uniform, they were used as lower edges.

# Need to adjust center location for ticks and labels.

ycl = np.copy(y) # y center labels

yc = np.copy(y[0:-1])

for i in range(0, y.shape[0]-1):

yc[i] = y[i] + (y[i+1]-y[i])/2

else:

# No adjustment needed. Indicate this with empty ycl.

ycl = np.array([])

return yc, ycl

def boundaryInfo(x, coord):

xlabels = None

xcl = None

N = Nx

if coord == "y":

N = Ny

if len(x.shape) == 1 and len(x) == N:

# Centers given. Calculate edges.

if iscategorical(x):

xlabels = x

x = np.linspace(0, x.shape[0]-1, x.shape[0], dtype='int32')

xedges = False

xc = x

x = calcEdges(x, coord)

xc, xcl = adjustCenters(x, xc)

else:

xc = np.array([])

xedges = True

return x, xc, xedges, xcl, xlabels

def iscategorical(x):

return isinstance(x[0], np.character)

def categoryinfo(x):

if len(x.shape) > 1:

raise ValueError('If x contains characters, it must have one column or one row.')

if not len(x) == Nx:

raise ValueError('If x contains characters, number of elements must match number of rows in z.')

xcategories = x

x = np.linspace(0, x.shape[0]-1, x.shape[0], dtype='int32')

return xcategories, x

def allint(x):

Ig = ~np.isnan(x)

if np.all(np.equal(x[Ig], np.int32(x[Ig]))):

return True

else:

return False

opts = {

'logging': False,

'title': '',

'xlabel': '',

'ylabel': '',

'logx': False,

'logy': False,

'backend': 'default',

'returnimage': False,

'transparent': False,

'ztitle': '',

'zlabel': '',

'logz': False,

'edgecolor': None,

'cmap': None,

'cmap.numcolors': 32,

'cmap.name': 'viridis',

'cmap.clim': None,

'nan.color': [0.95,0.95,0.95],

'nan.alpha': 1,

'nan.hatch': '',

'nan.hatch.color': [0.1,0.1,0.1] ,

'nan.legend': True,

'gap.color': [0,0,0],

'gap.alpha': 1,

'gap.hatch': '',

'gap.hatch.color': [0.95,0.95,0.95],

'gap.legend': True,

'logz0.color': [1,1,1],

'logz0.alpha': 1,

'logz0.hatch': '',

'logz0.hatch.color': [0.95,0.95,0.95],

'logz0.legend': True

}

for key, value in kwargs.items():

if key in opts:

opts[key] = value

else:

warnings.warn('Warning: Ignoring invalid keyword option "%s".' % key, SyntaxWarning)

if opts['returnimage']:

# When returnimage=True, the Matplotlib OO API is used b/c it is thread safe.

# Otherwise, the pyplot API is used. Ideally would always use the OO API,

# but this has problems with notebooks and showing figure when executing

# a script from the command line.

# TODO: Document issues.

#

# API differences description:

# https://www.saltycrane.com/blog/2007/01/how-to-use-pylab-api-vs-matplotlib-api_3134/

from matplotlib.figure import Figure

from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

else:

from matplotlib import pyplot as plt

if opts['logging']: print("heatmap(): Using Matplotlib back-end " + matplotlib.get_backend())

from matplotlib.ticker import MaxNLocator

if not opts['cmap.name'] in matplotlib.pyplot.colormaps():

warning('colormap name "'

+ opts['cmap.name']

+ '" is not in list of known names: '

+ str(matplotlib.pyplot.colormaps())

+ ". Using 'viridis'.")

opts['cmap.name'] = 'viridis'

if not opts['cmap']:

opts['cmap'] = matplotlib.pyplot.get_cmap(\

opts['cmap.name'], opts['cmap.numcolors'])

if opts['returnimage']:

fig = Figure()

# Calling FigureCanvas() attaches canvas to fig which is used later.

FigureCanvas(fig)

ax = fig.add_subplot(111)

else:

fig, ax = plt.subplots()

if opts['logging']:

print("timeseries(): Using Matplotlib back-end " + matplotlib.get_backend())

if type(x) == list:

x = np.array(x)

if type(y) == list:

y = np.array(y)

if type(z) == list:

z = np.array(z)

# Number of x values (columns)

Nx = z.shape[1] if len(z.shape) == 2 else 1

# Number of y values (rows)

Ny = z.shape[0] if len(z.shape) == 2 else 1

# If y is a matrix, assume it has two columns.

# First column is lower edge, second is upper edge.

# If given a two-row matrix, transpose.

if len(y.shape) > 1: # Bin edges given

if y.shape[1] != 2 and y.shape[0] == 2:

warning('If y is a matrix, it should have two columns. Two rows found. Transposing.')

y = np.transpose(y)

if y.shape[0] !=2 and y.shape[1] != 2:

raise ValueError('If y is a matrix, it must have two columns or two rows.')

if len(x.shape) > 1: # Bin edges given

if x.shape[1] != 2 and x.shape[0] == 2:

warning('If x is a matrix, it should have two columns. Two rows found. Transposing.')

x = np.transpose(x)

if x.shape[0] != 2 and x.shape[1] != 2:

raise ValueError('If x is a matrix, it must have two columns or two rows.')

if x.ndim == 1 and not (len(x) == Nx or len(x) == Nx+1):

raise ValueError('Required: len(x) == z.shape[1] or len(x) == z.shape[1] + 1.')

if y.ndim == 1 and not (len(y) == Ny or len(y) == Ny+1):

raise ValueError('Required: len(y) == z.shape[0] or len(y) == z.shape[0] + 1.')

inan = np.where(np.isnan(z))

havenans = False

allnan = False

if np.any(inan) > 0:

havenans = True

if np.all(np.isnan(z)):

allnan = True

categoricalx = iscategorical(x)

x, xc, xedges, xcl, xlabels = boundaryInfo(x,'x')

categoricaly = iscategorical(y)

y, yc, yedges, ycl, ylabels = boundaryInfo(y,'y')

xgaps = np.array([], dtype=np.int32)

if len(x.shape) == 2: # x is an matrix

x, z, xgaps = calcGaps(x, z, 'x')

ygaps = np.array([], dtype=np.int32)

if len(y.shape) == 2: # y is an matrix

y, z, ygaps = calcGaps(y, z, 'y')

legendh = []

havegaps = False

if len(xgaps) > 0 or len(ygaps) > 0:

havegaps = True

cmapg = colors.LinearSegmentedColormap.from_list('gap',[opts['gap.color'],opts['gap.color']],2)

zg = np.nan*np.copy(z)

if len(xgaps) > 0:

zg[:,xgaps] = 1

if len(ygaps) > 0:

zg[ygaps,:] = 1

if opts['gap.hatch'] == '':

im = ax.pcolormesh(x, y, zg, cmap=cmapg)

else:

# pcolormesh does not support hatch

# TODO: Must set hatch.color through rc params and context manager.

# opts['nan.hatch.color']

# https://stackoverflow.com/a/42672782/1491619

with rc_context(rc={'hatch.color': opts['gap.hatch.color']}):

im = ax.pcolor(x, y, zg, cmap=cmapg, hatch=opts['gap.hatch'])

if opts['gap.legend']:

with rc_context(rc={'hatch.color': opts['gap.hatch.color']}):

legendh.append(Patch(facecolor=opts['gap.color'],

hatch=opts['gap.hatch']+opts['gap.hatch'],

edgecolor=opts['edgecolor'], label='No data'))

if havenans:

edgecolor = opts['edgecolor']

if allnan and edgecolor is None:

edgecolor = 'k'

cmapn = colors.LinearSegmentedColormap.from_list('nan',[opts['nan.color'],opts['nan.color']],2)

zn = np.nan*np.copy(z)

if havegaps:

inan = np.where(np.logical_and(np.isnan(z),zg != 1))

zn[inan] = 1

if opts['nan.hatch'] == '':

im = ax.pcolormesh(x, y, zn, edgecolor=edgecolor, cmap=cmapn)

else:

# Must set hatch.color through rc params and context manager.

with rc_context(rc={'hatch.color': opts['nan.hatch.color']}):

im = ax.pcolor(x, y, zn, cmap=cmapn, edgecolor=edgecolor, hatch=opts['nan.hatch'])

if opts['nan.legend']:

with rc_context(rc={'hatch.color': opts['nan.hatch.color']}):

legendh.append(Patch(facecolor=opts['nan.color'],

hatch=opts['nan.hatch']+opts['nan.hatch'],

edgecolor=edgecolor, label='NaN'))

# TODO: Handle case where > 10.

# Label every Nth, etc. as needed

if xedges and x.size <= 10:

ax.set_xticks(x)

if xc.size > 0 and xc.size <= 10:

ax.set_xticks(xc)

if len(xcl) > 0:

# Relabel x-ticks b/c nonuniform center spacing.

ax.set_xticklabels(xcl[0:-1])

#for spine in ax.spines.values(): spine.set_edgecolor(None)

if yedges and y.size <= 10:

ax.set_yticks(y)

if yc.size > 0 and yc.size <= 10:

ax.set_yticks(yc)

if len(ycl) > 0:

# Relabel y-ticks b/c nonuniform center spacing.

ax.set_yticklabels(ycl[0:-1])

def setTicks(labels, coord):

categories, _ = categoryinfo(labels)

# TODO: This will create too many ticks if # of categories is large

if coord == 'x':

ax.set_xticklabels(categories)

ticks = ax.get_xticks()

ticklines = ax.get_xticklines()

ax.set_xticks(list(ticks) + [-0.5] + list(ticks+0.5))

else:

ax.set_yticklabels(categories)

ticks = ax.get_yticks()

ticklines = ax.get_yticklines()

ax.set_yticks(list(ticks) + [-0.5] + list(ticks+0.5))

k = 0

for l in ticklines:

if k < 2*len(categories):

l.set_markeredgewidth(0)

k = k+1

if categoricalx:

setTicks(xlabels, 'x')

if categoricaly:

setTicks(ylabels, 'y')

# TODO: categoricalz not implemented.

categoricalz = iscategorical(z)

if categoricalz:

zcategories, z = categoryinfo(z)

#ax.set_ylim(y[0], y[-1])

if opts['xlabel']:

ax.set_xlabel(opts['xlabel'])

if opts['logx']:

# TODO: Flip sign if needed.

ax.set_xscale('log')

if opts['ylabel']:

ax.set_ylabel(opts['ylabel'])

if opts['logy']:

# TODO: Flip sign if needed.

ax.set_yscale('log')

if opts['title']:

ax.set_title(opts['title'], fontsize=10)

if False and allint(x):

xa = ax.get_yaxis()

xa.set_major_locator(MaxNLocator(integer=True, steps=[1, 2, 4, 5, 10]))

if False and allint(y):

ya = ax.get_yaxis()

ya.set_major_locator(MaxNLocator(integer=True, steps=[1, 2, 4, 5, 10]))

# NaN, logz0, and gap legend

if len(legendh) > 0:

fig.legend(frameon=False, borderaxespad=0.25, borderpad=0.15,

handles=legendh, loc='upper right',

fontsize='x-small', ncol=len(legendh))

if allnan:

cb = None

else:

# Colorbar

zc = np.array([])

allintz = allint(z)

if not 'cmap' in kwargs and allintz:

Ig = ~np.isnan(z)

zc = np.unique(z[Ig])

if len(zc) == 1:

nc = len(zc)

else:

nc = np.int(zc[-1]-zc[0] + 1)

nc = np.min([1024, nc])

if 'cmap.numcolors' in kwargs:

if opts['cmap.numcolors'] != nc:

warning('Overriding requested number of colors. Using number of colors = ' + str(nc))

cmap_name = opts['cmap.name']

if len(Ig) == 2 and not 'cmap.name' in kwargs:

# Binary. Plot black and white.

cmap_name = 'gray'

opts['cmap'] = matplotlib.pyplot.get_cmap(cmap_name, nc)

zmin = np.nanmin(z)

zmax = np.nanmax(z)

if zmin < 0 and zmax > 0:

warning('Colorbar cannot have log scale when all values do not have the same sign.')

havelogz0 = False

if not opts['logz']:

im = ax.pcolormesh(x, y, z, cmap=opts['cmap'], edgecolor=opts['edgecolor'])

else:

# Log scale emits warning if data have NaNs.

warnings.filterwarnings(action='ignore',

message='invalid value encountered in less_equal')

flipsign = False

if zmin < 0 and zmax <= 0 and opts['logz']:

z = -z;

flipsign = True

if zmin == 0:

#warning('Log scale for z requested but min(z) = 0.')

havelogz0 = True

logz0idx = np.where(z==0)

z[z == 0] = np.nan

havelogz0 = True

zmin = np.nanmin(z)

norm = LogNorm(vmin=zmin, vmax=zmax)

im = ax.pcolormesh(x, y, z, cmap=opts['cmap'], norm=norm,

edgecolor=opts['edgecolor'])

if havelogz0:

cmapz = colors.LinearSegmentedColormap.from_list('logz0',[opts['logz0.color'],opts['logz0.color']],2)

z = np.nan*z

z[logz0idx] = 1

ax.pcolormesh(x, y, z, cmap=cmapz)

if opts['logz0.legend']:

legendh.append(Patch(facecolor=opts['logz0.color'],

edgecolor='k', label='0.0'))

cb = fig.colorbar(im, ax=ax, pad=0.01)

if allintz:

# Put tick label at center of color patch.

if not opts['cmap.clim']:

im.set_clim(zc[0]-0.5, zc[-1] + 0.5)

za = cb.ax.get_yaxis()

za.set_major_locator(MaxNLocator(integer=True, min_n_ticks=1, steps=[1, 2, 4, 5, 10]))

if opts['cmap.clim']:

im.set_clim(opts['cmap.clim'])

zt = cb.get_ticks()

# How to label zmin and zmax without overlap of default labels?

# Extend clim to next interval? But then one may think actual values in

# extended range. For now, just warn.

if zt[0] != zmin:

#warnings.warn('Lower z limit not labeled.')

zt = np.concatenate(([zmin],zt))

#cb.set_ticks(zt)

if zt[-1] != zmax:

#warnings.warn('Upper z limit not labeled.')

zt = np.concatenate((zt,[zmax]))

#cb.set_ticks(zt)

cb.set_label(opts['zlabel'], fontsize=10) # On side

cb.ax.set_title(opts['ztitle'], fontsize=10) # On top

if opts['logz'] and flipsign:

# Put negative sign in front of numbers on colorbar

zlabels = cb.ax.get_yticklabels()

for i in range(0, len(zlabels)):

text = zlabels[i].get_text()

if text != '':

zlabels[i].set_text('-'+text)

cb.ax.set_yticklabels(zlabels)

if isinstance(x[0], datetime.datetime):

datetick('x', axes=ax)

if isinstance(y[0], datetime.datetime):

datetick('y', axes=ax)

# The following two conditions will be replaced by more general

# code that calculates ax and cb position and dimensions based on

# computed text box sizes. Note that plt.tight_subplot() and

# bbox_inches='tight' have many issues that indicate they probably will

# require significant work-arounds to be developed:

# https://github.com/matplotlib/matplotlib/issues/12355/

# https://stackoverflow.com/questions/48128546/why-is-the-legend-not-present-in-the-generated-image-if-i-use-tight-for-bbox-i

# https://stackoverflow.com/questions/10101700/moving-matplotlib-legend-outside-of-the-axis-makes-it-cutoff-by-the-figure-box

ax_w = 0.8

if opts['zlabel'].count('\n') > 0:

ax_w = 0.75

ax_h = 0.73

if opts['title'].count('\n') > 0:

ax_h = 0.73

# Set positions of axes and colorbar

ax_x = 0.1 # Offset from left

ax_y = 0.14 # Offset from bottom

cb_x = ax_x+ax_w+0.005 # Offset from left

cb_y = ax_y # Offset from bottom

cb_w = 0.1 # Width

cb_h = ax_h # Height

ax.set_position([ax_x,ax_y,ax_w,ax_h])

if not allnan:

cb.ax.set_position([cb_x,cb_y,cb_w,cb_h])

if opts['transparent']:

fig.patch.set_alpha(0)

ax.patch.set_alpha(0)

if not opts['returnimage']:

plt.show()

return fig, cb