speed up likelihood function and switch to using fixed dx

This commit speeds up the likelihood function by about ~20% by using the
precomputed track positions, directions, times, etc. instead of interpolating
them on the fly.

It also switches to computing the number of points to integrate along the track
by dividing the track length by a specified distance, currently set to 1 cm.
This should hopefully speed things up for lower energies and result in more
stable fits at high energies.

1 files changed, 91 insertions, 41 deletions

diff --git a/src/likelihood.c b/src/likelihood.c
index 707ff38..8116261 100644
--- a/src/likelihood.c
+++ b/src/likelihood.c
@@ -320,7 +320,7 @@ double log_pt(double t, size_t n, double mu_noise, double mu_indirect, double *m
     return ln(n) + (n-1)*log1p(-F(t,mu_noise,mu_indirect,mu_direct,mu_shower,n2,ts,ts_shower,tmean,sigma,ts_sigma)) + log(f(t,mu_noise,mu_indirect,mu_direct,mu_shower,n2,ts,ts_shower,tmean,sigma,ts_sigma));
 }
 
-static void integrate_path_shower(double *x, double *pdf, double T0, double *pos0, double *dir0, int pmt, double a, double b, size_t n, double *mu_direct, double *mu_indirect, double *time, double *sigma, double rad, double pos_a, double pos_b)
+static void integrate_path_shower(double *x, double *pdf, double T0, double *pos0, double *dir0, int pmt, size_t n, double *mu_direct, double *mu_indirect, double *time, double *sigma, double rad)
 {
     size_t i;
     static double qs[MAX_NPOINTS];
@@ -357,10 +357,11 @@ static void integrate_path_shower(double *x, double *pdf, double T0, double *pos
         ts2[i] = t*t*qs[i];
     }
 
-    *mu_direct = trapz(qs,(b-a)/(n-1),n);
-    *mu_indirect = trapz(q_indirects,(b-a)/(n-1),n);
-    *time = trapz(ts,(b-a)/(n-1),n)/(*mu_direct);
-    double t2 = trapz(ts2,(b-a)/(n-1),n)/(*mu_direct);
+    double dx = x[1]-x[0];
+    *mu_direct = trapz(qs,dx,n);
+    *mu_indirect = trapz(q_indirects,dx,n);
+    *time = trapz(ts,dx,n)/(*mu_direct);
+    double t2 = trapz(ts2,dx,n)/(*mu_direct);
 
     *sigma = sqrt(t2-(*time)*(*time));
 
@@ -372,7 +373,7 @@ static void integrate_path_shower(double *x, double *pdf, double T0, double *pos
     }
 }
 
