plothist 1.3.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

plothist/plothist_style.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
-import matplotlib as mpl
-from importlib.resources import files
 
 
-def set_style(style="default"):
+def set_style(style: str = "default") -> None:
     """
     Set the plothist style.
 
@@ -29,22 +30,21 @@ def set_style(style="default"):
     available_styles = ["default"]
 
     if style in available_styles:
-        style_file = files("plothist").joinpath(f"{style}_style.mplstyle")
-        plt.style.use(style_file)
+        plt.style.use(f"plothist.{style}_style")
     else:
         raise ValueError(f"{style} not in the available styles: {available_styles}")
 
 
 def cubehelix_palette(
-    ncolors=7,
-    start=1.5,
-    rotation=1.5,
-    gamma=1.0,
-    hue=0.8,
-    lightest=0.8,
-    darkest=0.3,
-    reverse=True,
-):
+    ncolors: int = 7,
+    start: float = 1.5,
+    rotation: float = 1.5,
+    gamma: float = 1.0,
+    hue: float = 0.8,
+    lightest: float = 0.8,
+    darkest: float = 0.3,
+    reverse: bool = True,
+) -> list[tuple[float, float, float]]:
     """
     Make a sequential palette from the cubehelix system, in which the perceived brightness is linearly increasing.
     This code is adapted from seaborn, which implements equation (2) of reference [1] below.
