WIP: New application pctaddnoise

Author: acoussat

applications/pctaddnoise/pctaddnoise.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+import argparse
+import json
+import sys
+import itk
+from itk import PCT as pct
+import numpy as np
+import numpy.lib.recfunctions as rfn
+import hepunits
+
+
+def build_parser():
+    parser = pct.PCTArgumentParser(
+        description="Add noise to ROOT data"
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        help="Root phase space file of particles",
+        required=True,
+    )
+    parser.add_argument("-t","--tree", help="Root tree name", required=True)
+    parser.add_argument("-o", "--output", help="Output file name", required=True)
+
+    parser.add_argument('--seed', help="Random number generator seed", type=int)
+    parser.add_argument('--noise-energy', help="Standard deviation of the Gaussian noise on the energy", type=float)
+    parser.add_argument('--noise-time', help="Standard deviation of the Gaussian noise on the time", type=float)
+
+    parser.add_argument(
+        "--verbose", "-v", help="Verbose execution", default=False, action="store_true"
+    )
+
+    tracker_uncertainty_group = parser.add_argument_group("Tracker uncertainty")
+    tracker_uncertainty_group.add_argument("--material-budget", help="Material budget x/x0", type=float, default=5e-3)
+    tracker_uncertainty_group.add_argument("--noise-position", help="Standard deviation of the tracker uncertainty (mm)", type=float)
+    tracker_uncertainty_group.add_argument("--tracker-distance", help="Distance between trackers (cm)", type=float)
+
+    return parser
+
+def get_sigma_sc(energy, x_over_x0, sp, dt):
+    proton_mass_c2 = 938.272013 * hepunits.MeV
+    betap = (energy + 2 * proton_mass_c2) * energy / (energy + proton_mass_c2)
+
+    # Equation (25) and (26) from [Krah et al, PMB, 2018]:
+    T = np.zeros((2, 2))
+    T[0, 1] = 1
+    T[1, 0] = -1 / dt
+    T[1, 1] = 1 / dt
+
+    sigma = (
+        13.6 * hepunits.MeV / betap * np.sqrt(x_over_x0) * (1 + 0.038 * np.log(x_over_x0))
+    )
+    sigma_sc = np.zeros((energy.size, 2, 2))
+    sigma_sc[:, 1, 1] = sigma**2
+    return np.tile(sp**2 * T @ T.T, (energy.size, 1, 1)) + sigma_sc
+
+def add_tracker_uncertainty(data, rng, material_budget, noise_position, tracker_distance):
+    e = data["KineticEnergy"]
+    sigma = get_sigma_sc(e, material_budget, noise_position * hepunits.mm, tracker_distance * hepunits.cm)
+    w, q = np.linalg.eig(np.linalg.inv(sigma))
+    xr = rng.standard_normal((e.size, 2, 2))
+    W = np.zeros((e.size, 2, 2))
+    W[:, 0, 0] = 1.0 / np.sqrt(w[:, 0])
+    W[:, 1, 1] = 1.0 / np.sqrt(w[:, 1])
+    dy_uncert = np.matmul(np.matmul(q, W), xr)
+    data['Position_X'] += dy_uncert[:, 0, 0]
+    data['Position_Y'] += dy_uncert[:, 0, 1]
+    data['Direction_X'] += dy_uncert[:, 1, 0]
+    data['Direction_Y'] += dy_uncert[:, 1, 1]
+
+def add_gaussian_noise(data, branch, rng, noise):
+    if noise is not None:
+        try:
+            data[branch] += rng.normal(scale=noise, size=len(data["KineticEnergy"]))
+        except KeyError:
+            print(f"Warning: cannot apply noise of {noise} on branch {branch} as the branch does not exist in the Root file! Skipping.", file=sys.stderr)
+
+def process(args_info: argparse.Namespace):
+    import uproot
+
+    rng = np.random.default_rng(args_info.seed)
+
+    if args_info.verbose:
+        print("Reading input Root file")
+    tree = uproot.open(args_info.input)[args_info.tree]
+    data = tree.arrays(library="np")
+
+    if args_info.verbose:
+        print("Applying noise…")
+
+    add_tracker_uncertainty(data, rng, args_info.material_budget, args_info.noise_position, args_info.tracker_distance)
+
+    add_gaussian_noise(data, "KineticEnergy", rng, args_info.noise_energy)
+    add_gaussian_noise(data, "LocalTime", rng, args_info.noise_time)
+
+    if args_info.verbose:
+        print("Writing output Root file…")
+    with uproot.recreate(args_info.output) as output_file:
+        output_file[args_info.tree] = data
+
+
+def main(argv=None):
+    parser = build_parser()
+    args_info = parser.parse_args(argv)
+    process(args_info)
+
+
+if __name__ == "__main__":
+    main()

applications/pctpairprotons/pctpairprotons.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python
 import argparse
 import json
-import sys
 import itk
 from itk import PCT as pct
 import numpy as np
@@ -73,95 +72,8 @@ def build_parser():
         "--psout", help="Name of tree in output phase space", default="PhaseSpace"
     )
 
