-
Notifications
You must be signed in to change notification settings - Fork 17
/
utils.jl
106 lines (89 loc) · 3.31 KB
/
utils.jl
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
"""
macro stack(into, structs...)
Stack the fields of multiple structs and create a new one.
The first argument is the name of the new struct followed by the
ones to be stacked. Parametric types are not supported and
the fieldnames needs to be unique.
Example:
@stack Baz Foo Bar
Creates `Baz` with the concatenated fields of `Foo` and `Bar`
"""
macro stack(into, structs...)
fields = []
for _struct in structs
names = fieldnames(getfield(__module__, _struct))
types = [fieldtype(getfield(__module__, _struct), field) for field in names]
for (n, t) in zip(names, types)
push!(fields, :($n::$t))
end
end
esc(
quote
struct $into
$(fields...)
end
end
)
end
"""
unpack(x::CompressionHeader)
Return the following information:
- Name of compression algorithm
- Level of the compression
- compressedbytes and uncompressedbytes according to [uproot3](https://github.com/scikit-hep/uproot3/blob/54f5151fb7c686c3a161fbe44b9f299e482f346b/uproot3/source/compressed.py#L132)
"""
function unpack(x::CompressionHeader)
algname = String(x.algo)
# shift without casting to `Int` will give you 0x00 because we're shifting 0 bits into UInt8
compressedbytes = x.c1 + (Int(x.c2) << 8) + (Int(x.c3) << 16)
uncompressedbytes = x.u1 + (Int(x.u2) << 8) + (Int(x.u3) << 16)
return algname, x.method, compressedbytes, uncompressedbytes
end
abstract type JaggType end
struct Nojagg <:JaggType end
struct Nooffsetjagg<:JaggType end
struct Offsetjagg <:JaggType end
struct Offsetjaggjagg <:JaggType end
function JaggType(f, branch, leaf)
# https://github.com/scikit-hep/uproot3/blob/54f5151fb7c686c3a161fbe44b9f299e482f346b/uproot3/interp/auto.py#L144
(match(r"\[.*\]", leaf.fTitle) !== nothing) && return Nooffsetjagg
leaf isa TLeafElement && leaf.fLenType==0 && return Offsetjagg
!hasproperty(branch, :fClassName) && return Nojagg
try
streamer = streamerfor(f, branch.fClassName).streamer.fElements.elements[1]
(streamer.fSTLtype == Const.kSTLvector) && return Offsetjagg
catch
end
return Nojagg
end
"""
parseTH(th::Dict{Symbol, Any})
Parse the output of [`TH`](@ref) into a tuple of `counts`, `edges`, and `sumw2`.
A `StatsBase.Histogram` can then be constructed with `Histogram(edges, counts)`.
TH1 and TH2 inputs are supported.
"""
function parseTH(th::Dict{Symbol, Any})
xmin = th[:fXaxis_fXmin]
xmax = th[:fXaxis_fXmax]
xnbins = th[:fXaxis_fNbins]
xbins = isempty(th[:fXaxis_fXbins]) ? range(xmin, xmax, length=xnbins+1) : th[:fXaxis_fXbins];
counts = th[:fN]
sumw2 = th[:fSumw2]
if th[:fYaxis_fNbins] > 1
ymin = th[:fYaxis_fXmin]
ymax = th[:fYaxis_fXmax]
ynbins = th[:fYaxis_fNbins]
ybins = isempty(th[:fYaxis_fXbins]) ? range(ymin, ymax, length=ynbins+1) : th[:fYaxis_fXbins];
counts = reshape(counts, (xnbins+2, ynbins+2))[2:end-1, 2:end-1]
sumw2 = reshape(sumw2, (xnbins+2, ynbins+2))[2:end-1, 2:end-1]
edges = (xbins, ybins)
else
counts = counts[2:end-1]
sumw2 = sumw2[2:end-1]
edges = (xbins,)
end
return counts, edges, sumw2
end
function samplefile(filename::AbstractString)
return ROOTFile(normpath(joinpath(@__DIR__, "../test/samples", filename)))
end