switch to using fit_event2() by default

This commit updates the fit to use the fit_event2() function which can fit for
multi vertex hypotheses. It also uses the QUAD fitter and the Hough transform
of the event to seed the fit so the results for 1 particle fits will be
slightly different than before.

I also fixed a small bug in combinations_with_replacement().

3 files changed, 103 insertions, 30 deletions

diff --git a/src/fit.c b/src/fit.c
index 9bf0dda..cb5a11d 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -5765,6 +5765,7 @@ void usage(void)
     fprintf(stderr,"  -o      output file\n");
     fprintf(stderr,"  --max-time           maximum time in seconds per fit (default: 3600)\n");
     fprintf(stderr,"  --skip-second-event  only fit the first event after a MAST bank\n");
+    fprintf(stderr,"  --max-particles      maximum number of particles to fit for (default: 1)\n");
     fprintf(stderr,"  -h      display this help message\n");
     exit(1);
 }
@@ -5841,9 +5842,53 @@ size_t get_nhit(event *ev)
     return nhit;
 }
 
+void sprintf_particle_string(int *id, size_t n, char *str)
+{
+    /* Convert a list of particle id codes to a string.
+     *
+     * This function is used when writing out the results for a multi particle
+     * fit. Ideally we would use a set or something similar, but YAML doesn't
+     * have support for that and lists can't be used as dictionary keys so
+     * currently we just concatenate the integer representation of the
+     * particles. For example:
+     *
+     * electron -> "20"
+     * muon     -> "22"
+     * electron, muon -> "2022"
+     * etc.
+     *
+     */
+    size_t i;
+
+    sprintf(str,"%i",id[0]);
+
+    for (i = 1; i < n; i++)
+        sprintf(str+strlen(str),"%i",id[i]);
+}
+
+void sprintf_yaml_list(double *a, size_t n, size_t stride, char *str)
+{
+    /* Write out a yaml list of doubles to the string pointed to by `str`.
+     *
+     * The array elements are accessed as a[i*stride] for i = 0,1,...,`n`. */
+    size_t i;
+
+    if (n == 1) {
+        sprintf(str,"%.2f",*a);
+        return;
+    }
+
+    sprintf(str,"[%.2f",*a);
+
+    for (i = 1; i < n; i++)
+        sprintf(str+strlen(str),",%.2g",a[stride*i]);
+
+    sprintf(str+strlen(str),"]");
+}
+
 int main(int argc, char **argv)
 {
-    int i;
+    int i, j, k;
     zebraFile *f;
     bank b;
     int rv;
@@ -5861,6 +5906,11 @@ int main(int argc, char **argv)
     long long elapsed;
     double fmin_best;
     double maxtime = 3600.0;
+    int max_particles = 1;
+    int id[MAX_VERTICES];
+    size_t combos[100];
+    size_t len;
+    char tmp[256];
 
     for (i = 1; i < argc; i++) {
         if (strlen(argv[i]) >= 2 && !strncmp(argv[i], "--", 2)) {
@@ -5870,6 +5920,9 @@ int main(int argc, char **argv)
             } else if (!strcmp(argv[i]+2,"max-time")) {
                 maxtime = strtod(argv[++i],NULL);
                 continue;
+            } else if (!strcmp(argv[i]+2,"max-particles")) {
+                max_particles = atoi(argv[++i]);
+                continue;
             }
         } else if (argv[i][0] == '-') {
             switch (argv[i][1]) {
@@ -5890,6 +5943,11 @@ int main(int argc, char **argv)
     if (!filename)
         usage();
 
+    if (max_particles > MAX_VERTICES) {
+        fprintf(stderr, "max_particles must be less than or equal to %i\n", MAX_VERTICES);
+        exit(1);
+    }
+
     signal(SIGPIPE, SIG_IGN);
     signal(SIGINT, sigint_handler);
 
@@ -6024,33 +6082,48 @@ int main(int argc, char **argv)
                 fprintf(fout, "        fit:\n");
             }
 
-            for (i = 0; i < LEN(particles); i++) {
-                gettimeofday(&tv_start, NULL);
-                if (fit_event(&ev,xopt,&fmin,particles[i],maxtime) == NLOPT_FORCED_STOP) {
-                    printf("ctrl-c caught. quitting...\n");
-                    goto end;
-                }
-                gettimeofday(&tv_stop, NULL);
-
-                elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
-
-                fmin_best = nll_best(&ev);
-
-                if (fout) {
-                    fprintf(fout, "          %i:\n", particles[i]);
-                    fprintf(fout, "            energy: %.2f\n", xopt[0]);
-                    fprintf(fout, "            posx: %.2f\n", xopt[1]);
-                    fprintf(fout, "            posy: %.2f\n", xopt[2]);
-                    fprintf(fout, "            posz: %.2f\n", xopt[3]);
-                    fprintf(fout, "            theta: %.4f\n", xopt[4]);
-                    fprintf(fout, "            phi: %.4f\n", xopt[5]);
-                    fprintf(fout, "            t0: %.2f\n", xopt[6]);
-                    fprintf(fout, "            z1: %.2f\n", xopt[7]);
-                    fprintf(fout, "            z2: %.2f\n", xopt[8]);
-                    fprintf(fout, "            fmin: %.2f\n", fmin);
-                    fprintf(fout, "            time: %lld\n", elapsed);
-                    fprintf(fout, "            psi: %.2f\n", fmin-fmin_best);
-                    fflush(fout);
+            fmin_best = nll_best(&ev);
+
+            /* Loop over 1,2,...,max_particles particle hypotheses. */
+            for (i = 1; i <= max_particles; i++) {
+                /* Find all unique combinations of i particles. */
+                combinations_with_replacement(LEN(particles),i,combos,&len);
+                for (j = 0; j < len; j++) {
+                    /* Set up the id array with the particle id codes. */
+                    for (k = 0; k < i; k++) {
+                        id[k] = particles[combos[j*i+k]];
+                    }
+
+                    /* Fit the event. */
+                    gettimeofday(&tv_start, NULL);
+                    if (fit_event2(&ev,xopt,&fmin,id,i,maxtime) == NLOPT_FORCED_STOP) {
+                        printf("ctrl-c caught. quitting...\n");
+                        goto end;
+                    }
+                    gettimeofday(&tv_stop, NULL);
+
+                    elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
+
+                    if (fout) {
+                        sprintf_particle_string(id,i,tmp);
+                        fprintf(fout, "          %s:\n", tmp);
+                        fprintf(fout, "            posx: %.2f\n", xopt[0]);
+                        fprintf(fout, "            posy: %.2f\n", xopt[1]);
+                        fprintf(fout, "            posz: %.2f\n", xopt[2]);
+                        fprintf(fout, "            t0: %.2f\n", xopt[3]);
+                        sprintf_yaml_list(xopt+4,i,3,tmp);
+                        fprintf(fout, "            energy: %s\n", tmp);
+                        sprintf_yaml_list(xopt+5,i,3,tmp);
+                        fprintf(fout, "            theta: %s\n", tmp);
+                        sprintf_yaml_list(xopt+6,i,3,tmp);
+                        fprintf(fout, "            phi: %s\n", tmp);
+                        //fprintf(fout, "            z1: %.2f\n", xopt[7]);
+                        //fprintf(fout, "            z2: %.2f\n", xopt[8]);
+                        fprintf(fout, "            fmin: %.2f\n", fmin);
+                        fprintf(fout, "            time: %lld\n", elapsed);
+                        fprintf(fout, "            psi: %.2f\n", fmin-fmin_best);
+                        fflush(fout);
+                    }
                 }
             }
 
diff --git a/src/likelihood.h b/src/likelihood.h
index 4ce319a..84fa7ff 100644
--- a/src/likelihood.h
+++ b/src/likelihood.h
@@ -53,7 +53,7 @@
 #define PHOTONS_PER_MEV 400.0
 
 /* Maximum number of vertices to fit. */
-#define MAX_VERTICES 2
+#define MAX_VERTICES 10
 
 typedef struct vertex {
     int id;
diff --git a/src/misc.c b/src/misc.c
index 88e4b77..56760e4 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -726,7 +726,7 @@ void combinations_with_replacement(size_t n, size_t r, size_t *result, size_t *l
     size_t i, j;
     size_t *tmp;
 
-    tmp = malloc(sizeof(size_t)*n*ipow(r,n));
+    tmp = malloc(sizeof(size_t)*n*ipow(n,r));
 
     product(n,r,tmp);