plothist 1.4.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

plothist/plothist_style.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-from importlib.resources import files
-
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 
 
-def set_style(style="default"):
+def set_style(style: str = "default") -> None:
     """
     Set the plothist style.
 
@@ -32,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.
@@ -75,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
@@ -116,10 +112,12 @@ def cubehelix_palette(
     pal = cmap(x)[:, :3].tolist()
     if reverse:
         pal = pal[::-1]
-    return pal
+    return [tuple(c) for c in 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.
 
@@ -132,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
     ----------
@@ -171,10 +170,10 @@ def get_color_palette(cmap, N):
         )
 
     plt_cmap = plt.get_cmap(cmap)
-    return plt_cmap(np.linspace(0, 1, N))
+    return [tuple(k) for k in 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.
 
@@ -190,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()
@@ -214,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.
 
@@ -229,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.
@@ -279,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,
@@ -303,19 +301,19 @@ def add_text(
 
 
 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").
 
@@ -323,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.
@@ -383,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.
@@ -392,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
@@ -418,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