add ability to fit for multiple vertices

3 files changed, 175 insertions, 141 deletions

diff --git a/src/fit.c b/src/fit.c
index 0634116..714a1cc 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -4978,33 +4978,36 @@ static struct startingParameters {
 double nll(unsigned int n, const double *x, double *grad, void *params)
 {
     fitParams *fpars = (fitParams *) params;
-    double T, theta, phi, t0;
-    double pos[3], dir[3];
+    double theta, phi;
     double fval;
-    double z1[1], z2[1];
     struct timeval tv_start, tv_stop;
+    vertex v;
 
     if (stop) nlopt_force_stop(opt);
 
-    T = x[0];
+    v.id = fpars->id;
 
-    pos[0] = x[1];
-    pos[1] = x[2];
-    pos[2] = x[3];
+    v.T0 = x[0];
+
+    v.pos[0] = x[1];
+    v.pos[1] = x[2];
+    v.pos[2] = x[3];
 
     theta = x[4];
     phi = x[5];
-    dir[0] = sin(theta)*cos(phi);
-    dir[1] = sin(theta)*sin(phi);
-    dir[2] = cos(theta);
+    v.dir[0] = sin(theta)*cos(phi);
+    v.dir[1] = sin(theta)*sin(phi);
+    v.dir[2] = cos(theta);
+
+    v.t0 = x[6];
 
-    t0 = x[6];
+    v.z1[0] = x[7];
+    v.z2[0] = x[8];
 
-    z1[0] = x[7];
-    z2[0] = x[8];
+    v.n = 1;
 
     gettimeofday(&tv_start, NULL);
-    fval = nll_muon(fpars->ev, fpars->id, T, pos, dir, t0, z1, z2, 1, fpars->dx, fpars->dx_shower, fpars->fast);
+    fval = nll_muon(fpars->ev, &v, 1, fpars->dx, fpars->dx_shower, fpars->fast);
     gettimeofday(&tv_stop, NULL);
 
     long long elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
diff --git a/src/likelihood.c b/src/likelihood.c
index 93740a0..4e766b6 100644
--- a/src/likelihood.c
+++ b/src/likelihood.c
@@ -720,7 +720,7 @@ 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, double dx, double dx_shower, int fast)
+double nll_muon(event *ev, vertex *v, size_t n, double dx, double dx_shower, 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
@@ -742,13 +742,13 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     double mu_reflected;
     path *path;
 
-    static double mu_direct[MAX_PMTS];
-    static double mu_shower[MAX_PMTS];
-    static double ts[MAX_PMTS];
-    static double ts_shower[MAX_PMTS];
-    static double ts_sigma[MAX_PMTS];
+    static double mu_direct[MAX_PMTS][MAX_VERTICES];
+    static double mu_shower[MAX_PMTS][MAX_VERTICES];
+    static double ts[MAX_PMTS][MAX_VERTICES];
+    static double ts_shower[MAX_PMTS][MAX_VERTICES];
+    static double ts_sigma[MAX_PMTS][MAX_VERTICES];
     static double mu[MAX_PMTS];
-    double mu_noise, mu_indirect;
+    double mu_noise, mu_indirect[MAX_VERTICES], mu_indirect_total;
 
     double *x, *pdf;
 
@@ -764,148 +764,162 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
 
     double log_pt_fast;
 
+    if (n > MAX_VERTICES) {
+        fprintf(stderr, "maximum number of vertices is %i\n", MAX_VERTICES);
+        exit(1);
+    }
+
     /* 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;
-    for (i = 0; i < MAX_PMTS; i++) ts[i] = 0.0;
-    for (i = 0; i < MAX_PMTS; i++) ts_shower[i] = 0.0;
-    for (i = 0; i < MAX_PMTS; i++) ts_sigma[i] = PMT_TTS;
-    for (i = 0; i < MAX_PMTS; i++) mu[i] = 0.0;
+    for (i = 0; i < MAX_PMTS; i++) {
+        mu[i] = 0.0;
+        for (j = 0; j < n; j++) {
+            mu_direct[i][j] = 0.0;
+            mu_shower[i][j] = 0.0;
+            ts[i][j] = 0.0;
+            ts_shower[i][j] = 0.0;
+            ts_sigma[i][j] = PMT_TTS;
+        }
+    }
 
-    if (fast)
-        p = particle_init(id, T0, 1000);
-    else
-        p = particle_init(id, T0, 10000);
+    for (j = 0; j < n; j++) {
+        if (fast)
+            p = particle_init(v[j].id, v[j].T0, 1000);
+        else
+            p = particle_init(v[j].id, v[j].T0, 10000);
 
-    range = p->range;
+        range = p->range;
 
-    /* Number of points to sample along the particle track. */
-    npoints = (size_t) (range/dx + 0.5);
+        /* Number of points to sample along the particle track. */
+        npoints = (size_t) (range/dx + 0.5);
 
-    if (npoints < MIN_NPOINTS) npoints = MIN_NPOINTS;
+        if (npoints < MIN_NPOINTS) npoints = MIN_NPOINTS;
 
-    /* Calculate total energy */
-    E0 = T0 + p->mass;
-    p0 = sqrt(E0*E0 - p->mass*p->mass);
-    beta0 = p0/E0;
+        /* Calculate total energy */
+        E0 = v[j].T0 + p->mass;
+        p0 = sqrt(E0*E0 - p->mass*p->mass);
+        beta0 = p0/E0;
 
-    /* FIXME: is this formula valid for muons? */
-    theta0 = get_scattering_rms(range/2,p0,beta0,1.0)/sqrt(range/2);
+        /* 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, npoints, theta0, getKineticEnergy, p, z1, z2, n, p->mass);
+        if (!fast) {
+            path = path_init(v[j].pos, v[j].dir, v[j].T0, range, npoints, theta0, getKineticEnergy, p, v[j].z1, v[j].z2, v[j].n, p->mass);
 
-        if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
-            electron_get_position_distribution_parameters(T0, &pos_a, &pos_b);
+            if (v[j].id == IDP_E_MINUS || v[j].id == IDP_E_PLUS) {
+                electron_get_position_distribution_parameters(v[j].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);
+                /* 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_shower + 0.5);
+                /* Number of points to sample along the longitudinal direction for
+                 * the electromagnetic shower. */
+                npoints_shower = (size_t) ((xhi-xlo)/dx_shower + 0.5);
 
-            if (npoints_shower < MIN_NPOINTS) npoints_shower = MIN_NPOINTS;
+                if (npoints_shower < MIN_NPOINTS) npoints_shower = MIN_NPOINTS;
 
-            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);
-            }
+                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);
+                }
 
