python/plotting/plot_helpers.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 from collections import namedtuple
0010 import numpy as np
0011
0012 # ------------------------------------------------------------------------------
0013 # Common helper types to configure plots
0014 # ------------------------------------------------------------------------------
0015
0016 """ Pass plotting data between functions """
0017 plt_data = namedtuple(
0018     "plt_data",
0019     "fig axes lgd data bins mu rms errors",
0020     defaults=[None, None, None, None, 0, -1, -1, None],
0021 )
0022
0023 """ Configuration for plot axes """
0024 axis_options = namedtuple(
0025     "axis_options",
0026     "label min max log_scale tick_positions label_format",
0027     defaults=["x", None, None, None, None, None],
0028 )
0029
0030 """ Configuration for plot legends """
0031 legend_options = namedtuple(
0032     "legend_options",
0033     "title loc ncol colspacing handletextpad horiz_anchor vert_anchor",
0034     defaults=[None, "upper right", 1, 1, 1, 0.5, 0.5],
0035 )
0036
0037 # ------------------------------------------------------------------------------
0038 # Common helpers for plotting measurement data
0039 # ------------------------------------------------------------------------------
0040
0041 """ Filter the data in a data frame by a given prescription """
0042
0043
0044 def filter_data(data, filter=lambda df: [], variables=[]):
0045     data_coll = []
0046
0047     # Get global data
0048     if len(filter(data)) == 0:
0049         for var in variables:
0050             data_coll.append(data[var].to_numpy(dtype=np.double))
0051
0052     # Filtered data
0053     else:
0054         filtered = data.loc[filter]
0055         for var in variables:
0056             data_coll.append(filtered[var].to_numpy(dtype=np.double))
0057
0058     return data_coll[0] if len(data_coll) == 1 else tuple(data_coll)