make_figs.py

# -*- coding: utf-8 -*-
"""Generate teaser.png and framework.png for the SymboLLM-FE project page."""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
import os

plt.rcParams["font.family"] = "DejaVu Sans"

ASSET_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
os.makedirs(ASSET_DIR, exist_ok=True)

C_BLUE = "#1e5abb"
C_LIGHT = "#eef4ff"
C_GRAY = "#606f7b"
C_DARK = "#1f2933"
C_SOFT = "#f7f9fb"
C_GREEN = "#2e7d32"
C_ORANGE = "#b26a00"

# ============================================================
# 1. Teaser: three AutoFE paradigms (Figure 1 style)
# ============================================================
fig, axes = plt.subplots(1, 3, figsize=(13.5, 5.4), dpi=200)
fig.subplots_adjust(left=0.02, right=0.98, top=0.86, bottom=0.06, wspace=0.12)

panels = [
    dict(
        title="Traditional AutoFE",
        sub="blind operator stacks",
        color=C_GRAY,
        marker_pos=(0.78, 0.30),   # high score, low interpretability
        code=[
            "Feature_new_1 = x0^2 + 2*x0*x1 + x1^2",
            "Feature_new_2 = min(x0, x1)",
            "Feature_new_3 = ln(x0 / x2)",
        ],
        badge="Opaque math, no semantics",
    ),
    dict(
        title="LLM-based AutoFE",
        sub="unguided semantic generation",
        color=C_ORANGE,
        marker_pos=(0.30, 0.74),   # low (unstable) score, high interpretability
        code=[
            "Risk_Indicator =",
            "  IF (Age > 60 AND Fever='Yes')",
            "  THEN 'High' ELSE 'Low'",
            "~50 rounds of trial-and-error",
        ],
        badge="Costly iterations, hallucination",
    ),
    dict(
        title="SymboLLM-FE (Ours)",
        sub="symbolic regression + LLM refinement",
        color=C_BLUE,
        marker_pos=(0.80, 0.78),   # high score, high interpretability
        code=[
            "Infection_Exposure_Risk =",
            "  (Fever=='Yes')*0.4 + (Fatigue=='Yes')*0.3",
            "  + Smooth*0.2     # weighted symptom score",
            "Only ~4 LLM API calls",
        ],
        badge="High accuracy + interpretable",
    ),
]

for ax, p in zip(axes, panels):
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_edgecolor("#e4e7eb")

    # title
    ax.text(0.04, 0.93, p["title"], fontsize=13.5, fontweight="bold",
            color=C_DARK, va="center", transform=ax.transAxes)
    ax.text(0.04, 0.855, p["sub"], fontsize=9, color=C_GRAY,
            va="center", transform=ax.transAxes)

    # code box
    box = FancyBboxPatch((0.04, 0.10), 0.92, 0.52,
                         boxstyle="round,pad=0.012,rounding_size=0.02",
                         linewidth=1.2, edgecolor=p["color"], facecolor=C_SOFT)
    ax.add_patch(box)
    for i, line in enumerate(p["code"]):
        yy = 0.565 - i * 0.088
        ax.text(0.075, yy, line, fontsize=8.2, color="#333",
                family="DejaVu Sans Mono", va="center", transform=ax.transAxes)

    # badge
    ax.text(0.05, 0.035, p["badge"], fontsize=9.2, color=p["color"],
            fontweight="bold", va="center", transform=ax.transAxes)

    # scatter anchor (model score vs interpretability)
    ax.plot(p["marker_pos"][0], p["marker_pos"][1], "o", ms=16,
            color=p["color"], alpha=0.25, transform=ax.transAxes, zorder=5)
    ax.plot(p["marker_pos"][0], p["marker_pos"][1], "o", ms=9,
            color=p["color"], transform=ax.transAxes, zorder=6)
    ax.annotate("Model\nScore", xy=(p["marker_pos"][0] - 0.13, p["marker_pos"][1] + 0.07),
                fontsize=8, color=C_GRAY, ha="center", va="center",
                transform=ax.transAxes)
    ax.annotate("Interpret-\nability", xy=(p["marker_pos"][0] + 0.14, p["marker_pos"][1] + 0.07),
                fontsize=8, color=C_GRAY, ha="center", va="center",
                transform=ax.transAxes)

fig.suptitle("Feature Generation Mechanisms & Interpretability Paradigms across AutoFE",
             fontsize=14.5, fontweight="bold", color=C_DARK, y=0.98)
fig.savefig(os.path.join(ASSET_DIR, "teaser.png"), bbox_inches="tight", facecolor="white")
plt.close(fig)

# ============================================================
# 2. Framework: two-stage pipeline (Figure 2 style)
# ============================================================
fig, ax = plt.subplots(figsize=(12.8, 6.2), dpi=200)
ax.set_xlim(0, 100)
ax.set_ylim(0, 62)
ax.axis("off")


def box(x, y, w, h, text, fc=C_LIGHT, ec=C_BLUE, fs=8.6, bold=True, tc=C_DARK):
    b = FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.4,rounding_size=1.2",
                       linewidth=1.4, edgecolor=ec, facecolor=fc)
    ax.add_patch(b)
    ax.text(x + w / 2, y + h / 2, text, fontsize=fs, fontweight="bold" if bold else "normal",
            color=tc, ha="center", va="center", linespacing=1.45)


def arrow(x1, y1, x2, y2, color=C_BLUE, lw=1.6):
    a = FancyArrowPatch((x1, y1), (x2, y2), arrowstyle="-|>",
                        mutation_scale=16, linewidth=lw, color=color)
    ax.add_patch(a)


# ---- Stage 1 -------------------------------------------------
ax.text(2, 60.5, "Stage 1 — Formula Construction by Symbolic Regression",
        fontsize=11.5, fontweight="bold", color=C_DARK, va="center")

box(2, 44, 17, 9, "Raw Columns\nC = {c1, ..., cn}", fc="#ffffff")
arrow(19.5, 48.5, 25.5, 48.5)

box(26, 44, 20, 9, "Spearman Correlation\nPre-sorting", fc=C_SOFT, ec=C_GRAY)
arrow(46.5, 48.5, 52.5, 48.5)

box(53, 44, 24, 9, "Expanding-Sliding Window\nO(2^n)  →  O(n^2)", fc=C_SOFT, ec=C_GRAY)
arrow(77.5, 48.5, 83.5, 48.5)

box(84, 44, 15, 9, "Symbolic\nRegression", fc=C_LIGHT)

arrow(91.5, 44, 91.5, 28.5)
box(84, 19.5, 15, 8, "Formula\nRepository\nR = {(S_i, P, τ_i)}", fc=C_SOFT, ec=C_GRAY)

# ---- Stage 2 -------------------------------------------------
ax.text(2, 14, "Stage 2 — Feature Generation via LLMs",
        fontsize=11.5, fontweight="bold", color=C_DARK, va="center")

box(2, 2, 17, 8.5, "Task Description\n(+ Background)", fc="#ffffff")
arrow(19.5, 6.2, 25.5, 6.2)

box(26, 2, 20, 8.5, "LLM Refinement\n(chain-of-thought code)", fc=C_LIGHT)
arrow(46.5, 6.2, 52.5, 6.2)

box(53, 2, 17, 8.5, "Execute Code\n→ C_merge", fc=C_SOFT, ec=C_GRAY)
arrow(70.5, 6.2, 76.5, 6.2)

box(77, 2, 21, 8.5, "Downstream Predictor\nValidation (best score)", fc=C_SOFT, ec=C_GRAY)

# feedback loop from predictor back to LLM
arrow(87.5, 2, 87.5, -0.0)
ax.plot([87.5, 36], [0.8, 0.8], color=C_BLUE, linewidth=1.4)
arrow(36, 0.8, 36, 1.6, color=C_BLUE)
ax.text(60, 1.15, "iterative refinement loop (score feedback)", fontsize=7.8,
        color=C_GRAY, ha="center", va="bottom")

# from formula repository to LLM
arrow(91.5, 19.5, 91.5, 14.5)
ax.plot([91.5, 36], [14.8, 14.8], color=C_BLUE, linewidth=1.4)
arrow(36, 14.8, 36, 11.0, color=C_BLUE)

# final output
arrow(98.5, 6.2, 99.5, 6.2)
ax.text(99.8, 6.2, "", fontsize=8, color=C_GRAY)

fig.savefig(os.path.join(ASSET_DIR, "framework.png"), bbox_inches="tight", facecolor="white")
plt.close(fig)

print("OK:", os.listdir(ASSET_DIR))