update fit.c to fit multiple vertices

This commit adds a new function fit_event2() to fit multiple vertices. To seed
the fit, fit_event2() does the following:

  - use the QUAD fitter to find the position and initial time of the event
  - call find_peaks() to find possible directions for the particles
  - loop over all possible unique combinations of the particles and direction
    vectors and do a "fast" minimization

The best minimum found from the "fast" minimizations is then used to start the fit.

This commit has a few other updates:

- adds a hit_only parameter to the nll() function. This was necessary since
previously PMTs which weren't hit were always skipped for the fast
minimization, but when fitting for multiple vertices we need to include PMTs
which aren't hit since we float the energy.

- add the function guess_energy() to guess the energy of a particle given a
position and direction. This function estimates the energy by summing up the
QHS for all PMTs hit within the Cerenkov cone and dividing by 6.

- fixed a bug which caused the fit to freeze when hitting ctrl-c during the
fast minimization phase.

1 files changed, 361 insertions, 6 deletions

diff --git a/src/fit.c b/src/fit.c
index 7965441..ea062ba 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -23,6 +23,16 @@
 #include "release.h"
 #include "id_particles.h"
 #include "misc.h"
+#include "quad.h"
+#include "sno.h"
+#include "find_peaks.h"
+
+/* Maximum number of fit parameters. Should be at least 4 + 3*MAX_VERTICES. */
+#define MAX_PARS 100
+/* Maximum number of peaks to search for in Hough transform. */
+#define MAX_NPEAKS 10
+/* Maximum kinetic energy for any particle. */
+#define MAX_ENERGY 10000
 
 char *GitSHA1(void);
 char *GitDirty(void);
