feat: add hep.savelabels(), to save label variations

* fix: missing italics for cms.label

* docs: update savelabels

* chore: fix tests

README.md

@@ -59,7 +59,7 @@ import mplhep as hep
 plt.style.use(hep.style.ROOT)
 ```
 
-Experiment specific style are also available. **If the default styles are not what you need, I'd be happy to merge in new styles or modify the current ones.**
+Experiment specific style are also available. **If the default styles are not what you need, please open an issue.**
 
 Default experiment labels are also available.
 
@@ -91,6 +91,18 @@ hep.hist2dplot(H, xbins, ybins)
 
 # More Information
 
+### Save all labels at once
+- `hep.savelabels('test.png')` will produces 4 variation on experiment label
+  - "" -> "test_paper.png"
+  - "Preliminary" -> "test_pas.png"
+  - "Supplementary" -> "test_supp.png"
+  - "Work in Progress" -> "test_wip.png"
+- Options can also be specified manually
+  - `hep.savelabels('test', labels=["FOO", "BAR"], suffixes=["foo.pdf", "bar"])` will produce
+    - "FOO" -> "foo.pdf"
+    - "BAR" -> "test_bar.png"
+- Other components of `<experiment>.label()` will remain unchanged.
+
 ### Other styles:
 - `hep.set_style("fira")` - use Fira Sans
 - `hep.set_style("firamath")` - use Fira Math
@@ -130,7 +142,7 @@ http://www.gust.org.pl/projects/e-foundry/tex-gyre/heros
 You can compare yourself if the differences are meanigful below.
 
 <p float="center">
-  <img src="FontComp.png" width="400" />
+  <img src="img/FontComp.png" width="400" />
 </p>
 
 They are Tex Gyre Heros, Helvetica and Arial respecively.

src/mplhep/__init__.py

@@ -10,6 +10,7 @@
 from . import alice, atlas, cms, label, lhcb, plot
 from . import styles as style
 from ._tools import Config
+from .label import savelabels
 from .plot import (
     append_axes,
     box_aspect,
@@ -29,7 +30,7 @@
 rcParams = Config(
     label=Config(
         data=None,
-        paper=None,
+        kind=None,
         supplementary=None,
         year=None,
         lumi=None,
@@ -63,6 +64,7 @@
     "plot",
     "style",
     "label",
+    "savelabels",
     # Log plot functions
     "histplot",
     "hist2dplot",

src/mplhep/alice.py

@@ -24,9 +24,7 @@ def _alice_text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_text(
-        "ALICE", text=text, fontsize=28, loc=1, italic=(False, False), **kwargs
-    )
+    return label_base.exp_text("ALICE", text=text, fontsize=28, loc=1, **kwargs)
 
 
 def _alice_label(**kwargs):
@@ -37,9 +35,7 @@ def _alice_label(**kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_label(
-        exp="ALICE", fontsize=28, loc=1, italic=(False, False), rlabel="", **kwargs
-    )
+    return label_base.exp_label(exp="ALICE", fontsize=28, loc=1, rlabel="", **kwargs)
 
 
 @docstring.copy(label_base.exp_text)

src/mplhep/atlas.py

@@ -24,7 +24,9 @@ def _atlas_text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_text("ATLAS", text=text, italic=(True, False), **kwargs)
+        kwargs.setdefault("italic", (True, False))
+        kwargs.setdefault("loc", 4)
+    return label_base.exp_text("ATLAS", text=text, **kwargs)
 
 
 def _atlas_label(**kwargs):
@@ -35,7 +37,9 @@ def _atlas_label(**kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_label(exp="ATLAS", italic=(True, False), **kwargs)
+        kwargs.setdefault("italic", (True, False))
+        kwargs.setdefault("loc", 4)
+    return label_base.exp_label(exp="ATLAS", **kwargs)
 
 
 @mpl.docstring.copy(label_base.exp_text)

src/mplhep/cms.py

@@ -26,7 +26,8 @@ def _cms_text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_text("CMS", text=text, italic=(False, True), **kwargs)
+        kwargs.setdefault("italic", (False, True))
+    return label_base.exp_text("CMS", text=text, **kwargs)
 
 
 def _cms_label(**kwargs):
@@ -37,6 +38,7 @@ def _cms_label(**kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs[key] = value
+        kwargs.setdefault("italic", (False, True))
     return label_base.exp_label(exp="CMS", **kwargs)
 
 

src/mplhep/label.py

@@ -226,8 +226,7 @@ def exp_label(
     loc=0,
     *,
     data=False,
-    paper=False,
-    supplementary=False,
+    kind=None,
     year=None,
     lumi=None,
     llabel=None,
@@ -253,11 +252,11 @@ def exp_label(
             Axes object (if None, last one is fetched)
         data : bool, optional
             Prevents appending "Simulation" to experiment label.
-        paper : bool, optional
-            Prevents appending "Preliminary" to experiment label.
-        supplementary : bool, optional
-            Appends "Supplementary" to experiment label for public plots, included
-            in supplementary materials, but not paper body.
+        kind : str, optional
+            Options are:
+            - 'paper' : Without "Preliminary"
+            - 'supp' : Attaches "Supplementary"
+            - 'wip' : Attaches "Work in Progress"
         year : int, optional
             Year when data was collected
         lumi : float, optional
@@ -307,11 +306,12 @@ def exp_label(
         _label = ""
         if not data:
             _label = " ".join(["Simulation", _label])
-        if not paper:
-            _label = " ".join([_label, "Preliminary"])
-        if supplementary:
+        if kind == "supp":
             _label = " ".join([_label, "Supplementary"])
-
+        elif kind == "wip":
+            _label = " ".join([_label, "Work in Progress"])
+        elif kind != "paper":
+            _label = " ".join([_label, "Preliminary"])
         _label = " ".join(_label.split())
 
     exptext, expsuffix = exp_text(
@@ -331,7 +331,12 @@ def exp_label(
             units="inches",
             fig=ax.figure,
         )
-        _lumi = r"$\sqrt{s} = \mathrm{13\ TeV}, " + str(lumi) + r"\ \mathrm{fb}^{-1}$"
+        if lumi is not None:
+            _lumi = (
+                r"$\sqrt{s} = \mathrm{13\ TeV}, " + str(lumi) + r"\ \mathrm{fb}^{-1}$"
+            )
+        else:
+            _lumi = r"$\sqrt{s} = \mathrm{13\ TeV}$"
         explumi = ExpSuffix(
             *exptext.get_position(),
             text=rlabel if rlabel is not None else _lumi,
@@ -346,3 +351,74 @@ def exp_label(
         return exptext, expsuffix, explumi
 
     return exptext, expsuffix
+
+
+def savelabels(
+    fname: str,
+    ax: plt.Axes | None = None,
+    labels: list[str] | None = None,
+    suffixes: list[str] | None = None,
+    **kwargs,
+):
+    """
+    Save multiple copies of a figure on which axes is located with the
+    most commonly used label combinations (or user supplied).
+
+    Example:
+        >>> import mplhep as hep
+        >>> hep.cms.label(data=False)
+        >>> hep.savelabels('test.png')
+
+        This will produces 4 variation on experiment label:
+            - "" -> "test_paper.png"
+            - "Preliminary" -> "test_pas.png"
+            - "Supplementary" -> "test_supp.png"
+            - "Work in Progress" -> "test_wip.png"
+
+        Or the combination of labels and filenames can be set manually
+
+        >>> hep.savelabels('test', labels=["FOO", "BAR"], suffixes=["foo.pdf", "bar"])
+
+        Which will produce:
+        - "FOO" -> "foo.pdf"
+        - "BAR" -> "test_bar.png"
+
+    Parameters
+    ----------
+    fname : str
+        Primary filename to be passed to ``plt.savefig``.
+    ax : matplotlib.axes.Axes, optional
+            Axes object (if None, last one is fetched)
+    labels : list, optional
+        List of label versions to be produced.
+        By default ["", "Preliminary", "Supplementary", "Work in Progress"]
+    suffixes : list, optional
+        A matching list of suffixes to be appended to names of saved files. If
+        supplied strings contain suffixes such as ".png" the names will be assumed
+        to be absolute and will not incorporate ``fname``.
+        By default ["paper", "pas", "supp", "wip"]
+    """
+    if labels is None:
+        labels = ["", "Preliminary", "Supplementary", "Work in Progress"]
+    if suffixes is None:
+        suffixes = ["paper", "pas", "supp", "wip"]
+    if len(labels) != len(suffixes):
+        raise ValueError("Number of variations must match number of suffixes")
+    if ax is None:
+        ax = plt.gca()
+
+    label = [ch for ch in ax.get_children() if isinstance(ch, ExpSuffix)][0]
+    _sim = "Simulation" if "Simulation" in label.get_text() else ""
+
+    for label_text, suffix in zip(labels, suffixes):
+        label.set_text(" ".join([_sim, label_text]))
+
+        if "." in suffix:  # Expect full names
+            save_name = suffix
+        else:
+            if "." in fname:  # Expect
+                save_name = f"{fname.split('.')[0]}_{suffix}.{fname.split('.')[1]}"
+            else:
+                save_name = f"{fname}_{suffix}"
+
+        ax.figure.savefig(save_name, **kwargs)

src/mplhep/lhcb.py

@@ -39,15 +39,11 @@ def _lhcb_text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_text(
-        "LHCb",
-        text=text,
-        italic=(False, False),
-        fontsize=28,
-        fontname="Times New Roman",
-        loc=1,
-        **kwargs
-    )
+        kwargs.setdefault("italic", (False, False))
+        kwargs.setdefault("fontsize", 28)
+        kwargs.setdefault("fontname", "Times New Roman")
+        kwargs.setdefault("loc", 1)
+    return label_base.exp_text("LHCb", text=text, **kwargs)
 
 
 def _lhcb_label(**kwargs):
@@ -58,14 +54,11 @@ def _lhcb_label(**kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs[key] = value
-    return label_base.exp_label(
-        exp="LHCb",
-        italic=(False, False),
-        fontsize=28,
-        fontname="Times New Roman",
-        loc=1,
-        **kwargs
-    )
+        kwargs.setdefault("italic", (False, False))
+        kwargs.setdefault("fontsize", 28)
+        kwargs.setdefault("fontname", "Times New Roman")
+        kwargs.setdefault("loc", 1)
+    return label_base.exp_label("LHCb", **kwargs)
 
 
 @docstring.copy(label_base.exp_text)

tests/test_styles.py

@@ -31,7 +31,7 @@ def test_style_atlas():
     # Test suite does not have Helvetica
     plt.style.use([hep.style.ATLAS, {"font.sans-serif": ["Tex Gyre Heros"]}])
     fig, ax = plt.subplots()
-    hep.atlas.text()
+    hep.atlas.label()
 
     plt.rcParams.update(plt.rcParamsDefault)
     return fig
@@ -44,7 +44,7 @@ def test_style_cms():
 
     plt.style.use(hep.style.CMS)
     fig, ax = plt.subplots()
-    hep.cms.text()
+    hep.cms.label()
 
     plt.rcParams.update(plt.rcParamsDefault)
     return fig
@@ -57,7 +57,7 @@ def test_style_alice():
 
     plt.style.use(hep.style.ALICE)
     fig, ax = plt.subplots()
-    hep.alice.text()
+    hep.alice.label()
 
     plt.rcParams.update(plt.rcParamsDefault)
     return fig
@@ -231,11 +231,11 @@ def test_label_loc():
 def test_label_config(fig_test, fig_ref):
     hep.rcParams.label.data = True
     hep.rcParams.label.lumi = 30
-    hep.rcParams.label.paper = True
+    hep.rcParams.label.kind = "paper"
 
     test_ax = fig_test.subplots()
     hep.cms.label(data=False, ax=test_ax)
 
     ref_ax = fig_ref.subplots()
     hep.rcParams.clear()
-    hep.cms.label(data=False, lumi=30, paper=True, ax=ref_ax)
+    hep.cms.label(data=False, lumi=30, kind="paper", ax=ref_ax)