From e636475c6d7ff2aa8e594571c91bca90985c6e0a Mon Sep 17 00:00:00 2001 From: Yulei Zhang Date: Wed, 14 Jan 2026 12:32:55 +0100 Subject: [PATCH 1/2] Add configurable auxiliary grid panel --- paper_plot.py | 42 +++++++++++++++++++++ plot_styles/grid_unrolled.py | 73 ++++++++++++++++++++++++++++++------ 2 files changed, 104 insertions(+), 11 deletions(-) diff --git a/paper_plot.py b/paper_plot.py index 65981b6d..313391ec 100644 --- a/paper_plot.py +++ b/paper_plot.py @@ -331,6 +331,7 @@ "hspace": 0, "height_ratios": { "main": 2.0, + "aux": 1.0, "cutflow": 0.8, "winner": 0.5, "ratio": 1.0, @@ -388,6 +389,12 @@ "zorder": 1, "log": False, }, + "aux_panel": { + "enabled": True, + "metric_col": "effective_steps", + "ylabel": "Effective steps", + "y_log": False, + }, "ratios": [], "ratio_line": { "marker": "o", @@ -1842,6 +1849,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, @@ -1880,11 +1911,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, diff --git a/plot_styles/grid_unrolled.py b/plot_styles/grid_unrolled.py index eb786a7a..6f90eef8 100644 --- a/plot_styles/grid_unrolled.py +++ b/plot_styles/grid_unrolled.py @@ -170,6 +170,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, *, @@ -200,15 +201,26 @@ 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 aux_enabled else 0) + + (1 if cutflow_enabled else 0) + + (1 if winner_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 aux_enabled: + height_ratios.append(hr_aux) if cutflow_enabled: height_ratios.append(hr_cutflow) if winner_enabled: @@ -231,23 +243,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 - if winner_enabled: - ax_winner = axes[2] if cutflow_enabled else axes[1] - else: - ax_winner = None - ratio_start = 1 + axis_index = 0 + ax_main = 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 + ax_cutflow = axes[axis_index] + axis_index += 1 + ax_winner = None if winner_enabled: - ratio_start += 1 - ax_ratio_list = axes[ratio_start:] + ax_winner = 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: @@ -317,6 +335,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) + 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"] @@ -457,6 +506,8 @@ 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) _apply_axis_style(axis_list, cfg.get("apply_axis_style", False), style) From 704dfef495c13490dc0e93443964e07c38f0b393 Mon Sep 17 00:00:00 2001 From: AvencastF Date: Wed, 14 Jan 2026 14:42:15 +0100 Subject: [PATCH 2/2] fix better grid study plot --- paper_plot.py | 48 +++++++++++++++++++++++++++------- plot_styles/core/style_axis.py | 11 ++++++-- plot_styles/core/theme.py | 1 + plot_styles/grid_unrolled.py | 36 +++++++++++++++++++------ 4 files changed, 76 insertions(+), 20 deletions(-) diff --git a/paper_plot.py b/paper_plot.py index 313391ec..2ae89ab2 100644 --- a/paper_plot.py +++ b/paper_plot.py @@ -328,10 +328,10 @@ "base_plot_config": { "models": [], "figsize": (15, 7), - "hspace": 0, + "hspace": 0.05, "height_ratios": { "main": 2.0, - "aux": 1.0, + "aux": 1.2, "cutflow": 0.8, "winner": 0.5, "ratio": 1.0, @@ -376,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, @@ -392,7 +392,7 @@ "aux_panel": { "enabled": True, "metric_col": "effective_steps", - "ylabel": "Effective steps", + "ylabel": "Updates [K]", "y_log": False, }, "ratios": [], @@ -420,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": [ @@ -455,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}, }, @@ -494,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, + }, }, }, }, @@ -2242,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 diff --git a/plot_styles/core/style_axis.py b/plot_styles/core/style_axis.py index f665ebb9..36a72f34 100644 --- a/plot_styles/core/style_axis.py +++ b/plot_styles/core/style_axis.py @@ -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) @@ -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)) diff --git a/plot_styles/core/theme.py b/plot_styles/core/theme.py index 105312b3..022b6c91 100644 --- a/plot_styles/core/theme.py +++ b/plot_styles/core/theme.py @@ -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 diff --git a/plot_styles/grid_unrolled.py b/plot_styles/grid_unrolled.py index 6f90eef8..af31690d 100644 --- a/plot_styles/grid_unrolled.py +++ b/plot_styles/grid_unrolled.py @@ -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 @@ -179,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) @@ -208,9 +214,9 @@ def plot_unrolled_grid_with_winner_and_ratios( nrows = ( 1 + + (1 if winner_enabled else 0) + (1 if aux_enabled else 0) + (1 if cutflow_enabled else 0) - + (1 if winner_enabled else 0) + len(ratios_cfg) ) hr_main = cfg["height_ratios"]["main"] @@ -219,12 +225,12 @@ def plot_unrolled_grid_with_winner_and_ratios( hr_winner = cfg["height_ratios"]["winner"] hr_ratio = cfg["height_ratios"]["ratio"] height_ratios = [hr_main] + if winner_enabled: + height_ratios.append(hr_winner) if aux_enabled: height_ratios.append(hr_aux) if cutflow_enabled: height_ratios.append(hr_cutflow) - if winner_enabled: - height_ratios.append(hr_winner) height_ratios += [hr_ratio] * len(ratios_cfg) fig, axes = plt.subplots( @@ -246,6 +252,10 @@ def plot_unrolled_grid_with_winner_and_ratios( axis_index = 0 ax_main = axes[axis_index] axis_index += 1 + ax_winner = None + if winner_enabled: + ax_winner = axes[axis_index] + axis_index += 1 ax_aux = None if aux_enabled: ax_aux = axes[axis_index] @@ -254,10 +264,6 @@ def plot_unrolled_grid_with_winner_and_ratios( if cutflow_enabled: ax_cutflow = axes[axis_index] axis_index += 1 - ax_winner = None - if winner_enabled: - ax_winner = axes[axis_index] - axis_index += 1 ax_ratio_list = axes[axis_index:] if cfg.get("y_main_log", False): @@ -361,7 +367,7 @@ def plot_unrolled_grid_with_winner_and_ratios( color=model_colors[model], ) - ax_aux.set_ylabel(ylabel, fontsize=label_fontsize_y) + 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"]) @@ -510,8 +516,22 @@ def plot_unrolled_grid_with_winner_and_ratios( 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'],