Commit 53313549 authored by connor.hainje@pnnl.gov's avatar connor.hainje@pnnl.gov
Browse files

Rework bin computation, add cutting method by bins, use central contrib computation

parent 8a4646d9
......@@ -23,25 +23,14 @@ def make_labels(d):
def make_bins(d):
# TODO: consider using np.digitize or pd.cut (?)
p = d["p"].values
p_bin = np.ones_like(p, dtype=int) * -1
# for i, r in enumerate(const.P_RANGES):
# p_bin[(p >= r[0]) & (p <= r[1])] = i
for i in range(const.N_P_BINS):
lo, hi = const.P_BINS[[i, i + 1]]
p_bin[(p >= lo) & (p <= hi)] = i
d["p_bin"] = p_bin
theta = d["theta"].values
theta_bin = np.ones_like(theta, dtype=int) * -1
# for i, r in enumerate(const.THETA_RANGES):
# theta_bin[(theta >= r[0]) & (theta <= r[1])] = i
for i in range(const.N_THETA_BINS):
lo, hi = const.THETA_BINS[[i, i + 1]]
theta_bin[(theta >= lo) & (theta <= hi)] = i
d["theta_bin"] = theta_bin
d["p_bin"] = np.digitize(d["p"].values, const.P_BINS) - 1
d["theta_bin"] = np.digitize(d["theta"].values, const.THETA_BINS) - 1
def cut_outside_bins(d):
d = d.loc[(d["p_bin"] >= 0) & (d["p_bin"] < const.N_P_BINS)]
d = d.loc[(d["theta_bin"] >= 0) & (d["theta_bin"] < const.N_THETA_BINS)]
return d
def make_lrs(d):
......@@ -98,13 +87,10 @@ def make_pid_abl(d):
def make_contrib(d):
for det in const.DETECTORS:
abl_lrs = compute_abl_lrs(d, det)
reg_lrs = np.stack([d[f"lr_{p}"].values for p in const.PARTICLES], axis=-1)
pids = d["pid"].values.astype(int)
reg_lr = reg_lrs[np.arange(len(pids)), pids]
abl_lr = abl_lrs[np.arange(len(pids)), pids]
ctrb = reg_lr - abl_lr
from .compute import compute_contrib
ctrbs = compute_contrib(d)
for det, ctrb in ctrbs.items():
d[f"contrib_{det}"] = ctrb
......@@ -180,7 +166,7 @@ def read_h5(
df.dropna(inplace=True)
df = df[df["labels"] >= 0]
if drop_outside_bins:
df = df[(df["p_bin"] != -1) & (df["theta_bin"] != -1)]
df = cut_outside_bins(df)
return df
......@@ -217,7 +203,7 @@ def read_npz(
df.dropna(inplace=True)
df = df[df["labels"] >= 0]
if drop_outside_bins:
df = df[(df["p_bin"] != -1) & (df["theta_bin"] != -1)]
df = cut_outside_bins(df)
return df
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment