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

packages/pidplots/pidplots/data.py

@@ -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