plot_gpu.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

# Data organization
configs = [
    "Ti, Pyr.", "S, Pyr.", "M, Pyr.", "B, Pyr.",
    "Ti, Iso.", "S, Iso.", "B, Iso."
]
N_values = [256, 512, 1024, 2048]

data = np.array([
    [0.60318, 0.631048, 0.658676, 0.788872],    # Ti, Pyr.
    [0.68295, 0.738171, 0.750116, 0.836389],    # S, Pyr.
    [0.715045, 0.759782, 0.778339, 0.888872],   # M, Pyr.
    [0.70318, 0.798335, 0.882231, np.nan],      # B, Pyr. (no 2048 data)
    [0.665668, 0.676529, 0.697021, np.nan],     # Ti, Iso.
    [0.670633, 0.701255, 0.740179, np.nan],     # S, Iso.
    [0.726529, 0.775582, 0.800547, np.nan]      # B, Iso.
])

# Plot configuration
plt.figure(figsize=(14, 8))
ax = plt.gca()

# Style definitions
colors = {'Ti': '#E41A1C', 'S': '#4DAF4A', 'M': '#377EB8', 'B': '#984EA3'}
markers = {'Pyr.': 'o', 'Iso.': 's'}
linestyles = {'Pyr.': '-', 'Iso.': '--'}

# Plot each configuration, masking NaNs
for i, config in enumerate(configs):
    size, pyramid_type = map(str.strip, config.split(','))
    y = data[i]
    mask = ~np.isnan(y)
    x = np.array(N_values)[mask]
    y = y[mask]
    ax.plot(
        x, y,
        color=colors[size],
        marker=markers[pyramid_type],
        linestyle=linestyles[pyramid_type],
        linewidth=2.5,
        markersize=10,
        label=config
    )

# Axis formatting
ax.set_xscale('log', base=2)
ax.set_xticks(N_values)
ax.get_xaxis().set_major_formatter(ticker.ScalarFormatter())
ax.tick_params(axis='both', which='major', labelsize=17)
ax.grid(True, linestyle='--', alpha=0.7)

ax.set_title("Graph Construction Runtime Proportion for ViG on GPU", fontsize=25, pad=15)
ax.set_xlabel("Image Resolution (Height=Width)", fontsize=20, labelpad=10)
ax.set_ylabel("Runtime Proportion", fontsize=20, labelpad=10)


# Legend
handles, labels = ax.get_legend_handles_labels()
ax.legend(
    handles, labels,
    fontsize=15,
    title="Variants",
    title_fontsize=17,
    loc='upper left',
    bbox_to_anchor=(1, 1)
)

plt.tight_layout(rect=[0, 0, 0.85, 1])  # Make space for legend
plt.savefig('gpu_vig_profile.pdf', dpi=300, bbox_inches='tight')
plt.show()