docs/figures/plot_kvpath.py

"""Publication figure: vertical layers, one shared alpha, independent deltas."""
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, FancyArrowPatch

plt.rcParams.update({'font.family': 'STIXGeneral', 'font.size': 11,
    'mathtext.fontset': 'stix', 'pdf.fonttype': 42,
    'ps.fonttype': 42, 'svg.fonttype': 'none'})
OUT = Path(__file__).resolve().parent
INK, BLUE, ORANGE = '#242424', '#36647D', '#A66B3F'
fig, ax = plt.subplots(figsize=(6.5, 4.7))
ax.set(xlim=(-.15, 8.1), ylim=(.65, 7.2))
ax.axis('off')

def text(x, y, s, size=11, color=INK, **kwargs):
    ax.text(x, y, s, fontsize=size, color=color,
            ha='center', va='center', **kwargs)

def line(xs, ys, color=INK, lw=.8, **kwargs):
    ax.plot(xs, ys, color=color, lw=lw, **kwargs)

def arrow(a, b, color=INK, style='-'):
    ax.add_patch(FancyArrowPatch(a, b, arrowstyle='-|>',
        mutation_scale=8, linewidth=.85, color=color,
        linestyle=style, shrinkA=1, shrinkB=1))

# One tall alpha block denotes shared storage, with depth increasing upward.
alpha_x, alpha_w, delta_w = 1.65, 2.20, 1.35
delta_x = alpha_x + alpha_w
bottom, top = 1.40, 6.20
ax.add_patch(Rectangle((alpha_x, bottom), alpha_w, top-bottom,
    facecolor='#EDF3F6', edgecolor=BLUE, linewidth=.9))
text(alpha_x+alpha_w/2, 6.67, r'Alpha component ($d_\alpha$)', 11, BLUE)
text(delta_x+delta_w/2, 6.67, r'Delta ($d_\delta$)', 11, ORANGE)
text(7.05, 6.67, 'Layer input', 11)

rows = [(1.8 + idx * (4.0 / 3.0), str(idx)) for idx in range(4)]
for y, idx in rows:
    text(1.12, y, rf'${idx}$', 12)
    text(alpha_x+alpha_w/2, y, r'$(k^\alpha,v^\alpha)$', 14, BLUE)
    ax.add_patch(Rectangle((delta_x, y-.40), delta_w, .80,
        facecolor='#F8EEE4', edgecolor=ORANGE, linewidth=.9))
    text(delta_x+delta_w/2, y,
         rf'$(k^\delta_{{{idx}}},v^\delta_{{{idx}}})$', 12, ORANGE)
    arrow((6.75, y), (5.24, y), ORANGE)
    text(6.0, y+.27, rf'$W^\delta_{{{idx}}}$', 12, ORANGE)
    text(7.05, y, rf'$h_{{{idx}}}$', 14)

line([.65, .48, .48, .65], [bottom, bottom, top, top])
text(.08, 3.8, r'One path ($m=4$)', 11, rotation=90)
text(1.12, 6.67, 'Layer', 10)
text(2.75, 3.80, 'shared across layers', 10, BLUE, style='italic')
for (lower_y, _), (upper_y, idx) in zip(rows, rows[1:]):
    arrow((4.525, lower_y+.43), (4.525, upper_y-.43), ORANGE, '--')
    text(5.03, (lower_y+upper_y)/2, rf'$\Delta_{{{idx}}}$', 11, ORANGE)

# First-layer input generates the shared alpha once.
line([7.05, 7.05, 2.75], [1.52, .88, .88], BLUE)
arrow((2.75, .88), (2.75, 1.38), BLUE)
text(4.9, 1.12, r'$W^\alpha$', 12, BLUE)

fig.subplots_adjust(left=.015, right=.985, bottom=.02, top=.98)
for ext in ('png', 'pdf', 'svg'):
    path = OUT / f'kvpath_alpha_delta.{ext}'
    fig.savefig(path, dpi=240, facecolor='white', bbox_inches='tight', pad_inches=.04)
    print(path)
plt.close(fig)