@@ -72,10 +72,9 @@ def cubehelix_palette(
 
     Returns
     -------
-    list of RGB tuples
+    list[tuple[float, float, float]]
         The generated palette of colors represented as a list of RGB tuples.
 
-
     References
     ----------
     [1] Green, D. A. (2011). "A colour scheme for the display of astronomical
@@ -87,7 +86,7 @@ def cubehelix_palette(
         # Adapted from matplotlib
         def color(lambda_):
             # emphasise either low intensity values (gamma < 1),
-            # or high intensity values (γ > 1)
+            # or high intensity values (gamma > 1)
             lambda_gamma = lambda_**gamma
 
             # Angle and amplitude for the deviation
@@ -116,7 +115,9 @@ def cubehelix_palette(
     return pal
 
 
-def get_color_palette(cmap, N):
+def get_color_palette(
+    cmap: str, N: int
+) -> list[str] | list[tuple[float, float, float]]:
     """
     Get N different colors from a chosen colormap.
 
@@ -129,8 +130,9 @@ def get_color_palette(cmap, N):
 
     Returns
     -------
-    list
-        A list of RGB color tuples sampled from the colormap.
+    list[str] or list[tuple[float, float, float]]
+        A list of colors. If "ggplot" is selected, returns a list of hex color strings.
+        Otherwise, returns a list of RGB color tuples.
 
     References
     ----------
@@ -159,7 +161,7 @@ def get_color_palette(cmap, N):
             "#FFB5B8",
         ][0:N]
 
-    elif cmap == "cubehelix":
+    if cmap == "cubehelix":
         return cubehelix_palette(N)
 
     if N < 2:
@@ -171,7 +173,7 @@ def get_color_palette(cmap, N):
     return plt_cmap(np.linspace(0, 1, N))
 
 
-def set_fitting_ylabel_fontsize(ax):
+def set_fitting_ylabel_fontsize(ax: plt.Axes) -> float:
     """
     Get the suitable font size for a ylabel text that fits within the plot's y-axis limits.
 
@@ -187,9 +189,8 @@ def set_fitting_ylabel_fontsize(ax):
     """
     ylabel_fontsize = ax.yaxis.get_label().get_fontsize()
 
-    # Check if renderer is available
-    if ax.figure.canvas.get_renderer() is None:
-        ax.figure.canvas.draw()
+    # Force renderer to be initialized
+    ax.figure.canvas.draw()
 
     while (
         ax.yaxis.get_label()
@@ -211,14 +212,14 @@ def set_fitting_ylabel_fontsize(ax):
 
 
 def add_text(
-    text,
-    x="left",
-    y="top",
-    fontsize=12,
-    white_background=False,
-    ax=None,
+    text: str,
+    x: float | str = "left",
+    y: float | str = "top",
+    fontsize: int = 12,
+    white_background: bool = False,
+    ax: plt.Axes | None = None,
     **kwargs,
-):
+) -> None:
     """
     Add text to an axis.
 
@@ -226,9 +227,9 @@ def add_text(
     ----------
     text : str
         The text to add.
-    x : float, optional
+    x : float | str, optional
         Horizontal position of the text in unit of the normalized x-axis length. The default is value "left", which is an alias for 0.0. Other aliases are "right", "left_in", "right_in", "right_out".
-    y : float, optional
+    y : float | str, optional
         Vertical position of the text in unit of the normalized y-axis length. The default is value "top", which is an alias for 1.01. Other aliases are "top_in", "bottom_in", "top_out"="top", "bottom_out"="bottom".
     fontsize : int, optional
         Font size, by default 12.
@@ -276,14 +277,14 @@ def add_text(
     }
 
     if isinstance(x, str):
-        x = x_values.get(x)
-        if x is None:
-            raise ValueError(f"{x} not a float or a valid position")
+        if x not in x_values:
+            raise ValueError(f"{x!r} is not a valid x position.")
+        x = x_values[x]
 
     if isinstance(y, str):
-        y = y_values.get(y)
-        if y is None:
-            raise ValueError(f"{y} not a float or a valid position")
+        if y not in y_values:
+            raise ValueError(f"{y!r} is not a valid y position.")
+        y = y_values[y]
 
     t = ax.text(
         x,
@@ -296,23 +297,23 @@ def add_text(
 
     # Add background
     if white_background:
-        t.set_bbox(dict(facecolor="white", edgecolor="white"))
+        t.set_bbox({"facecolor": "white", "edgecolor": "white"})
 
 
 def add_luminosity(
-    collaboration,
-    x="right",
-    y="top",
-    fontsize=12,
-    is_data=True,
-    lumi="",
-    lumi_unit="fb",
-    preliminary=False,
-    two_lines=False,
-    white_background=False,
-    ax=None,
+    collaboration: str,
+    x: float | str = "right",
+    y: float | str = "top",
+    fontsize: int = 12,
+    is_data: bool = True,
+    lumi: int | str = "",
+    lumi_unit: str = "fb",
+    preliminary: bool = False,
+    two_lines: bool = False,
+    white_background: bool = False,
+    ax: plt.Axes | None = None,
     **kwargs,
-):
+) -> None:
     """
     Add the collaboration name and the integrated luminosity (or "Simulation").
 
@@ -320,17 +321,17 @@ def add_luminosity(
     ----------
     collaboration : str
         Collaboration name.
-    x : float, optional
+    x : float | str, optional
         Horizontal position of the text in unit of the normalized x-axis length. The default is value "right", which is an alias for 1.0. Can take other aliases such as "left", "left_in", "right_in", "right_out".
-    y : float, optional
+    y : float | str, optional
         Vertical position of the text in unit of the normalized y-axis length. The default is value "top", which is an alias for 1.01. Can take other aliases such as "top_in", "bottom_in", "top_out"="top", "bottom_out"="bottom".
     fontsize : int, optional
         Font size, by default 12.
     is_data : bool, optional
         If True, plot integrated luminosity. If False, plot "Simulation", by default True.
-    lumi : int/string, optional
+    lumi : int | str, optional
         Integrated luminosity. If empty, do not plot luminosity. Default value is empty.
-    lumi_unit : string, optional
+    lumi_unit : str, optional
         Integrated luminosity unit. Default value is fb. The exponent is automatically -1.
     preliminary : bool, optional
         If True, print "preliminary", by default False.
@@ -380,7 +381,7 @@ def add_luminosity(
     )
 
 
-def plot_reordered_legend(ax, order, **kwargs):
+def plot_reordered_legend(ax: plt.Axes, order: list[int], **kwargs) -> None:
     """
     Reorder the legend handlers and labels on the given Matplotlib axis based
     on the specified order and plot the reordered legend.
@@ -389,7 +390,7 @@ def plot_reordered_legend(ax, order, **kwargs):
     ----------
     ax : matplotlib.axes.Axes
         The Matplotlib Axes object on which the legend is to be reordered.
-    order : list of int
+    order : list[int]
         A list of integers specifying the new order of the legend items.
         The integers refer to the indices of the current legend items.
     kwargs : dict, optional
@@ -415,7 +416,6 @@ def plot_reordered_legend(ax, order, **kwargs):
     To reorder the legend so that 'Line 2' comes first, use:
 
     >>> plot_reordered_legend(ax, [1, 0])
-
     """
 
     # Extract the original handlers and labels
@@ -424,8 +424,7 @@ def plot_reordered_legend(ax, order, **kwargs):
     # Check if order is valid
     if not all(i in range(len(labels)) for i in order) or len(set(order)) < len(order):
         raise ValueError(
-            "The order list should contain all integers from 0 to "
-            f"{len(labels) - 1}."
+            f"The order list should contain all integers from 0 to {len(labels) - 1}."
         )
 
     # Reorder handlers and labels