-static void integrate_path(path *p, int pmt, double a, double b, size_t n, double *mu_direct, double *mu_indirect, double *time)
+static void integrate_path(path *p, int pmt, double *mu_direct, double *mu_indirect, double *time)
 {
     size_t i;
     static double qs[MAX_NPOINTS];
@@ -380,7 +381,7 @@ static void integrate_path(path *p, int pmt, double a, double b, size_t n, doubl
     static double ts[MAX_NPOINTS];
     double dir[3], pos[3], n_d2o, n_h2o, wavelength0, t, theta0, beta, l_d2o, l_h2o, x;
 
-    if (n > MAX_NPOINTS) {
+    if (p->len > MAX_NPOINTS) {
         fprintf(stderr, "number of points is greater than MAX_NPOINTS!\n");
         exit(1);
     }
@@ -390,10 +391,27 @@ static void integrate_path(path *p, int pmt, double a, double b, size_t n, doubl
     n_d2o = get_index_snoman_d2o(wavelength0);
     n_h2o = get_index_snoman_h2o(wavelength0);
 
-    for (i = 0; i < n; i++) {
-        x = a + i*(b-a)/(n-1);
+    for (i = 0; i < p->len; i++) {
+        x = p->s[i];
+
+        pos[0] = p->x[i];
+        pos[1] = p->y[i];
+        pos[2] = p->z[i];
 
-        path_eval(p, x, pos, dir, &beta, &t, &theta0);
+        dir[0] = p->dx[i];
+        dir[1] = p->dy[i];
+        dir[2] = p->dz[i];
+
+        beta = p->beta[i];
+
+        t = p->t[i];
+
+        if (p->n > 0) {
+            theta0 = p->theta0;
+        } else {
+            theta0 = fmax(MIN_THETA0,p->theta0*sqrt(x));
+            theta0 = fmin(MAX_THETA0,theta0);
+        }
 
         get_path_length(pos,pmts[pmt].pos,AV_RADIUS,&l_d2o,&l_h2o);
 
@@ -403,9 +421,10 @@ static void integrate_path(path *p, int pmt, double a, double b, size_t n, doubl
         ts[i] = t*qs[i];
     }
 
-    *mu_direct = trapz(qs,(b-a)/(n-1),n);
-    *mu_indirect = trapz(q_indirects,(b-a)/(n-1),n);
-    *time = trapz(ts,(b-a)/(n-1),n);
+    double dx = p->s[1]-p->s[0];
+    *mu_direct = trapz(qs,dx,p->len);
+    *mu_indirect = trapz(q_indirects,dx,p->len);
+    *time = trapz(ts,dx,p->len);
 }
 
 static double get_total_charge_approx(double T0, double *pos, double *dir, particle *p, int i, double smax, double theta0, double *t, double *mu_reflected, double n_d2o, double n_h2o, double cos_theta_cerenkov, double sin_theta_cerenkov)
@@ -628,8 +647,19 @@ static double getKineticEnergy(double x, void *p)
     return particle_get_energy(x, (particle *) p);
 }
 
-double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t0, double *z1, double *z2, size_t n, size_t npoints, int fast)
+double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t0, double *z1, double *z2, size_t n, double dx, int fast)
 {
+    /* Returns the negative log likelihood for event `ev` given a particle with
+     * id `id`, initial kinetic energy `T0`, position `pos`, direction `dir` and
+     * initial time `t0`.
+     *
+     * `dx` is the distance the track is sampled at to compute the likelihood.
+     *
+     * If `fast` is nonzero, returns an approximate faster version of the likelihood.
+     *
+     * `z1` and `z2` should be arrays of length `n` and represent coefficients
+     * in the Karhunen Loeve expansion of the track path. These are only used
+     * when `fast` is zero. */
     size_t i, j, nhit;
     static double logp[MAX_PE], nll[MAX_PMTS];
     double range, theta0, E0, p0, beta0, smax, log_mu, max_logp;
@@ -649,6 +679,18 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     static double mu[MAX_PMTS];
     double mu_noise, mu_indirect;
 
+    double *x, *pdf;
+
+    double result;
+
+    double min_ratio;
+
+    double pos_a, pos_b;
+
+    double xlo, xhi;
+
+    size_t npoints, npoints_shower;
+
     /* Initialize static arrays. */
     for (i = 0; i < MAX_PMTS; i++) mu_direct[i] = 0.0;
     for (i = 0; i < MAX_PMTS; i++) mu_shower[i] = 0.0;
@@ -657,28 +699,12 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     for (i = 0; i < MAX_PMTS; i++) ts_sigma[i] = PMT_TTS;
     for (i = 0; i < MAX_PMTS; i++) mu[i] = 0.0;
 
-    double pos_a, pos_b;
-
-    electron_get_position_distribution_parameters(T0, &pos_a, &pos_b);
-
-    /* Upper limit to integrate for the electromagnetic shower. */
-    double xlo = 0.0;
-    double xhi = gsl_cdf_gamma_Pinv(0.99,pos_a,pos_b);
-
-    double *x = malloc(sizeof(double)*npoints);
-    double *pdf = malloc(sizeof(double)*npoints);
-    for (i = 0; i < npoints; i++) {
-        x[i] = xlo + i*(xhi-xlo)/(npoints-1);
-        pdf[i] = gamma_pdf(x[i],pos_a,pos_b);
-    }
-
-    double result;
-
-    double min_ratio;
-
     p = particle_init(id, T0, 10000);
     range = p->range;
 
+    /* Number of points to sample along the particle track. */
+    npoints = (size_t) (range/dx + 0.5);
+
     /* Calculate total energy */
     E0 = T0 + p->mass;
     p0 = sqrt(E0*E0 - p->mass*p->mass);
@@ -687,8 +713,29 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     /* FIXME: is this formula valid for muons? */
     theta0 = get_scattering_rms(range/2,p0,beta0,1.0)/sqrt(range/2);
 
-    if (!fast)
-        path = path_init(pos, dir, T0, range, theta0, getKineticEnergy, p, z1, z2, n, p->mass);
+    if (!fast) {
+        path = path_init(pos, dir, T0, range, npoints, theta0, getKineticEnergy, p, z1, z2, n, p->mass);
+
+        if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
+            electron_get_position_distribution_parameters(T0, &pos_a, &pos_b);
+
+            /* Lower and upper limits to integrate for the electromagnetic
+             * shower. */
+            xlo = 0.0;
+            xhi = gsl_cdf_gamma_Pinv(0.99,pos_a,pos_b);
+
+            /* Number of points to sample along the longitudinal direction for
+             * the electromagnetic shower. */
+            npoints_shower = (size_t) ((xhi-xlo)/dx + 0.5);
+
+            x = malloc(sizeof(double)*npoints_shower);
+            pdf = malloc(sizeof(double)*npoints_shower);
+            for (i = 0; i < npoints_shower; i++) {
+                x[i] = xlo + i*(xhi-xlo)/(npoints_shower-1);
+                pdf[i] = gamma_pdf(x[i],pos_a,pos_b);
+            }
+        }
+    }
 
     if (beta0 > BETA_MIN)
         get_smax(p, BETA_MIN, range, &smax);
@@ -765,7 +812,7 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
         if (b > range) b = range;
 
         double q_indirect;
-        integrate_path(path, i, a, b, npoints, mu_direct+i, &q_indirect, &result);
+        integrate_path(path, i, mu_direct+i, &q_indirect, &result);
         mu_indirect += q_indirect;
 
         if (mu_direct[i] > 1e-9) {
@@ -792,15 +839,21 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
         }
 
         if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
-            integrate_path_shower(x,pdf,T0,pos,dir,i,0,xhi,npoints,mu_shower+i,&q_indirect,&result,&ts_sigma[i],p->rad,pos_a,pos_b);
+            integrate_path_shower(x,pdf,T0,pos,dir,i,npoints_shower,mu_shower+i,&q_indirect,&result,&ts_sigma[i],p->rad);
             mu_indirect += q_indirect;
             ts_shower[i] = t0 + result;
         }
     }
 
-    if (!fast)
+    if (!fast) {
         path_free(path);
 
+        if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
+            free(x);
+            free(pdf);
+        }
+    }
+
     particle_free(p);
 
     mu_noise = DARK_RATE*GTVALID*1e-9;
@@ -863,8 +916,5 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     for (i = 0; i < n; i++)
         logp_path += log_norm(z1[i],0,1) + log_norm(z2[i],0,1);
 
-    free(x);
-    free(pdf);
-
     return kahan_sum(nll,nhit) - logp_path;
 }