Skip to content

Add configurable auxiliary panel to unrolled grid plots #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AvencastF merged 2 commits into from
Jan 14, 2026
Merged

Add configurable auxiliary panel to unrolled grid plots #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e636475
Add configurable auxiliary grid panel
AvencastF Jan 14, 2026
704dfef
fix better grid study plot
AvencastF Jan 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 78 additions & 8 deletions paper_plot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,9 +328,10 @@
"base_plot_config": {
"models": [],
"figsize": (15, 7),
"hspace": 0,
"hspace": 0.05,
"height_ratios": {
"main": 2.0,
"aux": 1.2,
"cutflow": 0.8,
"winner": 0.5,
"ratio": 1.0,
Expand Down Expand Up @@ -375,11 +376,11 @@
"enabled": True,
"ylabel": "",
"ylabel_rotation": 90,
"ylabel_pad": 37,
# "ylabel_pad": 37,
},
"cutflow": {
"enabled": True,
"ylabel": "Stats. [K]",
"ylabel": "Signal \n Stats. [K]",
"color": "0.7",
"edgecolor": "0.2",
"alpha": 0.75,
Expand All @@ -388,6 +389,12 @@
"zorder": 1,
"log": False,
},
"aux_panel": {
"enabled": True,
"metric_col": "effective_steps",
"ylabel": "Updates [K]",
"y_log": False,
},
"ratios": [],
"ratio_line": {
"marker": "o",
Expand All @@ -413,12 +420,11 @@
# "tight_layout": {"pad": 0.2},
"subplot_adjust": {
"top": 0.83, "bottom": 0.08,
"left": 0.05, "right": 0.98,
"left": 0.1, "right": 0.98,
},
},
"plots": {
"individual": {
# "style": PlotStyle(base_font_size=10.0, tick_label_size=9.0, full_axis=True),
"output_name": "grid_sic_individual",
"title": "Grid SIC (individual)",
"series": [
Expand Down Expand Up @@ -448,11 +454,27 @@
},
],
"plot_config": {
"figsize": (15, 8),
"height_ratios": {
"main": 2.2,
"aux": 1.0,
"cutflow": 0.8,
"winner": 0.5,
"ratio": 1.0,
},
"subplot_adjust": {
"top": 0.865, "bottom": 0.08,
"left": 0.1, "right": 0.98,
},
"aux_panel": {
"enabled": True,
"ylabel": "Effective \n Steps [K]",
},
"apply_axis_style": True,
"ratios": [
# {"baseline": "XGBoost", "mode": "ratio", "ylabel": "/XGB", "reference_line": True},
# {"baseline": "TabPFN v2.5", "mode": "ratio", "ylabel": "/TabPFN", "reference_line": True},
{"baseline": "Full", "mode": "ratio", "ylabel": "Ratio to Full", "reference_line": True},
{"baseline": "Full", "mode": "ratio", "ylabel": "Ratio \n to Full", "reference_line": True},
],
"unc": {"enabled": True},
},
Expand Down Expand Up @@ -487,13 +509,25 @@
},
],
"plot_config": {
"figsize": (15, 6),
"hspace": 0.05,
"ratios": [
# {"baseline": "XGBoost (param)", "mode": "ratio", "ylabel": "/XGB", "reference_line": True, "y_log": True},
# {"baseline": "EveNet-Scratch (param)", "mode": "ratio", "ylabel": "/Scratch", "reference_line": True, "y_log": True},
{"baseline": "Full", "mode": "ratio", "ylabel": "Ratio to Full", "reference_line": True,
{"baseline": "Full", "mode": "ratio", "ylabel": "Ratio \n to Full", "reference_line": True,
"y_log": False},
],
"unc": {"enabled": True},
"aux_panel": {
"enabled": False,
},
"cutflow": {
"enabled": False,
},
"subplot_adjust": {
"top": 0.81, "bottom": 0.10,
"left": 0.09, "right": 0.98,
},
},
},
},
Expand Down Expand Up @@ -1842,6 +1876,30 @@ def _collect_grid_series(
return points, sic_by_model, unc_by_model, model_order, color_map


def _collect_grid_metric_map(
grid_data: pd.DataFrame,
*,
series_specs: list[dict],
metric_col: str,
points: list[tuple[float, float]],
):
metric_by_model = {}
for spec in series_specs:
label = spec.get("label", spec["model"])
df_model = grid_data[grid_data["model"].eq(spec["model"])]
if spec.get("type") is not None:
df_model = df_model[df_model["type"].eq(spec["type"])]

if metric_col in df_model.columns:
grouped_metric = df_model.groupby(["m_X", "m_Y"])[metric_col].mean()
else:
grouped_metric = pd.Series(dtype=float)
metric_map = {key: value for key, value in grouped_metric.items()}
metric_by_model[label] = {point: metric_map.get(point, np.nan) for point in points}

return metric_by_model


def plot_grid_results(
grid_data: pd.DataFrame,
cutflow_df: pd.DataFrame | None = None,
Expand Down Expand Up @@ -1880,11 +1938,22 @@ def plot_grid_results(
plot_config["unc"] = unc_cfg
plot_config['raw_model_series'] = [f"{m['model']}_{m['type']}" for m in plot_cfg.get("series", [])]

aux_cfg = plot_config.get("aux_panel", {})
aux_by_model = None
if aux_cfg.get("enabled", False):
aux_by_model = _collect_grid_metric_map(
grid_data,
series_specs=plot_cfg.get("series", []),
metric_col=aux_cfg.get("metric_col", "effective_steps"),
points=points,
)

with use_style(style):
fig, axes = plot_unrolled_grid_with_winner_and_ratios(
points,
sic_by_model,
unc_by_model=unc_by_model,
aux_by_model=aux_by_model,
cutflow_df=cutflow_df,
config=plot_config,
style=style,
Expand Down Expand Up @@ -2200,7 +2269,8 @@ def read_grid_data(file_path):
train_size=selected_grid_df["statistics"],
batch_size_per_GPU=selected_grid_df["effective_batch_size"],
GPUs=1,
)
) / 1000
selected_grid_df.loc[mask_par, "effective_steps"] /= len(cutflow_df)

return selected_grid_df, cutflow_df

Expand Down
11 changes: 9 additions & 2 deletions plot_styles/core/style_axis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ def apply_nature_axis_style(ax, *, style: PlotStyle | None = None):
if style is not None and style.tick_label_size is not None
else plt.rcParams.get("xtick.labelsize", 11)
)
unified_y_pad = (
style.unified_y_pad
if style is not None and style.unified_y_pad is not None
else None
)

ax.tick_params(reset=True)

Expand Down Expand Up @@ -64,10 +69,12 @@ def apply_nature_axis_style(ax, *, style: PlotStyle | None = None):
left=True,
right=False,
)

