From 35a2c11877c55019e63f239a36f42e6671aa0ee8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Coussat?= Date: Fri, 3 Jan 2025 10:23:51 +0100 Subject: [PATCH] ENH: Fit TOF to energy loss --- CMakeLists.txt | 2 +- pcttoftowepl.py => pcttoffit.py | 78 +++++++++++++++++++++++---------- 2 files changed, 55 insertions(+), 25 deletions(-) rename pcttoftowepl.py => pcttoffit.py (76%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 667e986..a58d7a2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -103,4 +103,4 @@ IF(PCT_WITH_GEANT4) ENDIF(PCT_WITH_GEANT4) CONFIGURE_FILE("pctpairprotons.py" "pctpairprotons.py" COPYONLY) -CONFIGURE_FILE("pcttoftowepl.py" "pcttoftowepl.py" COPYONLY) \ No newline at end of file +CONFIGURE_FILE("pcttoffit.py" "pcttoffit.py" COPYONLY) diff --git a/pcttoftowepl.py b/pcttoffit.py similarity index 76% rename from pcttoftowepl.py rename to pcttoffit.py index 3dde15b..b94daa0 100644 --- a/pcttoftowepl.py +++ b/pcttoffit.py @@ -9,8 +9,8 @@ import numpy as np -def tof_to_wepl_mc( - output='pcttoftowepl', +def tof_fit_mc( + output='pcttoffit', number_of_particles=1e4, visu=False, verbose=False @@ -90,7 +90,8 @@ def add_detector(name, translation): 'EventID', 'TrackID', 'Position', - 'PreGlobalTime' + 'PreGlobalTime', + 'KineticEnergy' ] particle_filter = sim.add_filter('ParticleFilter', 'Filter' + name) particle_filter.particle = 'proton' @@ -107,8 +108,8 @@ def add_detector(name, translation): sim.run() -def tof_to_wepl_fit( - output='pcttoftowepl', +def tof_fit( + output='pcttoffit', poly_deg=2, path_type='simple', display=False, @@ -138,6 +139,7 @@ def print_verbose(*args, **kwargs): tofs = [] wepls = [] + energy_losses = [] for n in np.unique(data['EventID']): event_mask = data['EventID'] == n @@ -148,6 +150,7 @@ def print_verbose(*args, **kwargs): vs = data['Position_Z'][event_mask] ws = data['Position_X'][event_mask] times = data['PreGlobalTime'][event_mask] + energies = data['KineticEnergy'][event_mask] tofs_event = [times[k] - times[0] for k in range(len(times))] @@ -169,36 +172,63 @@ def print_verbose(*args, **kwargs): else: sys.exit(f"Invalid path time {path_type}!") + energy_losses_event = [ + energies[0] - energies[k] + for k in range(len(energies)) + ] + tofs.extend(tofs_event) wepls.extend(wepls_event) + energy_losses.extend(energy_losses_event) - p = np.polyfit(tofs, wepls, deg=poly_deg) - print_verbose("Fitted coefficients:", p) + p_wepl = np.polyfit(tofs, wepls, deg=poly_deg) + print_verbose("Fitted coefficients for WEPL:", p_wepl) + + p_energy_loss = np.polyfit(tofs, energy_losses, deg=poly_deg) + print_verbose("Fitted coefficients for energy loss:", p_energy_loss) if display or savefig: - tof_fit = np.linspace(0, np.max(tofs), 100) - wepl_fit = np.polyval(p, tof_fit) + tof_samples = np.linspace(0, np.max(tofs), 100) + + wepl_fit = np.polyval(p_wepl, tof_samples) plt.figure() plt.plot(tofs, wepls, '+', label="Detected events") - plt.plot(tof_fit, wepl_fit, label=f"Polynomial fit (degree {poly_deg})") - plt.xlabel("TOF") + plt.plot(tof_samples, wepl_fit, label=f"Polynomial fit (degree {poly_deg})") + plt.xlabel("TOF [ns]") plt.ylabel("WEPL [mm]") plt.legend() if savefig: - plt.savefig(f'{output}/tof_to_wepl_fit.pdf') + plt.savefig(f'{output}/tof_wepl_fit.pdf') if display: plt.show() - with open(f'{output}/tof_to_wepl_fit.json', 'w', encoding='utf-8') as f: - json.dump(list(p), f) + eloss_fit = np.polyval(p_energy_loss, tof_samples) + + plt.figure() + plt.plot(tofs, energy_losses, '+', label="Detected events") + plt.plot(tof_samples, eloss_fit, label=f"Polynomial fit (degree {poly_deg})") + plt.xlabel("TOF [ns]") + plt.ylabel("Energy loss [MeV]") + plt.legend() + + if savefig: + plt.savefig(f'{output}/tof_eloss_fit.pdf') + if display: + plt.show() + + with open(f'{output}/tof_wepl_fit.json', 'w', encoding='utf-8') as f: + json.dump(list(p_wepl), f) + + with open(f'{output}/tof_eloss_fit.json', 'w', encoding='utf-8') as f: + json.dump(list(p_energy_loss), f) - return p + return p_wepl, p_energy_loss -def pcttoftowepl( - output='pcttoftowepl', +def pcttoffit( + output='pcttoffit', number_of_particles=1e4, poly_deg=2, path_type='simple', @@ -207,25 +237,25 @@ def pcttoftowepl( savefig=False, verbose=False ): - tof_to_wepl_mc(output, number_of_particles, visu, verbose) - p = tof_to_wepl_fit(output, poly_deg, path_type, display, savefig, verbose) - return p + tof_fit_mc(output, number_of_particles, visu, verbose) + ps = tof_fit(output, poly_deg, path_type, display, savefig, verbose) + return ps def main(): - parser = argparse.ArgumentParser(description='Convert TOF to WEPL using a fit on Monte Carlo data') - parser.add_argument('--output', help="Path of outputs", default='pcttoftowepl') + parser = argparse.ArgumentParser(description='Convert TOF to WEPL and energy loss using a fit on Monte Carlo data') + parser.add_argument('--output', help="Path of outputs", default='pcttoffit') parser.add_argument('-n', '--number-of-particles', help="Number of generated particles", default=1e4, type=int) parser.add_argument('--poly-deg', help="Degrees of polynomial fit", default=2, type=int) - parser.add_argument('--path-type', help="How to compute proton path", choices=['simple', 'realistic'], default='simple') + parser.add_argument('--path-type', help="How to compute proton path for WEPL", choices=['simple', 'realistic'], default='simple') parser.add_argument('--visu', help="Visualize Monte Carlo simulation", default=False, action='store_true') parser.add_argument('--display', help="Display polynomial fit plot", default=False, action='store_true') parser.add_argument('--savefig', help="Write polynomial fit plot to disk", default=False, action='store_true') parser.add_argument('--verbose', '-v', help="Verbose execution", default=False, action='store_true') args_info = parser.parse_args() - pcttoftowepl(**vars(args_info)) + pcttoffit(**vars(args_info)) if __name__ == '__main__': main()