Adding Functionality to Store Lines of Plots

examples/plot_basic_histogram.py

@@ -31,11 +31,34 @@
     atlas_second_tag="Example for more general plot",
     figsize=(6, 5),
     n_ratio_panels=1,
+    bin_array_path="plot_histo_example.pkl",
 )
 
 # Add histograms and plot
-plot_histo.add(expectation_hist, reference=True)
-plot_histo.add(measurement_hist)
-plot_histo.draw()
+plot_histo.add(expectation_hist, key="expectation", reference=True)
+plot_histo.add(measurement_hist, key="measurement")
 
+plot_histo.draw()
 plot_histo.savefig("histogram_basic_example.png", transparent=False)
+
+# Using the stored lines from plot_histo_example.pkl to plot again
+plot_hist_from_file = HistogramPlot(
+    ylabel="Number of events",
+    xlabel="Invariant mass $m$ [a.u.]",
+    logy=False,
+    bins=50,
+    bins_range=(1.1, 4),
+    norm=False,
+    atlas_first_tag="Simulation Internal",
+    atlas_second_tag="Example of a plot loaded from file",
+    figsize=(6, 5),
+    n_ratio_panels=1,
+    bin_array_path="plot_histo_example.pkl",
+)
+
+# Add histograms and plot
+plot_hist_from_file.add(key="expectation", reference=True)
+plot_hist_from_file.add(key="measurement")
+
+plot_hist_from_file.draw()
+plot_hist_from_file.savefig("histogram_basic_example_from_pickle.png", transparent=False)

examples/plot_basic_line_plot.py

@@ -16,10 +16,27 @@
     xlabel="Epoch",
     ylabel="Loss",
     atlas_second_tag="This is an example of a basic line plot",
+    bin_array_path="plot_line2d_example.pkl",
 )
 
 # Add and draw the lines
-line_plot.add(Line2D(epochs, training_loss, label="Training loss"))
-line_plot.add(Line2D(epochs, validation_loss, label="Validation loss"))
+line_plot.add(Line2D(epochs, training_loss, label="Training loss"), key="train")
+line_plot.add(Line2D(epochs, validation_loss, label="Validation loss"), key="val")
+
 line_plot.draw()
 line_plot.savefig("line_plot_example.png", transparent=False)
+
+# Plot from pickle file
+line_plot_from_file = Line2DPlot(
+    xlabel="Epoch",
+    ylabel="Loss",
+    atlas_second_tag="This is an example of a basic line plot from the pickle file",
+    bin_array_path="plot_line2d_example.pkl",
+)
+
+# Add and draw the lines
+line_plot_from_file.add(key="train")
+line_plot_from_file.add(key="val")
+
+line_plot_from_file.draw()
+line_plot_from_file.savefig("line_plot_example_from_file.png", transparent=False)

examples/plot_pt_vs_eff.py

@@ -92,9 +92,10 @@
     logy=False,
     atlas_second_tag="$\\sqrt{s}=13$ TeV, dummy jets \ndummy sample, $f_{c}=0.018$",
     n_ratio_panels=1,
+    bin_array_path="plot_VarVsEff_example.pkl",
 )
-plot_bkg_rej.add(rnnip_light, reference=True)
-plot_bkg_rej.add(dips_light)
+plot_bkg_rej.add(rnnip_light, key="rnnip_ujets", reference=True)
+plot_bkg_rej.add(dips_light, key="dips_ujets")
 
 plot_bkg_rej.draw()
 plot_bkg_rej.savefig("pt_light_rej.png")
@@ -107,8 +108,8 @@
     atlas_second_tag="$\\sqrt{s}=13$ TeV, dummy jets, \ndummy sample, $f_{c}=0.018$",
     n_ratio_panels=1,
 )
-plot_sig_eff.add(rnnip_light, reference=True)
-plot_sig_eff.add(dips_light)
+plot_sig_eff.add(rnnip_light, key="rnnip_bjets", reference=True)
+plot_sig_eff.add(dips_light, key="dips_bjets")
 
 plot_sig_eff.atlas_second_tag += "\nInclusive $\\epsilon_b=70\\%$"
 
@@ -118,3 +119,20 @@
 # Drawing a hline indicating inclusive efficiency
 plot_sig_eff.draw_hline(0.7)
 plot_sig_eff.savefig("pt_b_eff.png", transparent=False)
