add --plot-likelihood option to fit

2 files changed, 85 insertions, 0 deletions

diff --git a/src/fit.c b/src/fit.c
index 951f614..881f629 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -61,6 +61,8 @@ static nlopt_opt opt;
 #define MCTK_RECORD    0x4d43544b
 #define MCVX_RECORD    0x4d435658
 
+char *flikelihood;
+
 static size_t iter;
 
 typedef struct fitParams {
@@ -5465,6 +5467,29 @@ int fit_event2(event *ev, double *xopt, double *fmin, int *id, size_t n, double
 
     memcpy(x,xopt,sizeof(x));
 
+    if (flikelihood) {
+        FILE *f = fopen(flikelihood, "w");
+
+        for (i = 0; i < 50; i++) {
+            for (j = 0; j < 50; j++) {
+                x[5] = xopt[5] - 0.1 + 0.2*i/49;
+                x[6] = xopt[6] - 0.1 + 0.2*j/49;
+                fval = nlopt_nll2(4+3*n,x,NULL,&fpars);
+                fprintf(f, "%.18g %.18g %.18g\n", x[5], x[6], fval);
+
+                if (stop) goto close;
+            }
+        }
+
+close:
+
+        fclose(f);
+
+        if (stop) exit(1);
+
+        exit(0);
+    }
+
     gettimeofday(&tv_start, NULL);
     nlopt_optimize(opt,x,&fval);
     gettimeofday(&tv_stop, NULL);
@@ -5775,6 +5800,9 @@ void usage(void)
     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,"  --min-nhit           minimum nhit to fit an event (default: 100)\n");
+    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,"  -h      display this help message\n");
     exit(1);
 }
@@ -5935,6 +5963,7 @@ int main(int argc, char **argv)
     size_t nhit, min_nhit = 100;
     int last_run;
     char dqxx_file[256];
+    int32_t gtid = -1;
 
     for (i = 1; i < argc; i++) {
         if (strlen(argv[i]) >= 2 && !strncmp(argv[i], "--", 2)) {
@@ -5950,6 +5979,12 @@ int main(int argc, char **argv)
             } else if (!strcmp(argv[i]+2,"min-nhit")) {
                 min_nhit = atoi(argv[++i]);
                 continue;
+            } else if (!strcmp(argv[i]+2,"plot-likelihood")) {
+                flikelihood = argv[++i];
+                continue;
+            } else if (!strcmp(argv[i]+2,"gtid")) {
+                gtid = atoi(argv[++i]);
+                continue;
             }
         } else if (argv[i][0] == '-') {
             switch (argv[i][1]) {
@@ -6193,6 +6228,8 @@ skip_mc:
             ev.run = bev.run;
             ev.gtid = bev.gtr_id;
 
+            if (gtid > 0 && ev.gtid != gtid) goto skip_event;
+
             if (ev.run != last_run) {
                 printf("loading DQXX file for run %010i\n", ev.run);
 
diff --git a/utils/plot-likelihood b/utils/plot-likelihood
new file mode 100755
index 0000000..2442236
--- /dev/null
+++ b/utils/plot-likelihood
@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the Free
+# Software Foundation, either version 3 of the License, or (at your option)
+# any later version.
+#
+# This program is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
+# more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# this program. If not, see <https://www.gnu.org/licenses/>.
+
+from __future__ import print_function, division
+import numpy as np
+
+# on retina screens, the default plots are way too small
+# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
+# Qt5 will scale everything using the dpi in ~/.Xresources
+import matplotlib
+matplotlib.use("Qt5Agg")
+
+if __name__ == '__main__':
+    import argparse
+    from mpl_toolkits.mplot3d import axes3d
+    import matplotlib.pyplot as plt
+
+    parser = argparse.ArgumentParser("plot likelihood function")
+    parser.add_argument("filenames", nargs='+', help="input files")
+    args = parser.parse_args()
+
+    for filename in args.filenames:
+        print(filename)
+        data = np.genfromtxt(filename)
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+
+        X = data[:,0].reshape((50,50))
+        Y = data[:,1].reshape((50,50))
+        Z = data[:,2].reshape((50,50))
+
+        ax.plot_wireframe(X, Y, Z)
+
+    plt.show()