heatmap.png

heatmap_coding.py

from pandas import DataFrame
from numpy.random import default_rng

from matplotlib.pyplot import subplot_mosaic, rcdefaults, rc, show
from matplotlib.ticker import MultipleLocator
from matplotlib.patches import Circle, RegularPolygon
from matplotlib.collections import PatchCollection
from matplotlib.colors import Normalize
from matplotlib.transforms import blended_transform_factory

from itertools import product
from functools import partial
from math import pi


rng = default_rng(0)
data = DataFrame(
    data=rng.uniform(-1, 1, size=(5, 6)),
    index=[*'ABCDE']
)

rcdefaults()
rc('font', size=14)

mosaic = [
    ['colormesh',   'cax',            'cax'           ],
    ['circle - no', 'circle - minor', 'circle - major'],
    ['square - no', 'square - minor', 'square - major'],
]

fig, axd = subplot_mosaic(
    mosaic, figsize=(20, 16), dpi=100,
    gridspec_kw={
        'hspace': .3, 'wspace': .1,
        'left': .1, 'right': .9, 'top': .92, 'bottom': .05
    },
)

# Plot heatmap using colormesh
ax = axd['colormesh']
img = ax.pcolormesh(
    data.columns, data.index, data.values,
    shading='nearest', vmin=-1, vmax=1, cmap='RdBu'
)
ax.set_title('Default - Colormesh', fontsize='x-large')
ax.xaxis.set_minor_locator(MultipleLocator(.5))
ax.yaxis.set_minor_locator(MultipleLocator(.5))
ax.xaxis.set_major_locator(MultipleLocator(1))
ax.yaxis.set_major_locator(MultipleLocator(1))
ax.tick_params(which='minor', length=0)

# Plot colorbar
cax = axd['cax']
cax.set_aspect(.1, anchor='S')
fig.colorbar(img, cax=cax, orientation='horizontal')
cax.tick_params(bottom=False, labelbottom=False, top=True, labeltop=True)

# Plot product of these markers and grid lines
markers = {
    'circle': Circle,
    'square': partial(RegularPolygon, numVertices=4, orientation=pi/4)
}
grids = {
    'no': lambda ax: ax.grid(visible=False),
    'major': lambda ax: ax.grid(which='major'),
    'minor': lambda ax: ax.grid(which='minor')
}

# Transform data so it is easily plottable with `patches`
scatter_data = data.stack()
y = axd['colormesh'].convert_yunits(scatter_data.index.get_level_values(0))
x = axd['colormesh'].convert_xunits(scatter_data.index.get_level_values(1))
sizes = img.norm(abs(scatter_data)) * .4 # sizes exist on linear scale, max raidus=.4

for (patch_label, patch), (grid_label, grid_cb) in product(markers.items(), grids.items()):
    label = f'{patch_label} - {grid_label}'
    ax = axd[label]

    # Could use `ax.scatter` setting s and c parameters,
    #   but we want shapes to be easily tied to the data coordinates
    collection = PatchCollection(
        [patch((xi, yi), radius=si) for xi, yi, si in zip(x, y, sizes)],
        cmap=img.cmap, norm=img.norm, edgecolors='k', zorder=2
    )
    collection.set_array(scatter_data)
    ax.add_collection(collection)

    grid_cb(ax)

    ax.set_title(f'{patch_label} - {grid_label} Grid'.title(), fontsize='x-large')
    ax.set_facecolor('gainsboro')

    ax.sharex(axd['colormesh'])
    ax.sharey(axd['colormesh'])
    ax.tick_params(which='minor', length=0)


# axd['colormesh'].invert_yaxis()

fig.suptitle('Redundant Coding Heat Maps', fontsize='xx-large')

text = '''
You may have seen heatmaps presented like the one on the left.

However colors can be tricky to compare, especially when they are not adjacent to one
another. Moreover, interpreting the magnitude of these values across 0 is near
impossible as it requires one to compare the intensity of different colors.

To simplify this comparison in a heatmap we can redundantly code the data we are
representing: mapping a value to both color and size as seen below.

Note that without a proper size legend, the sizes are primarily useful for relative comparison.
'''.strip()

axd['cax'].text(
    0, 1, s=text, va='top', fontsize='large',
    transform=blended_transform_factory(
        axd['cax'].transAxes, axd['colormesh'].transAxes
    ),
)

fig.savefig('heatmap.png')
show()