.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "starterkit/Ex13_plot_invariant_mass.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_starterkit_Ex13_plot_invariant_mass.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_starterkit_Ex13_plot_invariant_mass.py:


.. _exercise13_label:

Common Analysis Workflow Continued
==================================

.. hint::
    This example demonstrates another post ``basf2`` task. Examplary, it uses
    the merged ntuple files from the previous task to perform an invariant mass
    fit. For demonstartion purposes, the task does not contain the actual run
    logic but calls a fitting module that is defined at a different location. In
    this way, the pipeline gains more readability since each task only contains
    a method call instead of the actual, potentianally long logic.

In a separate module, we define the ``PlotInvariantMass`` class that
performs the actual plotting of the invariant mass.

.. GENERATED FROM PYTHON SOURCE LINES 18-79

.. code-block:: Python

    import numpy as np
    import matplotlib.pyplot as plt
    import pyarrow.parquet as pq
    from plothist import make_hist, plot_hist

    variable_dict = {
        "deltaE": {"label": "$\\Delta E$ [GeV]", "range": (-0.1, 0.1)},
        "Mbc": {"label": "$M_\\mathrm{{bc}}$ [GeV/c$^2$]", "range": (5.25, 5.29)},
    }


    class PlotMetaData:
        _label_store = {
            "mumu": "MC $\\mu^+ \\mu^-$",
            "ee": "MC $e^+ e^-$",
            "mumu_treeFit": "MC $\\mu^+ \\mu^-$ (TreeFit)",
            "ee_treeFit": "MC $e^+ e^-$ (TreeFit)",
        }

        def __init__(self, filename, event_type, treeFit) -> None:
            self.filename = filename
            self.event_type = event_type
            self.treeFit = treeFit

        @staticmethod
        def read_parquet(filename):
            table = pq.read_table(filename, use_pandas_metadata=True)
            return table.to_pandas(strings_to_categorical=True, deduplicate_objects=True, self_destruct=False)

        @property
        def df(self):
            return self.read_parquet(self.filename)

        @property
        def plot_label(self):
            return self._label_store[self.event_type + ("_treeFit" if self.treeFit else "")]


    class PlotInvariantMass:
        def __init__(self, input_data, resultpaths):
            # Expect the output of the following structure:
            # {"event_type": "mumu", "treeFit": False, "filename": "results/Merged/merged.parquet"}
            self.input_data = input_data
            self.resultpaths = resultpaths
            self.plotmetadata = [PlotMetaData(**data) for data in input_data]

        def plot(self):
            for resultpath in self.resultpaths:
                fig, ax = plt.subplots()
                variable = resultpath.split("/")[-1].split(".")[0]
                bins = np.linspace(*variable_dict[variable]["range"], 100)
                for data in self.plotmetadata:
                    hist = make_hist(data.df[variable], bins=bins)
                    plot_hist(hist, ax=ax, histtype="step", linewidth=1.2, label=data.plot_label)

                ax.set_xlabel(variable_dict[variable]["label"])
                ax.set_ylabel("Entries")
                ax.legend()
                fig.savefig(resultpath, bbox_inches="tight")



.. GENERATED FROM PYTHON SOURCE LINES 80-85

In the workflow pipeline, we only define the ``PlotInvariantMassTask`` which handles the
input and output files. The actual plotting is outsourced to the ``PlotInvariantMass``
class defined in a separate location and can potentially be reused in other tasks.
The task inherits from the ``MergerTask`` and uses the output files of the
``MergerTask`` as input files.

.. GENERATED FROM PYTHON SOURCE LINES 87-131

.. code-block:: Python

    from itertools import product
    import b2luigi
    from Ex12_merge_files import MergerTask


    @b2luigi.inherits(MergerTask, without=["event_type", "treeFit"])
    class PlotInvariantMassTask(b2luigi.Task):
        result_dir = "results/Plot"

        def requires(self):
            for event_type, tree_fit in product(["mumu", "ee"], [True, False]):
                yield MergerTask(n_events=self.n_events, n_loops=self.n_loops, event_type=event_type, treeFit=tree_fit)

        def output(self):
            yield self.add_to_output("deltaE.png")
            yield self.add_to_output("Mbc.png")

        def run(self):
            input_data = []
            for input in self.get_all_input_file_names():
                components = input.split("/")

                # Extract event_type and treeFit values
                event_type = [comp.split("=")[1] for comp in components if comp.startswith("event_type=")][0]
                tree_fit = [comp.split("=")[1] == "True" for comp in components if comp.startswith("treeFit=")][0]

                # Append the dictionary to the result list
                input_data.append({"event_type": event_type, "treeFit": tree_fit, "filename": input})

            # Call "general" plotting module that contains the actual task and
            # produces files according to the prefined `output` method
            plotter = PlotInvariantMass(
                input_data=input_data,
                resultpaths=[
                    self.get_output_file_name("deltaE.png"),
                    self.get_output_file_name("Mbc.png"),
                ],
            )

            plotter.plot()


    if __name__ == "__main__":
        b2luigi.process(PlotInvariantMassTask(n_events=100, n_loops=40), workers=40)


.. _sphx_glr_download_starterkit_Ex13_plot_invariant_mass.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: Ex13_plot_invariant_mass.py <Ex13_plot_invariant_mass.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: Ex13_plot_invariant_mass.zip <Ex13_plot_invariant_mass.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_