[benchmarker/artifacts] Improve matplotlib artifact capabilities and rendering

monitoring/benchmarker/artifacts/matplotlib/matplotlib_figure.py

@@ -54,6 +54,12 @@ def _configure_axis(
         else:
             ax_obj.set_ylabel(axis_spec.label)
 
+    if "margin" in axis_spec and axis_spec.margin is not None:
+        if is_x:
+            ax_obj.margins(x=axis_spec.margin)
+        else:
+            ax_obj.margins(y=axis_spec.margin)
+
     min_val = None
     if "min_value" in axis_spec and axis_spec.min_value is not None:
         min_val = axis_spec.min_value
@@ -101,12 +107,56 @@ def generate_matplotlib_figure(
         f"Generating Matplotlib figure artifact '{fig_spec.name}' -> {out_path}"
     )
 
+    max_sub_rows = max((s.n_subplot_rows for s in fig_spec.subfigures), default=1)
+    max_sub_cols = max((s.n_subplot_cols for s in fig_spec.subfigures), default=1)
+    fig_width = (
+        fig_spec.width
+        if "width" in fig_spec and fig_spec.width is not None
+        else 8 * fig_spec.n_subfigure_cols * max_sub_cols
+    )
+    fig_height = (
+        fig_spec.height
+        if "height" in fig_spec and fig_spec.height is not None
+        else 5 * fig_spec.n_subfigure_rows * max_sub_rows
+    )
+
     fig = plt.figure(
-        figsize=(8 * fig_spec.n_subfigure_cols, 5 * fig_spec.n_subfigure_rows)
+        figsize=(fig_width, fig_height),
+        layout="constrained",
     )
+
+    engine = fig.get_layout_engine()
+    if engine is not None:
+        engine_kwargs: dict[str, Any] = {
+            "w_pad": fig_spec.w_pad
+            if "w_pad" in fig_spec and fig_spec.w_pad is not None
+            else 0.08,
+            "h_pad": fig_spec.h_pad
+            if "h_pad" in fig_spec and fig_spec.h_pad is not None
+            else 0.12,
+        }
+        if "wspace" in fig_spec and fig_spec.wspace is not None:
+            engine_kwargs["wspace"] = fig_spec.wspace
+        if "hspace" in fig_spec and fig_spec.hspace is not None:
+            engine_kwargs["hspace"] = fig_spec.hspace
+        engine.set(**engine_kwargs)
+
     if "title" in fig_spec and fig_spec.title:
         fig.suptitle(fig_spec.title)
-    subfigs_res = fig.subfigures(fig_spec.n_subfigure_rows, fig_spec.n_subfigure_cols)
+
+    subfig_kwargs: dict[str, Any] = {}
+    if "wspace" in fig_spec and fig_spec.wspace is not None:
+        subfig_kwargs["wspace"] = fig_spec.wspace
+    elif fig_spec.n_subfigure_cols > 1:
+        subfig_kwargs["wspace"] = 0.08
+    if "hspace" in fig_spec and fig_spec.hspace is not None:
+        subfig_kwargs["hspace"] = fig_spec.hspace
+    elif fig_spec.n_subfigure_rows > 1:
+        subfig_kwargs["hspace"] = 0.08
+
+    subfigs_res = fig.subfigures(
+        fig_spec.n_subfigure_rows, fig_spec.n_subfigure_cols, **subfig_kwargs
+    )
     if isinstance(subfigs_res, np.ndarray):
         subfigs = list(subfigs_res.flatten())
     elif isinstance(subfigs_res, (list, tuple)):
@@ -156,8 +206,16 @@ def generate_matplotlib_figure(
         if "title" in subfig_spec and subfig_spec.title:
             subfig.suptitle(subfig_spec.title)
 
+        gridspec_kw: dict[str, Any] = {}
+        if "wspace" in subfig_spec and subfig_spec.wspace is not None:
+            gridspec_kw["wspace"] = subfig_spec.wspace
+        if "hspace" in subfig_spec and subfig_spec.hspace is not None:
+            gridspec_kw["hspace"] = subfig_spec.hspace
+
         axes_res = subfig.subplots(
-            subfig_spec.n_subplot_rows, subfig_spec.n_subplot_cols
+            subfig_spec.n_subplot_rows,
+            subfig_spec.n_subplot_cols,
+            gridspec_kw=gridspec_kw if gridspec_kw else None,
         )
         if isinstance(axes_res, np.ndarray):
             axes = list(axes_res.flatten())
@@ -192,32 +250,12 @@ def generate_matplotlib_figure(
             if "title" in subplot_spec and subplot_spec.title:
                 ax.set_title(subplot_spec.title)
 
-            if "x_axis" in subplot_spec and subplot_spec.x_axis:
-                _configure_axis(
-                    ax, subplot_spec.x_axis, True, subplot_interpreter, "x_axis"
-                )
-
-            if "y_axis" in subplot_spec and subplot_spec.y_axis:
-                _configure_axis(
-                    subplot_y_axes[0],
-                    subplot_spec.y_axis,
-                    False,
-                    subplot_interpreter,
-                    "y_axis",
-                )
-
             if "y_axes" in subplot_spec and subplot_spec.y_axes:
-                for i, sec_y_spec in enumerate(subplot_spec.y_axes):
-                    y_idx = i + 1
-                    while len(subplot_y_axes) <= y_idx:
-                        subplot_y_axes.append(ax.twinx())
-                    _configure_axis(
-                        subplot_y_axes[y_idx],
-                        sec_y_spec,
-                        False,
-                        subplot_interpreter,
-                        f"y_axes[{i}]",
-                    )
+                for i in range(len(subplot_spec.y_axes)):
+                    sec_ax = ax.twinx()
+                    if i >= 1:
+                        sec_ax.spines["right"].set_position(("outward", 60 * i))
+                    subplot_y_axes.append(sec_ax)
 
             for xy_plot in subplot_spec.xy_plots:
                 _, xyplot_interpreter = get_updated_context(
@@ -271,7 +309,13 @@ def generate_matplotlib_figure(
                 if "y_axis" in xy_plot and xy_plot.y_axis is not None:
                     y_axis_idx = xy_plot.y_axis
                 while len(subplot_y_axes) <= y_axis_idx:
-                    subplot_y_axes.append(ax.twinx())
+                    sec_ax = ax.twinx()
+                    offset_idx = len(subplot_y_axes) - 1
+                    if offset_idx >= 1:
+                        sec_ax.spines["right"].set_position(
+                            ("outward", 60 * offset_idx)
+                        )
+                    subplot_y_axes.append(sec_ax)
                 target_ax = subplot_y_axes[y_axis_idx]
 
                 if xy_plot.type == XYPlotType.Scatter:
@@ -283,6 +327,30 @@ def generate_matplotlib_figure(
                         f"XYPlotType '{xy_plot.type}' not implemented"
                     )
 
+            if "x_axis" in subplot_spec and subplot_spec.x_axis:
+                _configure_axis(
+                    ax, subplot_spec.x_axis, True, subplot_interpreter, "x_axis"
+                )
+
+            if "y_axis" in subplot_spec and subplot_spec.y_axis:
+                _configure_axis(
+                    subplot_y_axes[0],
+                    subplot_spec.y_axis,
+                    False,
+                    subplot_interpreter,
+                    "y_axis",
+                )
+
+            if "y_axes" in subplot_spec and subplot_spec.y_axes:
+                for i, sec_y_spec in enumerate(subplot_spec.y_axes):
+                    _configure_axis(
+                        subplot_y_axes[i + 1],
+                        sec_y_spec,
+                        False,
+                        subplot_interpreter,
+                        f"y_axes[{i}]",
+                    )
+
             if "legend" in subplot_spec and subplot_spec.legend:
                 legend_spec = subplot_spec.legend
                 legend_kwargs: dict[str, Any] = {}
@@ -311,6 +379,8 @@ def generate_matplotlib_figure(
                 else:
                     subplot_y_axes[-1].legend(**legend_kwargs)
 
-    plt.tight_layout()
-    fig.savefig(out_path, bbox_inches="tight")
+    save_kwargs: dict[str, Any] = {"bbox_inches": "tight"}
+    if "dpi" in fig_spec and fig_spec.dpi is not None:
+        save_kwargs["dpi"] = fig_spec.dpi
+    fig.savefig(out_path, **save_kwargs)
     plt.close(fig)

monitoring/benchmarker/configurations/artifacts/matplotlib_figure.py

@@ -9,6 +9,9 @@
 class AxisSpecification(ImplicitDict):
     label: Optional[str]
 
+    margin: Optional[float]
+    """If specified, margin fraction (e.g. 0.05 for 5%) to apply to this axis."""
+
     min_value: Optional[float]
     """If specified, the minimum value of this axis."""
 
@@ -139,6 +142,12 @@ class SubfigureSpecification(ImplicitDict):
     n_subplot_rows: int = 1
     n_subplot_cols: int = 1
 
+    wspace: Optional[float]
+    """If specified, width space between subplots within this subfigure."""
+
+    hspace: Optional[float]
+    """If specified, height space between subplots within this subfigure."""
+
     render_expr: Optional[ASTExpression]
     """If specified and this expression evaluates to false, skip rendering this subfigure."""
 
@@ -158,6 +167,27 @@ class MatplotlibFigureSpecification(ImplicitDict):
     n_subfigure_rows: int = 1
     n_subfigure_cols: int = 1
 
+    width: Optional[float]
+    """Figure width in inches. Defaults based on columns and subplots."""
+
+    height: Optional[float]
+    """Figure height in inches. Defaults based on rows and subplots."""
+
+    wspace: Optional[float]
+    """If specified, width space between subfigures."""
+
+    hspace: Optional[float]
+    """If specified, height space between subfigures."""
+
+    w_pad: Optional[float]
+    """Width padding for layout engine (in inches)."""
+
+    h_pad: Optional[float]
+    """Height padding for layout engine (in inches)."""
+
+    dpi: Optional[int]
+    """Figure DPI when saving to disk."""
+
     evaluation_context: Optional[list[SymbolExpression]]
     """Symbols available to other expressions in this figure specification."""
 

schemas/monitoring/benchmarker/configurations/artifacts/matplotlib_figure/AxisSpecification.json

@@ -13,6 +13,13 @@
         "null"
       ]
     },
+    "margin": {
+      "description": "If specified, margin fraction (e.g. 0.05 for 5%) to apply to this axis.",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
     "max_value": {
       "description": "If specified, the maximum value of this axis.",
       "type": [

schemas/monitoring/benchmarker/configurations/artifacts/matplotlib_figure/MatplotlibFigureSpecification.json

@@ -7,6 +7,13 @@
       "description": "Path to content that replaces the $ref",
       "type": "string"
     },
+    "dpi": {
+      "description": "Figure DPI when saving to disk.",
+      "type": [
+        "integer",
+        "null"
+      ]
+    },
     "evaluation_context": {
       "description": "Symbols available to other expressions in this figure specification.",
       "items": {
@@ -17,6 +24,27 @@
         "null"
       ]
     },
+    "h_pad": {
+      "description": "Height padding for layout engine (in inches).",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
+    "height": {
+      "description": "Figure height in inches. Defaults based on rows and subplots.",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
+    "hspace": {
+      "description": "If specified, height space between subfigures.",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
     "n_subfigure_cols": {
       "type": "integer"
     },
@@ -39,6 +67,27 @@
         "string",
         "null"
       ]
+    },
+    "w_pad": {
+      "description": "Width padding for layout engine (in inches).",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
+    "width": {
+      "description": "Figure width in inches. Defaults based on columns and subplots.",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
+    "wspace": {
+      "description": "If specified, width space between subfigures.",
+      "type": [
+        "number",
+        "null"
+      ]
     }
   },
   "required": [

schemas/monitoring/benchmarker/configurations/artifacts/matplotlib_figure/SubfigureSpecification.json

@@ -17,6 +17,13 @@
         "null"
       ]
     },
+    "hspace": {
+      "description": "If specified, height space between subplots within this subfigure.",
+      "type": [
+        "number",
+        "null"
+      ]
+    },
     "n_subplot_cols": {
       "type": "integer"
     },
@@ -41,6 +48,13 @@
         "string",
         "null"
       ]
+    },
+    "wspace": {
+      "description": "If specified, width space between subplots within this subfigure.",
+      "type": [
+        "number",
+        "null"
+      ]
     }
   },
   "required": [