add ability to specify a particle combo on the command line

This commit updates the fit program to accept a particle combo from the command
line so you can fit for a single particle combination hypothesis. For example
running:

    $ ./fit ~/zdabs/mu_minus_700_1000.hdf5 -p 2020

would just fit for the 2 electron hypothesis.

The reason for adding this ability is that my grid jobs were getting evicted
when fitting muons in run 10,000 since it takes 10s of hours to fit for all the
particle hypothesis. With this change, and a small update to the
submit-grid-jobs script we now submit a single grid job per particle
combination hypothesis which should make each grid job run approximately 4
times faster.

2 files changed, 140 insertions, 64 deletions

diff --git a/src/fit.c b/src/fit.c
index 1018c28..9efee54 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -5821,6 +5821,7 @@ void usage(void)
     fprintf(stderr,"  --plot-likelihood FILENAME\n");
     fprintf(stderr,"                       scan the likelihood space and write out the results to FILENAME\n");
     fprintf(stderr,"  --gtid               only fit a single GTID\n");
+    fprintf(stderr,"  -p                   specify particle combo to fit for\n");
     fprintf(stderr,"  -h      display this help message\n");
     exit(1);
 }
@@ -5874,6 +5875,100 @@ void sprintf_yaml_list(double *a, size_t n, size_t stride, char *str)
     sprintf(str+strlen(str),"]");
 }
 
