demo_foxualize.ipynb

foxualizer.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as p
from matplotlib.collections import PatchCollection
from itertools import chain
import descartes, shapely.geometry as sg  # install: !pip3 install --user shapely descartes
deg_sin = lambda a: np.sin(a * (2 * np.pi / 360.0))
deg_cos = lambda a: np.cos(a * (2 * np.pi / 360.0))


def draw_sectors(x, y, proportions, colors, captions=None, r=1.0, offset_angle=225):
    if captions is None:
        captions = [''] * len(proportions)
    assert np.min(proportions) >= 0 and len(proportions) == len(colors) == len(captions)
    assert np.shape(colors) in {(len(proportions), 3), (len(proportions), 4)}, 'colors must be rgb or rgba'
    canonic_arc = lambda angle1, angle2: sorted([(offset_angle - angle1), (offset_angle - angle2)])
    denominator = np.sum(proportions) if any(proportions > 0) else 1.0
    d_angles = proportions / denominator * 360.0
        
    angle = 0.0
    shell = sg.Polygon([[x - r, y - r], [x + r, y - r], [x + r, y + r], [x - r, y + r]])
    patches, annotations = [], []

    for d_angle, color, text in zip(d_angles, colors, captions):
        if d_angle == 0: continue
        theta1, theta2 = canonic_arc(angle, angle + d_angle)
        shell = sg.Polygon([[x - r, y - r], [x + r, y - r], [x + r, y + r], [x - r, y + r]])
        shifted_theta2 = theta2 if theta2 >= theta1 else theta2 + 360
        dx, dy = deg_cos((theta1 + shifted_theta2) / 2), deg_sin((theta1 + shifted_theta2) / 2)
        if np.allclose(shifted_theta2 - theta1, 360):
            wedge = shell
        else:
            wedge = sg.Polygon([
                        (x, y), 
                        (x + 2 * r * deg_cos(theta1), y + 2 * r * deg_sin(theta1)),
                        (x + 2 * r * (deg_cos(theta1)) + dx,
                         y + 2 * r * (deg_sin(theta1)) + dy),
                        (x + 2 * r * (deg_cos(theta2)) + dx,
                         y + 2 * r * (deg_sin(theta2)) + dy),
                        (x + 2 * r * deg_cos(theta2), y + 2 * r * deg_sin(theta2)),
                    ])
        mesh = shell.intersection(wedge)
        patches.append(descartes.PolygonPatch(mesh, color=color))
        
        annotations.append(dict(s=text, xy=(mesh.centroid.x, mesh.centroid.y)))
        angle += d_angle
        
    
    return patches, annotations

def draw_foxtable(data, cmaps, show=True, show_zeros=False, colorbars=True,
                  annotation_params=None, ticks=False, caption_format='{:0.4}',
                  **kwargs):
    """
    :param data: a 3d array [height, width, n_components] of values of each component
        Values should be in [0, 1] and at least one value must be positive
    :param cmaps: a list[n_components] of colormaps - mappings from values in data to rgba colors
        colormap is any function that takes x \in [0, 1]
            and returns RGBA vector with each component also \in [0, 1]
        All default colormaps: 
            https://matplotlib.org/examples/color/colormaps_reference.html
        Custom example:
            lambda x: [1.0, 1.0 - x, 1.0 - x, 1.0]  # white-to-red cmap
    """
    assert np.min(data) >= 0.0 and np.max(data) <= 1.0 and np.max(data) > 0
    assert len(cmaps) == np.shape(data)[-1]

    nrow, ncol, n_components = np.shape(data)
    i_grid, j_grid = map(np.ravel, np.meshgrid(np.arange(nrow), np.arange(ncol)))

    get_colors = lambda values: np.array([cmap(val) for cmap, val in zip(cmaps, values)])
    get_proportions = lambda values: np.ones_like(values) if show_zeros else (values > 0)
    get_captions = lambda values: [caption_format.format(value) for value in values]
    
    patches, annotations = zip(*chain(*(
        iter(zip(*draw_sectors(j, nrow - i - 1, get_proportions(data[i, j]), get_colors(data[i, j]),
                               captions=get_captions(data[i, j]), r=0.5)))
        for i, j in zip(i_grid, j_grid)
    )))

    fig, ax = plt.subplots(**kwargs)
    ax.add_collection(PatchCollection(patches, match_original=True))
    for ann in annotations:
        ax.annotate(**ann, **dict(ha='center', va='center', **(annotation_params or {})))
    
    if show:
        ax.set_ylim(-0.5, nrow - 0.5)
        ax.set_xlim(-0.5, ncol - 0.5)
        if ticks:
            ax.set_xticks(np.arange(ncol))
            ax.set_yticks(np.arange(nrow))
            ax.set_yticklabels(np.arange(nrow)[::-1])
        else:
            ax.set_xticks([])
            ax.set_yticks([])
        
        plt.show()
    
    return ax