Raptor/validation/tools/pim/plot_style.py

"""Shared appearance and PNG/SVG export for PIM result plots."""

from pathlib import Path
from math import ceil

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.path import Path as PlotPath
from matplotlib.patheffects import AbstractPathEffect
from matplotlib.transforms import Affine2D, IdentityTransform, TransformedPath


TEXT_COLOR = "#243449"
HATCH_COLOR = "#C62828"
COLORS = ("#2B65A1", "#168C8B", "#D18D32", "#8062A5", "#B65D60")


class CenteredHatch(AbstractPathEffect):
    """Anchor dot and diagonal hatch patterns to each rectangle's center."""

    def draw_path(self, renderer, gc, tpath, affine, rgbFace=None):
        hatch = gc.get_hatch() or ""
        base_gc = renderer.new_gc()
        base_gc.copy_properties(gc)
        base_gc.set_hatch("".join(symbol for symbol in hatch if symbol not in "./\\"))
        renderer.draw_path(base_gc, tpath, affine, rgbFace)
        base_gc.restore()

        bounds = affine.transform_path(tpath).get_extents()
        center_x, center_y = (bounds.x0 + bounds.x1) / 2, (bounds.y0 + bounds.y1) / 2
        vertices = []
        diagonal = 2 ** -0.5
        for symbols, dx, dy in (("/", diagonal, diagonal), ("\\", diagonal, -diagonal)):
            count = sum(hatch.count(symbol) for symbol in symbols)
            if count == 0:
                continue
            spacing = renderer.points_to_pixels(12 / count)
            half_span = (abs(dy) * bounds.width + abs(dx) * bounds.height) / 2
            half_length = (abs(dx) * bounds.width + abs(dy) * bounds.height) / 2
            extent = ceil(half_span / spacing)
            for step in range(-extent, extent + 1):
                x = center_x - dy * step * spacing
                y = center_y + dx * step * spacing
                vertices.extend(((x - dx * half_length, y - dy * half_length),
                                 (x + dx * half_length, y + dy * half_length)))
        line_gc = renderer.new_gc()
        line_gc.copy_properties(gc)
        line_gc.set_hatch(None)
        line_gc.set_foreground(gc.get_hatch_color())
        line_gc.set_linewidth(gc.get_hatch_linewidth())
        line_gc.set_clip_path(TransformedPath(tpath, affine))
        if vertices:
            renderer.draw_path(line_gc, PlotPath(vertices, [PlotPath.MOVETO, PlotPath.LINETO] * (len(vertices) // 2)),
                               IdentityTransform())
        if "." in hatch:
            spacing = renderer.points_to_pixels(12 / hatch.count("."))
            x_extent = ceil(bounds.width / (2 * spacing))
            y_extent = ceil(bounds.height / (2 * spacing))
            centers = [(center_x + x * spacing, center_y + y * spacing)
                       for x in range(-x_extent, x_extent + 1)
                       for y in range(-y_extent, y_extent + 1)]
            line_gc.set_linewidth(0)
            renderer.draw_markers(line_gc, PlotPath.unit_circle(),
                                  Affine2D().scale(renderer.points_to_pixels(0.9)),
                                  PlotPath(centers), IdentityTransform(), gc.get_hatch_color())
        line_gc.restore()


def configure_style() -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.labelcolor": TEXT_COLOR,
        "text.color": TEXT_COLOR,
        "xtick.color": TEXT_COLOR,
        "ytick.color": TEXT_COLOR,
        "axes.edgecolor": "#CFD6DF",
        "savefig.facecolor": "white",
        "figure.facecolor": "white",
    })


def value_label(ax, x, y, label: str, status: str = "PASS", direction: int = 1,
                fontsize: float = 8) -> None:
    alignment = "bottom" if direction > 0 else "top"
    rotation = 90 if direction > 0 else -90
    text = ax.text(x, y, label, rotation=rotation, ha="center", va=alignment,
                   fontsize=fontsize, color=TEXT_COLOR)
    if status != "PASS":
        height = text.get_window_extent(ax.figure.canvas.get_renderer()).height
        offset = direction * (height * 72 / ax.figure.dpi + 3)
        ax.annotate("INVALID" if status == "FAIL" else status, (x, y),
                    xytext=(0, offset), textcoords="offset points",
                    rotation=rotation, ha="center", va=alignment, fontsize=fontsize,
                    color=HATCH_COLOR)


def fit_value_labels(fig) -> None:
    """Reserve vertical space for rotated labels using their rendered size."""
    for ax in fig.axes:
        for text in ax.texts:
            text.set_in_layout(False)
    fig.canvas.draw()
    renderer = fig.canvas.get_renderer()
    limits = []
    for ax in fig.axes:
        bounds = ax.get_window_extent(renderer)
        labels = []
        for text in ax.texts:
            x, y = text.xy if hasattr(text, "xy") else text.get_position()
            anchor = ax.transData.transform((x, y))[1]
            extent = text.get_window_extent(renderer)
            labels.append((y, max(0, anchor - extent.y0 + 3) / bounds.height,
                           max(0, extent.y1 - anchor + 3) / bounds.height))
        lower, upper = ax.get_ylim()
        for _ in range(12):
            lower = min([lower, *((y - below * upper) / (1 - below) for y, below, _ in labels)])
            upper = max([upper, *((y - above * lower) / (1 - above) for y, _, above in labels)])
        limits.append((lower, upper))
    for ax in fig.axes:
        ax.set_ylim(min(lower for lower, _ in limits), max(upper for _, upper in limits))


def save_figure(fig, basename: str, output_dir: Path) -> None:
    for patch in fig.findobj(Patch):
        hatch = patch.get_hatch() or ""
        if any(symbol in hatch for symbol in "./\\"):
            patch.set_path_effects([CenteredHatch()])
    for extension in ("png", "svg"):
        target = output_dir / extension / f"{basename}.{extension}"
        target.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(target, dpi=220, bbox_inches="tight", pad_inches=0.04)
        print(target)
    plt.close(fig)