179 lines
3.4 KiB
Python
179 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from scipy.interpolate import interp1d
|
|
import argparse
|
|
|
|
#python3 eloss.py E_vs_X_table.txt --Ei <starting energy> --Ef <final energy>
|
|
#python3 eloss.py E_vs_X_table.txt --Ei <starting energy> --dx <distance traveled>
|
|
#python3 eloss.py E_vs_X_table.txt --Ef <final energy> --dx <distance traveled>
|
|
|
|
def load_table(filename):
|
|
"""
|
|
Load table with columns:
|
|
x(cm) E(MeV)
|
|
|
|
Returns:
|
|
x_array, E_array
|
|
"""
|
|
|
|
data = pd.read_csv(
|
|
filename,
|
|
sep='\s+',
|
|
comment="#",
|
|
header=None,
|
|
skiprows=1
|
|
)
|
|
|
|
x = data.iloc[:, 0].values
|
|
E = data.iloc[:, 1].values
|
|
|
|
return x, E
|
|
|
|
|
|
def build_interpolators(x, E):
|
|
|
|
# E(x)
|
|
E_of_x = interp1d(
|
|
x,
|
|
E,
|
|
bounds_error=False,
|
|
fill_value=0.0
|
|
)
|
|
|
|
# x(E)
|
|
x_of_E = interp1d(
|
|
E[::-1], # reverse so energy increases
|
|
x[::-1],
|
|
bounds_error=False,
|
|
fill_value="extrapolate"
|
|
)
|
|
|
|
return E_of_x, x_of_E
|
|
|
|
def distance_traveled(Ei, Ef, x_of_E):
|
|
|
|
xi = float(x_of_E(Ei))
|
|
xf = float(x_of_E(Ef))
|
|
|
|
return abs(xf - xi)
|
|
|
|
def final_energy(Ei, dx, x_of_E, E_of_x):
|
|
|
|
xi = float(x_of_E(Ei))
|
|
|
|
xf = xi + dx
|
|
|
|
Ef = float(E_of_x(xf))
|
|
|
|
return max(Ef, 0.0)
|
|
|
|
|
|
def initial_energy(Ef, dx, x_of_E, E_of_x):
|
|
|
|
xf = float(x_of_E(Ef))
|
|
|
|
xi = xf - dx
|
|
|
|
Ei = float(E_of_x(xi))
|
|
|
|
return max(Ei, 0.0)
|
|
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Energy-distance calculator using E vs X tables"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"table",
|
|
help="E vs X table file"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--Ei",
|
|
type=float,
|
|
help="Initial energy (MeV)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--Ef",
|
|
type=float,
|
|
help="Final energy (MeV)"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dx",
|
|
type=float,
|
|
help="Distance traveled (cm)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Count supplied variables
|
|
supplied = [
|
|
args.Ei is not None,
|
|
args.Ef is not None,
|
|
args.dx is not None
|
|
]
|
|
|
|
if sum(supplied) != 2:
|
|
print("\nERROR:")
|
|
print("Supply exactly TWO of: Ei, Ef, dx\n")
|
|
return
|
|
|
|
# Load table, skip the first row
|
|
x, E = load_table(args.table)
|
|
|
|
# Build interpolators
|
|
E_of_x, x_of_E = build_interpolators(x, E)
|
|
|
|
if args.Ei is not None and args.Ef is not None:
|
|
|
|
dx = distance_traveled(
|
|
args.Ei,
|
|
args.Ef,
|
|
x_of_E
|
|
)
|
|
|
|
print(f"\nDistance traveled: {dx:.6f} cm")
|
|
|
|
elif args.Ei is not None and args.dx is not None:
|
|
|
|
Ef = final_energy(
|
|
args.Ei,
|
|
args.dx,
|
|
x_of_E,
|
|
E_of_x
|
|
)
|
|
|
|
print(f"\nFinal energy: {Ef:.6f} MeV")
|
|
|
|
elif args.Ef is not None and args.dx is not None:
|
|
|
|
Ei = initial_energy(
|
|
args.Ef,
|
|
args.dx,
|
|
x_of_E,
|
|
E_of_x
|
|
)
|
|
|
|
print(f"\nInitial energy: {Ei:.6f} MeV")
|
|
|
|
elif args.Ei is not None and args.Ef is not None and args.dx is not None:
|
|
print("\nERROR: Supply exactly TWO of: Ei, Ef, dx\n")
|
|
return
|
|
|
|
elif args.help:
|
|
parser.print_help()
|
|
return
|
|
|
|
|
|
# ============================================================
|
|
# Run
|
|
# ============================================================
|
|
|
|
if __name__ == "__main__":
|
|
main() |