+
+
+# Plot again the background rejection, but this time using the pickle file
+plot_bkg_rej_from_file = VarVsEffPlot(
+    mode="bkg_rej",
+    ylabel="Light-flavour jets rejection",
+    xlabel=r"$p_{T}$ [GeV]",
+    logy=False,
+    atlas_second_tag="Example plot from pickle file",
+    n_ratio_panels=1,
+    bin_array_path="plot_VarVsEff_example.pkl",
+)
+plot_bkg_rej_from_file.add(key="rnnip_ujets", reference=True)
+plot_bkg_rej_from_file.add(key="dips_ujets")
+
+plot_bkg_rej_from_file.draw()
+plot_bkg_rej_from_file.savefig("pt_light_rej_from_pickle_file.png")

examples/plot_rocs.py

@@ -62,6 +62,7 @@
     atlas_second_tag="$\\sqrt{s}=13$ TeV, dummy jets \ndummy sample, $f_{c}=0.018$",
     figsize=(6.5, 6),
     y_scale=1.4,
+    bin_array_path="plot_roc_example.pkl",
 )
 plot_roc.add_roc(
     Roc(
@@ -72,6 +73,7 @@
         signal_class="bjets",
         label="RNNIP",
     ),
+    key="rnnip_ujets",
     reference=True,
 )
 plot_roc.add_roc(
@@ -83,6 +85,7 @@
         signal_class="bjets",
         label="DIPS r22",
     ),
+    key="dips_ujets",
 )
 plot_roc.add_roc(
     Roc(
@@ -93,6 +96,7 @@
         signal_class="bjets",
         label="RNNIP",
     ),
+    key="rnnip_cjets",
     reference=True,
 )
 plot_roc.add_roc(
@@ -104,10 +108,34 @@
         signal_class="bjets",
         label="DIPS r22",
     ),
+    key="dips_cjets",
 )
 # setting which flavour rejection ratio is drawn in which ratio panel
 plot_roc.set_ratio_class(1, "ujets")
 plot_roc.set_ratio_class(2, "cjets")
 
 plot_roc.draw()
 plot_roc.savefig("roc.png", transparent=False)
+
+# Using the stored lines from plot_roc_example.pkl to plot again
+plot_roc_from_file = RocPlot(
+    n_ratio_panels=2,
+    ylabel="Background rejection",
+    xlabel="$b$-jet efficiency",
+    atlas_second_tag="Example of a plot loaded from file",
+    figsize=(6.5, 6),
+    y_scale=1.4,
+    bin_array_path="plot_roc_example.pkl",
+)
+
+plot_roc_from_file.add_roc(key="rnnip_ujets", reference=True)
+plot_roc_from_file.add_roc(key="dips_ujets")
+plot_roc_from_file.add_roc(key="rnnip_cjets", reference=True)
+plot_roc_from_file.add_roc(key="dips_cjets")
+
+# setting which flavour rejection ratio is drawn in which ratio panel
+plot_roc_from_file.set_ratio_class(1, "ujets")
+plot_roc_from_file.set_ratio_class(2, "cjets")
+
+plot_roc_from_file.draw()
+plot_roc_from_file.savefig("roc_from_pickle.png", transparent=False)

puma/histogram.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import pickle
+
 import matplotlib as mpl
 import numpy as np
 import pandas as pd
@@ -314,33 +316,51 @@
             raise ValueError("Not more than one ratio panel supported.")
         self.initialise_figure()
 
+        # Check if a pickle file is given for the plot
+        if self.bin_array_path and self.bin_array_path.is_file():
+            # Load the pickle file
+            with open(self.bin_array_path, "rb") as f:
+                self.plot_objects = pickle.load(f)
+
     def add(
         self,
-        histogram: Histogram,
+        histogram: Histogram | None = None,
         key: str | None = None,
         reference: bool = False,
     ):
         """Adding histogram object to figure.
 
         Parameters
         ----------
-        histogram : Histogram class
-            Histogram curve
-        key : str, optional
-            Unique identifier for histogram, by default None
+        histogram : Histogram class, optional
+            Histogram curve. Need to be provided if no bin array file was given when the
+            plot was initialised. By default None
+        key : str
+            Unique identifier for histogram. Should be set when saving/loading the bin arrays.
+            Otherwise, the order is not changeable later. By default None
         reference : bool, optional
             If this histogram is used as reference for ratio calculation, by default
             False
 
         Raises
         ------
+        ValueError
+            If no Histogram is provided and no bin array file was used in the creation
+            of the plot instance
         KeyError
             If unique identifier key is used twice
         """
         if key is None:
             key = len(self.plot_objects) + 1
