Source code for lib.sty_functions

import matplotlib
import importlib.util
import plotly.express as px

from typing import Optional
from rich import print as rprint
from matplotlib import pyplot as plt

from .io_functions import check_key
from .fig_config import figure_features, add_grid

matplotlib.use("Qt5Agg")


[docs]def style_selector(OPT: Optional[dict] = None): if OPT is None: OPT = {} OPT["STYLE"] = "CIEMAT" elif isinstance(OPT, dict) == True and "STYLE" not in OPT.keys(): OPT["STYLE"] = "CIEMAT" styles = ["CIEMAT"] if check_key(OPT, "STYLE") == False: pass if OPT["STYLE"] == "CIEMAT": figure_features() if OPT["STYLE"] == "DUNE": import dunestyle.matplotlib as dune try: importlib.util.find_spec("dune") styles.append("DUNE") except: rprint("DUNE style not found") if OPT["STYLE"] == "HEP" or OPT["STYLE"] == "ATLAS": import mplhep as hep try: importlib.util.find_spec("hep") styles.append("HEP", "ATLAS") except: rprint("HEP style not found") if OPT["STYLE"] == "ATLAS": hep.style.use("ATLAS") # Include watermark to the top left of the plot # plt.text(0.02, 0.98, "Preliminary", fontsize=14, color='gray', alpha=0.5, transform=plt.gca().transAxes, ha='left', va='top') plt.rcParams.update({"font.size": 14}) matplotlib.rcParams["axes.prop_cycle"] = matplotlib.cycler(color=get_prism_colors()) rprint(f"* You can change your plotting style with [green]OPT[STYLE]={styles}![/green]")
[docs]def get_prism_colors(): prism = [c.split("rgb")[-1] for c in px.colors.qualitative.Prism] # Convert strng of tuple of int to tuple of int prism = [tuple(int(c) / 255 for c in s[1:-1].split(",")) for s in prism] # Change the order of the colors: insert the first color at the second to last position prism.insert(-1, prism.pop(0)) # rprint(prism) return prism
[docs]def get_color(number, even=True, style="prism", debug=False): if style == "prism": colors = get_prism_colors() else: # Use standard matplotlib colors colors = [c for c in plt.rcParams["axes.prop_cycle"].by_key()["color"]] if even: number = get_even_color(number) else: number = 1 + get_even_color(number) return colors[number]
[docs]def get_even_color(number): try: number = int(number) except ValueError: number = int(number.split("+")[0]) number = number - 2 if number % 2 == 0: pass else: number = number - 1 if number > 0 else number + 1 return number