+/* Convert a particle combo integer to an array of particle ids.
+ *
+ * Example:
+ *
+ *     int id[MAX_VERTICES];
+ *     particle_combo_to_array(2020,id);
+ *
+ * and then id should be equal to {20,20}. */
+size_t particle_combo_to_array(int particle_combo, int *id)
+{
+    size_t i;
+
+    i = 0;
+    while (particle_combo) {
+        id[i++] = particle_combo % 100;
+        particle_combo /= 100;
+    }
+
+    return i;
+}
+
+/* Helper function to perform a fit for a single particle combination hypothesis.
+ *
+ * Arguments:
+ *
+ *     ev         - event
+ *     id         - array of particle id's to fit for (see id_particle.h)
+ *     n          - number of particles (maximum: 3)
+ *     fit_result - pointer to HDF5Fit result object to store fit results in
+ *     maxtime    - maximum time in seconds to perform the minimization (after the
+ *                  "quick" minimization phase)
+ *     fmin_best  - the return value from nll_best()
+ *
+ * Returns the return value from fit_event2(). */
+int do_fit(event *ev, int *id, size_t n, HDF5Fit *fit_result, double maxtime, double fmin_best)
+{
+    int rv;
+    struct timeval tv_start, tv_stop;
+    long long elapsed;
+    double xopt[MAX_PARS];
+    double fmin;
+
+    /* Fit the event. */
+    gettimeofday(&tv_start, NULL);
+    rv = fit_event2(ev,xopt,&fmin,id,n,maxtime);
+    gettimeofday(&tv_stop, NULL);
+
+    if (rv == NLOPT_FORCED_STOP) return rv;
+
+    elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;
+
+    fit_result->run = ev->run;
+    fit_result->gtid = ev->gtid;
+    fit_result->n = n;
+    fit_result->x = xopt[0];
+    fit_result->y = xopt[1];
+    fit_result->z = xopt[2];
+    fit_result->t0 = xopt[3];
+    fit_result->id1 = id[0];
+    fit_result->energy1 = xopt[4];
+    fit_result->theta1 = xopt[5];
+    fit_result->phi1 = xopt[6];
+
+    if (n > 1) {
+        fit_result->id2 = id[1];
+        fit_result->energy2 = xopt[7];
+        fit_result->theta2 = xopt[8];
+        fit_result->phi2 = xopt[9];
+    } else {
+        fit_result->id2 = 0;
+        fit_result->energy2 = NAN;
+        fit_result->theta2 = NAN;
+        fit_result->phi2 = NAN;
+    }
+
+    if (n > 2) {
+        fit_result->id3 = id[2];
+        fit_result->energy3 = xopt[10];
+        fit_result->theta3 = xopt[11];
+        fit_result->phi3 = xopt[12];
+    } else {
+        fit_result->id3 = 0;
+        fit_result->energy3 = NAN;
+        fit_result->theta3 = NAN;
+        fit_result->phi3 = NAN;
+    }
+
+    fit_result->fmin = fmin;
+    fit_result->time = elapsed;
+    fit_result->psi = fmin-fmin_best;
+
+    return rv;
+}
+
 int main(int argc, char **argv)
 {
     int i, j, k;
@@ -5888,13 +5983,9 @@ int main(int argc, char **argv)
     MCTKBank bmctk;
     MCVXBank bmcvx;
     event ev = {0};
-    double fmin;
-    double xopt[MAX_PARS];
     char *filename = NULL;
     char *output = NULL;
     int skip_second_event = 0;
-    struct timeval tv_start, tv_stop;
-    long long elapsed;
     double fmin_best;
     double maxtime = 3600.0;
     int max_particles = 1;
@@ -5905,6 +5996,7 @@ int main(int argc, char **argv)
     int last_run;
     char dqxx_file[256];
     int32_t gtid = -1;
+    int particle_combo = 0;
 
     int nevents = 0;
     /* Array of events to write out to HDF5 file.
@@ -5950,6 +6042,9 @@ int main(int argc, char **argv)
             case 'o':
                 output = argv[++i];
                 break;
+            case 'p':
+                particle_combo = atoi(argv[++i]);
+                break;
             case 'h':
                 usage();
             default:
@@ -6223,70 +6318,37 @@ skip_mc:
 
             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]];
-                    }
+            if (particle_combo) {
+                i = particle_combo_to_array(particle_combo,id);
 
-                    /* 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 (output) {
-                        hdf5_fits[nfits].run = ev.run;
-                        hdf5_fits[nfits].gtid = ev.gtid;
-                        hdf5_fits[nfits].n = i;
-                        hdf5_fits[nfits].x = xopt[0];
-                        hdf5_fits[nfits].y = xopt[1];
-                        hdf5_fits[nfits].z = xopt[2];
-                        hdf5_fits[nfits].t0 = xopt[3];
-                        hdf5_fits[nfits].id1 = id[0];
-                        hdf5_fits[nfits].energy1 = xopt[4];
-                        hdf5_fits[nfits].theta1 = xopt[5];
-                        hdf5_fits[nfits].phi1 = xopt[6];
-
-                        if (i > 1) {
-                            hdf5_fits[nfits].id2 = id[1];
-                            hdf5_fits[nfits].energy2 = xopt[7];
-                            hdf5_fits[nfits].theta2 = xopt[8];
-                            hdf5_fits[nfits].phi2 = xopt[9];
-                        } else {
-                            hdf5_fits[nfits].id2 = 0;
-                            hdf5_fits[nfits].energy2 = NAN;
-                            hdf5_fits[nfits].theta2 = NAN;
-                            hdf5_fits[nfits].phi2 = NAN;
+                if (do_fit(&ev,id,i,hdf5_fits+nfits,maxtime,fmin_best) == NLOPT_FORCED_STOP) {
+                    printf("ctrl-c caught. quitting...\n");
+                    goto end;
+                }
+
+                nfits++;
+
+                if (output) save_output(output, hdf5_events, nevents, hdf5_mcgn, nmcgn, hdf5_fits, nfits);
+            } else {
+                /* 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]];
                         }
 
-                        if (i > 2) {
-                            hdf5_fits[nfits].id3 = id[2];
-                            hdf5_fits[nfits].energy3 = xopt[10];
-                            hdf5_fits[nfits].theta3 = xopt[11];
-                            hdf5_fits[nfits].phi3 = xopt[12];
-                        } else {
-                            hdf5_fits[nfits].id3 = 0;
-                            hdf5_fits[nfits].energy3 = NAN;
-                            hdf5_fits[nfits].theta3 = NAN;
-                            hdf5_fits[nfits].phi3 = NAN;
+                        if (do_fit(&ev,id,i,hdf5_fits+nfits,maxtime,fmin_best) == NLOPT_FORCED_STOP) {
+                            printf("ctrl-c caught. quitting...\n");
+                            goto end;
                         }
 
-                        hdf5_fits[nfits].fmin = fmin;
-                        hdf5_fits[nfits].time = elapsed;
-                        hdf5_fits[nfits].psi = fmin-fmin_best;
                         nfits++;
-                    }
 
-                    if (output) save_output(output, hdf5_events, nevents, hdf5_mcgn, nmcgn, hdf5_fits, nfits);
+                        if (output) save_output(output, hdf5_events, nevents, hdf5_mcgn, nmcgn, hdf5_fits, nfits);
+                    }
                 }
             }
 
diff --git a/utils/submit-grid-jobs b/utils/submit-grid-jobs
index 63cbcb7..a55a723 100755
--- a/utils/submit-grid-jobs
+++ b/utils/submit-grid-jobs
@@ -140,13 +140,16 @@ ID = uuid.uuid1()
 class MyTemplate(string.Template):
     delimiter = '@'
 
-def submit_job(filename, run, gtid, dir, dqxx_dir, min_nhit, max_particles):
+def submit_job(filename, run, gtid, dir, dqxx_dir, min_nhit, max_particles, particle_combo=None):
     print("submitting job for %s gtid %i" % (filename, gtid))
     head, tail = split(filename)
     root, ext = splitext(tail)
 
     # all output files are prefixed with FILENAME_GTID_UUID
-    prefix = "%s_%08i_%s" % (root,gtid,ID.hex)
+    if particle_combo:
+        prefix = "%s_%08i_%i_%s" % (root,gtid,particle_combo,ID.hex)
+    else:
+        prefix = "%s_%08i_%s" % (root,gtid,ID.hex)
 
     # fit output filename
     output = "%s.hdf5" % prefix
@@ -166,6 +169,8 @@ def submit_job(filename, run, gtid, dir, dqxx_dir, min_nhit, max_particles):
         sys.exit(1)
 
     args = [tail,"-o",output,"--gtid",gtid,"--min-nhit",min_nhit,"--max-particles",max_particles]
+    if particle_combo:
+        args += ["-p",particle_combo]
     transfer_input_files = ",".join([executable,filename,join(dqxx_dir,"DQXX_%010i.dat" % run)] + [join(dir,filename) for filename in INPUT_FILES])
     transfer_output_files = ",".join([output])
 
@@ -191,12 +196,19 @@ def submit_job(filename, run, gtid, dir, dqxx_dir, min_nhit, max_particles):
     # submit the job
     check_call(["condor_submit",condor_submit])
 
+def array_to_particle_combo(combo):
+    particle_combo = 0
+    for i, id in enumerate(combo[::-1]):
+        particle_combo += id*100**i
+    return particle_combo
+
 if __name__ == '__main__':
     import argparse
     from subprocess import check_call
     import os
     import tempfile
     import h5py
+    from itertools import combinations_with_replacement
 
     parser = argparse.ArgumentParser("submit grid jobs")
     parser.add_argument("filenames", nargs='+', help="input files")
@@ -275,7 +287,9 @@ if __name__ == '__main__':
         with h5py.File(output.name) as f:
             for ev in f['ev']:
                 if ev['nhit'] >= args.min_nhit:
-                    submit_job(filename, ev['run'], ev['gtid'], dir, dqxx_dir, args.min_nhit, args.max_particles)
+                    for i in range(1,args.max_particles+1):
+                        for particle_combo in map(array_to_particle_combo,combinations_with_replacement([20,22],i)):
+                            submit_job(filename, ev['run'], ev['gtid'], dir, dqxx_dir, args.min_nhit, args.max_particles, particle_combo)
 
         # Delete temporary HDF5 file
         os.unlink(output.name)