-
Notifications
You must be signed in to change notification settings - Fork 68
/
TProfile3D.py
117 lines (102 loc) · 4.31 KB
/
TProfile3D.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
"""
This module defines the behavior of ``TProfile3D``.
"""
from __future__ import annotations
import numpy
import uproot
import uproot.behaviors.TH3
import uproot.behaviors.TProfile
from uproot.behaviors.TH1 import boost_axis_metadata, boost_metadata
class TProfile3D(uproot.behaviors.TProfile.Profile):
"""
Behaviors for three-dimensional profiles: ROOT's ``TProfile3D``.
"""
no_inherit = (uproot.behaviors.TH3.TH3,)
@property
def axes(self):
return (self.member("fXaxis"), self.member("fYaxis"), self.member("fZaxis"))
def axis(self, axis):
if axis == 0 or axis == -3 or axis == "x":
return self.member("fXaxis")
elif axis == 1 or axis == -2 or axis == "y":
return self.member("fYaxis")
elif axis == 2 or axis == -1 or axis == "z":
return self.member("fZaxis")
else:
raise ValueError(
"axis must be 0 (-3), 1 (-2), 2 (-1) or 'x', 'y', 'z' for a TProfile3D"
)
@property
def weighted(self):
fBinSumw2 = self.member("fBinSumw2", none_if_missing=True)
return fBinSumw2 is None or len(fBinSumw2) != len(self.member("fNcells"))
def counts(self, flow=False):
fBinEntries = numpy.asarray(self.member("fBinEntries"))
out = uproot.behaviors.TProfile._effective_counts_1d(
fBinEntries.reshape(-1),
numpy.asarray(self.member("fBinSumw2")).reshape(-1),
self.member("fNcells"),
)
out = out.reshape(fBinEntries.shape)
if flow:
return out
else:
return out[1:-1, 1:-1, 1:-1]
def values(self, flow=False):
if hasattr(self, "_values"):
values = self._values
else:
(root_cont,) = self.base(uproot.models.TArray.Model_TArray)
root_cont = numpy.asarray(root_cont, dtype=numpy.float64)
values = uproot.behaviors.TProfile._values_1d(
numpy.asarray(self.member("fBinEntries")).reshape(-1),
root_cont.reshape(-1),
)
xaxis_fNbins = self.member("fXaxis").member("fNbins")
yaxis_fNbins = self.member("fYaxis").member("fNbins")
zaxis_fNbins = self.member("fZaxis").member("fNbins")
values = numpy.transpose(
values.reshape(zaxis_fNbins + 2, yaxis_fNbins + 2, xaxis_fNbins + 2)
)
self._values = values
if flow:
return values
else:
return values[1:-1, 1:-1, 1:-1]
def _values_errors(self, flow, error_mode):
attr = "_errors" + uproot.behaviors.TProfile._error_mode_str(error_mode)
if hasattr(self, attr):
values = self._values
errors = getattr(self, attr)
else:
(root_cont,) = self.base(uproot.models.TArray.Model_TArray)
root_cont = numpy.asarray(root_cont, dtype=numpy.float64)
fSumw2 = self.member("fSumw2", none_if_missing=True)
if fSumw2 is not None:
fSumw2 = numpy.asarray(fSumw2).reshape(-1)
values, errors = uproot.behaviors.TProfile._values_errors_1d(
error_mode,
numpy.asarray(self.member("fBinEntries")).reshape(-1),
root_cont.reshape(-1),
fSumw2,
self.member("fNcells"),
numpy.asarray(self.member("fBinSumw2")).reshape(-1),
)
xaxis_fNbins = self.member("fXaxis").member("fNbins")
yaxis_fNbins = self.member("fYaxis").member("fNbins")
zaxis_fNbins = self.member("fZaxis").member("fNbins")
values = numpy.transpose(
values.reshape(zaxis_fNbins + 2, yaxis_fNbins + 2, xaxis_fNbins + 2)
)
errors = numpy.transpose(
errors.reshape(zaxis_fNbins + 2, yaxis_fNbins + 2, xaxis_fNbins + 2)
)
self._values = values
setattr(self, attr, errors)
if flow:
return values, errors
else:
return values[1:-1, 1:-1, 1:-1], errors[1:-1, 1:-1, 1:-1]
def to_boost(self, metadata=boost_metadata, axis_metadata=boost_axis_metadata):
raise NotImplementedError("FIXME @henryiii: this one kinda doesn't exist")