Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-05-27 07:24:18

0001 # This file is part of the ACTS project.
0002 #
0003 # Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 #
0005 # This Source Code Form is subject to the terms of the Mozilla Public
0006 # License, v. 2.0. If a copy of the MPL was not distributed with this
0007 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008 
0009 # detray includes
0010 import plotting
0011 
0012 # python includes
0013 import numpy as np
0014 import pandas as pd
0015 import math
0016 import sys
0017 
0018 # Common options
0019 lgd_loc = "upper right"
0020 
0021 """ Read track position data """
0022 
0023 
0024 def read_track_data(file, logging):
0025     if file:
0026         # Preserve floating point precision
0027         df = pd.read_csv(file, float_precision="round_trip")
0028         logging.debug(df)
0029 
0030         return df
0031     else:
0032         logging.error("Could not find navigation data file: " + file)
0033         sys.exit(1)
0034 
0035 
0036 """ Plot the distributions of track parameter data """
0037 
0038 
0039 def plot_track_params(opts, detector, track_type, plot_factory, out_format, df):
0040 
0041     detector_name = detector.replace(" ", "_")
0042 
0043     fig_size = (8.5, 8.5)
0044 
0045     # Configure the plot legend
0046     lgd_ops = plotting.legend_options(
0047         loc=lgd_loc,
0048         ncol=4,
0049         colspacing=0.8,
0050         handletextpad=0.005,
0051         horiz_anchor=1.02,
0052         vert_anchor=1.12,
0053     )
0054 
0055     # Configure the y axes for all hists
0056     y_axis_opts = plotting.axis_options(label="")
0057 
0058     # Plot the track data
0059     def __create_plot(x, bins, suffix, x_axis, y_axis=y_axis_opts):
0060         hist_data = plot_factory.hist1D(
0061             x=x,
0062             bins=bins,
0063             x_axis=x_axis,
0064             y_axis=y_axis,
0065             lgd_ops=lgd_ops,
0066             show_stats=False,
0067             figsize=fig_size,
0068         )
0069 
0070         plot_factory.write_plot(
0071             hist_data, f"{detector_name}_{track_type}_{suffix}", out_format
0072         )
0073 
0074     # Plot the charge
0075     x_axis_opts = plotting.axis_options(label=r"$q\,\mathrm{[e]}$")
0076     __create_plot(x=df["q"], bins=4, x_axis=x_axis_opts, suffix="charge_dist")
0077 
0078     # Plot the total momentum
0079     p = np.sqrt(np.square(df["px"]) + np.square(df["py"]) + np.square(df["pz"]))
0080     x_axis_opts = plotting.axis_options(label=r"$p_{tot}\,\mathrm{[GeV]}$")
0081 
0082     __create_plot(
0083         x=p,
0084         bins=1 if np.std(p) < 1e-10 else 100,
0085         x_axis=x_axis_opts,
0086         y_axis=y_axis_opts._replace(log_scale=10),
0087         suffix="p_dist",
0088     )
0089 
0090     # Plot the transverse momentum
0091     pT = np.sqrt(np.square(df["px"]) + np.square(df["py"]))
0092     x_axis_opts = plotting.axis_options(label=r"$p_{T}\,\mathrm{[GeV]}$")
0093 
0094     __create_plot(
0095         x=pT,
0096         bins=1 if np.std(pT) < 1e-10 else 100,
0097         x_axis=x_axis_opts,
0098         suffix="pT_dist",
0099     )
0100 
0101     # Plot the x-origin
0102     x_axis_opts = plotting.axis_options(label=r"$x\,\mathrm{[mm]}$")
0103     __create_plot(x=df["x"], bins=100, x_axis=x_axis_opts, suffix="x_origin")
0104 
0105     # Plot the y-origin
0106     x_axis_opts = plotting.axis_options(label=r"$y\,\mathrm{[mm]}$")
0107     __create_plot(x=df["y"], bins=100, x_axis=x_axis_opts, suffix="y_origin")
0108 
0109     # Plot the z-origin
0110     x_axis_opts = plotting.axis_options(label=r"$z\,\mathrm{[mm]}$")
0111     __create_plot(x=df["z"], bins=100, x_axis=x_axis_opts, suffix="z_origin")
0112 
0113     # Plot the phi angle of the track direction
0114     phi = np.arctan2(df["py"], df["px"])
0115     x_axis_opts = plotting.axis_options(label=r"$\varphi\,\mathrm{[rad]}$")
0116 
0117     __create_plot(x=phi, bins=100, x_axis=x_axis_opts, suffix="dir_phi")
0118 
0119     # Plot the theta value of the track direction
0120     theta = np.arctan2(pT, df["pz"])
0121     x_axis_opts = plotting.axis_options(label=r"$\theta\,\mathrm{[rad]}$")
0122 
0123     __create_plot(x=theta, bins=100, x_axis=x_axis_opts, suffix="dir_theta")
0124 
0125     # Plot the eta value of the track direction
0126     eta = np.arctanh(df["pz"] / p)
0127     x_axis_opts = plotting.axis_options(label=r"$\eta$")
0128 
0129     __create_plot(x=eta, bins=100, x_axis=x_axis_opts, suffix="dir_eta")
0130 
0131 
0132 """ Plot the track positions of two data sources - rz view """
0133 
0134 
0135 def compare_track_pos_xy(
0136     opts,
0137     detector,
0138     scan_type,
0139     plot_factory,
0140     out_format,
0141     df1,
0142     label1,
0143     color1,
0144     df2,
0145     label2,
0146     color2,
0147 ):
0148 
0149     n_rays = np.max(df1["track_id"]) + 1
0150     tracks = "rays" if scan_type == "ray" else "helices"
0151 
0152     # Reduce data to the requested z-range (50mm tolerance)
0153     min_z = opts.z_range[0]
0154     max_z = opts.z_range[1]
0155     assert min_z < max_z, "xy plotting range: min z must be smaller that max z"
0156     pos_range = lambda data: ((data["z"] > min_z) & (data["z"] < max_z))
0157 
0158     first_x, first_y = plotting.filter_data(
0159         data=df1, filter=pos_range, variables=["x", "y"]
0160     )
0161 
0162     second_x, second_y = plotting.filter_data(
0163         data=df2, filter=pos_range, variables=["x", "y"]
0164     )
0165 
0166     # Plot the xy coordinates of the filtered track positions
0167     lgd_ops = plotting.legend_options(
0168         loc="upper center",
0169         ncol=4,
0170         colspacing=0.4,
0171         handletextpad=0.005,
0172         horiz_anchor=0.5,
0173         vert_anchor=1.095,
0174     )
0175 
0176     hist_data = plot_factory.scatter(
0177         figsize=(10, 10),
0178         x=first_x,
0179         y=first_y,
0180         x_axis=plotting.axis_options(label=r"$x\,\mathrm{[mm]}$"),
0181         y_axis=plotting.axis_options(label=r"$y\,\mathrm{[mm]}$"),
0182         label=label1,
0183         color=color1,
0184         alpha=1.0,
0185         show_stats=lambda x, y: f"{n_rays} {tracks}",
0186         lgd_ops=lgd_ops,
0187     )
0188 
0189     # Compare against second data set
0190     plot_factory.highlight_region(hist_data, second_x, second_y, color2, label2)
0191 
0192     detector_name = detector.replace(" ", "_")
0193     l1 = label1.replace(" ", "_").replace("(", "").replace(")", "")
0194     l2 = label2.replace(" ", "_").replace("(", "").replace(")", "")
0195 
0196     # Need a very high dpi to reach a good coverage of the individual points
0197     plot_factory.write_plot(
0198         hist_data,
0199         f"{detector_name}_{scan_type}_track_pos_{l1}_{l2}_xy",
0200         out_format,
0201         dpi=600,
0202     )
0203 
0204 
0205 """ Plot the track positions of two data sources - rz view """
0206 
0207 
0208 def compare_track_pos_rz(
0209     opts,
0210     detector,
0211     scan_type,
0212     plot_factory,
0213     out_format,
0214     df1,
0215     label1,
0216     color1,
0217     df2,
0218     label2,
0219     color2,
0220 ):
0221 
0222     n_rays = np.max(df1["track_id"]) + 1
0223     tracks = "rays" if scan_type == "ray" else "helices"
0224 
0225     first_x, first_y, first_z = plotting.filter_data(
0226         data=df1, variables=["x", "y", "z"]
0227     )
0228 
0229     second_x, second_y, second_z = plotting.filter_data(
0230         data=df2, variables=["x", "y", "z"]
0231     )
0232 
0233     # Plot the xy coordinates of the filtered intersections points
0234     lgd_ops = plotting.legend_options(
0235         loc="upper center",
0236         ncol=4,
0237         colspacing=0.8,
0238         handletextpad=0.005,
0239         horiz_anchor=0.5,
0240         vert_anchor=1.168,
0241     )
0242 
0243     hist_data = plot_factory.scatter(
0244         figsize=(12, 6),
0245         x=first_z,
0246         y=np.hypot(first_x, first_y),
0247         x_axis=plotting.axis_options(label=r"$z\,\mathrm{[mm]}$"),
0248         y_axis=plotting.axis_options(label=r"$r\,\mathrm{[mm]}$"),
0249         label=label1,
0250         color=color1,
0251         alpha=1.0,
0252         show_stats=lambda x, y: f"{n_rays} {tracks}",
0253         lgd_ops=lgd_ops,
0254     )
0255 
0256     # Compare against second data set
0257     plot_factory.highlight_region(
0258         hist_data, second_z, np.hypot(second_x, second_y), color2, label2
0259     )
0260 
0261     detector_name = detector.replace(" ", "_")
0262     l1 = label1.replace(" ", "_").replace("(", "").replace(")", "")
0263     l2 = label2.replace(" ", "_").replace("(", "").replace(")", "")
0264 
0265     # Need a very high dpi to reach a good coverage of the individual points
0266     plot_factory.write_plot(
0267         hist_data,
0268         f"{detector_name}_{scan_type}_track_pos_{l1}_{l2}_rz",
0269         out_format,
0270         dpi=600,
0271     )
0272 
0273 
0274 """ Plot the absolute track positions distance """
0275 
0276 
0277 def plot_track_pos_dist(
0278     opts, detector, scan_type, plot_factory, out_format, df1, label1, df2, label2
0279 ):
0280 
0281     dist = np.sqrt(
0282         np.square(df1["x"] - df2["x"])
0283         + np.square(df1["y"] - df2["y"])
0284         + np.square(df1["z"] - df2["z"])
0285     )
0286 
0287     dist_outlier = math.sqrt(3 * opts.outlier**2)
0288 
0289     # Remove outliers
0290     filter_dist = np.absolute(dist) < dist_outlier
0291     filtered_dist = dist[filter_dist]
0292 
0293     if not np.all(filter_dist == True):
0294         print("\nRemoved outliers (dist):")
0295         for i, d in enumerate(dist):
0296             if math.fabs(d) > dist_outlier:
0297                 track_id = (df1["track_id"].to_numpy())[i]
0298                 print(f"track {track_id}: {d}")
0299 
0300     # Where to place the legend box
0301     lgd_ops = plotting.legend_options(
0302         loc=lgd_loc,
0303         ncol=4,
0304         colspacing=0.8,
0305         handletextpad=0.005,
0306         horiz_anchor=1.02,
0307         vert_anchor=1.24,
0308     )
0309 
0310     # Plot the xy coordinates of the filtered intersections points
0311     hist_data = plot_factory.hist1D(
0312         x=filtered_dist,
0313         bins=100,
0314         x_axis=plotting.axis_options(label=r"$d\,\mathrm{[mm]}$"),
0315         y_axis=plotting.axis_options(label="", log_scale=10),
0316         figsize=(8.5, 8.5),
0317         lgd_ops=lgd_ops,
0318     )
0319 
0320     detector_name = detector.replace(" ", "_")
0321     l1 = label1.replace(" ", "_").replace("(", "").replace(")", "")
0322     l2 = label2.replace(" ", "_").replace("(", "").replace(")", "")
0323     plot_factory.write_plot(
0324         hist_data, f"{detector_name}_{scan_type}_dist_{l1}_{l2}", out_format
0325     )
0326 
0327 
0328 """ Plot the track position residual for the given variable """
0329 
0330 
0331 def plot_track_pos_res(
0332     opts, detector, scan_type, plot_factory, out_format, df1, label1, df2, label2, var
0333 ):
0334 
0335     tracks = "rays" if scan_type == "ray" else "helices"
0336 
0337     assert len(df1[var]) == len(df2[var])
0338     res = df1[var] - df2[var]
0339 
0340     # Remove outliers
0341     filter_res = np.absolute(res) < opts.outlier
0342     filtered_res = res[filter_res]
0343 
0344     u_out = o_out = int(0)
0345     if not np.all(filter_res == True):
0346         print(f"\nRemoved outliers ({var}):")
0347         for i, r in enumerate(res):
0348             if math.fabs(r) > opts.outlier:
0349                 track_id = (df1["track_id"].to_numpy())[i]
0350                 print(f"track {track_id}: {df1[var][i]} - {df2[var][i]} = {r}")
0351 
0352                 if r < 0.0:
0353                     u_out = u_out + 1
0354                 else:
0355                     o_out = o_out + 1
0356 
0357     lgd_ops = plotting.legend_options(
0358         loc=lgd_loc,
0359         ncol=4,
0360         colspacing=0.01,
0361         handletextpad=0.0005,
0362         horiz_anchor=1.02,
0363         vert_anchor=1.28,
0364     )
0365 
0366     # Plot the residuals as a histogram and fit a gaussian to it
0367     hist_data = plot_factory.hist1D(
0368         x=filtered_res,
0369         figsize=(9, 9),
0370         bins=100,
0371         x_axis=plotting.axis_options(
0372             label=r"$\mathrm{res}" + rf"\,{var}" + r"\,\mathrm{[mm]}$"
0373         ),
0374         lgd_ops=lgd_ops,
0375         u_outlier=u_out,
0376         o_outlier=o_out,
0377     )
0378 
0379     mu, sig = plot_factory.fit_gaussian(hist_data)
0380     if mu is None or sig is None:
0381         print(rf"WARNING: fit failed (res ({tracks}): {label1} - {label2} )")
0382 
0383     detector_name = detector.replace(" ", "_")
0384     l1 = label1.replace(" ", "_").replace("(", "").replace(")", "")
0385     l2 = label2.replace(" ", "_").replace("(", "").replace(")", "")
0386 
0387     plot_factory.write_plot(
0388         hist_data, f"{detector_name}_{scan_type}_track_res_{var}_{l1}_{l2}", out_format
0389     )