-    data_noise_group = parser.add_argument_group("Data noise")
-    data_noise_group.add_argument('--noise-measurements', help="Standard deviation of the Gaussian noise on the measurements (energy or time)", type=float)
-    data_noise_group.add_argument('--noise-position', help="Standard deviation of the Gaussian noise on the position", type=float)
-    data_noise_group.add_argument('--tracker-distance', help="Distance between the two trackers of the upstream and downstream detectors, used to calculate noise on directions", type=float)
-    data_noise_group.add_argument('--seed', help="Random number generator seed", type=int)
-
     return parser
 
-def add_noise(pairs, measurement_column, noise_measurements, noise_position, tracker_distance, seed):
-
-    rng = np.random.default_rng(seed)
-
-    if noise_measurements is not None:
-        measurement_in = pairs[measurement_column + "_in"]
-        measurement_out = pairs[measurement_column + "_out"]
-        measurement_in += rng.normal(scale=noise_measurements, size=len(measurement_in))
-        measurement_out += rng.normal(scale=noise_measurements, size=len(measurement_out))
-
-    if noise_position is not None:
-
-        u_in = pairs["u_in"]
-        u_out = pairs["u_out"]
-        v_in = pairs["v_in"]
-        v_out = pairs["v_out"]
-
-        noise_u_in = u_in + rng.normal(scale=noise_position, size=len(u_in))
-        noise_u_out = u_out + rng.normal(scale=noise_position, size=len(u_out))
-        noise_v_in = v_in + rng.normal(scale=noise_position, size=len(v_in))
-        noise_v_out = v_out + rng.normal(scale=noise_position, size=len(v_out))
-
-        if tracker_distance is None:
-            print("Warning: noise on position was provided, but tracker distance is unspecified. No noise on directions will be applied.", file=sys.stderr)
-        else:
-            # Recover the point on the second tracker in each detector
-            w_in = pairs["w_in"]
-            w_out = pairs["w_out"]
-            du_in = pairs["du_in"]
-            du_out = pairs["du_out"]
-            dv_in = pairs["dv_in"]
-            dv_out = pairs["dv_out"]
-            dw_in = pairs["dw_in"]
-            dw_out = pairs["dw_out"]
-
-            w_in_2 = w_in - tracker_distance
-            slope_in = (w_in_2 - w_in) / dw_in
-            u_in_2 = u_in + slope_in * du_in
-            v_in_2 = v_in + slope_in * dv_in
-
-            w_out_2 = w_out + tracker_distance
-            slope_out = (w_out_2 - w_out) / dw_out
-            u_out_2 = u_out + slope_out * du_out
-            v_out_2 = v_out + slope_out * dv_out
-
-            # Add noise to the point on the second tracker
-            noise_u_in_2 = u_in_2 + rng.normal(scale=noise_position, size=len(u_in_2))
-            noise_v_in_2 = v_in_2 + rng.normal(scale=noise_position, size=len(v_in_2))
-            noise_u_out_2 = u_out_2 + rng.normal(scale=noise_position, size=len(u_out_2))
-            noise_v_out_2 = v_out_2 + rng.normal(scale=noise_position, size=len(v_out_2))
-
-            # Build a direction from these noisy points
-            noise_du_in = noise_u_in - noise_u_in_2
-            noise_dv_in = noise_v_in - noise_v_in_2
-            noise_dw_in = tracker_distance
-            noise_du_out = noise_u_out_2 - noise_u_out
-            noise_dv_out = noise_v_out_2 - noise_v_out
-            noise_dw_out = tracker_distance
-
-            norm_in = np.sqrt(noise_du_in**2 + noise_dv_in**2 + noise_dw_in**2)
-            norm_out = np.sqrt(noise_du_out**2 + noise_dv_out**2 + noise_dw_out**2)
-
-            noise_du_in /= norm_in
-            noise_dv_in /= norm_in
-            noise_dw_in /= norm_in
-            noise_du_out /= norm_out
-            noise_dv_out /= norm_out
-            noise_dw_out /= norm_out
-
-            pairs["du_in"] = noise_du_in
-            pairs["du_out"] = noise_du_out
-            pairs["dv_in"] = noise_dv_in
-            pairs["dv_out"] = noise_dv_out
-            pairs["dw_in"] = noise_dw_in
-            pairs["dw_out"] = noise_dw_out
-
-        pairs["u_in"] = noise_u_in
-        pairs["v_in"] = noise_v_in
-        pairs["u_out"] = noise_u_out
-        pairs["v_out"] = noise_v_out
-
 
 def process(args_info: argparse.Namespace):
     import uproot
@@ -176,7 +88,7 @@ def verbose(message):
         def verbose(message):
             pass
 
-    measurement_column = "PreGlobalTime" if args_info.store_time else "KineticEnergy"
+    measurement_column = "LocalTime" if args_info.store_time else "KineticEnergy"
 
     def load_tree_as_df(root_file, tree_name):
 
@@ -274,10 +186,6 @@ def load_tree_as_df(root_file, tree_name):
         np.recarray.sort(pairs, order=["RunID", "EventID", "TrackID_in", "TrackID_out"])
     verbose("Merged input and output phase spaces.")
 
-    if args_info.noise_measurements is not None or args_info.noise_position is not None:
-        verbose("Adding noise…")
-        add_noise(pairs, measurement_column, args_info.noise_measurements, args_info.noise_position, args_info.tracker_distance, args_info.seed)
-
     if args_info.fit is not None:
         verbose("Converting energy loss or TOF to WEPL…")
         with open(args_info.fit, encoding="utf-8") as f: