Core/python/histogram.py

0001 # This file is part of the ACTS project.
0002 #
0003 # Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 #
0005 # This Source Code Form is subject to the terms of the Mozilla Public
0006 # License, v. 2.0. If a copy of the MPL was not distributed with this
0007 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009
0010 def plot_histogram(self, ax=None, **kwargs):
0011     import matplotlib.pyplot as plt
0012     import mplhep
0013
0014     ax = ax or plt.subplots()[1]
0015     artists = mplhep.histplot(self._to_boost_histogram_(), ax=ax, **kwargs)
0016     if label := self.histogram.axis(0).label:
0017         ax.set_xlabel(label)
0018     if self.title:
0019         ax.set_title(self.title)
0020     return artists
0021
0022
0023 def plot_profile(self, ax=None, **kwargs):
0024     import matplotlib.pyplot as plt
0025     import mplhep
0026
0027     ax = ax or plt.subplots()[1]
0028     artists = mplhep.histplot(self._to_boost_histogram_(), ax=ax, **kwargs)
0029     if label := self.histogram.axis(0).label:
0030         ax.set_xlabel(label)
0031     if self.sampleAxisTitle:
0032         ax.set_ylabel(self.sampleAxisTitle)
0033     if self.title:
0034         ax.set_title(self.title)
0035     return artists
0036
0037
0038 def plot_efficiency(self, ax=None, **kwargs):
0039     import matplotlib.pyplot as plt
0040     import mplhep
0041
0042     ax = ax or plt.subplots()[1]
0043     xlabel = kwargs.pop("xlabel", self.total.axis(0).label)
0044     mplhep.comp.comparison(
0045         self.accepted._to_boost_histogram_(),
0046         self.total._to_boost_histogram_(),
0047         ax=ax,
0048         xlabel=xlabel,
0049         comparison="efficiency",
0050         **kwargs,
0051     )
0052     if self.title:
0053         ax.set_title(self.title)
0054
0055
0056 def _patch_plot_methods(m):
0057     m.Histogram1.plot = plot_histogram
0058     m.ProfileHistogram1.plot = plot_profile
0059     m.Efficiency1.plot = plot_efficiency
0060
0061
0062 def _patch_histogram_types(m):
0063     """Patch ACTS histogram types with the boost-histogram protocol.
0064
0065     Adds _to_boost_histogram_() to BoostHistogram, BoostProfileHistogram,
0066     Histogram1/2/3, ProfileHistogram1, and Efficiency1/2 so that
0067     bh.Histogram(acts_obj) works.
0068
0069     boost_histogram is imported lazily on first conversion; an ImportError
0070     is printed and re-raised if the package is not installed.
0071     """
0072
0073     def _bh_axes(bh, self):
0074         return [
0075             bh.axis.Variable(list(self.axis(i).edges), metadata=self.axis(i).label)
0076             for i in range(self.rank)
0077         ]
0078
0079     # BoostHistogram -> bh.Histogram with Double storage
0080     def _boost_hist_to_bh(self):
0081         import boost_histogram as bh
0082
0083         h = bh.Histogram(*_bh_axes(bh, self), storage=bh.storage.Double())
0084         h.view(flow=False)[:] = self.values()
0085         return h
0086
0087     m.BoostHistogram._to_boost_histogram_ = _boost_hist_to_bh
0088
0089     # BoostProfileHistogram -> bh.Histogram with Mean storage
0090     def _profile_to_bh(self):
0091         import boost_histogram as bh
0092
0093         h = bh.Histogram(*_bh_axes(bh, self), storage=bh.storage.Mean())
0094         view = h.view()
0095         view["count"] = self.counts()
0096         view["value"] = self.means()
0097         view["_sum_of_deltas_squared"] = self.sum_of_deltas_squared()
0098         return h
0099
0100     m.BoostProfileHistogram._to_boost_histogram_ = _profile_to_bh
0101
0102     # Histogram<Dim>: delegate to .histogram which is a BoostHistogram
0103     for cls in (m.Histogram1, m.Histogram2, m.Histogram3):
0104         cls._to_boost_histogram_ = lambda self: self.histogram._to_boost_histogram_()
0105
0106     # ProfileHistogram<1>: delegate to .histogram which is a BoostProfileHistogram
0107     m.ProfileHistogram1._to_boost_histogram_ = (
0108         lambda self: self.histogram._to_boost_histogram_()
0109     )
0110
0111     _patch_plot_methods(m)