-
Notifications
You must be signed in to change notification settings - Fork 4
/
mu.py
73 lines (63 loc) · 2.34 KB
/
mu.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
import argparse
from numpy.random.mtrand import shuffle
import pandas as pd
import numpy as np
import h5py
from scipy.optimize import minimize_scalar
from scipy.special import logsumexp
psr = argparse.ArgumentParser()
psr.add_argument("-o", dest="opt", type=str, help="output file")
psr.add_argument("ipt", type=str, help="input file")
psr.add_argument("--sparse", type=str, help="LucyDDM file")
psr.add_argument("--ref", type=str, help="truth file")
args = psr.parse_args()
sample = pd.read_hdf(args.ipt, "sample").set_index(["TriggerNo", "ChannelID"])
index = pd.read_hdf(args.sparse, "index").set_index(["TriggerNo", "ChannelID"])
mu0 = index["mu0"]
pe = pd.read_hdf(args.ref, "SimTriggerInfo/PEList").set_index(["TriggerNo", "PMTId"])
pe_count = pe.groupby(level=[0, 1])["Charge"].count()
t0_truth = pd.read_hdf(args.ref, "SimTruth/T").set_index(["TriggerNo", "ChannelID"])
def rescale(ent):
"""
rescale mu0 to a better mu.
"""
eid, cid = ent.index[0]
mu_t = mu0.loc[(eid, cid)]
NPE0 = int(mu_t + 0.5)
NPE_truth = pe_count.loc[(eid, cid)]
size = len(ent)
burn = size // 5
steps = ent["flip"].values
steps[np.abs(steps) == 2] = 0
NPE_evo = np.cumsum(np.insert(steps, 0, NPE0))[burn:]
NPE, counts = np.unique(NPE_evo, return_counts=True)
freq = counts / (size - burn)
loggN = -NPE * np.log(mu_t) + np.log(freq)
is_bracket = True
try:
assert NPE[0] <= mu_t and mu_t <= NPE[-1], "not a bracket"
rst = minimize_scalar(
lambda μ: μ - logsumexp(loggN + NPE * np.log(μ)),
bracket=(NPE[0], mu_t, NPE[-1]),
)
except (ValueError, AssertionError):
is_bracket = False
rst = minimize_scalar(
lambda μ: μ - logsumexp(loggN + NPE * np.log(μ)), bounds=(NPE[0], NPE[-1])
)
t0 = np.average(ent["t0"].values[burn:])
return pd.Series(
{
"mu": rst.x,
"is_bracket": is_bracket,
"mu0": mu_t,
"NPE_truth": NPE_truth,
"NPE_mean": np.average(NPE_evo),
"NPE_median": np.median(NPE_evo),
"t0": t0,
"t0_truth": t0_truth.loc[(eid, cid)]["T0"],
}
)
mu_fit = sample.groupby(level=[0, 1]).apply(rescale)
with h5py.File(args.opt, "w") as opt:
opt.create_dataset("mu", data=mu_fit.to_records(), compression="gzip", shuffle=True)