label_pad = 4 * scale
ax.xaxis.labelpad = label_pad
ax.yaxis.labelpad = label_pad
if unified_y_pad is None:
ax.yaxis.labelpad = label_pad
else:
ax.yaxis.labelpad = unified_y_pad

if style is not None and style.nbins is not None:
ax.yaxis.set_major_locator(MaxNLocator(nbins=style.nbins))
Expand Down
1 change: 1 addition & 0 deletions plot_styles/core/theme.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ class PlotStyle:
legend_loc: str | None = None
cms_label_fontsize: float | None = None
cms_label_y_start: float | None = None
unified_y_pad: float | None = None

def _resolved_sizes(self) -> dict:
tick_size = self.tick_label_size or self.base_font_size * 0.9
Expand Down
97 changes: 84 additions & 13 deletions plot_styles/grid_unrolled.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from plot_styles.core.legend import plot_legend
from plot_styles.core.style_axis import apply_nature_axis_style
from plot_styles.core.theme import PlotStyle
from plot_styles.style import MODEL_COLORS


Expand Down Expand Up @@ -170,6 +171,7 @@ def plot_unrolled_grid_with_winner_and_ratios(
points,
sic_by_model,
unc_by_model=None,
aux_by_model=None,
cutflow_df=None,
config=None,
*,
Expand All @@ -178,6 +180,11 @@ def plot_unrolled_grid_with_winner_and_ratios(
"""Plot an unrolled grid with optional winner strip and ratio panels."""
cfg = _merge_configs({}, config or {})
style = cfg.get("style", style)

# # hotfitx
# if style:
# style.unified_y_pad = getattr(cfg.get('style',{}), "unified_y_pad", None)

label_fontsize = cfg.get("label_fontsize", None)
label_fontsize_y = cfg.get("label_fontsize_y", label_fontsize)
label_fontsize_top = cfg.get("label_fontsize_top", label_fontsize)
Expand All @@ -200,19 +207,30 @@ def plot_unrolled_grid_with_winner_and_ratios(

ratios_cfg = cfg.get("ratios", [])
winner_enabled = cfg.get("winner", {}).get("enabled", True)
aux_cfg = cfg.get("aux_panel", {})
aux_enabled = aux_cfg.get("enabled", False) and aux_by_model is not None
cutflow_cfg = cfg.get("cutflow", {})
cutflow_enabled = cutflow_cfg.get("enabled", False) and cutflow_df is not None

nrows = 1 + (1 if cutflow_enabled else 0) + (1 if winner_enabled else 0) + len(ratios_cfg)
nrows = (
1
+ (1 if winner_enabled else 0)
+ (1 if aux_enabled else 0)
+ (1 if cutflow_enabled else 0)
+ len(ratios_cfg)
)
hr_main = cfg["height_ratios"]["main"]
hr_aux = cfg["height_ratios"].get("aux", 1.0)
hr_cutflow = cfg["height_ratios"].get("cutflow", 0.5)
hr_winner = cfg["height_ratios"]["winner"]
hr_ratio = cfg["height_ratios"]["ratio"]
height_ratios = [hr_main]
if cutflow_enabled:
height_ratios.append(hr_cutflow)
if winner_enabled:
height_ratios.append(hr_winner)
if aux_enabled:
height_ratios.append(hr_aux)
if cutflow_enabled:
height_ratios.append(hr_cutflow)
height_ratios += [hr_ratio] * len(ratios_cfg)

fig, axes = plt.subplots(
Expand All @@ -231,23 +249,29 @@ def plot_unrolled_grid_with_winner_and_ratios(
for ax in axes:
ax.set_xlim(-0.5, n_points - 0.5)

ax_main = axes[0]
ax_cutflow = axes[1] if cutflow_enabled else None
axis_index = 0
ax_main = axes[axis_index]
axis_index += 1
ax_winner = None
if winner_enabled:
ax_winner = axes[2] if cutflow_enabled else axes[1]
else:
ax_winner = None
ratio_start = 1
ax_winner = axes[axis_index]
axis_index += 1
ax_aux = None
if aux_enabled:
ax_aux = axes[axis_index]
axis_index += 1
ax_cutflow = None
if cutflow_enabled:
ratio_start += 1
if winner_enabled:
ratio_start += 1
ax_ratio_list = axes[ratio_start:]
ax_cutflow = axes[axis_index]
axis_index += 1
ax_ratio_list = axes[axis_index:]

if cfg.get("y_main_log", False):
ax_main.set_yscale("log")

boost_axes = [ax_main]
if ax_aux is not None:
boost_axes.append(ax_aux)
# if ax_winner is not None:
# boost_axes.append(ax_winner)
if ax_cutflow is not None:
Expand Down Expand Up @@ -317,6 +341,37 @@ def plot_unrolled_grid_with_winner_and_ratios(

draw_block_separators(ax_main, block_edges, cfg["block_separator"])

if aux_enabled:
aux_vals = series_from_dict(aux_by_model, models, ordered)
aux_line_cfg = dict(line_cfg)
aux_line_cfg.update(aux_cfg.get("line", {}))
metric_col = aux_cfg.get("metric_col", "effective_steps")
ylabel = aux_cfg.get("ylabel")
if ylabel is None:
if metric_col == "effective_steps":
ylabel = "Effective steps"
elif metric_col == "min_val_loss":
ylabel = "Min val loss"
else:
ylabel = metric_col.replace("_", " ")
for model in models:
values = np.array(aux_vals[model], dtype=float)
values[~np.isfinite(values)] = np.nan
values[values <= 0.0] = np.nan
ax_aux.plot(
x,
values,
marker=aux_line_cfg.get("marker", "o"),
markersize=aux_line_cfg.get("markersize", 2.5),
linewidth=aux_line_cfg.get("linewidth", 1.2),
color=model_colors[model],
)

ax_aux.set_ylabel(ylabel, fontsize=label_fontsize_y,labelpad = aux_cfg.get("ylabel_pad", 15))
if aux_cfg.get("y_log", False):
ax_aux.set_yscale("log")
draw_block_separators(ax_aux, block_edges, cfg["block_separator"])

if cutflow_enabled:
passed_lookup = {
(row["m_X"], row["m_Y"]): row["passed"]
Expand Down Expand Up @@ -457,10 +512,26 @@ def plot_unrolled_grid_with_winner_and_ratios(
)

axis_list = [ax_main, *ax_ratio_list]
if ax_aux is not None:
axis_list.append(ax_aux)
if ax_cutflow is not None:
axis_list.append(ax_cutflow)
if ax_winner is not None:
axis_list.append(ax_winner)
_apply_axis_style(axis_list, cfg.get("apply_axis_style", False), style)

fig.align_ylabels(axis_list)

for ax in axes[:-1]:
ax.tick_params(
axis="x",
which="both",
bottom=False, # remove tick marks
top=False,
labelbottom=False # remove tick labels
)
ax.set_xlabel("") # optional: remove x-label text

plot_legend(
fig,
active_models=cfg['raw_model_series'],
Expand Down
Loading