update muon kinetic energy calculation

This commit updates the calculation of the muon kinetic energy as a function of
distance along the track. Previously I was using an approximation from the PDG,
but it doesn't seem to be very accurate and won't generalize to the case of
electrons. The kinetic energy is now calculated using the tabulated values of
dE/dx as a function of energy.

7 files changed, 109 insertions, 60 deletions

diff --git a/src/likelihood.c b/src/likelihood.c
index 76364fd..6881e8f 100644
--- a/src/likelihood.c
+++ b/src/likelihood.c
@@ -122,7 +122,7 @@ static double gsl_muon_charge(double x, void *params)
     return get_expected_charge(x, T, theta0, pos, dir, pmts[pars->i].pos, pmts[pars->i].normal, PMT_RADIUS);
 }
 
-double get_total_charge_approx(double T0, double *pos, double *dir, int i, double smax, double theta0, double *t)
+double get_total_charge_approx(double T0, double *pos, double *dir, muon_energy *m, int i, double smax, double theta0, double *t)
 {
     /* Returns the approximate expected number of photons seen by PMT `i` using
      * an analytic formula.
@@ -218,7 +218,7 @@ double get_total_charge_approx(double T0, double *pos, double *dir, int i, doubl
     b = (s-z)/a;
 
     /* Compute the kinetic energy at the point `s` along the track. */
-    T = get_T(T0,s,HEAVY_WATER_DENSITY);
+    T = muon_get_energy(s,m);
 
     /* Calculate the particle velocity at the point `s`. */
     E = T + MUON_MASS;
@@ -258,36 +258,36 @@ double get_total_charge_approx(double T0, double *pos, double *dir, int i, doubl
     return f*exp(-a/absorption_length)*n*x*beta0*prob*(1/sin_theta)*omega*(erf((a+b*(smax-s)+n*(smax-z)*beta0)/frac) + erf((-a+b*s+n*z*beta0)/frac))/(b+n*beta0)/(4*M_PI);
 }
 
-double getKineticEnergy(double x, double T0)
-{
-    return get_T(T0, x, HEAVY_WATER_DENSITY);
-}
+typedef struct betaRootParams {
+    muon_energy *m;
+    double beta_min;
+} betaRootParams;
 
 static double beta_root(double x, void *params)
 {
     /* Function used to find at what point along a track a particle crosses
      * some threshold in beta. */
-    double T, E, p, beta, T0, beta_min;
+    double T, E, p, beta;
+    betaRootParams *pars;
 
-    T0 = ((double *) params)[0];
-    beta_min = ((double *) params)[1];
+    pars = (betaRootParams *) params;
 
-    T = get_T(T0, x, HEAVY_WATER_DENSITY);
+    T = muon_get_energy(x, pars->m);
 
     /* Calculate total energy */
     E = T + MUON_MASS;
     p = sqrt(E*E - MUON_MASS*MUON_MASS);
     beta = p/E;
 
-    return beta - beta_min;
+    return beta - pars->beta_min;
 }
 
-static int get_smax(double T0, double beta_min, double range, double *smax)
+static int get_smax(muon_energy *m, double beta_min, double range, double *smax)
 {
     /* Find the point along the track at which the particle's velocity drops to
      * `beta_min`. */
     int status;
-    double params[2];
+    betaRootParams pars;
     gsl_root_fsolver *s;
     gsl_function F;
     int iter = 0, max_iter = 100;
@@ -295,11 +295,11 @@ static int get_smax(double T0, double beta_min, double range, double *smax)
 
     s = gsl_root_fsolver_alloc(gsl_root_fsolver_brent);
 
-    params[0] = T0;
-    params[1] = beta_min;
+    pars.m = m;
+    pars.beta_min = beta_min;
 
     F.function = &beta_root;
-    F.params = params;
+    F.params = &pars;
 
     gsl_root_fsolver_set(s, &F, 0.0, range);
 
@@ -323,6 +323,11 @@ static int get_smax(double T0, double beta_min, double range, double *smax)
     return status;
 }
 
+double getKineticEnergy(double x, void *params)
+{
+    return muon_get_energy(x,(muon_energy *) params);
+}
+
 double nll_muon(event *ev, double T0, double *pos, double *dir, double t0, double *z1, double *z2, size_t n, double epsrel, int fast)
 {
     size_t i, j, nhit;
@@ -330,6 +335,7 @@ double nll_muon(event *ev, double T0, double *pos, double *dir, double t0, doubl
     double total_charge;
     double logp[MAX_PE], nll[MAX_PMTS], range, theta0, E0, p0, beta0, smax, log_mu, max_logp;
     double tmean = 0.0;
+    muon_energy *m;
 
     double mu_direct[MAX_PMTS];
     double ts[MAX_PMTS];
@@ -354,10 +360,12 @@ double nll_muon(event *ev, double T0, double *pos, double *dir, double t0, doubl
     /* FIXME: is this formula valid for muons? */
     theta0 = get_scattering_rms(range/2,p0,beta0,1.0,HEAVY_WATER_DENSITY)/sqrt(range/2);
 
-    params.p = path_init(pos, dir, T0, range, theta0, getKineticEnergy, z1, z2, n, MUON_MASS);
+    m = muon_init_energy(T0,HEAVY_WATER_DENSITY,10000);
+
+    params.p = path_init(pos, dir, T0, range, theta0, getKineticEnergy, m, z1, z2, n, MUON_MASS);
 
     if (beta0 > BETA_MIN)
-        get_smax(T0, BETA_MIN, range, &smax);
+        get_smax(m, BETA_MIN, range, &smax);
     else
         smax = 0.0;
 
@@ -368,7 +376,7 @@ double nll_muon(event *ev, double T0, double *pos, double *dir, double t0, doubl
         params.i = i;
 
         if (fast) {
-            mu_direct[i] = get_total_charge_approx(T0, pos, dir, i, smax, theta0, &ts[i]);
+            mu_direct[i] = get_total_charge_approx(T0, pos, dir, m, i, smax, theta0, &ts[i]);
             ts[i] += t0;
         } else {
             F.function = &gsl_muon_charge;
diff --git a/src/muon.c b/src/muon.c
index f2a7733..5eb1f54 100644
--- a/src/muon.c
+++ b/src/muon.c
@@ -14,6 +14,7 @@
 #include "sno.h"
 #include "scattering.h"
 #include "pmt_response.h"
+#include "misc.h"
 
 static int initialized = 0;
 
@@ -171,41 +172,69 @@ double get_range(double T, double rho)
     return gsl_spline_eval(spline_range, T, acc_range)/rho;
 }
 
-double get_T(double T0, double x, double rho)
+muon_energy *muon_init_energy(double T0, double rho, size_t n)
 {
-    /* Returns the approximate kinetic energy of a muon in water after
-     * travelling `x` cm with an initial kinetic energy `T`.
+    /* Returns a struct with arrays of the muon position and kinetic energy.
+     * This struct can then be passed to muon_get_energy() to interpolate the
+     * muon's kinetic energy at any point along the track. For example:
      *
-     * `T` should be in MeV, `x` in cm, and `rho` in g/cm^3.
+     *     muon_energy *m = muon_init_energy(1000.0, HEAVY_WATER_DENSITY, 100);
+     *     double T = muon_get_energy(x, m);
      *
-     * Return value is in MeV.
+     * To compute the kinetic energy as a function of distance we need to solve
+     * the following integral equation:
      *
-     * See http://pdg.lbl.gov/2018/AtomicNuclearProperties/adndt.pdf. */
-    double a, b, range, T;
+     *     T(x) = T0 - \int_0^x dT(T(x))/dx
+     *
+     * which we solve by dividing the track up into `n` segments and then
+     * numerically computing the energy at each step. */
+    size_t i;
+    double range, dEdx;
 
-    if (!initialized) {
-        if (init()) {
-            exit(1);
-        }
+    /* Get the range of the muon. */
+    range = get_range(T0, rho);
+
+    muon_energy *m = malloc(sizeof(muon_energy));
+    m->x = malloc(sizeof(double)*n);
+    m->T = malloc(sizeof(double)*n);
+    m->n = n;
+
+    m->x[0] = 0;
+    m->T[0] = T0;
+
+    /* Now we loop over the points along the track and assume dE/dx is constant
+     * between points. */
+    for (i = 1; i < n; i++) {
+        m->x[i] = range*i/(n-1);
+        dEdx = get_dEdx(m->T[i-1], rho);
+        m->T[i] = m->T[i-1] - dEdx*(m->x[i]-m->x[i-1]);
     }
 
-    range = get_range(T0, rho);
+    return m;
+}
 
-    /* This comes from Equation 33.42 in the PDG Passage of Particles Through
-     * Matter article. */
-    b = log(1 + T0/MUON_CRITICAL_ENERGY_D2O)/range;
-    /* Now we compute the ionization energy loss from the known range and b. */
-    a = b*T0/(exp(b*range)-1.0);
+double muon_get_energy(double x, muon_energy *m)
+{
+    /* Returns the approximate kinetic energy of a muon in water after
+     * travelling `x` cm with an initial kinetic energy `T`.
+     *
+     * Return value is in MeV. */
+    double T;
 
-    /* Compute the kinetic energy after travelling a distance `x` in the
-     * continuous slowing down approximation. */
-    T = -a/b + (T0+a/b)*exp(-b*x);
+    T = interp1d(x,m->x,m->T,m->n);
 
     if (T < 0) return 0;
 
     return T;
 }
 
+void muon_free_energy(muon_energy *m)
+{
+    free(m->x);
+    free(m->T);
+    free(m);
+}
+
 double get_dEdx(double T, double rho)
 {
     /* Returns the approximate dE/dx for a muon in water with kinetic energy
@@ -222,6 +251,8 @@ double get_dEdx(double T, double rho)
         }
     }
 
+    if (T < spline_dEdx->x[0]) return spline_dEdx->y[0];
+
     return gsl_spline_eval(spline_dEdx, T, acc_dEdx)*rho;
 }
 
diff --git a/src/muon.h b/src/muon.h
index e95cb3c..253a9df 100644
--- a/src/muon.h
+++ b/src/muon.h
@@ -1,10 +1,20 @@
 #ifndef MUON_H
 #define MUON_H
 
+#include <stddef.h> /* for size_t */
+
 #define EULER_CONSTANT 0.57721
 
+typedef struct muon_energy {
+    double *x;
+    double *T;
+    size_t n;
+} muon_energy;
+
+muon_energy *muon_init_energy(double T0, double rho, size_t n);
+double muon_get_energy(double x, muon_energy *m);
+void muon_free_energy(muon_energy *m);
 double get_range(double T, double rho);
-double get_T(double T0, double x, double rho);
 double get_dEdx(double T, double rho);
 double get_expected_charge(double x, double T, double T0, double *pos, double *dir, double *pmt_pos, double *pmt_normal, double r);
 
diff --git a/src/path.c b/src/path.c
index 478a609..2083dc1 100644
--- a/src/path.c
+++ b/src/path.c
@@ -31,7 +31,7 @@ double path_get_coefficient(unsigned int k, double *s, double *x, double theta0,
     return sum*pow(k-0.5,2)*pow(M_PI,2)/pow(theta0*range,2);
 }
 
-path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, double *z1, double *z2, size_t n, double m)
+path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, void *params, double *z1, double *z2, size_t n, double m)
 {
     size_t i, j;
     double E, mom, beta, theta, phi;
@@ -68,7 +68,7 @@ path *path_init(double *pos, double *dir, double T0, double range, double theta0
             theta1[i] += z1[j]*foo(s[i],range,theta0,j+1);
             theta2[i] += z2[j]*foo(s[i],range,theta0,j+1);
         }
-        T[i] = fun(s[i],T0);
+        T[i] = fun(s[i],params);
         if (i > 0) {
             theta = sqrt(theta1[i]*theta1[i] + theta2[i]*theta2[i]);
             phi = atan2(theta2[i],theta1[i]);
diff --git a/src/path.h b/src/path.h
index ed9591b..449088f 100644
--- a/src/path.h
+++ b/src/path.h
@@ -14,7 +14,7 @@
  * FIXME: Should do some tests to figure out what is the best value. */
 #define MIN_THETA0 0.01
 
-typedef double getKineticEnergyFunc(double x, double T0);
+typedef double getKineticEnergyFunc(double x, void *params);
 
 typedef struct path {
     double pos[3];
@@ -35,7 +35,7 @@ typedef struct path {
 } path;
 
 double path_get_coefficient(unsigned int k, double *s, double *x, double theta0, size_t n);
-path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, double *z1, double *z2, size_t n, double m);
+path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, void *params, double *z1, double *z2, size_t n, double m);
 void path_eval(path *p, double s, double *pos, double *dir, double *T, double *t, double *theta0);
 void path_free(path *p);
 
diff --git a/src/test-path.c b/src/test-path.c
index 7b0a9cd..bca82a4 100644
--- a/src/test-path.c
+++ b/src/test-path.c
@@ -13,7 +13,7 @@ typedef struct sim {
     path *p;
 } sim;
 
-static double getKineticEnergy(double x, double T0)
+static double getKineticEnergy(double x, void *params)
 {
     return 1.0;
 }
@@ -85,7 +85,7 @@ sim *simulate_path(double *pos0, double *dir0, double range, double theta0, size
         z2[i] = path_get_coefficient(i+1,s,theta2,theta0,N);
     }
 
-    simvalue->p = path_init(pos0,dir0,1.0,range,theta0,getKineticEnergy,z1,z2,n,1.0);
+    simvalue->p = path_init(pos0,dir0,1.0,range,theta0,getKineticEnergy,NULL,z1,z2,n,1.0);
 
     free(s);
     free(theta1);
diff --git a/src/test.c b/src/test.c
index 6d7d934..b74728d 100644
--- a/src/test.c
+++ b/src/test.c
@@ -113,26 +113,26 @@ int test_muon_get_T(char *err)
 {
     /* A very simple test to make sure the energy as a function of distance
      * along the track makes sense. Should include more detailed tests later. */
-    double T, E, range;
+    double T0, T, range;
+    muon_energy *m;
 
     /* Assume initial kinetic energy is 1 GeV. */
-    T = 1000.0;
-    E = get_T(T,1e-9,1.0);
+    T0 = 1000.0;
+    m = muon_init_energy(T0,1.0,10000);
+    T = muon_get_energy(1e-9,m);
 
     /* At the beginning of the track we should have roughly the same energy. */
-    if (!isclose(E, T, 1e-5, 0)) {
-        sprintf(err, "KE = %.5f, but expected %.5f", E, T);
+    if (!isclose(T, T0, 1e-5, 0)) {
+        sprintf(err, "KE = %.5f, but expected %.5f", T, T0);
         return 1;
     }
 
-    /* Assume initial kinetic energy is 1 GeV. */
-    T = 1000.0;
-    range = get_range(T,1.0);
-    E = get_T(T,range,1.0);
+    range = get_range(T0,1.0);
+    T = muon_get_energy(range,m);
 
-    /* At the end of the track we should have roughly the same energy. */
-    if (!isclose(E, 0, 1e-5, 1e-5)) {
-        sprintf(err, "KE = %.5f, but expected %.5f", E, 0.0);
+    /* At the end of the track the energy should be approximately 0. */
+    if (!isclose(T, 0, 1e-5, 1e-5)) {
+        sprintf(err, "KE = %.5f, but expected %.5f", T, 0.0);
         return 1;
     }
 
@@ -319,7 +319,7 @@ int test_logsumexp(char *err)
     return 0;
 }
 
-double getKineticEnergy(double x, double T0)
+double getKineticEnergy(double x, void *params)
 {
     return 1.0;
 }
@@ -349,7 +349,7 @@ int test_path(char *err)
 
         rand_sphere(dir0);
 
-        path *p = path_init(pos0,dir0,T0,range,0.1,getKineticEnergy,NULL,NULL,0,m);
+        path *p = path_init(pos0,dir0,T0,range,0.1,getKineticEnergy,NULL,NULL,NULL,0,m);
 
         for (j = 0; j < 100; j++) {
             s = range*j/99;