const.py 3.9 KB
Newer Older
1
2
3
4
5
6
7
"""
const.py
Connor Hainje (connor.hainje@pnnl.gov)

Defines various constants that are useful for the PID studies.
"""

8
9
import numpy as np

10
PARTICLES = ["e", "mu", "pi", "K", "p", "d"]
11
"""The abbreviations of the six primary particle hypotheses"""
12
13

PDG_CODES = [11, 13, 211, 321, 2212, 1000010020]
14
"""The PDG codes associated with the six primary particle hypotheses"""
15

16
PDG_TO_PART = dict(zip(PDG_CODES, PARTICLES))
17
"""A map from PDG code to particle abbreviation"""
18

19
PART_TO_PDG = dict(zip(PARTICLES, PDG_CODES))
20
"""A map from particle abbreviation to PDG code"""
21
22

PART_TO_LABEL = dict(zip(PARTICLES, range(len(PARTICLES))))
23
"""A map from particle abbreviation to learning label"""
24
25

PDG_TO_LABEL = dict(zip(PDG_CODES, range(len(PDG_CODES))))
26
"""A map from PDG code to learning label"""
27

28
PART_PLOT_LABELS = [r"$e$", r"$\mu$", r"$\pi$", r"$K$", r"$p$", r"$d$"]
29
"""TeX labels for the six primary particle hypotheses"""
30

31
32
33
34
35
36
37
38
LOGL_PLOT_LABELS = [
    r"$\log \mathcal{L}_e$",
    r"$\log \mathcal{L}_\mu$",
    r"$\log \mathcal{L}_\pi$",
    r"$\log \mathcal{L}_K$",
    r"$\log \mathcal{L}_p$",
    r"$\log \mathcal{L}_d$",
]
39
"""TeX labels for the six hypothesis likelihoods"""
40
41

N_PARTICLES = len(PARTICLES)
42
"""Number of particles"""
43

44
DETECTORS = ["SVD", "CDC", "TOP", "ARICH", "ECL", "KLM"]
45
"""The names of the six PID detectors"""
46
47

N_DETECTORS = len(DETECTORS)
48
"""Number of detectors"""
49

50
51
52
53
54
55

def _compute_bin_centers(bins):
    return 0.5 * (bins[:-1] + bins[1:])


P_BINS = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.5])
56
57
"""Momentum bin edges."""

58
P_BIN_CENTERS = _compute_bin_centers(P_BINS)
59
60
"""Momentum bin centers."""

61
N_P_BINS = len(P_BIN_CENTERS)
62
63
"""Number of bins in momentum."""

64
P_BIN_LABELS = [f"{x:.1f}" for x in P_BINS]
65
66
"""Labels for momentum bin edges, for use in plotting routines."""

67
P_BIN_UNIT = "GeV"
68
"""Unit for momentum bins, for use in plotting routines."""
69
70

THETA_BINS = np.radians(np.array([17, 28, 40, 60, 77, 96, 115, 133, 150]))
71
72
"""Theta bin edges, specified in radians."""

73
THETA_BIN_CENTERS = _compute_bin_centers(THETA_BINS)
74
75
"""Theta bin centers, specified in radians."""

76
N_THETA_BINS = len(THETA_BIN_CENTERS)
77
78
"""Number of bins in theta."""

79
THETA_BIN_LABELS = [f"{x:.0f}" for x in np.degrees(THETA_BINS)]
80
81
"""Labels for theta bin edges, for use in plotting routines. Can be given in degrees or radians. Set the units to match."""

82
THETA_BIN_UNIT = "deg"
83
"""Unit for theta bins, for use in plotting routines."""
84
85
86


def update_p_bins(bins, labels=None, unit=None):
87
88
89
90
91
92
93
94
95
96
97
    """Update the global variables for momentum bins.

    Parameters
    ----------
    bins : array
        The momentum bin edges.
    labels : list of strings or None, optional
        Labels for the edges (used in plotting methods), by default None
    unit : string or None, optional
        Unit for the bin labels (used in plotting methods), by default None
    """
98
99
100
101
102
103
104
105
106
    global P_BINS, P_BIN_CENTERS, N_P_BINS, P_BIN_LABELS, P_BIN_UNIT
    P_BINS = bins
    P_BIN_CENTERS = _compute_bin_centers(P_BINS)
    N_P_BINS = len(P_BIN_CENTERS)
    P_BIN_LABELS = labels if labels is not None else [f"{x:.1f}" for x in P_BINS]
    P_BIN_UNIT = unit if unit is not None else "GeV"


def update_theta_bins(bins, labels=None, unit=None):
107
108
109
110
111
112
113
114
115
116
117
    """Update the global variables for theta bins.

    Parameters
    ----------
    bins : array
        The theta bin edges in radians.
    labels : list of strings or None, optional
        Labels for the bin edges (used in plotting methods), by default None
    unit : string or None, optional
        Unit for the bin labels (used in plotting methods), by default None
    """
118
119
120
121
    global THETA_BINS, THETA_BIN_CENTERS, N_THETA_BINS, THETA_BIN_LABELS, THETA_BIN_UNIT
    THETA_BINS = bins
    THETA_BIN_CENTERS = _compute_bin_centers(THETA_BINS)
    N_THETA_BINS = len(THETA_BIN_CENTERS)
122
123
124
125
    THETA_BIN_LABELS = (
        labels if labels is not None else [f"{x:.0f}" for x in np.degrees(THETA_BINS)]
    )
    THETA_BIN_UNIT = unit if unit is not None else "degrees"