eloss calculator

ELoss/eloss.py

@@ -0,0 +1,179 @@
+#!/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()