优化二三章结构

scripts/plot_loss_heatmap.py

@@ -7,9 +7,12 @@ from pathlib import Path
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib import font_manager
+from matplotlib.lines import Line2D
 
 
 OVERSEAS_CITIES = {"迪拜", "法兰克福", "雅加达", "开普敦"}
+OVERSEAS_COLOR = "#4285F4"
+DOMESTIC_COLOR = "#EA4335"
 
 
 def leading_group_size(labels: list[str], group: set[str]) -> int:
@@ -94,8 +97,29 @@ def plot_heatmap(csv_path: Path, output_path: Path) -> None:
     ax.set_xlabel("接收端")
     ax.set_ylabel("发送端")
     # ax.set_title("公网链路平均丢包率")
+    ax.xaxis.tick_top()
+    ax.xaxis.set_label_position("top")
+    ax.tick_params(axis="x", top=True, labeltop=True, bottom=False, labelbottom=False, pad=6)
 
-    plt.setp(ax.get_xticklabels(), rotation=35, ha="right", rotation_mode="anchor")
+    plt.setp(ax.get_xticklabels(), rotation=-35, ha="right", rotation_mode="anchor")
+    for label in ax.get_xticklabels():
+        label.set_color(OVERSEAS_COLOR if label.get_text() in OVERSEAS_CITIES else DOMESTIC_COLOR)
+    for label in ax.get_yticklabels():
+        label.set_color(OVERSEAS_COLOR if label.get_text() in OVERSEAS_CITIES else DOMESTIC_COLOR)
+
+    legend_handles = [
+        Line2D([0], [0], marker="s", color="none", markerfacecolor=OVERSEAS_COLOR, markeredgecolor=OVERSEAS_COLOR, markersize=9, label="海外"),
+        Line2D([0], [0], marker="s", color="none", markerfacecolor=DOMESTIC_COLOR, markeredgecolor=DOMESTIC_COLOR, markersize=9, label="国内"),
+    ]
+    ax.legend(
+        handles=legend_handles,
+        loc="lower center",
+        bbox_to_anchor=(0.5, 1.15),
+        ncol=2,
+        frameon=False,
+        columnspacing=1.4,
+        handletextpad=0.4,
+    )
 
     for row_index in range(len(senders)):
         for col_index in range(len(receivers)):
@@ -120,7 +144,7 @@ def plot_heatmap(csv_path: Path, output_path: Path) -> None:
     ax.tick_params(which="minor", bottom=False, left=False)
 
     colorbar = fig.colorbar(image, ax=ax, shrink=0.88)
-    colorbar.set_label("平均丢包率（%）", fontsize=14)
+    colorbar.set_label("平均丢包率", fontsize=14)
     colorbar.ax.tick_params(labelsize=12)
 
     output_path.parent.mkdir(parents=True, exist_ok=True)