Отображение изменений в древовидной карте в matplotlib

Эта задача состоит из двух частей.

1. Вычисление макета прямоугольников.
2. Рисуем прямоугольники.

Первая часть может оказаться весьма увлекательной: люди публикуют научные работы по теме. Здесь нецелесообразно изобретать велосипед. Однако вторая часть довольно проста и может быть выполнена в matplotlib.

В приведенном ниже решении используется squrify для вычисления макета с использованием большего значения для каждой пары значений, а затем matplotlib для рисования двух прямоугольников друг над другом.

import numpy as np
import matplotlib.pyplot as plt
import squarify

from matplotlib import colormaps
from matplotlib.colors import to_rgba

DEFAULT_COLORS = list(zip(colormaps["tab20"].colors[::2],
                          colormaps["tab20"].colors[1::2]))


def color_to_grayscale(color):
    # Adapted from: https://stackoverflow.com/a/689547/2912349
    r, g, b, a = to_rgba(color)
    return (0.299 * r + 0.587 * g + 0.114 * b) * a


class PairedTreeMap:

    def __init__(self, values, colors=DEFAULT_COLORS, labels=None, ax=None, bbox=(0, 0, 200, 100)):
        """
        Draw a treemap of value pairs.

        values : list[tuple[float, float]]
            A list of value pairs.

        colors : list[tuple[RGBA, RGBA]]
            The corresponding color pairs. Defaults to light/dark tab20 matplotlib color pairs.

        labels : list[str]
            The labels, one for each pair.

        ax : matplotlib.axes._axes.Axes
            The matplotlib axis instance to draw on.

        bbox : tuple[float, float, float, float]
            The (x, y) origin and (width, height) extent of the treemap.

        """

        self.ax = self.initialize_axis(ax)
        self.rects = self.get_layout(values, bbox)
        self.artists = list(self.draw(self.rects, values, colors, self.ax))

        if labels:
            self.labels = list(self.add_labels(self.rects, labels, values, colors, self.ax))


    def get_layout(self, values, bbox):
        maxima = np.max(values, axis=1)
        order = np.argsort(maxima)[::-1]
        normalized_maxima = squarify.normalize_sizes(maxima[order], *bbox[2:])
        rects = squarify.padded_squarify(normalized_maxima, *bbox)
        reorder = np.argsort(order)
        return [rects[ii] for ii in reorder]


    def initialize_axis(self, ax=None):
        if ax is None:
            fig, ax = plt.subplots()
        ax.set_aspect("equal")
        ax.axis("off")
        return ax


    def _get_artist_pair(self, rect, value_pair, color_pair):
        x, y, w, h = rect["x"], rect["y"], rect["dx"], rect["dy"]
        (small, large), (color_small, color_large) = zip(*sorted(zip(value_pair, color_pair)))
        ratio = np.sqrt(small / large)
        return (plt.Rectangle((x, y), w,         h,         color=color_large, zorder=1),
                plt.Rectangle((x, y), w * ratio, h * ratio, color=color_small, zorder=2))


    def draw(self, rects, values, colors, ax):
        for rect, value_pair, color_pair in zip(rects, values, colors):
            large_patch, small_patch = self._get_artist_pair(rect, value_pair, color_pair)
            ax.add_patch(large_patch)
            ax.add_patch(small_patch)
            yield(large_patch, small_patch)
        ax.autoscale_view()


    def add_labels(self, rects, labels, values, colors, ax):
        for rect, label, value_pair, color_pair in zip(rects, labels, values, colors):
            x, y, w, h = rect["x"], rect["y"], rect["dx"], rect["dy"]
            # decide a fontcolor based on background brightness
            (small, large), (color_small, color_large) = zip(*sorted(zip(value_pair, color_pair)))
            ratio = small / large
            background_brightness = color_to_grayscale(color_large) if ratio < 0.33 else color_to_grayscale(color_small) # i.e. 0.25 + some fudge
            fontcolor = "white" if background_brightness < 0.5 else "black"
            yield ax.text(x + w/2, y + h/2, label, va = "center", ha = "center", color=fontcolor)


if __name__ == "__main__":

    values = [
        (4, 10),
        (13, 15),
        (5, 8),
    ]

    colors = [
        ("red", "coral"),
        ("royalblue", "cornflowerblue"),
        ("darkslategrey", "gray"),
    ]

    labels = [
        "Denmark",
        "Sweden",
        "Norway"
    ]

    PairedTreeMap(values, colors=colors, labels=labels, bbox=(0, 0, 100, 100))
    plt.show()