-        if key in self.plot_objects:
-            raise KeyError(f"Duplicated key {key} already used for unique identifier.")
+
+        if self.bin_array_path and self.bin_array_path.is_file():
+            histogram = self.plot_objects[key]
+
+        elif histogram is None:
+            raise ValueError("No Histogram provided for addition!")
+
+        elif key in self.plot_objects:
+            raise KeyError(f"Duplicated key! {key} already used as unique identifier.")
 
         # Add key to histogram object
         histogram.key = key
@@ -617,6 +637,11 @@
                 )
             )
 
+        # If a pickle filepath is given and the file doesn't exist, create it and dump the data
+        if self.bin_array_path and not self.bin_array_path.is_file():
+            with open(self.bin_array_path, "wb") as f:
+                pickle.dump(self.plot_objects, f)
+
         if self.discrete_vals is not None:
             self.bins = bins
 

puma/line_plot_2d.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import pickle
+
 import matplotlib as mpl
 import numpy as np
 import pandas as pd
@@ -123,32 +125,49 @@
 
         self.initialise_figure()
 
+        # Check if a pickle file is given for the plot
+        if self.bin_array_path and self.bin_array_path.is_file():
+            # Load the pickle file
+            with open(self.bin_array_path, "rb") as f:
+                self.plot_objects = pickle.load(f)
+
     def add(
         self,
-        curve: object,
+        curve: Line2D | None = None,
         key: str | None = None,
         is_marker: bool = False,
     ):
         """Adding `puma.Line2D` object to figure.
 
         Parameters
         ----------
-        curve : puma.line_plot_2D.Line2D
-            Line2D object
+        curve : puma.line_plot_2D.Line2D, , optional
+            Line2D object. Need to be provided if no bin array file was given when the
+            plot was initialised. By default None
         key : str, optional
             Unique identifier for the curve, by default None
         is_marker : bool, optional
             Defines if this is a marker (True) or a line (False). By default False.
 
         Raises
         ------
+        ValueError
+            If no Line2D is provided and no bin array file was used in the creation
+            of the plot instance
         KeyError
             If unique identifier key is used twice
         """
         # If key not defined, set it to a numerical value
         if key is None:
             key = len(self.plot_objects) + 1
 
+        # Check if the Line2D objects are loaded from file
+        if self.bin_array_path and self.bin_array_path.is_file():
+            curve = self.plot_objects[key]
+
+        elif curve is None:
+            raise ValueError("No Line2D provided for addition!")
+
         # Check that key is not double used
         if key in self.plot_objects:
             raise KeyError(f"Duplicated key {key} already used for unique identifier.")
@@ -241,6 +260,12 @@
             )
 
         self.plotting_done = True
+
+        # If a pickle filepath is given and the file doesn't exist, create it and dump the data
+        if self.bin_array_path and not self.bin_array_path.is_file():
+            with open(self.bin_array_path, "wb") as f:
+                pickle.dump(self.plot_objects, f)
+
         return plt_handles
 
     def draw(self):

puma/plot_base.py

@@ -4,6 +4,7 @@
 
 import tkinter as tk
 from dataclasses import dataclass
+from pathlib import Path
 
 import atlasify
 from IPython import get_ipython
@@ -147,6 +148,10 @@
     plotting_done : bool
         Bool that indicates if plotting is done. Only then `atlasify()` can be called,
         by default False
+    bin_array_path : str
+        File in which the bin edges, values etc. are stored. If given, it is checked
+        that the file exists and then these values are taken for plotting instead of
+        recalculating everything. By default None
     """
 
     title: str = ""
@@ -196,6 +201,8 @@
 
     plotting_done: bool = False
 
+    bin_array_path: str = None
+
     def __post_init__(self):
         """Check for allowed values.
 
@@ -239,6 +246,13 @@
                 "You specified an ATLAS tag, but `apply_atlas_style` is set to false. "
                 "Tag will therefore not be shown on plot."
             )
+        if self.bin_array_path is not None:
+            # Check that the given value will be a Path object
+            if isinstance(self.bin_array_path, str):
+                self.bin_array_path = Path(self.bin_array_path)
+
+            elif not isinstance(self.bin_array_path, Path):
+                raise ValueError("Given bin_array_path is neither a string nor a Path object")
 
     def __check_figsize(self):
         """Check `figsize`.