-            /* Make sure the PDF is normalized. */
-            double norm = trapz(pdf,x[1]-x[0],npoints_shower);
-            for (i = 0; i < npoints_shower; i++) {
-                pdf[i] /= norm;
+                /* Make sure the PDF is normalized. */
+                double norm = trapz(pdf,x[1]-x[0],npoints_shower);
+                for (i = 0; i < npoints_shower; i++) {
+                    pdf[i] /= norm;
+                }
             }
-
         }
-    }
 
-    if (beta0 > BETA_MIN)
-        get_smax(p, BETA_MIN, range, &smax);
-    else
-        smax = 0.0;
+        if (beta0 > BETA_MIN)
+            get_smax(p, BETA_MIN, range, &smax);
+        else
+            smax = 0.0;
 
-    /* Compute the Cerenkov angle at the start of the track.
-     *
-     * These are computed at the beginning of this function and then passed to
-     * the different functions to avoid recomputing them on the fly. */
-    wavelength0 = 400.0;
-    n_d2o = get_index_snoman_d2o(wavelength0);
-    n_h2o = get_index_snoman_h2o(wavelength0);
-    cos_theta_cerenkov = 1/(n_d2o*beta0);
-    sin_theta_cerenkov = sqrt(1-pow(cos_theta_cerenkov,2));
+        /* Compute the Cerenkov angle at the start of the track.
+         *
+         * These are computed at the beginning of this function and then passed to
+         * the different functions to avoid recomputing them on the fly. */
+        wavelength0 = 400.0;
+        n_d2o = get_index_snoman_d2o(wavelength0);
+        n_h2o = get_index_snoman_h2o(wavelength0);
+        cos_theta_cerenkov = 1/(n_d2o*beta0);
+        sin_theta_cerenkov = sqrt(1-pow(cos_theta_cerenkov,2));
+
+        mu_indirect[j] = 0.0;
+        for (i = 0; i < MAX_PMTS; i++) {
+            if (pmts[i].pmt_type != PMT_NORMAL) continue;
+
+            /* Skip PMTs which weren't hit when doing the "fast" likelihood
+             * calculation. */
+            if (fast && !ev->pmt_hits[i].hit) continue;
 
-    mu_indirect = 0.0;
-    for (i = 0; i < MAX_PMTS; i++) {
-        if (pmts[i].pmt_type != PMT_NORMAL) continue;
+            if (fast) {
+                mu_direct[i][j] = get_total_charge_approx(beta0, v[j].pos, v[j].dir, p, i, smax, theta0, &ts[i][j], &mu_reflected, n_d2o, n_h2o, cos_theta_cerenkov, sin_theta_cerenkov);
+                ts[i][j] += v[j].t0;
+                mu_indirect[j] += mu_reflected;
+                continue;
+            }
 
-        /* Skip PMTs which weren't hit when doing the "fast" likelihood
-         * calculation. */
-        if (fast && !ev->pmt_hits[i].hit) continue;
+            double q_indirect;
+            integrate_path(path, i, &mu_direct[i][j], &q_indirect, &result);
+            mu_indirect[j] += q_indirect;
+
+            if (mu_direct[i][j] > 1e-9) {
+                ts[i][j] = v[j].t0 + result/mu_direct[i][j];
+            } else {
+                /* If the expected number of PE is very small, our estimate of
+                 * the time can be crazy since the error in the total charge is
+                 * in the denominator. Therefore, we estimate the time using
+                 * the same technique as in get_total_charge_approx(). See that
+                 * function for more details.
+                 *
+                 * Note: I don't really know how much this affects the likelihood.
+                 * I should test it to see if we can get away with something even
+                 * simpler (like just computing the distance from the start of the
+                 * track to the PMT). */
+                ts[i][j] = v[j].t0 + guess_time(v[j].pos,v[j].dir,pmts[i].pos,smax,n_d2o,n_h2o,cos_theta_cerenkov,sin_theta_cerenkov);
+            }
 
-        if (fast) {
-            mu_direct[i] = get_total_charge_approx(beta0, pos, dir, p, i, smax, theta0, &ts[i], &mu_reflected, n_d2o, n_h2o, cos_theta_cerenkov, sin_theta_cerenkov);
-            ts[i] += t0;
-            mu_indirect += mu_reflected;
-            continue;
+            if (v[j].id == IDP_E_MINUS || v[j].id == IDP_E_PLUS) {
+                /* If we are fitting for an electron or positron, we need to
+                 * estimate the expected number of photons produced by the
+                 * electromagnetic shower.
+                 *
+                 * Technically we should also have a term for the electromagnetic
+                 * shower produced by muons, but for the energy range I'm
+                 * considering right now, it is very small. */
+                integrate_path_shower(p,x,pdf,v[j].T0,v[j].pos,v[j].dir,i,npoints_shower,&mu_shower[i][j],&q_indirect,&result,&ts_sigma[i][j],p->rad);
+                mu_indirect[j] += q_indirect;
+                ts_shower[i][j] = v[j].t0 + result;
+            }
         }
 
-        double q_indirect;
-        integrate_path(path, i, mu_direct+i, &q_indirect, &result);
-        mu_indirect += q_indirect;
+        if (!fast) {
+            path_free(path);
 
-        if (mu_direct[i] > 1e-9) {
-            ts[i] = t0 + result/mu_direct[i];
-        } else {
-            /* If the expected number of PE is very small, our estimate of
-             * the time can be crazy since the error in the total charge is
-             * in the denominator. Therefore, we estimate the time using
-             * the same technique as in get_total_charge_approx(). See that
-             * function for more details.
-             *
-             * Note: I don't really know how much this affects the likelihood.
-             * I should test it to see if we can get away with something even
-             * simpler (like just computing the distance from the start of the
-             * track to the PMT). */
-            ts[i] = t0 + guess_time(pos,dir,pmts[i].pos,smax,n_d2o,n_h2o,cos_theta_cerenkov,sin_theta_cerenkov);
+            if (v[j].id == IDP_E_MINUS || v[j].id == IDP_E_PLUS) {
+                free(x);
+                free(pdf);
+            }
         }
 
-        if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
-            /* If we are fitting for an electron or positron, we need to
-             * estimate the expected number of photons produced by the
-             * electromagnetic shower.
-             *
-             * Technically we should also have a term for the electromagnetic
-             * shower produced by muons, but for the energy range I'm
-             * considering right now, it is very small. */
-            integrate_path_shower(p,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;
-        }
+        particle_free(p);
     }
 
-    if (!fast) {
-        path_free(path);
+    mu_noise = DARK_RATE*GTVALID*1e-9;
 
-        if (id == IDP_E_MINUS || id == IDP_E_PLUS) {
-            free(x);
-            free(pdf);
-        }
+    mu_indirect_total = 0.0;
+    for (j = 0; j < n; j++) {
+        if (v[j].id == IDP_E_MINUS || v[j].id == IDP_E_PLUS)
+            mu_indirect_total += mu_indirect[j]*CHARGE_FRACTION_ELECTRON/10000.0;
+        else
+            mu_indirect_total += mu_indirect[j]*CHARGE_FRACTION_MUON/10000.0;
     }
 
-    particle_free(p);
-
-    mu_noise = DARK_RATE*GTVALID*1e-9;
-    if (id == IDP_E_MINUS || id == IDP_E_PLUS)
-        mu_indirect *= CHARGE_FRACTION_ELECTRON/10000.0;
-    else
-        mu_indirect *= CHARGE_FRACTION_MUON/10000.0;
-
     /* Compute the expected number of photons reaching each PMT by adding up
      * the contributions from the noise hits and the direct, indirect, and
      * shower light. */
@@ -916,10 +930,12 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
          * calculation. */
         if (fast && !ev->pmt_hits[i].hit) continue;
 
-        if (fast)
-            mu[i] = mu_direct[i] + mu_noise;
-        else
-            mu[i] = mu_direct[i] + mu_shower[i] + mu_indirect + mu_noise;
+        for (j = 0; j < n; j++) {
+            if (fast)
+                mu[i] += mu_direct[i][j] + mu_noise;
+            else
+                mu[i] += mu_direct[i][j] + mu_shower[i][j] + mu_indirect_total + mu_noise;
+        }
     }
 
     if (fast)
@@ -939,14 +955,14 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
         log_mu = log(mu[i]);
 
         if (fast)
-            log_pt_fast = log_pt(ev->pmt_hits[i].t, 1, mu_noise, mu_indirect, &mu_direct[i], &mu_shower[i], 1, &ts[i], &ts_shower[i], ts[i], PMT_TTS, &ts_sigma[i]);
+            log_pt_fast = log_pt(ev->pmt_hits[i].t, 1, mu_noise, mu_indirect_total, &mu_direct[i][0], &mu_shower[i][0], n, &ts[i][0], &ts_shower[i][0], ts[i][0], PMT_TTS, &ts_sigma[i][0]);
 
         if (ev->pmt_hits[i].hit) {
             for (j = 1; j < MAX_PE; j++) {
                 if (fast)
                     logp[j] = get_log_pq(ev->pmt_hits[i].qhs,j) - mu[i] + j*log_mu - lnfact(j) + log_pt_fast;
                 else
-                    logp[j] = get_log_pq(ev->pmt_hits[i].qhs,j) - mu[i] + j*log_mu - lnfact(j) + log_pt(ev->pmt_hits[i].t, j, mu_noise, mu_indirect, &mu_direct[i], &mu_shower[i], 1, &ts[i], &ts_shower[i], ts[i], PMT_TTS, &ts_sigma[i]);
+                    logp[j] = get_log_pq(ev->pmt_hits[i].qhs,j) - mu[i] + j*log_mu - lnfact(j) + log_pt(ev->pmt_hits[i].t, j, mu_noise, mu_indirect_total, &mu_direct[i][0], &mu_shower[i][0], n, &ts[i][0], &ts_shower[i][0], ts[i][0], PMT_TTS, &ts_sigma[i][0]);
 
                 if (j == 1 || logp[j] > max_logp) max_logp = logp[j];
 
@@ -971,8 +987,9 @@ double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t
     }
 
     logp_path = 0.0;
-    for (i = 0; i < n; i++)
-        logp_path += log_norm(z1[i],0,1) + log_norm(z2[i],0,1);
+    for (j = 0; j < n; j++)
+        for (i = 0; i < v[j].n; i++)
+            logp_path += log_norm(v[j].z1[i],0,1) + log_norm(v[j].z2[i],0,1);
 
     /* Add up the negative log likelihood terms from each PMT. I use a kahan
      * sum here instead of just summing all the values to help prevent round
diff --git a/src/likelihood.h b/src/likelihood.h
index be003d3..921a97d 100644
--- a/src/likelihood.h
+++ b/src/likelihood.h
@@ -52,6 +52,20 @@
  * quantity from simulation. */
 #define PHOTONS_PER_MEV 400.0
 
+/* Maximum number of vertices to fit. */
+#define MAX_VERTICES 2
+
+typedef struct vertex {
+    int id;
+    double T0;
+    double pos[3];
+    double dir[3];
+    double t0;
+    double z1[10];
+    double z2[10];
+    size_t n;
+} vertex;
+
 typedef struct particle {
     int id;
     double mass;
@@ -69,6 +83,6 @@ double particle_get_energy(double x, particle *p);
 void particle_free(particle *p);
 double time_pdf(double t, double mu_noise, double mu_indirect, double *mu_direct, double *mu_shower, size_t n, double *ts, double *ts_shower, double tmean, double sigma, double *ts_sigma);
 double time_cdf(double t, double mu_noise, double mu_indirect, double *mu_direct, double *mu_shower, size_t n, double *ts, double *ts_shower, double tmean, double sigma, double *ts_sigma);
-double nll_muon(event *ev, int id, double T0, double *pos, double *dir, double t0, double *z1, double *z2, size_t n, double dx, double dx_shower, int fast);
+double nll_muon(event *ev, vertex *v, size_t n, double dx, double dx_shower, int fast);
 
 #endif