feat: added subplot wrapper

Author: 0ctagon

src/mplhep/__init__.py

@@ -72,6 +72,7 @@
     sort_legend,
     yscale_anchored_text,
     yscale_legend,
+    subplots,
 )
 
 # Configs
@@ -148,4 +149,5 @@
     "style",
     "yscale_anchored_text",
     "yscale_legend",
+    "subplots",
 ]

src/mplhep/comparison_plotters.py

@@ -28,6 +28,7 @@
 from .utils import (
     _get_model_type,
     set_fitting_ylabel_fontsize,
+    subplots,
 )
 
 
@@ -111,15 +112,7 @@ def hists(
     _check_counting_histogram(h2_plottable)
 
     if fig is None and ax_main is None and ax_comparison is None:
-        figsize = plt.rcParams["figure.figsize"]
-        fig, (ax_main, ax_comparison) = plt.subplots(
-            nrows=2,
-            figsize=(figsize[0], figsize[1] * 1.25),
-            gridspec_kw={"height_ratios": [4, 1]},
-        )
-        fig.subplots_adjust(hspace=0.15)
-        ax_main.xaxis.set_ticklabels([])
-        ax_main.set_xlabel(" ")
+        fig, (ax_main, ax_comparison) = subplots(nrows=2)
     elif fig is None or ax_main is None or ax_comparison is None:
         msg = "Need to provide fig, ax_main and ax_comparison (or none of them)."
         raise ValueError(msg)
@@ -496,15 +489,7 @@ def data_model(
 
     if fig is None and ax_main is None and ax_comparison is None:
         if plot_only is None:
-            figsize = plt.rcParams["figure.figsize"]
-            fig, (ax_main, ax_comparison) = plt.subplots(
-                nrows=2,
-                figsize=(figsize[0], figsize[1] * 1.25),
-                gridspec_kw={"height_ratios": [4, 1]},
-            )
-            fig.subplots_adjust(hspace=0.15)
-            ax_main.xaxis.set_ticklabels([])
-            ax_main.set_xlabel(" ")
+            fig, (ax_main, ax_comparison) = subplots(nrows=2)
         elif plot_only == "ax_main":
             _, ax_comparison = plt.subplots()
             fig, ax_main = plt.subplots()

src/mplhep/utils.py

@@ -533,3 +533,49 @@ def _get_model_type(components):
     if all(callable(x) for x in components):
         return "functions"
     return "histograms"
+
+
+def subplots(
+    figsize: tuple[float, float] | None = None,
+    nrows: int = 1,
+    gridspec_kw: dict | None = None,
+    hspace: float = 0.15,
+    *args,
+    **kwargs,
+) -> tuple[plt.Figure, np.ndarray]:
+    """
+    Wrapper around plt.subplots to create a figure with multiple subplots
+
+    Parameters
+    ----------
+    figsize : tuple[float, float], optional
+        Figure size in inches.
+    nrows : int, optional
+        Number of rows in the subplot grid. Default is 2.
+    gridspec_kw : dict | None, optional
+        Additional keyword arguments for the GridSpec. Default is None.
+        If None is provided, this is set to {"height_ratios": [4, 1]}.
+    hspace : float, optional
+        Height spacing between subplots. Default is 0.15.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The created figure.
+    axes : np.ndarray
+        Array of Axes objects representing the subplots.
+    """
+    if gridspec_kw is None and nrows > 1:
+        gridspec_kw = {"height_ratios": [4*(1+0.25*(nrows-2)), *(1 for _ in range(nrows-1))]}
+    if figsize is None:
+        figsize = (plt.rcParams["figure.figsize"][0], plt.rcParams["figure.figsize"][1] * (1 + 0.25 * (nrows - 1)))
+
+    fig, axes = plt.subplots(nrows=nrows, figsize=figsize, gridspec_kw=gridspec_kw)
+    if nrows > 1:
+        fig.subplots_adjust(hspace=hspace)
+
+        for ax in axes[:-1]:
+            _ = ax.xaxis.set_ticklabels([])
+            ax.set_xlabel(" ")
+
+    return fig, axes