feat: Adding publication/supplemental info to label (#466)

* First running version (additional tweaking for text offset is required)

* Added test case for publication information

* Changing canvas size

* Fixed styling for loc 1,3

* Updating styling for loc=0 and 4

* Tighting the distance for loc=1

* Removed unneeded variable

* Tweaking positions

* Changed direction

* Adding reference image

---------

Co-authored-by: Andrzej Novak <[email protected]>

Author: yimuchen

src/mplhep/atlas.py

@@ -25,7 +25,7 @@ def text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (True, False))
+    kwargs.setdefault("italic", (True, False, True))
     kwargs.setdefault("loc", 4)
     return label_base.exp_text("ATLAS", text=text, **kwargs)
 
@@ -39,7 +39,7 @@ def label(label=None, **kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (True, False))
+    kwargs.setdefault("italic", (True, False, True))
     kwargs.setdefault("loc", 4)
     if label is not None:
         kwargs["label"] = label

src/mplhep/cms.py

@@ -25,7 +25,7 @@ def text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (False, True))
+    kwargs.setdefault("italic", (False, True, False))
     kwargs.setdefault("exp", "CMS")
     return label_base.exp_text(text=text, **kwargs)
 
@@ -39,7 +39,7 @@ def label(label=None, **kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (False, True))
+    kwargs.setdefault("italic", (False, True, False))
     if label is not None:
         kwargs["label"] = label
     kwargs.setdefault("exp", "CMS")

src/mplhep/label.py

@@ -25,16 +25,24 @@ def __repr__(self):
         )
 
 
+class SuppText(mtext.Text):
+    def __repr__(self):
+        return "supptext: Custom Text({}, {}, {})".format(
+            self._x, self._y, repr(self._text)
+        )
+
+
 def exp_text(
     exp="",
     text="",
+    supp="",
     loc=0,
     *,
     ax=None,
     fontname=None,
     fontsize=None,
     exp_weight="bold",
-    italic=(False, False),
+    italic=(False, False, False),
     pad=0,
 ):
     """Add typical LHC experiment primary label to the axes.
@@ -57,7 +65,7 @@ def exp_text(
             Defines size of "secondary label". Experiment label is 1.3x larger.
         exp_weight : string, optional
             Set fontweight of <exp> label. Default "bold".
-        italic : (bool, bool), optional
+        italic : (bool, bool, bool), optional
             Tuple of bools to switch which label is italicized
         pad : float, optional
             Additional padding from axes border in units of axes fraction size.
@@ -87,6 +95,14 @@ def exp_text(
         4: {"xy": (0.05, 0.9550 - pad), "va": "bottom"},
     }
 
+    loc3_dict = {
+        0: {"xy": (1.012, 1 + pad), "va": "top", "ha": "left"},
+        1: {"xy": (0.05, 0.945 - pad), "va": "top"},
+        2: {"xy": (0.05, 0.935 - pad), "va": "top"},
+        3: {"xy": (0.05, 0.940 - pad), "va": "top"},
+        4: {"xy": (0.05, 0.9450 - pad), "va": "top"},
+    }
+
     if loc not in [0, 1, 2, 3, 4]:
         raise ValueError(
             "loc must be in {0, 1, 2}:\n"
@@ -182,7 +198,52 @@ def dist(tup):
     )
     ax._add_text(expsuffix)
 
-    return exptext, expsuffix
+    if loc == 0:
+        # No transformation, fixed location
+        _t = mtransforms.offset_copy(exptext._transform, units="inches", fig=ax.figure)
+    elif loc == 1:
+        _t = mtransforms.offset_copy(
+            exptext._transform,
+            y=-exptext.get_window_extent().height / _dpi,
+            units="inches",
+            fig=ax.figure,
+        )
+    elif loc == 2:
+        _t = mtransforms.offset_copy(
+            expsuffix._transform,
+            y=-expsuffix.get_window_extent().height / _dpi,
+            units="inches",
+            fig=ax.figure,
+        )
+    elif loc == 3:
+        _t = mtransforms.offset_copy(
+            expsuffix._transform,
+            y=-expsuffix.get_window_extent().height / _dpi,
+            units="inches",
+            fig=ax.figure,
+        )
+    elif loc == 4:
+        _t = mtransforms.offset_copy(
+            exptext._transform,
+            y=-exptext.get_window_extent().height / _dpi,
+            units="inches",
+            fig=ax.figure,
+        )
+
+    supptext = SuppText(
+        *loc3_dict[loc]["xy"],
+        text=supp,
+        transform=_t,
+        ha=loc3_dict[loc].get("ha", "left"),
+        va=loc3_dict[loc]["va"],
+        fontsize=_font_size / 1.3,
+        fontname=fontname,
+        rotation=0 if loc != 0 else 90,
+        fontstyle="italic" if italic[2] else "normal",
+    )
+    ax._add_text(supptext)
+
+    return exptext, expsuffix, supptext
 
 
 # Lumi text
@@ -234,6 +295,7 @@ def exp_label(
     *,
     data=False,
     label="",
+    pub="",
     year=None,
     lumi=None,
     lumi_format="{0}",
@@ -244,7 +306,7 @@ def exp_label(
     fontsize=None,
     exp_weight="bold",
     pad=0,
-    italic=(False, False),
+    italic=(False, False, False),
     ax=None,
 ):
     """A convenience wrapper combining ``<exp>.text`` and ``lumitext`` providing for