@@ -42,8 +52,10 @@ typedef struct fitParams {
     double dx_shower;
     int fast;
     int print;
-    int id;
+    int id[MAX_VERTICES];
+    size_t n;
     int charge_only;
+    int hit_only;
 } fitParams;
 
 int particles[] = {
@@ -4976,6 +4988,62 @@ static struct startingParameters {
     { 800.0, 800.0, 800.0},
 };
 
+static double nlopt_nll2(unsigned int n, const double *x, double *grad, void *params)
+{
+    size_t i;
+    fitParams *fpars = (fitParams *) params;
+    double theta, phi;
+    double fval;
+    struct timeval tv_start, tv_stop;
+    vertex v[MAX_VERTICES];
+
+    if (stop) nlopt_force_stop(opt);
+
+    for (i = 0; i < fpars->n; i++) {
+        v[i].id = fpars->id[i];
+
+        v[i].T0 = x[4+3*i];
+
+        v[i].pos[0] = x[0];
+        v[i].pos[1] = x[1];
+        v[i].pos[2] = x[2];
+        v[i].t0 = x[3];
+
+        theta = x[5+3*i];
+        phi = x[6+3*i];
+        v[i].dir[0] = sin(theta)*cos(phi);
+        v[i].dir[1] = sin(theta)*sin(phi);
+        v[i].dir[2] = cos(theta);
+
+        v[i].n = 0;
+    }
+
+    gettimeofday(&tv_start, NULL);
+    fval = nll(fpars->ev, v, fpars->n, fpars->dx, fpars->dx_shower, fpars->fast, fpars->charge_only, fpars->hit_only);
+    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;
+
+    if (fpars->print) {
+        printf("%5zu %7.2f %7.2f %7.2f %6.2f ",
+               iter++,
+               x[0],
+               x[1],
+               x[2],
+               x[3]);
+
+        for (i = 0; i < fpars->n; i++)
+            printf("%10.2f %7.2f %7.2f ",
+                   x[4+3*i],
+                   x[5+3*i],
+                   x[6+3*i]);
+
+        printf("f() = %7.3e took %lld ms\n", fval, elapsed);
+    }
+
+    return fval;
+}
+
 static double nlopt_nll(unsigned int n, const double *x, double *grad, void *params)
 {
     fitParams *fpars = (fitParams *) params;
@@ -4986,7 +5054,7 @@ static double nlopt_nll(unsigned int n, const double *x, double *grad, void *par
 
     if (stop) nlopt_force_stop(opt);
 
-    v.id = fpars->id;
+    v.id = fpars->id[0];
 
     v.T0 = x[0];
 
@@ -5008,7 +5076,7 @@ static double nlopt_nll(unsigned int n, const double *x, double *grad, void *par
     v.n = 1;
 
     gettimeofday(&tv_start, NULL);
-    fval = nll(fpars->ev, &v, 1, fpars->dx, fpars->dx_shower, fpars->fast, fpars->charge_only);
+    fval = nll(fpars->ev, &v, 1, fpars->dx, fpars->dx_shower, fpars->fast, fpars->charge_only, fpars->hit_only);
     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;
@@ -5111,6 +5179,290 @@ void guess_direction(event *ev, double *pos, double *theta, double *phi)
     *phi = atan2(dir[1],dir[0]);
 }
 
+double guess_energy(event *ev, double *pos, double *dir)
+{
+    /* Returns the approximate energy for a particle at position `pos` in
+     * direction `dir`. This guess is made by summing up all the charge within
+     * a Cerenkov cone and making the approximation that the particle creates
+     * approximately 6 photons/MeV.
+     *
+     * FIXME: Should update this to something better. */
+    size_t i;
+    double qhs_sum, n_d2o;
+    double pmt_dir[3];
+
+    n_d2o = get_index_snoman_d2o(400.0);
+
+    qhs_sum = 0.0;
+    for (i = 0; i < MAX_PMTS; i++) {
+        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL || !ev->pmt_hits[i].hit) continue;
+        SUB(pmt_dir,pmts[i].pos,pos);
+        normalize(pmt_dir);
+        if (DOT(pmt_dir,dir) > 1/n_d2o)
+            qhs_sum += ev->pmt_hits[i].qhs;
+    }
+
+    return qhs_sum/6.0;
+}
+
+int fit_event2(event *ev, double *xopt, double *fmin, int *id, size_t n, double maxtime)
+{
+    size_t i, j;
+    fitParams fpars;
+    double x[100], ss[100], lb[100], ub[100], fval, n_d2o, x0[100], T0, Tmin, mass, pos[3], t0, dir[3];
+    struct timeval tv_start, tv_stop;
+    double time_elapsed;
+
+    // x
+    // y
+    // z
+    // t0
+    // E
+    // theta
+    // phi
+
+    opt = nlopt_create(NLOPT_LN_BOBYQA, 4+3*n);
+    nlopt_set_min_objective(opt,nlopt_nll2,&fpars);
+
+    /* Guess the position and t0 of the event using the QUAD fitter. */
+    quad(ev,pos,&t0,10000);
+
+    if (t0 < 0 || t0 > GTVALID || NORM(pos) > PSUP_RADIUS) {
+        /* If the QUAD fitter returns something outside the PSUP or with an
+         * invalid time we just assume it's at the center. */
+        fprintf(stderr, "quad returned pos = %.2f, %.2f, %.2f t0 = %.2f. Assuming vertex is at the center!\n",
+                pos[0], pos[1], pos[2], t0);
+        pos[0] = 0.0;
+        pos[1] = 0.0;
+        pos[2] = 0.0;
+        t0 = guess_t0(ev,pos);
+    }
+
+    x0[0] = pos[0];
+    x0[1] = pos[1];
+    x0[2] = pos[2];
+    x0[3] = t0;
+
+    ss[0] = 10.0;
+    ss[1] = 10.0;
+    ss[2] = 10.0;
+    ss[3] = 1.0;
+
+    lb[0] = -1000.0;
+    lb[1] = -1000.0;
+    lb[2] = -1000.0;
+    lb[3] = 0.0;
+
+    ub[0] = +1000.0;
+    ub[1] = +1000.0;
+    ub[2] = +1000.0;
+    ub[3] = GTVALID;
+
+    double peak_theta[MAX_NPEAKS];
+    double peak_phi[MAX_NPEAKS];
+    size_t npeaks;
+
+    find_peaks(ev, pos, 100, 100, peak_theta, peak_phi, &npeaks, LEN(peak_theta), 0.5);
+
+    size_t *result = malloc(sizeof(size_t)*n*ipow(npeaks,n));
+
+    size_t nvertices;
+
+    unique_vertices(id,n,npeaks,result,&nvertices);
+
+    n_d2o = get_index_snoman_d2o(400.0);
+
+    fpars.ev = ev;
+
+    iter = 0;
+
+    /* First we do a set of "quick" minimizations to try and start the
+     * minimizer close to the minimum. To make these function evaluations
+     * faster, we set the absolute tolerance on the likelihood to 1.0, the
+     * maximum number of function evaluations to 100, and the relative
+     * tolerance on the numerical integration to 10%. */
+    fpars.dx = 1.0;
+    fpars.dx_shower = 10.0;
+    fpars.fast = 1;
+    fpars.hit_only = 0;
+    fpars.print = 0;
+    fpars.charge_only = 0;
+    nlopt_set_ftol_abs(opt, 1e-2);
+    nlopt_set_maxeval(opt, 10000);
+
+    fpars.n = n;
+    for (i = 0; i < n; i++)
+        fpars.id[i] = id[i];
+
+    for (i = 0; i < nvertices; i++) {
+        memcpy(x,x0,sizeof(x));
+        for (j = 0; j < n; j++) {
+            x[5+3*j] = peak_theta[result[i*n+j]];
+            x[6+3*j] = peak_phi[result[i*n+j]];
+            dir[0] = sin(peak_theta[result[i*n+j]])*cos(peak_phi[result[i*n+j]]);
+            dir[1] = sin(peak_theta[result[i*n+j]])*sin(peak_phi[result[i*n+j]]);
+            dir[2] = cos(peak_theta[result[i*n+j]]);
+            T0 = guess_energy(ev,x0,dir);
+
+            switch (id[j]) {
+            case IDP_E_MINUS:
+            case IDP_E_PLUS:
+                mass = ELECTRON_MASS;
+                break;
+            case IDP_MU_MINUS:
+            case IDP_MU_PLUS:
+                mass = MUON_MASS;
+                break;
+            case IDP_PROTON:
+                mass = PROTON_MASS;
+                break;
+            default:
+                fprintf(stderr, "unknown particle id %i\n", id[j]);
+                exit(1);
+            }
+
+            /* Calculate the Cerenkov threshold for a muon. */
+            Tmin = mass/sqrt(1.0-1.0/(n_d2o*n_d2o)) - mass;
+
+            /* If our guess is below the Cerenkov threshold, start at the Cerenkov
+             * threshold. */
+            double Tmin2 = sqrt(mass*mass/(1-pow(0.9,2)))-mass;
+            if (T0 < Tmin2) T0 = Tmin2;
+
+            x[4+3*j] = T0;
+            ss[4+3*j] = x[4+3*j]*0.02;
+            ss[5+3*j] = 0.01;
+            ss[6+3*j] = 0.01;
+            lb[4+3*j] = Tmin;
+            lb[5+3*j] = -INFINITY;
+            lb[6+3*j] = -INFINITY;
+            ub[4+3*j] = MAX_ENERGY;
+            ub[5+3*j] = +INFINITY;
+            ub[6+3*j] = +INFINITY;
+        }
+
+        /* Fix the position. */
+        lb[0] = pos[0];
+        lb[1] = pos[1];
+        lb[2] = pos[2];
+
+        ub[0] = pos[0];
+        ub[1] = pos[1];
+        ub[2] = pos[2];
+
+        nlopt_set_lower_bounds(opt, lb);
+        nlopt_set_upper_bounds(opt, ub);
+
+        nlopt_set_initial_step(opt, ss);
+
+        gettimeofday(&tv_start, NULL);
+        nlopt_optimize(opt,x,&fval);
+        gettimeofday(&tv_stop, NULL);
+
+        if (stop) goto stop;
+
+        long long elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
+
+        printf("%4zu/%4zu %7.2f %7.2f %7.2f %6.2f ",
+               i+1,
+               nvertices,
+               x[0],
+               x[1],
+               x[2],
+               x[3]);
+
+        for (j = 0; j < n; j++)
+            printf("%10.2f %7.2f %7.2f ",
+                   x[4+3*j],
+                   x[5+3*j],
+                   x[6+3*j]);
+
+        printf("f() = %7.3e took %lld ms\n", fval, elapsed);
+
+        if (i == 0 || fval < *fmin) {
+            *fmin = fval;
+            memcpy(xopt,x,sizeof(x));
+        }
+    }
+
+    /* Reset the lower and upper bounds. */
+    lb[0] = -1000.0;
+    lb[1] = -1000.0;
+    lb[2] = -1000.0;
+    lb[3] = 0.0;
+
+    ub[0] = +1000.0;
+    ub[1] = +1000.0;
+    ub[2] = +1000.0;
+    ub[3] = GTVALID;
+
+    for (i = 0; i < nvertices; i++) {
+        for (j = 0; j < n; j++) {
+            ss[4+3*j] = xopt[4+3*j]*0.02;
+        }
+    }
+
+    nlopt_set_lower_bounds(opt, lb);
+    nlopt_set_upper_bounds(opt, ub);
+
+    nlopt_set_initial_step(opt, ss);
+
+    /* Now, we do the "real" minimization. */
+    fpars.dx = 1.0;
+    fpars.dx_shower = 10.0;
+    fpars.fast = 0;
+    fpars.print = 1;
+    fpars.charge_only = 0;
+    fpars.hit_only = 0;
+    nlopt_set_ftol_abs(opt, 1e-5);
+    nlopt_set_xtol_rel(opt, 1e-2);
+    nlopt_set_maxeval(opt, 1000);
+    nlopt_set_maxtime(opt, maxtime);
+
+    memcpy(x,xopt,sizeof(x));
+
+    gettimeofday(&tv_start, NULL);
+    nlopt_optimize(opt,x,&fval);
+    gettimeofday(&tv_stop, NULL);
+
+    time_elapsed = tv_stop.tv_sec - tv_start.tv_sec + (tv_stop.tv_usec - tv_start.tv_usec)/1e6;
+
+    if (time_elapsed > maxtime) goto end;
+
+    if (stop) goto stop;
+
+    do {
+        *fmin = fval;
+        memcpy(xopt,x,sizeof(x));
+
+        nlopt_set_maxtime(opt, maxtime-time_elapsed);
+
+        gettimeofday(&tv_start, NULL);
+        nlopt_optimize(opt,x,&fval);
+        gettimeofday(&tv_stop, NULL);
+
+        time_elapsed += tv_stop.tv_sec - tv_start.tv_sec + (tv_stop.tv_usec - tv_start.tv_usec)/1e6;
+
+        if (stop) goto stop;
+    } while (fval < *fmin && fabs(fval-*fmin) > 1e-2 && time_elapsed < maxtime);
+
+    if (fval < *fmin) {
+        *fmin = fval;
+        memcpy(xopt,x,sizeof(x));
+    }
+
+end:
+
+    nlopt_destroy(opt);
+
+    return 0;
+
+stop:
+    nlopt_destroy(opt);
+
+    return NLOPT_FORCED_STOP;
+}
+
 int fit_event(event *ev, double *xopt, double *fmin, int id, double maxtime)
 {
     size_t i;
@@ -5208,7 +5560,8 @@ int fit_event(event *ev, double *xopt, double *fmin, int id, double maxtime)
     nlopt_set_initial_step(opt, ss);
 
     fpars.ev = ev;
-    fpars.id = id;
+    fpars.id[0] = id;
+    fpars.n = 1;
 
     iter = 0;
 
@@ -5222,6 +5575,7 @@ int fit_event(event *ev, double *xopt, double *fmin, int id, double maxtime)
     fpars.fast = 1;
     fpars.print = 0;
     fpars.charge_only = 0;
+    fpars.hit_only = 1;
     nlopt_set_ftol_abs(opt, 1.0);
     nlopt_set_maxeval(opt, 20);
 
@@ -5262,7 +5616,7 @@ int fit_event(event *ev, double *xopt, double *fmin, int id, double maxtime)
         nlopt_optimize(opt,x,&fval);
         gettimeofday(&tv_stop, NULL);
 
-        if (stop) goto end;
+        if (stop) goto stop;
 
         long long elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
 
@@ -5317,6 +5671,7 @@ int fit_event(event *ev, double *xopt, double *fmin, int id, double maxtime)
     fpars.fast = 0;
     fpars.print = 1;
     fpars.charge_only = 0;
+    fpars.hit_only = 0;
     nlopt_set_ftol_abs(opt, 1e-5);
     nlopt_set_xtol_rel(opt, 1e-2);
     nlopt_set_maxeval(opt, 1000);
@@ -5461,7 +5816,7 @@ int main(int argc, char **argv)
     MCVXBank bmcvx;
     event ev = {0};
     double fmin;
-    double xopt[9];
+    double xopt[MAX_PARS];
     char *filename = NULL;
     char *output = NULL;
     FILE *fout = NULL;