""" AgentEA-Studio: Visualization helpers for Gradio. """ import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.patches as mpatches import numpy as np import io import base64 from src.config import ALIGNMENT_RESULTS, ABLATION_RESULTS, AGENTS, PIPELINE_STEPS def fig_to_base64(fig, dpi=150): """Convert matplotlib figure to base64 encoded PNG.""" buf = io.BytesIO() fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight", facecolor="white") buf.seek(0) img_base64 = base64.b64encode(buf.read()).decode("utf-8") plt.close(fig) return img_base64 def plot_alignment_comparison(dataset_name): """Plot bar chart comparing alignment methods on a dataset.""" results = ALIGNMENT_RESULTS.get(dataset_name) if not results: return None methods = list(results.keys()) hits1 = [results[m]["Hits@1"] * 100 for m in methods] hits10 = [results[m]["Hits@10"] * 100 for m in methods] mrr = [results[m]["MRR"] * 100 for m in methods] x = np.arange(len(methods)) width = 0.25 fig, ax = plt.subplots(figsize=(12, 5)) bars1 = ax.bar(x - width, hits1, width, label="Hits@1", color="#2196F3", edgecolor="white") bars2 = ax.bar(x, hits10, width, label="Hits@10", color="#4CAF50", edgecolor="white") bars3 = ax.bar(x + width, mrr, width, label="MRR", color="#FF9800", edgecolor="white") ax.set_ylabel("Score (%)", fontsize=12) ax.set_title(f"Alignment Performance on {dataset_name}", fontsize=14, fontweight="bold") ax.set_xticks(x) ax.set_xticklabels(methods, rotation=30, ha="right", fontsize=9) ax.legend(loc="lower right", fontsize=10) ax.set_ylim(0, 105) ax.grid(axis="y", alpha=0.3) # Highlight our methods for i, m in enumerate(methods): if "AdaCoAgentEA" in m: for bar in [bars1[i], bars2[i], bars3[i]]: bar.set_edgecolor("#D32F2F") bar.set_linewidth(2) plt.tight_layout() return fig_to_base64(fig) def plot_ablation_study(dataset_name): """Plot ablation study results.""" results = ABLATION_RESULTS.get(dataset_name) if not results: return None configs = list(results.keys()) hits1 = [results[c]["Hits@1"] * 100 for c in configs] mrr = [results[c]["MRR"] * 100 for c in configs] fig, ax = plt.subplots(figsize=(10, 5)) x = np.arange(len(configs)) width = 0.35 bars1 = ax.bar(x - width/2, hits1, width, label="Hits@1", color="#1565C0", edgecolor="white") bars2 = ax.bar(x + width/2, mrr, width, label="MRR", color="#E65100", edgecolor="white") ax.set_ylabel("Score (%)", fontsize=12) ax.set_title(f"Ablation Study on {dataset_name}", fontsize=14, fontweight="bold") ax.set_xticks(x) ax.set_xticklabels(configs, rotation=25, ha="right", fontsize=9) ax.legend(fontsize=10) ax.set_ylim(0, 105) ax.grid(axis="y", alpha=0.3) plt.tight_layout() return fig_to_base64(fig) def plot_heterogeneity_radar(): """Plot radar chart showing heterogeneity dimensions across datasets.""" from src.config import DATASETS categories = ["Fact Disparity", "Density Disparity", "Entity Overlap", "Struct. Similarity"] datasets_to_plot = ["ICEWS-WIKI", "ICEWS-YAGO", "DBP15K-EN-FR", "DBP-WIKI"] N = len(categories) angles = [n / float(N) * 2 * np.pi for n in range(N)] angles += angles[:1] fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(polar=True)) colors = ["#D32F2F", "#F44336", "#2196F3", "#4CAF50"] for i, ds_name in enumerate(datasets_to_plot): ds = DATASETS[ds_name] values = [ min(ds["fact_diff_pct"] / 40, 100), min(ds["density_diff_pct"] / 35, 100), ds["overlapping"], ds["structural_similarity"], ] values += values[:1] ax.fill(angles, values, alpha=0.15, color=colors[i]) ax.plot(angles, values, "o-", linewidth=2, label=ds_name, color=colors[i]) ax.set_xticks(angles[:-1]) ax.set_xticklabels(categories, fontsize=10) ax.set_ylim(0, 100) ax.set_title("KG Heterogeneity Radar", fontsize=14, fontweight="bold", pad=20) ax.legend(loc="upper right", bbox_to_anchor=(1.3, 1.1), fontsize=9) plt.tight_layout() return fig_to_base64(fig) def plot_pipeline_flow(): """Plot the multi-agent pipeline flow diagram.""" fig, ax = plt.subplots(figsize=(16, 8)) ax.set_xlim(0, 16) ax.set_ylim(0, 8) ax.axis("off") # Colors color_input = "#E3F2FD" color_area1 = "#C8E6C9" color_area2 = "#FFF9C4" color_area3 = "#FFCCBC" color_gating = "#E1BEE7" color_output = "#BBDEFB" # Draw boxes boxes = [ (0.5, 5.5, 2.5, 1.5, "Input\nHHKGs", color_input), (3.5, 5.5, 4.0, 1.5, "Area 1\nStructural-Semantic\nDecoupling", color_area1), (3.5, 3.0, 4.0, 0.8, "Gating G1", color_gating), (8.0, 5.5, 4.0, 1.5, "Area 2\nDeep Localization", color_area2), (8.0, 3.0, 4.0, 0.8, "Gating G2", color_gating), (12.5, 5.5, 3.0, 1.5, "Area 3\nKnowledge\nAggregation", color_area3), (12.5, 3.0, 3.0, 0.8, "Gating G3", color_gating), (0.5, 1.0, 2.5, 1.0, "Output\nAligned Pairs", color_output), ] for x, y, w, h, label, color in boxes: rect = mpatches.FancyBboxPatch( (x, y), w, h, boxstyle="round,pad=0.1", facecolor=color, edgecolor="gray", linewidth=1.5 ) ax.add_patch(rect) ax.text(x + w/2, y + h/2, label, ha="center", va="center", fontsize=9, fontweight="bold") # Draw arrows arrows = [ (3.0, 6.25, 3.5, 6.25), # Input → Area 1 (7.5, 6.25, 8.0, 6.25), # Area 1 → Area 2 (12.0, 6.25, 12.5, 6.25), # Area 2 → Area 3 (14.0, 5.5, 14.0, 3.8), # Area 3 → G3 (10.0, 3.0, 10.0, 3.8), # G2 → Area 2 (5.5, 3.0, 5.5, 3.8), # G1 → Area 1 (7.5, 3.4, 12.5, 3.4), # G2 → G3 (3.5, 3.4, 0.5, 3.4), # G1 → Output (10.0, 3.4, 0.5, 3.4), # G2 → Output ] for x1, y1, x2, y2 in arrows: ax.annotate("", xy=(x2, y2), xytext=(x1, y1), arrowprops=dict(arrowstyle="->", color="gray", lw=1.5)) ax.set_title("AdaCoAgentEA Multi-Agent Pipeline Architecture", fontsize=16, fontweight="bold", pad=10) # Legend legend_elements = [ mpatches.Patch(facecolor=color_input, edgecolor="gray", label="Input/Output"), mpatches.Patch(facecolor=color_area1, edgecolor="gray", label="Area 1: Structural-Semantic"), mpatches.Patch(facecolor=color_area2, edgecolor="gray", label="Area 2: Deep Localization"), mpatches.Patch(facecolor=color_area3, edgecolor="gray", label="Area 3: Knowledge Aggregation"), mpatches.Patch(facecolor=color_gating, edgecolor="gray", label="Gating Mechanism"), ] ax.legend(handles=legend_elements, loc="lower center", ncol=5, fontsize=8, bbox_to_anchor=(0.5, -0.08)) plt.tight_layout() return fig_to_base64(fig) def plot_agent_communication(): """Plot the agent communication protocol diagram.""" fig, ax = plt.subplots(figsize=(14, 7)) ax.set_xlim(0, 14) ax.set_ylim(0, 7) ax.axis("off") # Shared Message Pool pool = mpatches.FancyBboxPatch( (4, 2), 6, 4, boxstyle="round,pad=0.15", facecolor="#E8EAF6", edgecolor="#3F51B5", linewidth=2 ) ax.add_patch(pool) ax.text(7, 5.5, "Shared Message Pool", ha="center", va="center", fontsize=13, fontweight="bold", color="#1A237E") # Agents agents = [ (1, 5, "S1-S3", "Small Models\nArea 1", "#C8E6C9"), (1, 3, "L1", "LLM\nArea 2", "#FFF9C4"), (1, 1, "S4", "Small Model\nArea 2", "#C8E6C9"), (12, 5, "L2", "LLM\nArea 2", "#FFF9C4"), (12, 3, "L3", "Master\nArea 3", "#FFCCBC"), (12, 1, "L4+", "Domain Experts\nArea 3", "#FFCCBC"), ] for x, y, label, desc, color in agents: rect = mpatches.FancyBboxPatch( (x - 0.8, y - 0.5), 1.6, 1.0, boxstyle="round,pad=0.05", facecolor=color, edgecolor="gray", linewidth=1 ) ax.add_patch(rect) ax.text(x, y + 0.1, label, ha="center", va="center", fontsize=10, fontweight="bold") ax.text(x, y - 0.25, desc, ha="center", va="center", fontsize=7, color="gray") # Communication arrows for x, y, label, desc, color in agents: ax.annotate("", xy=(4.5, y), xytext=(x + 0.8, y), arrowprops=dict(arrowstyle="<->", color="#7986CB", lw=1.2, ls="--")) ax.text(7, 1.2, "Structured Agent Communication + Publish-Subscribe Mechanism", ha="center", fontsize=10, fontstyle="italic", color="gray") ax.set_title("Agent Communication Protocol", fontsize=14, fontweight="bold", pad=10) plt.tight_layout() return fig_to_base64(fig) def plot_improvement_waterfall(dataset_name): """Plot waterfall chart showing improvement contributions.""" results = ABLATION_RESULTS.get(dataset_name) if not results: return None base = results["w/o LLM Agents"]["Hits@1"] * 100 ada = results["AdaCoAgentEA+"]["Hits@1"] * 100 components = [ ("Base (w/o LLM)", base, base), ("+ LLM Agents", results["w/o Communication Protocols"]["Hits@1"] * 100 - base, results["w/o Communication Protocols"]["Hits@1"] * 100), ("+ Communication Protocols", results["w/o Expert Role-Playing"]["Hits@1"] * 100 - results["w/o Communication Protocols"]["Hits@1"] * 100, results["w/o Expert Role-Playing"]["Hits@1"] * 100), ("+ Expert Role-Playing", results["w/o Meta-Logic Rules"]["Hits@1"] * 100 - results["w/o Expert Role-Playing"]["Hits@1"] * 100, results["w/o Meta-Logic Rules"]["Hits@1"] * 100), ("+ Meta-Logic Rules", ada - results["w/o Meta-Logic Rules"]["Hits@1"] * 100, ada), ] fig, ax = plt.subplots(figsize=(10, 5)) labels = [c[0] for c in components] values = [c[1] for c in components] cumulative = [c[2] for c in components] colors = ["#90A4AE", "#42A5F5", "#66BB6A", "#FFA726", "#EF5350"] # Plot bars bottom = base for i, (label, val, cum) in enumerate(components): if i == 0: ax.bar(i, val, color=colors[i], edgecolor="white") else: ax.bar(i, val, bottom=bottom, color=colors[i], edgecolor="white") bottom += val ax.set_xticks(range(len(components))) ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9) ax.set_ylabel("Hits@1 (%)", fontsize=12) ax.set_title(f"Component Contribution Waterfall ({dataset_name})", fontsize=14, fontweight="bold") ax.grid(axis="y", alpha=0.3) # Add value labels for i, (label, val, cum) in enumerate(components): ax.text(i, cum + 1, f"{cum:.1f}%", ha="center", fontsize=9, fontweight="bold") plt.tight_layout() return fig_to_base64(fig)