switch from YAML output to HDF5 to speed things up

1 files changed, 110 insertions, 84 deletions

diff --git a/src/fit.c b/src/fit.c
index adc6ebc..0360a8f 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -47,12 +47,21 @@
 #include "electron.h"
 #include "muon.h"
 #include "proton.h"
+#include "hdf5.h"
+#include "hdf5_utils.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 number of events, primary seed particle tracks, and fits.
+ *
+ * FIXME: I should write out the events, tracks, and fits as they come up
+ * instead of storing them until the end. */
+#define MAX_NEVENTS 1000000
+
 char *GitSHA1(void);
 char *GitDirty(void);
 
@@ -5869,12 +5878,13 @@ int main(int argc, char **argv)
 {
     int i, j, k;
     zebraFile *f;
-    zebraBank bmast, bmc, bmcgn, mctk, b;
+    zebraBank bmast, mc, bmcgn, mctk, b;
     int rv;
     EVBank bev;
     FTPVBank bftpv;
     FTXKBank bftxk;
     RSPBank bftxr;
+    MCBank bmc;
     MCTKBank bmctk;
     MCVXBank bmcvx;
     event ev = {0};
@@ -5882,7 +5892,6 @@ int main(int argc, char **argv)
     double xopt[MAX_PARS];
     char *filename = NULL;
     char *output = NULL;
-    FILE *fout = NULL;
     int skip_second_event = 0;
     struct timeval tv_start, tv_stop;
     long long elapsed;
@@ -5892,12 +5901,29 @@ int main(int argc, char **argv)
     int id[MAX_VERTICES];
     size_t combos[100];
     size_t len;
-    char tmp[256];
-    size_t nhit, min_nhit = 100;
+    size_t min_nhit = 100;
     int last_run;
     char dqxx_file[256];
     int32_t gtid = -1;
 
+    int nevents = 0;
+    /* Array of events to write out to HDF5 file.
+     *
+     * Note: Declared static since otherwise it will cause a stack overflow. */
+    static HDF5Event hdf5_events[MAX_NEVENTS];
+
+    int nmcgn = 0;
+    /* Array of primary seed tracks to write out to HDF5 file.
+     *
+     * Note: Declared static since otherwise it will cause a stack overflow. */
+    static HDF5MCGN hdf5_mcgn[MAX_NEVENTS];
+
+    int nfits = 0;
+    /* Array of fit results to write out to HDF5 file.
+     *
+     * Note: Declared static since otherwise it will cause a stack overflow. */
+    static HDF5Fit hdf5_fits[MAX_NEVENTS];
+
     for (i = 1; i < argc; i++) {
         if (strlen(argv[i]) >= 2 && !strncmp(argv[i], "--", 2)) {
             if (!strcmp(argv[i]+2,"skip-second-event")) {
@@ -5953,22 +5979,8 @@ int main(int argc, char **argv)
         return 1;
     }
 
-    if (output) {
-        fout = fopen(output, "w");
-
-        if (!fout) {
-            fprintf(stderr, "failed to open '%s': %s\n", output, strerror(errno));
-            return 1;
-        }
-
-        fprintf(fout, "git_sha1: %s\n", GitSHA1());
-        fprintf(fout, "git_dirty: %s\n", GitDirty());
-        fprintf(fout, "data:\n");
-    }
-
     if (load_pmt_info()) {
         zebra_close(f);
-        if (output) fclose(fout);
         return 1;
     }
 
@@ -6038,30 +6050,29 @@ int main(int argc, char **argv)
             continue;
         }
 
-        if (fout) fprintf(fout, "  -\n");
-
         if (bmast.links[KMAST_MC-1] == 0) goto skip_mc;
 
-        rv = zebra_get_bank(f,&bmc,bmast.links[KMAST_MC-1]);
+        rv = zebra_get_bank(f,&mc,bmast.links[KMAST_MC-1]);
 
         if (rv) {
             fprintf(stderr, "error getting MC bank: %s\n", zebra_err);
             goto err;
         }
 
+        unpack_mc(mc.data, &bmc);
+
         if (bmast.links[KMC_MCGN-1] == 0) {
             fprintf(stderr, "MCGN link is zero!\n");
             goto err;
         }
 
-        rv = zebra_get_bank(f,&bmcgn,bmc.links[KMC_MCGN-1]);
+        rv = zebra_get_bank(f,&bmcgn,mc.links[KMC_MCGN-1]);
 
         if (rv) {
             fprintf(stderr, "error getting MCGN bank: %s\n", zebra_err);
             goto err;
         }
 
-        if (fout) fprintf(fout, "    mcgn:\n");
         while (1) {
             if (bmcgn.links[KMCGN_MCTK-1] == 0) {
                 fprintf(stderr, "MCTK link is zero!\n");
@@ -6110,17 +6121,18 @@ int main(int argc, char **argv)
 
             unpack_mcvx(b.data, &bmcvx);
 
-            if (fout) {
-                fprintf(fout, "      -\n");
-                fprintf(fout, "        id: %" PRIu32 "\n", bmctk.idp);
-                fprintf(fout, "        energy: %.2f\n", bmctk.ene);
-                fprintf(fout, "        posx: %.2f\n", bmcvx.x);
-                fprintf(fout, "        posy: %.2f\n", bmcvx.y);
-                fprintf(fout, "        posz: %.2f\n", bmcvx.z);
-                fprintf(fout, "        dirx: %.4f\n", bmctk.drx);
-                fprintf(fout, "        diry: %.4f\n", bmctk.dry);
-                fprintf(fout, "        dirz: %.4f\n", bmctk.drz);
-                fprintf(fout, "        time: %.4f\n", bmcvx.tim);
+            if (output) {
+                hdf5_mcgn[nmcgn].evn = bmc.evn;
+                hdf5_mcgn[nmcgn].id = bmctk.idp;
+                hdf5_mcgn[nmcgn].energy = bmctk.ene;
+                hdf5_mcgn[nmcgn].x = bmcvx.x;
+                hdf5_mcgn[nmcgn].y = bmcvx.y;
+                hdf5_mcgn[nmcgn].z = bmcvx.z;
+                hdf5_mcgn[nmcgn].dirx = bmctk.drx;
+                hdf5_mcgn[nmcgn].diry = bmctk.dry;
+                hdf5_mcgn[nmcgn].dirz = bmctk.drz;
+                hdf5_mcgn[nmcgn].time = bmcvx.tim;
+                nmcgn++;
             }
 
             if (bmcgn.next) {
@@ -6156,7 +6168,6 @@ skip_mc:
             }
         }
 
-        if (fout) fprintf(fout, "    ev:\n");
         while (1) {
             unpack_ev(b.data, &bev);
             ev.run = bev.run;
@@ -6185,49 +6196,44 @@ skip_mc:
 
             if (get_event(f,&ev,&b)) goto skip_event;
 
-            nhit = get_nhit(&ev);
+            if (output) {
+                hdf5_events[nevents].run = ev.run;
+                hdf5_events[nevents].evn = bmc.evn;
+                hdf5_events[nevents].gtr = ev.trigger_time;
+                hdf5_events[nevents].nhit = ev.nhit;
+                hdf5_events[nevents].gtid = ev.gtid;
+                hdf5_events[nevents].trg_type = ev.trigger_type;
+                hdf5_events[nevents].dc = get_dc_word(&ev, f, &bmast, &b);
 
-            if (fout) {
-                fprintf(fout, "      - run: %i\n", ev.run);
-                fprintf(fout, "        gtr: %.0f\n", ev.trigger_time);
-                fprintf(fout, "        nhit: %zu\n", nhit);
-                fprintf(fout, "        gtid: %i\n", ev.gtid);
-                fprintf(fout, "        trg_type: 0x%08x\n", ev.trigger_type);
-                fprintf(fout, "        dc: 0x%08x\n", get_dc_word(&ev, f, &bmast, &b));
-            }
-
-            if (fout) {
                 if (get_ftpv(f,&b,&bftpv)) {
                     fprintf(stderr, "%s\n", zdab_err);
+                    hdf5_events[nevents].ftp_x = NAN;
+                    hdf5_events[nevents].ftp_y = NAN;
+                    hdf5_events[nevents].ftp_z = NAN;
                 } else {
-                    fprintf(fout, "        ftp:\n");
-                    fprintf(fout, "            x: %.2f\n", bftpv.x);
-                    fprintf(fout, "            y: %.2f\n", bftpv.y);
-                    fprintf(fout, "            z: %.2f\n", bftpv.z);
+                    hdf5_events[nevents].ftp_x = bftpv.x;
+                    hdf5_events[nevents].ftp_y = bftpv.y;
+                    hdf5_events[nevents].ftp_z = bftpv.z;
                 }
-            }
 
-            if (fout) {
                 if (get_ftxk(f,&b,&bftxk)) {
                     fprintf(stderr, "%s\n", zdab_err);
+                    hdf5_events[nevents].ftk_energy = NAN;
                 } else {
-                    fprintf(fout, "        ftk:\n");
-                    fprintf(fout, "            energy: %.2f\n", bftxk.energy);
+                    hdf5_events[nevents].ftk_energy = bftxk.energy;
                 }
-            }
 
-            if (fout) {
                 if (get_rsp(f,&b,&bftxr)) {
                     fprintf(stderr, "%s\n", zdab_err);
+                    hdf5_events[nevents].rsp_energy = NAN;
                 } else {
-                    fprintf(fout, "        rsp:\n");
-                    fprintf(fout, "            energy: %.2f\n", bftxr.ene);
+                    hdf5_events[nevents].rsp_energy = bftxr.ene;
                 }
-            }
 
-            if (nhit < min_nhit) goto skip_event;
+                nevents++;
+            }
 
-            if (fout) fprintf(fout, "        fit:\n");
+            if (ev.nhit < min_nhit) goto skip_event;
 
             fmin_best = nll_best(&ev);
 
@@ -6251,25 +6257,47 @@ skip_mc:
 
                     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);
+                    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 (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;
+                        }
+
+                        hdf5_fits[nfits].fmin = fmin;
+                        hdf5_fits[nfits].time = elapsed;
+                        hdf5_fits[nfits].psi = fmin-fmin_best;
+                        nfits++;
                     }
                 }
             }
@@ -6292,13 +6320,13 @@ skip_event:
 
 end:
 
+    if (output) save_output(output, hdf5_events, nevents, hdf5_mcgn, nmcgn, hdf5_fits, nfits);
+
     free_interpolation();
     pmt_response_free();
 
     db_free(db);
 
-    if (fout) fclose(fout);
-
     zebra_close(f);
 
     return 0;
@@ -6309,8 +6337,6 @@ err:
 
     db_free(db);
 
-    if (fout) fclose(fout);
-
     zebra_close(f);
 
     return 1;