From ee1a017ef2851d16d9d38b30907c0d0f63d41325 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Coussat?= Date: Fri, 21 Aug 2026 11:10:22 +0200 Subject: [PATCH] ENH: Add pctaddnoise application --- .github/workflows/build-test-package.yml | 2 +- AddTrackerUncertainty.py | 90 ---------- applications/pctaddnoise/pctaddnoise.py | 199 ++++++++++++++++++++++ documentation/docs/getting_started.md | 18 ++ examples/Reconstruction/Reconstruction.py | 17 +- pyproject.toml | 1 + test/pct_application_test.py | 17 ++ wrapping/__init_pct__.py | 1 + 8 files changed, 251 insertions(+), 94 deletions(-) delete mode 100755 AddTrackerUncertainty.py create mode 100644 applications/pctaddnoise/pctaddnoise.py diff --git a/.github/workflows/build-test-package.yml b/.github/workflows/build-test-package.yml index 1baa5557..4aa5d914 100644 --- a/.github/workflows/build-test-package.yml +++ b/.github/workflows/build-test-package.yml @@ -51,7 +51,7 @@ jobs: echo "Installing wheel: $wheel" pip install $wheel - pip install pytest uproot opengate + pip install pytest uproot opengate hepunits # Force the installation of Geant4 data, required by opengate opengate_info diff --git a/AddTrackerUncertainty.py b/AddTrackerUncertainty.py deleted file mode 100755 index d23daefc..00000000 --- a/AddTrackerUncertainty.py +++ /dev/null @@ -1,90 +0,0 @@ -#!/usr/bin/env python - -import itk -import hepunits -import numpy as np -import click - - -# Equation (27) and (28) from [Krah et al, PMB, 2018] -def GetSigmaSc(energy, xOverX0, 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 - - sigmaSc = ( - 13.6 * hepunits.MeV / betap * np.sqrt(xOverX0) * (1 + 0.038 * np.log(xOverX0)) - ) - SigmaSc = np.zeros((energy.size, 2, 2)) - SigmaSc[:, 1, 1] = sigmaSc**2 - return np.tile(sp**2 * T @ T.T, (energy.size, 1, 1)) + SigmaSc - - -@click.command() -@click.option("xOverX0", "--xOverX0", default=5e-3, help="Material budget (unitless)") -@click.option( - "--sp", default=0.15, help="Standard deviation of the tracker uncertainty (mm)" -) -@click.option("--dt", default=10, help="Distance between trackers (cm)") -@click.option("-i", "--input", required=True, help="Input file name") -@click.option("-o", "--output", required=True, help="Output file name") -@click.option( - "entryTranslation", - "--entryTranslation", - default=0.0, - help="Amount of translation of the entry detector position in the direction perpendicular to the detector (mm)", -) -@click.option( - "exitTranslation", - "--exitTranslation", - default=0.0, - help="Amount of translation of the exit detector position in the direction perpendicular to the detector (mm)", -) -def AddTrackerUncertainty( - xOverX0, sp, dt, input, output, entryTranslation, exitTranslation -): - dt = dt * hepunits.cm - sp = sp * hepunits.mm - - pairs = itk.imread(input) - pairs = itk.GetArrayFromImage(pairs) - - # Move to entrance and exit detector (new) positions - pairs[:, 0, :] += ((entryTranslation / pairs[:, 2, 2]) * pairs[:, 2, :].T).T - pairs[:, 1, :] += ((exitTranslation / pairs[:, 3, 2]) * pairs[:, 3, :].T).T - - eEntry = pairs[:, 4, 0] - eExit = pairs[:, 4, 1] - SigmaEntry = GetSigmaSc(eEntry, xOverX0, sp, dt) - SigmaExit = GetSigmaSc(eExit, xOverX0, sp, dt) - wEntry, Qentry = np.linalg.eig(np.linalg.inv(SigmaEntry)) - wExit, Qexit = np.linalg.eig(np.linalg.inv(SigmaExit)) - xrEntry = np.random.randn(eEntry.size, 2, 2) - xrExit = np.random.randn(eExit.size, 2, 2) - # Wentry = np.diag(1./np.sqrt(wEntry[i,:])) - Wentry = np.zeros((eEntry.size, 2, 2)) - Wentry[:, 0, 0] = 1.0 / np.sqrt(wEntry[:, 0]) - Wentry[:, 1, 1] = 1.0 / np.sqrt(wEntry[:, 1]) - # Wexit = np.diag(1./np.sqrt(wExit[i,:])) - Wexit = np.zeros((eExit.size, 2, 2)) - Wexit[:, 0, 0] = 1.0 / np.sqrt(wExit[:, 0]) - Wexit[:, 1, 1] = 1.0 / np.sqrt(wExit[:, 1]) - # dYuncertEntry = Qentry[i,:,:].dot(Wentry).dot(xrEntry[i,:,:]).T - dYuncertEntry = np.matmul(np.matmul(Qentry, Wentry), xrEntry) - # dYuncertExit = Qexit[i,:,:].dot(Wexit).dot(xrExit[i,:,:]).T - dYuncertExit = np.matmul(np.matmul(Qexit, Wexit), xrExit) - pairs[:, 0, 0:2] += dYuncertEntry[:, 0, :] # Entrance position X/Y - pairs[:, 2, 0:2] += dYuncertEntry[:, 1, :] # Entrance direction X/Y - pairs[:, 1, 0:2] += dYuncertExit[:, 0, :] # Exit position X/Y - pairs[:, 3, 0:2] += dYuncertExit[:, 1, :] # Exit direction X/Y - - itk.imwrite(itk.GetImageFromArray(pairs, is_vector=True), output) - - -if __name__ == "__main__": - AddTrackerUncertainty() diff --git a/applications/pctaddnoise/pctaddnoise.py b/applications/pctaddnoise/pctaddnoise.py new file mode 100644 index 00000000..0b347699 --- /dev/null +++ b/applications/pctaddnoise/pctaddnoise.py @@ -0,0 +1,199 @@ +""" +Add uncertainties to proton CT tracker data. + +Proton CT detectors are typically simulated in GATE using a single PhaseSpaceActor. However, proton CT trackers actually consist of two detector planes. The position of the proton is detected twice, which allows to additionally detect the proton direction. However, several effects affect the measured positions and directions: protons may undergo scattering in the detector, and proton CT trackers are usually made out of strips that discretize the detection area. + +This application takes as input a ROOT file generated from a PhaseSpaceActor representing a proton CT tracker. The following uncertainties are taken into account: +- if the ROOT file contains position and direction branches: realistic position and direction uncertainties based on Krah et al. (PMB, 2018) +- if the ROOT file contains an energy branch: Gaussian energy uncertainty +- if the ROOT file contains a time branch: Gaussian time uncertainty + +This script is adapted from "AddTrackerUncertainty.py" (https://github.com/RTKConsortium/PCT/blob/6afe8ee0a0c25fd3e892761a237f120fd19198e4/AddTrackerUncertainty.py) from earlier versions of PCT. +""" + +#!/usr/bin/env python +import sys +import numpy as np +import hepunits +import uproot +import argparse +import itk +from itk import PCT as pct + + +def build_parser(): + parser = pct.PCTArgumentParser( + description="Add uncertainties to proton CT tracker data" + ) + parser.add_argument( + "--material-budget", default=5e-3, type=float, help="Material budget (unitless)" + ) + parser.add_argument( + "--tracker-distance", + default=10.0, + type=float, + help="Distance between trackers (cm)", + ) + parser.add_argument("-i", "--input", required=True, help="Input file name") + parser.add_argument("--tree", required=True, help="Name of tree in ROOT file") + parser.add_argument("-o", "--output", required=True, help="Output file name") + parser.add_argument( + "--translation", + default=0.0, + help="Translation of the detector position in the direction perpendicular to the detector (mm)", + ) + parser.add_argument( + "--noise-position", + help="Standard deviation of the Gaussian noise on the position", + type=float, + ) + 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("--seed", help="Random seed", type=int) + parser.add_argument( + "--verbose", "-v", help="Verbose execution", default=False, action="store_true" + ) + return parser + + +def get_sigma_sc(energy, x_over_x0, sp, dt): + """ + Get the Σ_sc matrix as defined by Equation (27) of Krah et al. (PMB, 2018). + + Args: + - energy: energies of the protons (MeV). + - x_over_x0: material budget. + - sp: standard deviation of the tracker uncertainty (mm). + - dt: distance between trackers (cm). + """ + if sp is None: + return + + 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 +): + """ + Add tracker uncertainties as detailed in Krah et al. (PMB, 2018), section 2.5. + + Args: + - data: data from the ROOT file, in NumPy format. + - rng: random number generator. + - material_budget: material budget. + - noise_position: standard deviation of the tracker uncertainty (mm). + - tracker_distance: distance between trackers (cm). + """ + 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)) + q = np.real(q) + 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, clamp=None): + """ + Add Gaussian noise to an arbitrary branch of ROOT data. + + Args: + - data: data from the ROOT file, in NumPy format. + - branch: what branch to consider. + - rng: random number generator. + - noise: standard deviation. + - clamp: values below this number will be clamped to this number. + """ + if noise is None: + return + + try: + data[branch] += rng.normal(scale=noise, size=len(data["KineticEnergy"])) + if clamp is not None: + data[branch] = np.where(data[branch] < clamp, clamp, data[branch]) + 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): + + rng = np.random.default_rng(args_info.seed) + + if args_info.verbose: + print("Reading input ROOT file") + data = uproot.open(args_info.input)[args_info.tree].arrays(library="np") + + # Move to entrance and exit detector (new) positions + if args_info.verbose and args_info.translation is not None: + print("Applying translations…") + for pos in ["X", "Y", "Z"]: + data[f"Position_{pos}"] += ( + args_info.translation / data["Direction_Z"] + ) * data[f"Direction_{pos}"] + + if args_info.noise_position is not None and args_info.noise_position > 0.0: + 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, clamp=0.0) + 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() diff --git a/documentation/docs/getting_started.md b/documentation/docs/getting_started.md index 1db1dca9..fdcfc2ff 100644 --- a/documentation/docs/getting_started.md +++ b/documentation/docs/getting_started.md @@ -33,6 +33,24 @@ python gate/protonct.py --help ``` Of course, feel free to explore the content of `protonct.py` directly to adapt it to your needs. +### Optional: add realistic noise to the GATE output + +PCT provides the `pctaddnoise` application that adds realistic noise to the generated data. This application takes as input a ROOT file generated from a GATE simulation (such as `protonct.py`) that represents either the upstream or the downstream detector. The following uncertainties are taken into account: +- if the ROOT file contains position and direction branches: realistic position and direction uncertainties based on [Krah et al. (PMB, 2018)](https://doi.org/10.1088/1361-6560/aaca1f) +- if the ROOT file contains an energy branch: Gaussian energy uncertainty +- if the ROOT file contains a time branch: Gaussian time uncertainty + +For instance, noise can be added using the following command: +```bash +pctaddnoise -i gate_simulation/PhaseSpaceIn.root -o gate_simulation/PhaseSpaceIn_noisy.root --material-budget .01 --tracker-distance 10 --noise-energy 10 +``` +which needs to be repeated for `PhaseSpaceOut.root` as well. + +As usual, you can explore additional options provided by `pctaddnoise` by running +```bash +pctaddnoise --help +``` + (pairing)= ## Protons pairing diff --git a/examples/Reconstruction/Reconstruction.py b/examples/Reconstruction/Reconstruction.py index abfb139f..2ca380d7 100644 --- a/examples/Reconstruction/Reconstruction.py +++ b/examples/Reconstruction/Reconstruction.py @@ -19,14 +19,25 @@ gate_folder = os.path.join(output_folder, "gate") protonct(gate_folder, projections=number_of_projections, verbose=False) -# TODO example on how to make data noisy +# Add noise to generated data +for file in ["PhaseSpaceIn", "PhaseSpaceOut"]: + pct.pctaddnoise( + input=os.path.join(gate_folder, f"{file}.root"), + output=os.path.join(gate_folder, f"{file}_noisy.root"), + tree=file, + material_budget=0.01, + tracker_distance=10.0, + noise_position=1.0, + noise_energy=1.0, + seed=1234, + ) # Convert GATE data to PCT list-mode pairs_folder = os.path.join(output_folder, "pairs") os.makedirs(pairs_folder, exist_ok=True) pct.pctpairprotons( - input_in=os.path.join(gate_folder, "PhaseSpaceIn.root"), - input_out=os.path.join(gate_folder, "PhaseSpaceOut.root"), + input_in=os.path.join(gate_folder, "PhaseSpaceIn_noisy.root"), + input_out=os.path.join(gate_folder, "PhaseSpaceOut_noisy.root"), output=os.path.join(pairs_folder, "pairs.mhd"), psin="PhaseSpaceIn", psout="PhaseSpaceOut", diff --git a/pyproject.toml b/pyproject.toml index bffa5e48..cbc5503f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ pctfdk = "itk.pctfdk:main" pctpairprotons = "itk.pctpairprotons:main" pctweplfit = "itk.pctweplfit:main" pctlomalinda = "itk.pctlomalinda:main" +pctaddnoise = "itk.pctaddnoise:main" [project.urls] Download = "https://github.com/RTKConsortium/PCT" diff --git a/test/pct_application_test.py b/test/pct_application_test.py index 2e871770..a7f0918c 100644 --- a/test/pct_application_test.py +++ b/test/pct_application_test.py @@ -3,6 +3,7 @@ import itk import urllib.request import numpy as np +import uproot from itk import PCT as pct @@ -100,3 +101,19 @@ def test_lomalinda_application( reference_lomalinda = itk.array_from_image(itk.imread(baseline_lomalinda_mhd)) assert np.array_equal(test_lomalinda, reference_lomalinda) return output0000 + + +baseline_addnoise = download_file_fixture( + "6a8561052688ba21262c390a", "baseline_addnoise.root" +) + + +def test_addnoise_application(tmp_path, phasespacein_root, baseline_addnoise): + output = tmp_path / "noise_test.root" + tree = "PhaseSpaceIn" + pct.pctaddnoise( + f"-i {phasespacein_root} -o {output} --tree {tree} --material-budget .01 --tracker-distance 10 --translation -5 --noise-position 10 --noise-energy 10 --seed 1234" + ) + root_test = uproot.open(output)[tree].arrays(library="np") + root_baseline = uproot.open(baseline_addnoise)[tree].arrays(library="np") + assert np.array_equal(root_test, root_baseline) diff --git a/wrapping/__init_pct__.py b/wrapping/__init_pct__.py index 3466b211..191c94ed 100644 --- a/wrapping/__init_pct__.py +++ b/wrapping/__init_pct__.py @@ -26,6 +26,7 @@ "pctpairprotons", "pctweplfit", "pctlomalinda", + "pctaddnoise", ] # Dynamically access make_application_func from pctExtras