1 files changed, 50 insertions, 0 deletions
diff --git a/utils/sddm/utils.py b/utils/sddm/utils.py
new file mode 100644
index 0000000..2b67e63
--- /dev/null
+++ b/utils/sddm/utils.py
@@ -0,0 +1,50 @@
+import numpy as np
+from scipy.stats import norm
+import pandas as pd
+
+_fast_cdf_x = np.linspace(-10,10,100000)
+_fast_cdf_y = norm.cdf(_fast_cdf_x)
+def fast_cdf(x,loc,scale):
+    """
+    A faster version of norm.cdf() which just uses interpolation based on a
+    lookup table.
+    """
+    return np.interp((x-loc)/scale,_fast_cdf_x,_fast_cdf_y)
+
+# Energy bias of reconstruction relative to Monte Carlo.
+#
+# Note: You can recreate this array using:
+#
+#     ./plot-fit-results -o [filename]
+ENERGY_BIAS = np.array([
+    (100.0,  0.050140,  0.078106),
+    (200.0,  0.058666,  0.078106),
+    (300.0,  0.049239, -0.000318),
+    (400.0,  0.045581, -0.020932),
+    (500.0,  0.050757, -0.028394),
+    (600.0,  0.048310, -0.029017),
+    (700.0,  0.052434, -0.020770),
+    (800.0,  0.032920, -0.019298),
+    (900.0,  0.040963, -0.015354)],
+    dtype=[('T',np.float32),('e_bias',np.float32),('u_bias',np.float32)])
+
+def correct_energy_bias(df):
+    """
+    Corrects for the energy bias of the reconstruction relative to the true
+    Monte Carlo energy.
+    """
+    # Note: We divide here since we define the bias as:
+    #
+    #     bias = (T_reco - T_true)/T_true
+    #
+    # Therefore,
+    #
+    #     T_true = T_reco/(1+bias)
+    df.loc[df['id1'] == 20,'energy1'] /= (1+np.interp(df.loc[df['id1'] == 20,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['e_bias']))
+    df.loc[df['id2'] == 20,'energy2'] /= (1+np.interp(df.loc[df['id2'] == 20,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['e_bias']))
+    df.loc[df['id3'] == 20,'energy3'] /= (1+np.interp(df.loc[df['id3'] == 20,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['e_bias']))
+    df.loc[df['id1'] == 22,'energy1'] /= (1+np.interp(df.loc[df['id1'] == 22,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['u_bias']))
+    df.loc[df['id2'] == 22,'energy2'] /= (1+np.interp(df.loc[df['id2'] == 22,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['u_bias']))
+    df.loc[df['id3'] == 22,'energy3'] /= (1+np.interp(df.loc[df['id3'] == 22,'energy1'],ENERGY_BIAS['T'],ENERGY_BIAS['u_bias']))
+    df['ke'] = df['energy1'].fillna(0) + df['energy2'].fillna(0) + df['energy3'].fillna(0)
+    return df