@@ -283,7 +345,7 @@ def exp_label(
             Defines size of "secondary label". Experiment label is 1.3x larger.
         exp_weight : string, optional
             Set fontweight of <exp> label. Default "bold".
-        italic : (bool, bool), optional
+        italic : (bool, bool, bool), optional
             Tuple of bools to switch which label is italicized
         pad : float, optional
             Additional padding from axes border in units of axes fraction size.
@@ -321,16 +383,17 @@ def exp_label(
     if llabel is not None:
         _label = llabel
     else:
-        _label = ""
+        _label = label
+        if pub:
+            _label = " ".join(["Supplementary", _label])
         if not data:
-            _label = " ".join(["Simulation", label])
-        else:
-            _label = label
+            _label = " ".join(["Simulation", _label])
         _label = " ".join(_label.split())
 
-    exptext, expsuffix = exp_text(
+    exptext, expsuffix, supptext = exp_text(
         exp=exp,
         text=_label,
+        supp=pub if loc != 4 else "",  # Special handling for loc4
         loc=loc,
         ax=ax,
         fontname=fontname,
@@ -341,8 +404,8 @@ def exp_label(
     )
     if loc == 4:
         _t = mtransforms.offset_copy(
-            exptext._transform,
-            y=-exptext.get_window_extent().height / ax.figure.dpi,
+            supptext._transform,
+            y=-supptext.get_window_extent().height / ax.figure.dpi,
             units="inches",
             fig=ax.figure,
         )
@@ -356,16 +419,35 @@ def exp_label(
             *exptext.get_position(),
             text=rlabel if rlabel is not None else _lumi,
             transform=_t,
-            ha=exptext.get_ha(),
+            ha=supptext.get_ha(),
             va="top",
             fontsize=fontsize,
             fontname=fontname,
             fontstyle="normal",
         )
         ax._add_text(explumi)
-        return exptext, expsuffix, explumi
 
-    return exptext, expsuffix
+        _t = mtransforms.offset_copy(
+            explumi._transform,
+            y=-explumi.get_window_extent().height / ax.figure.dpi,
+            units="inches",
+            fig=ax.figure,
+        )
+        _font_size = rcParams["font.size"] if fontsize is None else fontsize
+        supptext = SuppText(
+            *explumi.get_position(),
+            text=pub,
+            transform=_t,
+            ha=explumi.get_ha(),
+            va="top",
+            fontsize=_font_size / 1.3,
+            fontname=fontname,
+            fontstyle="italic" if italic[2] else "normal",
+        )
+        ax._add_text(supptext)
+        return exptext, expsuffix, supptext, explumi
+
+    return exptext, expsuffix, supptext
 
 
 def savelabels(

src/mplhep/lhcb.py

@@ -39,7 +39,7 @@ def text(text="", **kwargs):
             and key in inspect.getfullargspec(label_base.exp_text).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (False, False))
+    kwargs.setdefault("italic", (False, False, False))
     kwargs.setdefault("fontsize", 28)
     kwargs.setdefault("fontname", "Times New Roman")
     kwargs.setdefault("loc", 1)
@@ -56,7 +56,7 @@ def label(label=None, **kwargs):
             and key in inspect.getfullargspec(label_base.exp_label).kwonlyargs
         ):
             kwargs.setdefault(key, value)
-    kwargs.setdefault("italic", (False, False))
+    kwargs.setdefault("italic", (False, False, False))
     kwargs.setdefault("fontsize", 28)
     kwargs.setdefault("fontname", "Times New Roman")
     kwargs.setdefault("exp_weight", "normal")

tests/baseline/test_pub_loc.png

@@ -220,6 +220,14 @@ def test_label_loc():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style="default")
+def test_pub_loc():
+    fig, axs = plt.subplots(2, 5, figsize=(20, 8))
+    for i, ax in enumerate(axs.flatten()):
+        hep.cms.label(loc=i % 5, ax=ax, lumi=50, pub="arXiv:aaaa.bbbbb", data=(i >= 5))
+    return fig
+
+
 @check_figures_equal(extensions=["pdf"])
 def test_label_config(fig_test, fig_ref):
     hep.rcParams.label.data = True