eng-lab/hwe/report/charts.py

"""Charts for the report, drawn with matplotlib and returned as inline SVG text.

Style rules (the same ones the book's figures follow):
- one colour per series, in a fixed order (blue, then orange, then aqua);
  a single series needs no legend, because the title names it;
- 2-pixel lines, markers at least 8 pixels with a thin ring in the background colour;
- solid hairline gridlines, recessive axes; text in ink, never in a series colour;
- label selectively: the current design and the extreme, not every point.
"""
from __future__ import annotations

import io

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
INK, MUTED, GRID, SURFACE = "#173446", "#52514e", "#e4e6e8", "#ffffff"

plt.rcParams.update({
    "font.family": "sans-serif", "font.size": 9, "axes.edgecolor": GRID, "axes.labelcolor": MUTED,
    "xtick.color": MUTED, "ytick.color": MUTED, "axes.grid": True, "grid.color": GRID,
    "grid.linewidth": 0.6, "grid.linestyle": "-", "axes.axisbelow": True, "svg.fonttype": "none",
    "axes.spines.top": False, "axes.spines.right": False,
})


def _svg(fig) -> str:
    buf = io.StringIO()
    fig.savefig(buf, format="svg", bbox_inches="tight", transparent=True)
    plt.close(fig)
    text = buf.getvalue()
    return text[text.index("<svg"):]                      # drop the XML prolog for inline use


def mass_bars(by_subsystem: dict) -> str:
    """Horizontal bars: predicted mass of each subsystem, largest at the top."""
    items = sorted(((k, v.to("g").m) for k, v in by_subsystem.items()), key=lambda kv: kv[1])
    fig, ax = plt.subplots(figsize=(5.6, 0.42 * len(items) + 0.6))
    names, grams = zip(*items)
    bars = ax.barh(names, grams, height=0.5, color=BLUE)
    for bar, g in zip(bars, grams):
        ax.text(bar.get_width(), bar.get_y() + bar.get_height() / 2, f"  {g:.0f} g",
                va="center", ha="left", color=INK, fontsize=8.5)
    ax.set_xlabel("predicted mass (g), growth allowance included")
    ax.grid(axis="y", visible=False)
    ax.set_xlim(0, max(grams) * 1.22)
    return _svg(fig)


def mass_spiral(sweep: dict) -> str:
    """Flight time against battery capacity: current design, longest flight, and the T/W limit."""
    caps = [c.to("mA*h").m for c in sweep["capacity"]]
    times = [t.to("min").m for t in sweep["flight_time"]]
    fig, ax = plt.subplots(figsize=(5.6, 2.9))
    failing = [c for c, r in zip(caps, sweep.get("thrust_to_weight", [])) if r.m < sweep.get("min_thrust_to_weight", 0)]
    if failing:
        ax.axvspan(min(failing), max(caps), color=GRID, alpha=0.7, linewidth=0)
        ax.text(min(failing), max(times) * 0.06, f"  T/W below {sweep['min_thrust_to_weight']:g}:\n  fails the requirement",
                color=MUTED, fontsize=7.5, va="bottom")
    ax.plot(caps, times, color=BLUE, linewidth=1.6, solid_capstyle="round")
    now_c, now_t = sweep["current_capacity"].to("mA*h").m, sweep["current_time"].to("min").m
    best = max(range(len(times)), key=lambda i: times[i])
    ax.plot([now_c], [now_t], "o", color=BLUE, markersize=7, markeredgecolor=SURFACE, markeredgewidth=1.5)
    ax.annotate(f"your design\n{now_c:,.0f} mAh, {now_t:.1f} min", (now_c, now_t), textcoords="offset points",
                xytext=(8, -26), color=INK, fontsize=8)
    if best < len(times) - 1:
        ax.plot([caps[best]], [times[best]], "o", color=ORANGE, markersize=7, markeredgecolor=SURFACE,
                markeredgewidth=1.5)
        ax.annotate(f"longest flight\n{caps[best]:,.0f} mAh, {times[best]:.1f} min", (caps[best], times[best]),
                    textcoords="offset points", xytext=(0, -32), ha="center", color=INK, fontsize=8)
    ax.set_xlabel("battery capacity (mAh), same cells and same energy per gram")
    ax.set_ylabel("hover flight time (min)")
    ax.set_ylim(0, max(times) * 1.18)
    return _svg(fig)


def phone_traces(phone: dict) -> str:
    """Three stacked panels sharing time: raw reading, velocity, height."""
    t = phone["t"].to("s").m
    fig, axes = plt.subplots(3, 1, figsize=(5.6, 4.6), sharex=True)
    series = [(phone["a"].to("m/s^2").m, "reading (m/s²)"), (phone["v"].to("m/s").m, "velocity (m/s)"),
              (phone["h"].to("cm").m, "height (cm)")]
    still = phone["still_until"].to("s").m
    for ax, (y, label) in zip(axes, series):
        ax.axvspan(t[0], still, color=GRID, alpha=0.6, linewidth=0)
        ax.plot(t, y, color=BLUE, linewidth=1.3)
        ax.set_ylabel(label)
    axes[0].text(t[0], axes[0].get_ylim()[1], " still window (bias estimate)", va="top", color=MUTED, fontsize=7.5)
    axes[-1].set_xlabel("time (s)")
    fig.align_ylabels(axes)
    return _svg(fig)