update zebra library to be able to use links

This commit updates the zebra library files zebra.{c,h} so that it's now
possible to traverse the data structure using links! This was originally
motivated by wanting to figure out which MC particles were generated from the
MCGN bank (from which it's only possible to access the tracks and vertices
using structural links).

I've also added a new test to test-zebra which checks the consistency of all of
the next/up/orig, structural, and reference links in a zebra file.

9 files changed, 835 insertions, 221 deletions

diff --git a/src/fit.c b/src/fit.c
index 79d8c8e..64a2a5a 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -5775,52 +5775,51 @@ void sigint_handler(int dummy)
     stop = 1;
 }
 
-int get_event(zebraFile *f, event *ev, bank *b)
+int get_event(zebraFile *f, event *ev, zebraBank *bev)
 {
     /* Read all the PMT banks from the zebra file and update `ev`.
      *
-     * Returns the last return value from get_bank(). */
+     * Returns 0 on success, -1 on error. */
     int i, rv;
     PMTBank bpmt;
+    zebraBank b;
     int id, crate, card, channel;
 
     for (i = 0; i < MAX_PMTS; i++) {
         ev->pmt_hits[i].hit = 0;
     }
 
-    while (1) {
-        rv = next_bank(f, b);
+    rv = zebra_get_bank(f,&b,bev->links[KEV_PMT-1]);
 
-        if (rv == -1) {
-            break;
-        } else if (rv == 1) {
-            /* EOF */
-            break;
-        }
+    if (rv) {
+        fprintf(stderr, "error getting PMT bank: %s\n", zebra_err);
+        return -1;
+    }
 
-        switch (b->name) {
-        case PMT_RECORD:
-            unpack_pmt(b->data, &bpmt);
-            card = bpmt.pin/1024;
-            crate = (bpmt.pin % 1024)/32;
-            channel = bpmt.pin % 32;
-            id = crate*512 + card*32 + channel;
-            ev->pmt_hits[id].hit = 1;
-            ev->pmt_hits[id].t = bpmt.pt;
-            ev->pmt_hits[id].qhl = bpmt.phl;
-            ev->pmt_hits[id].qhs = bpmt.phs;
-            ev->pmt_hits[id].qlx = bpmt.plx;
-            ev->pmt_hits[id].flags |= bpmt.pf & KPF_DIS;
-            break;
-        case EV_RECORD:
-        case MAST_RECORD:
-            goto end;
+    while (1) {
+        unpack_pmt(b.data, &bpmt);
+        card = bpmt.pin/1024;
+        crate = (bpmt.pin % 1024)/32;
+        channel = bpmt.pin % 32;
+        id = crate*512 + card*32 + channel;
+        ev->pmt_hits[id].hit = 1;
+        ev->pmt_hits[id].t = bpmt.pt;
+        ev->pmt_hits[id].qhl = bpmt.phl;
+        ev->pmt_hits[id].qhs = bpmt.phs;
+        ev->pmt_hits[id].qlx = bpmt.plx;
+        ev->pmt_hits[id].flags |= bpmt.pf & KPF_DIS;
+
+        if (!b.next) break;
+
+        rv = zebra_get_bank(f,&b,b.next);
+
+        if (rv) {
+            fprintf(stderr, "error getting PMT bank: %s\n", zebra_err);
+            return -1;
         }
     }
 
-end:
-
-    return rv;
+    return 0;
 }
 
 size_t get_nhit(event *ev)
@@ -5890,7 +5889,7 @@ int main(int argc, char **argv)
 {
     int i, j, k;
     zebraFile *f;
-    bank b;
+    zebraBank bmast, bmc, bmcgn, mctk, b;
     int rv;
     EVBank bev;
     MCTKBank bmctk;
@@ -6014,69 +6013,139 @@ int main(int argc, char **argv)
         exit(1);
     }
 
-    int first_ev = 1;
-    int first_mctk = 1;
-    int first_mcvx = 1;
-    int skip = 0;
     while (1) {
-        if (!skip)
-            rv = next_bank(f, &b);
-
-        skip = 0;
+        rv = zebra_read_next_logical_record(f);
 
         if (rv == -1) {
-            fprintf(stderr, "error getting bank: %s\n", zebra_err);
+            fprintf(stderr, "error getting logical record: %s\n", zebra_err);
             goto err;
         } else if (rv == 1) {
             /* EOF */
             break;
         }
 
-        switch (b.name) {
-        case MAST_RECORD:
-            /* New event. */
-            if (fout) fprintf(fout, "  -\n");
-            first_ev = 1;
-            first_mctk = 1;
-            first_mcvx = 1;
-            break;
-        case MCVX_RECORD:
-            /* New MC vertex. */
-            unpack_mcvx(b.data, &bmcvx);
-            if (fout && b.status & KMCVX_CRE) {
-                if (first_mcvx) fprintf(fout, "    mcvx:\n");
-                fprintf(fout, "      -\n");
-                fprintf(fout, "        posx: %.2f\n", bmcvx.x);
-                fprintf(fout, "        posy: %.2f\n", bmcvx.y);
-                fprintf(fout, "        posz: %.2f\n", bmcvx.z);
+        rv = zebra_get_bank(f, &bmast, f->first_bank);
+
+        if (rv) {
+            fprintf(stderr, "error getting MAST bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        if (fout) fprintf(fout, "  -\n");
+
+        if (bmast.links[KMAST_MC-1] == 0) {
+            fprintf(stderr, "MAST link is zero!\n");
+            goto err;
+        }
+
+        rv = zebra_get_bank(f,&bmc,bmast.links[KMAST_MC-1]);
+
+        if (rv) {
+            fprintf(stderr, "error getting MC bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        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]);
+
+        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");
+                goto err;
             }
-            first_mcvx = 0;
-            break;
-        case MCTK_RECORD:
-            /* New MC track. */
+
+            rv = zebra_get_bank(f,&mctk,bmcgn.links[KMCGN_MCTK-1]);
+
+            if (rv) {
+                fprintf(stderr, "error getting MCTK bank: %s\n", zebra_err);
+                goto err;
+            }
+
+            rv = zebra_get_bank(f,&b,mctk.orig);
+
+            if (rv) {
+                fprintf(stderr, "error getting MCTK bank: %s\n", zebra_err);
+                goto err;
+            }
+
             unpack_mctk(b.data, &bmctk);
-            if (fout && bmctk.idp != 0) {
-                if (first_mctk) fprintf(fout, "    mctk:\n");
+
+            if (mctk.up == 0) {
+                fprintf(stderr, "MCVX link is zero!\n");
+                goto err;
+            }
+
+            rv = zebra_get_bank(f,&b,mctk.up);
+
+            if (rv) {
+                fprintf(stderr, "error getting MCVX bank: %s\n", zebra_err);
+                goto err;
+            }
+
+            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);
             }
-            first_mctk = 0;
-            break;
-        case EV_RECORD:
+
+            if (bmcgn.next) {
+                rv = zebra_get_bank(f,&bmcgn,bmcgn.next);
+
+                if (rv) {
+                    fprintf(stderr, "error getting MCGN bank: %s\n", zebra_err);
+                    goto err;
+                }
+            } else {
+                break;
+            }
+        }
+
+        rv = zebra_get_bank(f,&b,bmast.links[KMAST_EV-1]);
+
+        if (rv) {
+            fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        /* Skip to the last event so we can traverse them in reverse order. The
+         * reason for this is that for some reason SNOMAN puts the events in
+         * reverse order within each logical record. */
+        while (b.next) {
+            rv = zebra_get_bank(f,&b,b.next);
+
+            if (rv) {
+                fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
+                goto err;
+            }
+        }
+
+        if (fout) fprintf(fout, "    ev:\n");
+        while (1) {
             unpack_ev(b.data, &bev);
             ev.run = bev.run;
             ev.gtid = bev.gtr_id;
 
-            if (!first_ev && skip_second_event) continue;
-
             rv = get_event(f,&ev,&b);
 
             if (fout) {
-                if (first_ev) fprintf(fout, "    ev:\n");
                 fprintf(fout, "      - gtid: %i\n", ev.gtid);
                 fprintf(fout, "        nhit: %zu\n", get_nhit(&ev));
                 fprintf(fout, "        fit:\n");
@@ -6127,11 +6196,17 @@ int main(int argc, char **argv)
                 }
             }
 
-            first_ev = 0;
+            /* Note the origin link for the first EV bank points back to the
+             * structural link location in the MAST bank. These links are super
+             * confusing! */
+            if ((b.orig == f->first_bank-KMAST_EV) || skip_second_event) break;
 
-            /* Skip reading in the next bank since get_event() already read in
-             * the next bank. */
-            skip = 1;
+            rv = zebra_get_bank(f,&b,b.orig);
+
+            if (rv) {
+                fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
+                goto err;
+            }
         }
     }
 
diff --git a/src/test-find-peaks.c b/src/test-find-peaks.c
index c82f9ec..ccd581c 100644
--- a/src/test-find-peaks.c
+++ b/src/test-find-peaks.c
@@ -95,52 +95,51 @@ void plot_find_peaks(event *ev, double *peak_theta, double *peak_phi, size_t n)
     }
 }
 
-int get_event(zebraFile *f, event *ev, bank *b)
+int get_event(zebraFile *f, event *ev, zebraBank *bev)
 {
     /* Read all the PMT banks from the zebra file and update `ev`.
      *
-     * Returns the last return value from get_bank(). */
+     * Returns 0 on success, -1 on error. */
     int i, rv;
     PMTBank bpmt;
+    zebraBank b;
     int id, crate, card, channel;
 
     for (i = 0; i < MAX_PMTS; i++) {
         ev->pmt_hits[i].hit = 0;
     }
 
-    while (1) {
-        rv = next_bank(f, b);
+    rv = zebra_get_bank(f,&b,bev->links[KEV_PMT-1]);
 
-        if (rv == -1) {
-            break;
-        } else if (rv == 1) {
-            /* EOF */
-            break;
-        }
+    if (rv) {
+        fprintf(stderr, "error getting PMT bank: %s\n", zebra_err);
+        return -1;
+    }
 
-        switch (b->name) {
-        case PMT_RECORD:
-            unpack_pmt(b->data, &bpmt);
-            card = bpmt.pin/1024;
-            crate = (bpmt.pin % 1024)/32;
-            channel = bpmt.pin % 32;
-            id = crate*512 + card*32 + channel;
-            ev->pmt_hits[id].hit = 1;
-            ev->pmt_hits[id].t = bpmt.pt;
-            ev->pmt_hits[id].qhl = bpmt.phl;
-            ev->pmt_hits[id].qhs = bpmt.phs;
-            ev->pmt_hits[id].qlx = bpmt.plx;
-            ev->pmt_hits[id].flags |= bpmt.pf & KPF_DIS;
-            break;
-        case EV_RECORD:
-        case MAST_RECORD:
-            goto end;
+    while (1) {
+        unpack_pmt(b.data, &bpmt);
+        card = bpmt.pin/1024;
+        crate = (bpmt.pin % 1024)/32;
+        channel = bpmt.pin % 32;
+        id = crate*512 + card*32 + channel;
+        ev->pmt_hits[id].hit = 1;
+        ev->pmt_hits[id].t = bpmt.pt;
+        ev->pmt_hits[id].qhl = bpmt.phl;
+        ev->pmt_hits[id].qhs = bpmt.phs;
+        ev->pmt_hits[id].qlx = bpmt.plx;
+        ev->pmt_hits[id].flags |= bpmt.pf & KPF_DIS;
+
+        if (!b.next) break;
+
+        rv = zebra_get_bank(f,&b,b.next);
+
+        if (rv) {
+            fprintf(stderr, "error getting PMT bank: %s\n", zebra_err);
+            return -1;
         }
     }
 
-end:
-
-    return rv;
+    return 0;
 }
 
 void usage(void)
@@ -156,7 +155,7 @@ int main(int argc, char **argv)
     int i;
     char *filename = NULL;
     zebraFile *f;
-    bank b;
+    zebraBank bmast, bmc, bmcgn, b;
     EVBank bev;
     event ev = {0};
     int rv;
@@ -198,70 +197,117 @@ int main(int argc, char **argv)
         return 1;
     }
 
-    int skip = 0;
     while (1) {
-        if (!skip)
-            rv = next_bank(f, &b);
-
-        skip = 0;
+        rv = zebra_read_next_logical_record(f);
 
         if (rv == -1) {
-            fprintf(stderr, "error getting bank: %s\n", zebra_err);
+            fprintf(stderr, "error getting MAST bank: %s\n", zebra_err);
             goto err;
         } else if (rv == 1) {
             /* EOF */
             break;
         }
 
-        switch (b.name) {
-        case MCVX_RECORD:
-            /* New MC vertex. */
-            if (b.status & KMCVX_CRE) {
-                unpack_mcvx(b.data, &bmcvx);
-            }
-            break;
-        case EV_RECORD:
-            unpack_ev(b.data, &bev);
-            ev.run = bev.run;
-            ev.gtid = bev.gtr_id;
+        rv = zebra_get_bank(f, &bmast, f->first_bank);
 
-            rv = get_event(f,&ev,&b);
+        if (rv) {
+            fprintf(stderr, "error getting MAST bank: %s\n", zebra_err);
+            goto err;
+        }
 
-            pos[0] = bmcvx.x;
-            pos[1] = bmcvx.y;
-            pos[2] = bmcvx.z;
+        if (bmast.links[KMAST_MC-1] == 0) {
+            fprintf(stderr, "MAST link is zero!\n");
+            goto err;
+        }
+
+        rv = zebra_get_bank(f,&bmc,bmast.links[KMAST_MC-1]);
 
-            size_t n = 100;
-            size_t m = 100;
+        if (rv) {
+            fprintf(stderr, "error getting MC bank: %s\n", zebra_err);
+            goto err;
+        }
 
-            double *x = calloc(n,sizeof(double));
-            double *y = calloc(m,sizeof(double));
-            double *result = calloc(n*m,sizeof(double));
+        if (bmast.links[KMC_MCGN-1] == 0) {
+            fprintf(stderr, "MCGN link is zero!\n");
+            goto err;
+        }
 
-            for (i = 0; i < n; i++) {
-                x[i] = i*M_PI/(n-1);
-            }
+        rv = zebra_get_bank(f,&bmcgn,bmc.links[KMC_MCGN-1]);
 
-            for (i = 0; i < m; i++) {
-                y[i] = i*2*M_PI/(m-1);
-            }
+        if (rv) {
+            fprintf(stderr, "error getting MCGN bank: %s\n", zebra_err);
+            goto err;
+        }
 
-            get_hough_transform(&ev,pos,x,y,n,m,result,0,0);
+        if (bmcgn.links[KMCGN_MCTK-1] == 0) {
+            fprintf(stderr, "MCTK link is zero!\n");
+            goto err;
+        }
 
-            find_peaks(&ev,pos,n,m,peak_theta,peak_phi,&npeaks,LEN(peak_theta));
+        rv = zebra_get_bank(f,&b,bmcgn.links[KMCGN_MCTK-1]);
 
-            printf("gtid %i\n", ev.gtid);
-            for (i = 0; i < npeaks; i++) {
-                printf("%.2f %.2f\n", peak_theta[i], peak_phi[i]);
-            }
+        if (rv) {
+            fprintf(stderr, "error getting MCTK bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        if (b.up == 0) {
+            fprintf(stderr, "MCVX link is zero!\n");
+            goto err;
+        }
+
+        rv = zebra_get_bank(f,&b,b.up);
+
+        if (rv) {
+            fprintf(stderr, "error getting MCVX bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        unpack_mcvx(b.data, &bmcvx);
+
+        rv = zebra_get_bank(f,&b,bmast.links[KMAST_EV-1]);
 
-            if (plot)
-                plot_find_peaks(&ev,peak_theta,peak_phi,npeaks);
+        if (rv) {
+            fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        unpack_ev(b.data, &bev);
+        ev.run = bev.run;
+        ev.gtid = bev.gtr_id;
+
+        rv = get_event(f,&ev,&b);
+
+        pos[0] = bmcvx.x;
+        pos[1] = bmcvx.y;
+        pos[2] = bmcvx.z;
+
+        size_t n = 100;
+        size_t m = 100;
+
+        double *x = calloc(n,sizeof(double));
+        double *y = calloc(m,sizeof(double));
+        double *result = calloc(n*m,sizeof(double));
 
-            /* Skip reading in the next bank since get_event() already read in
-             * the next bank. */
-            skip = 1;
+        for (i = 0; i < n; i++) {
+            x[i] = i*M_PI/(n-1);
         }
+
+        for (i = 0; i < m; i++) {
+            y[i] = i*2*M_PI/(m-1);
+        }
+
+        get_hough_transform(&ev,pos,x,y,n,m,result,0,0);
+
+        find_peaks(&ev,pos,n,m,peak_theta,peak_phi,&npeaks,LEN(peak_theta));
+
+        printf("gtid %i\n", ev.gtid);
+        for (i = 0; i < npeaks; i++) {
+            printf("%.2f %.2f\n", peak_theta[i], peak_phi[i]);
+        }
+
+        if (plot)
+            plot_find_peaks(&ev,peak_theta,peak_phi,npeaks);
     }
 
     zebra_close(f);
diff --git a/src/test-zebra.c b/src/test-zebra.c
index ffd0c50..86aff3c 100644
--- a/src/test-zebra.c
+++ b/src/test-zebra.c
@@ -3,12 +3,15 @@
 #include "Record_Info.h"
 #include "misc.h"
 
+char *filename;
+
 typedef int testFunction(char *err);
 
 int test_zdab(char *err)
 {
+    /* Simple test to make sure we can read the 5 events from Muons.zdab. */
     int rv;
-    bank b;
+    zebraBank b;
     int nzdabs = 0;
 
     zebraFile *z = zebra_open("Muons.zdab");
@@ -19,16 +22,23 @@ int test_zdab(char *err)
     }
 
     while (1) {
-        rv = next_bank(z, &b);
+        rv = zebra_read_next_logical_record(z);
 
         if (rv == -1) {
-            sprintf(err, "error getting bank: %s\n", zebra_err);
+            sprintf(err, "error getting logical record: %s", zebra_err);
             goto err;
         } else if (rv == 1) {
             break;
         }
 
-        if (b.name == ZDAB_RECORD) nzdabs += 1;
+        rv = zebra_get_bank(z, &b, z->first_bank);
+
+        if (rv) {
+            sprintf(err, "error getting first bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        if (b.idh == ZDAB_RECORD) nzdabs += 1;
     }
 
     if (nzdabs != 5) {
@@ -46,11 +56,110 @@ err:
     return -1;
 }
 
+int test_bank_links(zebraFile *z, uint32_t link, char *err)
+{
+    int i, rv;
+    zebraBank b1, b2;
+
+    rv = zebra_get_bank(z, &b1, link);
+
+    if (rv) {
+        sprintf(err, "error getting bank: %s\n", zebra_err);
+        goto err;
+    }
+
+    /* Test that all the structural links point back to the bank. */
+    for (i = 0; i < b1.num_structural_links; i++) {
+        if (b1.links[i] == 0) continue;
+
+        rv = zebra_get_bank(z, &b2, b1.links[i]);
+
+        if (rv) {
+            sprintf(err, "error getting bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        /* Check that the up pointer points back to the original bank. */
+        if (b2.up != link) {
+            sprintf(err, "bank up pointer from bank %.4s does not point back to %.4s bank!",
+                    b2.name, b1.name);
+            goto err;
+        }
+
+        rv = test_bank_links(z,b1.links[i],err);
+
+        if (rv) goto err;
+    }
+
+    if (b1.next) {
+        /* Test that the next/orig pointers are consistent. */
+        rv = zebra_get_bank(z, &b2, b1.next);
+
+        if (rv) {
+            sprintf(err, "error getting bank: %s\n", zebra_err);
+            goto err;
+        }
+
+        /* Check that the orig pointer points back to the original bank. */
+        if (b2.orig != link) {
+            sprintf(err, "bank orig pointer from bank %.4s does not point back to %.4s bank!",
+                    b2.name, b1.name);
+            goto err;
+        }
+
+        rv = test_bank_links(z,b1.next,err);
+
+        if (rv) goto err;
+    }
+
+    return 0;
+
+err:
+    return -1;
+}
+
+int test_links(char *err)
+{
+    /* Tests that all the next/up/orig bank pointers in the zebra file are
+     * consistent. */
+    int rv;
+
+    zebraFile *z = zebra_open(filename);
+
+    if (!z) {
+        sprintf(err, "%s", zebra_err);
+        return 1;
+    }
+
+    while (1) {
+        rv = zebra_read_next_logical_record(z);
+
+        if (rv == -1) {
+            sprintf(err, "error getting logical record: %s", zebra_err);
+            goto err;
+        } else if (rv == 1) {
+            break;
+        }
+
+        if (test_bank_links(z,z->first_bank,err)) goto err;
+    }
+
+    zebra_close(z);
+
+    return 0;
+
+err:
+    zebra_close(z);
+
+    return -1;
+}
+
 struct tests {
     testFunction *test;
     char *name;
 } tests[] = {
-    {test_zdab, "test_zdab"}
+    {test_zdab, "test_zdab"},
+    {test_links, "test_links"}
 };
 
 int main(int argc, char **argv)
@@ -60,6 +169,13 @@ int main(int argc, char **argv)
     int retval = 0;
     struct tests test;
 
+    if (argc < 2) {
+        fprintf(stderr, "usage: test-zebra FILENAME\n");
+        return 1;
+    }
+
+    filename = argv[1];
+
     for (i = 0; i < LEN(tests); i++) {
         test = tests[i];
 
diff --git a/src/zdab_utils.c b/src/zdab_utils.c
index 1cde62a..1f2ed03 100644
--- a/src/zdab_utils.c
+++ b/src/zdab_utils.c
@@ -49,6 +49,30 @@ void swap_int16(int16_t *val_pt, int count)
     return;
 }
 
+void unpack_mcgn(uint32_t *data, MCGNBank *b)
+{
+    unpack((uint8_t *) data,      "l",&b->id);
+    unpack((uint8_t *) (data+1),  "l",&b->num);
+    unpack((uint8_t *) (data+2),  "l",&b->radcor_proc);
+    unpack((uint8_t *) (data+3),  "l",&b->radcor_made_gamma);
+    unpack((uint8_t *) (data+4),  "l",&b->spare5);
+    unpack((uint8_t *) (data+5),  "l",&b->spare6);
+    unpack((uint8_t *) (data+7),  "l",&b->spare7);
+    unpack((uint8_t *) (data+8),  "l",&b->spare8);
+    unpack((uint8_t *) (data+9),  "l",&b->spare9);
+    unpack((uint8_t *) (data+10), "l",&b->spare10);
+    unpack((uint8_t *) (data+12), "f",&b->radcor_xtot);
+    unpack((uint8_t *) (data+13), "f",&b->radcor_xdif);
+    unpack((uint8_t *) (data+14), "f",&b->gentim_xpar);
+    unpack((uint8_t *) (data+15), "f",&b->spare14);
+    unpack((uint8_t *) (data+16), "f",&b->spare15);
+    unpack((uint8_t *) (data+17), "f",&b->spare16);
+    unpack((uint8_t *) (data+18), "f",&b->spare17);
+    unpack((uint8_t *) (data+19), "f",&b->spare18);
+    unpack((uint8_t *) (data+20), "f",&b->spare19);
+    unpack((uint8_t *) (data+21), "f",&b->spare20);
+}
+
 void unpack_mcvx(uint32_t *data, MCVXBank *b)
 {
     unpack((uint8_t *) data,      "l",&b->cls);
diff --git a/src/zdab_utils.h b/src/zdab_utils.h
index f70482a..8d8badc 100644
--- a/src/zdab_utils.h
+++ b/src/zdab_utils.h
@@ -18,6 +18,74 @@
 #define SWAP_INT16(a,b)
 #endif
 
+/* Links for the MCVX bank. */
+#define KMCVX_MCVXG 4
+#define KMCVX_MCPM  3
+#define KMCVX_MCTKI 2
+#define KMCVX_MCTK  1
+
+/* Links for the MC bank. */
+#define KMC_MCNS 6
+#define KMC_MCNH 5
+#define KMC_MCPM 4
+#define KMC_EGS4 3
+#define KMC_MCVX 2
+#define KMC_MCGN 1
+
+/* Links for the EV bank. */
+#define KEV_NPFA     32
+#define KEV_NECL     31
+#define KEV_NPA      30
+#define KEV_NEMG_PAR 29
+#define KEV_NEMG     28
+#define KEV_NES      27
+#define KEV_NESG     26
+#define KEV_RSP      25
+#define KEV_NECK     24
+#define KEV_BUTT     23
+#define KEV_ANAL     22
+#define KEV_FECD     21
+#define KEV_LG 	     20
+#define KEV_OWL      19
+#define KEV_PMT      18
+#define KEV_CL 	     17
+#define KEV_PIF      16
+#define KEV_PIT0     15
+#define KEV_PIT      14
+#define KEV_PIHS     13
+#define KEV_PIHL     12
+#define KEV_PILX     11
+#define KEV_PIN      10
+#define KEV_PT0      9
+#define KEV_PLX      8
+#define KEV_PHS      7
+#define KEV_ZDAB     6
+#define KEV_PF       5
+#define KEV_PN       4
+#define KEV_PHL      3
+#define KEV_PT       2
+#define KEV_FT       1
+
+/* Links for the MAST bank. */
+#define KMAST_RLOG 25 /* Runlog data. */
+#define KMAST_RLAI 24 /* Runlog analysis data. */
+#define KMAST_LABL 12 /* Tape label information. */
+#define KMAST_CAAC 11 /* AV Status. */
+#define KMAST_CAST 10 /* Manipulator status. */
+#define KMAST_SOSF 9  /* LED status. */
+#define KMAST_SOSX 8  /* Manipulator source status. */
+#define KMAST_TRIG 7  /* Trigger information bank. */
+#define KMAST_EPED 6  /* Electronics calibration pedestal bank. */
+#define KMAST_RHDR 5  /* Top level run header bank. */
+#define KMAST_ANAL 4  /* Analysis top level bank. */
+#define KMAST_DSML 3  /* Data structure manager log bank. */
+#define KMAST_MC   2  /* Monte carlo top level bank. */
+#define KMAST_EV   1  /* First event trigger bank. */
+
+/* MAST bank data words. */
+#define KMAST_SNOMAN_VER 1
+#define KMAST_SNOMAN_VER_ORG 2
+
 /* Status bits for the MC Vertex bank.
  *
  * See http://www.lip.pt/~barros/SNO/doc/snoman/companion_frames.html. */
@@ -63,6 +131,46 @@
 #define KPF_NO_CAL   0x08000000
 #define KPF_BAD_CAL  0x10000000
 
+/* Links for the MCGN bank. */
+#define KMCGN_MCTK 1
+
+/* Links for the MCTK bank. */
+#define KMCTK_MCVX 1
+
+typedef struct MCGNBank {
+    /* Particle id. */
+    uint32_t id;
+    /* Number of outgoing particles generated. */
+    uint32_t num;
+    /* Radiative process if any. */
+    uint32_t radcor_proc;
+    /* +1 particle brems and produces a gamma.
+     * -1 particle is a gamma produced by brem. */
+    uint32_t radcor_made_gamma;
+    /* Spares. */
+    uint32_t spare5;
+    uint32_t spare6;
+    uint32_t spare7;
+    uint32_t spare8;
+    uint32_t spare9;
+    uint32_t spare10;
+    /* Correction to total cross sections. */
+    float radcor_xtot;
+    /* Correction to differential cross sections. */
+    float radcor_xdif;
+    /* Value of exponential parameter used for generating particle start times,
+     * if time type IDTEX1 was selected. */
+    float gentim_xpar;
+    /* More spares. */
+    float spare14;
+    float spare15;
+    float spare16;
+    float spare17;
+    float spare18;
+    float spare19;
+    float spare20;
+} MCGNBank;
+
 typedef struct MCVXBank {
     /* Class:  = 1 Source, = 2 Boundary, = 3 Interaction, = 4 Sink
 	       = 5 Pre-source. */
@@ -282,6 +390,7 @@ typedef struct PMTBank {
     uint32_t qrc;
 } PMTBank;
 
+void unpack_mcgn(uint32_t *data, MCGNBank *b);
 void unpack_mcvx(uint32_t *data, MCVXBank *b);
 void unpack_mctk(uint32_t *data, MCTKBank *b);
 void unpack_ev(uint32_t *data, EVBank *b);
diff --git a/src/zebra.c b/src/zebra.c
index 4ce51a2..24973d7 100644
--- a/src/zebra.c
+++ b/src/zebra.c
@@ -10,6 +10,10 @@ char zebra_err[256];
 
 zebraFile *zebra_open(const char *filename)
 {
+    /* Returns a pointer to a new zebra file object.
+     *
+     * Returns the pointer to the zebraFile object on success, or NULL on
+     * error. */
     FILE *f = fopen(filename, "r");
 
     if (!f) {
@@ -18,16 +22,23 @@ zebraFile *zebra_open(const char *filename)
     }
 
     zebraFile *z = (zebraFile *) malloc(sizeof(zebraFile));
+
+    if (!z) {
+        sprintf(zebra_err, "failed to allocate space for zebraFile struct: %s", strerror(errno));
+        fclose(f);
+        return NULL;
+    }
+
     z->f = f;
-    z->offset = 0;
     z->lr_size = 0;
     z->buf = NULL;
     z->buf_size = 0;
+    z->tab = NULL;
 
     return z;
 }
 
-int read_next_physical_record(zebraFile *z)
+static int read_next_physical_record(zebraFile *z)
 {
     /* Read the next physical record from the zebra file and append it to the
      * buffer.
@@ -94,7 +105,7 @@ int read_next_physical_record(zebraFile *z)
     return 0;
 }
 
-int get_bytes(zebraFile *z, size_t size)
+static int get_bytes(zebraFile *z, size_t size)
 {
     /* Read bytes from the zdab file until there are at least `size` words in
      * the buffer.
@@ -102,7 +113,7 @@ int get_bytes(zebraFile *z, size_t size)
      * Returns 0 on success, -1 on error, and 1 on EOF. */
     int rv;
 
-    while ((z->buf_size - z->offset) < size*4) {
+    while (z->buf_size < size*4) {
         rv = read_next_physical_record(z);
 
         if (rv == -1) {
@@ -116,65 +127,189 @@ int get_bytes(zebraFile *z, size_t size)
     return 0;
 }
 
-int read_next_logical_record(zebraFile *z)
+static uint32_t get_link(zebraFile *z, uint32_t link)
+{
+    /* Returns the offset for a zebra bank in the local buffer from a link
+     * pointing to the original bank. */
+    int i;
+    uint32_t sum;
+
+    if (z->nwtab == 0) return link;
+
+    sum = 0;
+    for (i = 0; i < z->nwtab/2; i++) {
+        if (link <= z->tab[i*2+1]) break;
+        sum += z->tab[i*2+1] - z->tab[i*2];
+    }
+
+    return sum + link - z->tab[i*2] + z->lr_offset;
+}
+
+static int rewrite_links(zebraFile *z)
+{
+    /* Rewrite the next/up/orig, structural, and reference links for all banks
+     * in the current logical record.
+     *
+     * Returns 0 on success, -1 on error. */
+    int i, rv;
+    uint32_t offset;
+    uint32_t io, noff;
+    zebraBank b;
+
+    offset = z->lr_offset*4;
+
+    while (offset < z->lr_size) {
+        io = unpacki32(z->buf+offset);
+        offset += 4;
+
+        noff = io & 0xFFFF;
+
+        if (noff > 12) {
+            offset += (noff - 12)*4;
+        }
+
+        rv = zebra_get_bank(z, &b, offset/4);
+
+        if (rv) return rv;
+
+        /* Rewrite the next, up, and orig pointers. */
+        if (b.next)
+            packi32(z->buf+offset,get_link(z,b.next));
+        if (b.up)
+            packi32(z->buf+offset+4,get_link(z,b.up));
+        if (b.orig)
+            packi32(z->buf+offset+8,get_link(z,b.orig));
+
+        /* Rewrite the structural and reference links. */
+        for (i = 0; i < b.num_links; i++) {
+            if (b.links[i])
+                packi32(z->buf+offset-(i+1)*4,get_link(z,b.links[i]));
+        }
+
+        offset += 9*4;
+
+        offset += b.num_data_words*4;
+    }
+
+    return 0;
+}
+
+int zebra_read_next_logical_record(zebraFile *z)
 {
     /* Read the next logical record into the buffer.
      *
      * Returns 0 on success, -1 on error, and 1 on EOF. */
+    int i;
     uint32_t size, type, nwtx, nwseg, nwtab, nwuh;
+    uint32_t io, noff;
+    uint32_t offset;
 
     while (1) {
-        memmove(z->buf,z->buf+z->offset,z->buf_size-z->offset);
-        z->buf = realloc(z->buf, z->buf_size-z->offset);
-        z->buf_size -= z->offset;
-        z->offset = 0;
+        /* Move the next logical record to the start of the buffer. */
+        memmove(z->buf,z->buf+z->lr_size,z->buf_size-z->lr_size);
+        z->buf = realloc(z->buf, z->buf_size-z->lr_size);
+        z->buf_size -= z->lr_size;
+        offset = 0;
 
-        switch (get_bytes(z,1)) {
+        switch (get_bytes(z,offset+1)) {
         case 1:
             return 1;
         case -1:
             return -1;
         }
-        size = unpacki32(z->buf+z->offset);
-        z->offset += 4;
-
+        size = unpacki32(z->buf+offset*4);
+        offset += 1;
+
+        /* If the size of the logical record is zero, we just skip it.
+         *
+         * From page 142 of the ZEBRA document at
+         * https://cdsweb.cern.ch/record/2296399/files/zebra.pdf:
+         *
+         * "The total size of a padding record is NWLR+1, including the LR
+         * length and the LR type; thus a padding record with NWLR=1 occupies
+         * two words. To pad exactly one word, NWLR=0 should be stored, in this
+         * case the LR type 5 is implied without being present in the data." */
         if (size == 0) continue;
 
-        switch (get_bytes(z,1)) {
+        switch (get_bytes(z,offset+1)) {
         case 1:
             return 1;
         case -1:
             return -1;
         }
-        type = unpacki32(z->buf+z->offset);
-        z->offset += 4;
+
+        /* Get the logical record type. */
+        type = unpacki32(z->buf+offset*4);
+        offset += 1;
 
         switch (type) {
+        /* Padding records. */
         case 5:
         case 6:
-            get_bytes(z,size-1);
-            z->offset += (size-1)*4;
+            get_bytes(z,offset+size-1);
+            z->lr_size = (size + 1)*4;
             continue;
+        /* Start or end of run. */
         case 1:
-            get_bytes(z,size);
-            z->offset += size*4;
+            get_bytes(z,offset+size);
+            z->lr_size = (size + 2)*4;
             continue;
+        /* Normal records. */
+        case 2:
+        case 3:
+        case 4:
+            break;
         }
 
+        /* For normal records the size of the logical record is NWLR + the two
+         * words for the size and type. */
         z->lr_size = (size + 2)*4;
-        get_bytes(z,size);
+        get_bytes(z,offset+size);
+
+        nwtx = unpacki32(z->buf+(offset+4)*4);
+        nwseg = unpacki32(z->buf+(offset+5)*4);
+        nwtab = unpacki32(z->buf+(offset+6)*4);
+        nwuh = unpacki32(z->buf+(offset+9)*4);
 
-        nwtx = unpacki32(z->buf+z->offset+4*4);
-        nwseg = unpacki32(z->buf+z->offset+5*4);
-        nwtab = unpacki32(z->buf+z->offset+6*4);
-        nwuh = unpacki32(z->buf+z->offset+9*4);
+        if (nwtab % 2) {
+            sprintf(zebra_err, "number of words in relocation table is not a multiple of 2!");
+            return -1;
+        }
 
-        z->offset += 10*4;
+        offset += 10;
 
-        z->offset += nwtx*4;
-        z->offset += nwseg*4;
-        z->offset += nwtab*4;
-        z->offset += nwuh*4;
+        /* Currently we don't use the user header, segment table, or the text
+         * vector. */
+        offset += nwuh;
+        offset += nwseg;
+        offset += nwtx;
+
+        /* Read the relocation table. */
+        if (z->tab) free(z->tab);
+        z->tab = malloc(sizeof(double)*nwtab);
+        for (i = 0; i < nwtab; i++)
+            z->tab[i] = unpacki32(z->buf+(offset+i)*4);
+        z->nwtab = nwtab;
+        offset += nwtab;
+
+        z->lr_offset = offset;
+
+        /* Find the pointer to the first bank. */
+        io = unpacki32(z->buf+offset*4);
+
+        noff = io & 0xFFFF;
+
+        /* First bank is just past the end of the logical record (one word for
+         * the io control word). */
+        z->first_bank = offset + 1;
+
+        /* Skip the links to get to the bank's next pointer. */
+        if (noff > 12) {
+            z->first_bank += noff - 12;
+        }
+
+        /* Rewrite the next/up/orig, structural, and reference links. */
+        if (rewrite_links(z)) return -1;
 
         break;
     }
@@ -182,31 +317,47 @@ int read_next_logical_record(zebraFile *z)
     return 0;
 }
 
-int next_bank(zebraFile *z, bank *b)
+int zebra_get_bank(zebraFile *z, zebraBank *b, uint32_t link)
 {
-    int rv;
-    uint32_t io, noff;
+    /* Get the bank located at `link`.
+     *
+     * On success 0 is returned and the zebraBank struct `b` is filled with the
+     * bank information. On error, -1 is returned and zebra_err is set. */
+    int i;
 
-    if (z->offset == z->lr_size) {
-        rv = read_next_logical_record(z);
-        if (rv) return rv;
+    if ((link+9)*4 > z->buf_size) {
+        sprintf(zebra_err, "link location of %i is outside buffer!", link);
+        return -1;
     }
 
-    io = unpacki32(z->buf+z->offset);
-    z->offset += 4;
+    unpack(z->buf+link*4,"lllllllll",&b->next, &b->up, &b->orig, &b->number, &b->idh, &b->num_links, &b->num_structural_links, &b->num_data_words, &b->status);
+
+    unpack(z->buf+link*4+16,"cccc",b->name,b->name+1,b->name+2,b->name+3);
 
-    noff = io & 0xFFFF;
+    b->name[4] = '\0';
 
-    if (noff > 12) {
-        z->offset += (noff-12)*4;
+    if (b->num_links > MAX_LINKS) {
+        sprintf(zebra_err, "number of links in bank is greater than %i!", MAX_LINKS);
+        return -1;
+    }
+
+    if (link < b->num_links) {
+        sprintf(zebra_err, "buffer only has %i words before bank, not enough for %i links!", link, b->num_links);
+        return -1;
     }
 
-    unpack(z->buf+z->offset,"lllllllll",&b->next, &b->up, &b->orig, &b->number, &b->name, &b->num_links, &b->num_structural_links, &b->num_data_words, &b->status);
-    z->offset += 9*4;
+    /* Read in reference and structural links. */
+    for (i = 0; i < b->num_links; i++)
+        b->links[i] = unpacki32(z->buf+(link-1-i)*4);
 
-    b->data = (uint32_t *) (z->buf+z->offset);
+    link += 9;
+
+    if ((link+b->num_data_words)*4 > z->buf_size) {
+        sprintf(zebra_err, "bank contains %i data bytes, but there are only %zu bytes left in buffer!", b->num_data_words*4, z->buf_size-link*4);
+        return -1;
+    }
 
-    z->offset += b->num_data_words*4;
+    b->data = (uint32_t *) (z->buf+link*4);
 
     return 0;
 }
@@ -214,6 +365,7 @@ int next_bank(zebraFile *z, bank *b)
 void zebra_close(zebraFile *z)
 {
     fclose(z->f);
+    if (z->tab) free(z->tab);
     if (z->buf) free(z->buf);
     free(z);
 }
diff --git a/src/zebra.h b/src/zebra.h
index 2f24939..c05f62d 100644
--- a/src/zebra.h
+++ b/src/zebra.h
@@ -1,3 +1,37 @@
+/* Library for reading ZEBRA files.
+ *
+ * The ZEBRA file format is an old file format used by FORTRAN programs for
+ * memory management and storing data structures. In particular, it's the
+ * default file format used by SNOMAN.
+ *
+ * This library is fairly simple and almost certainly does *not* include all
+ * the details needed to read an arbitrary ZEBRA file. Instead I implemented as
+ * much as was necessary to read in SNOMAN files and fixed bugs along the way.
+ *
+ * Example usage:
+ *
+ *     zebraBank b;
+ *     zebraFile *z = zebra_open("Muons.zdab")
+ *
+ *     while (1) {
+ *         rv = read_next_logical_record(z);
+ *
+ *         switch (rv) {
+ *             case 1:
+ *                 // EOF
+ *                 goto end;
+ *             case -1:
+ *                 fprintf(stderr, "error getting logical record: %s\n", zebra_err);
+ *                 goto err;
+ *          }
+ *
+ *          zebra_get_bank(z,&b,z->first_bank);
+ *
+ *          // etc.
+ *      }
+ *
+ */
+
 #ifndef ZEBRA_H
 #define ZEBRA_H
 
@@ -5,44 +39,102 @@
 #include <stdlib.h> /* for size_t */
 #include <stdint.h> /* for uint8_t, etc. */
 
+/* Maximum number of links in a bank.
+ *
+ * Technically we could malloc() this but that means we would have to free
+ * banks which would be annoying. */
+#define MAX_LINKS 100
+
+/* Global error string when any function returns -1. */
 extern char zebra_err[256];
 
+/* Physical record markers. */
 #define ZEBRA_SIG0 0x0123cdefUL
 #define ZEBRA_SIG1 0x80708070UL
 #define ZEBRA_SIG2 0x4321abcdUL
 #define ZEBRA_SIG3 0x80618061UL
 
+/* Bitmasks for the physical record block size control word. */
 #define ZEBRA_BLOCK_SIZE_MASK 0x00ffffffUL
 #define ZEBRA_EMERGENCY_STOP  0x80000000UL
 #define ZEBRA_END_OF_RUN      0x20000000UL
 #define ZEBRA_START_OF_RUN    0x40000000UL
 
-typedef struct bank {
+typedef struct zebraBank {
+    /* Pointer to the next bank in the chain. */
     uint32_t next;
+    /* Pointer to the supporting bank. */
     uint32_t up;
+    /* Pointer to the previous bank in the chain (or to the referencing link in
+     * the supporting bank if it's the first). */
     uint32_t orig;
+    /* Bank number. From the SNOMAN FAQ:
+     *
+     * "In general, ZEBRA attaches no special significance to the bank number,
+     * and it is perfectly O.K. to change it at any time." */
     uint32_t number;
-    uint32_t name;
+    /* Hollerith bank ID as an integer. */
+    uint32_t idh;
+    /* Total number of links in the bank.
+     *
+     * Note: In SNOMAN reference links aren't guaranteed to point to the start
+     * of a bank. According to the SNOMAN docs all reference links in SNOMAN
+     * should point to the bank's status word instead of the start of the bank,
+     * but I tried the following and it didn't seem to work:
+     *
+     *     zebra_get_bank(z,&b,mctk.links[KMCTK_MCVX-1]-8);
+     *
+     * So, I'm not really sure how reference links work in SNOMAN. */
     uint32_t num_links;
+    /* Number of structural links. The order of the links in the bank is:
+     *
+     *     structural link 1
+     *     structural link 2
+     *     ...
+     *     structural link num_structural_links
+     *     reference link num_structural_links+1
+     *     reference link num_structural_links+2
+     *     ...
+     *     reference link num_links
+     *
+     */
     uint32_t num_structural_links;
+    /* Number of data words. */
     uint32_t num_data_words;
+    /* Status word. */
     uint32_t status;
+    /* Pointer to the bank data in zebraFile->buf. Note that after calling
+     * zebra_read_next_logical_record(), this pointer will no longer point to
+     * the bank data. */
     uint32_t *data;
-} bank;
+    /* Reference and structural links. */
+    uint32_t links[MAX_LINKS];
+    /* The Hollerith bank name as a string instead of an integer. This can be
+     * useful when printing error messages or for debugging. */
+    char name[5];
+} zebraBank;
 
 typedef struct zebraFile {
     FILE *f;
-    size_t offset;
+    /* Size of the current logical record in bytes. */
     size_t lr_size;
+    /* Buffer used to read in the zebra file. */
     uint8_t *buf;
+    /* Total size of the current buffer. */
     size_t buf_size;
+    /* Relocation table. */
+    uint32_t *tab;
+    /* Number of words in the relocation table. */
+    size_t nwtab;
+    /* Link to first bank. */
+    uint32_t first_bank;
+    /* Number of words from the start of the buffer to the first bank word. */
+    uint32_t lr_offset;
 } zebraFile;
 
 zebraFile *zebra_open(const char *filename);
-int read_next_physical_record(zebraFile *z);
-int get_bytes(zebraFile *z, size_t size);
-int read_next_logical_record(zebraFile *z);
-int next_bank(zebraFile *z, bank *b);
+int zebra_read_next_logical_record(zebraFile *z);
+int zebra_get_bank(zebraFile *z, zebraBank *b, uint32_t link);
 void zebra_close(zebraFile *z);
 
 #endif
diff --git a/utils/plot b/utils/plot
index 59eb4c2..d671016 100755
--- a/utils/plot
+++ b/utils/plot
@@ -75,7 +75,7 @@ if __name__ == '__main__':
         psi = []
         for event in data['data']:
             # get the particle ID
-            id = event['mctk'][-1]['id']
+            id = event['mcgn'][0]['id']
 
             if 'ev' not in event:
                 continue
@@ -88,24 +88,24 @@ if __name__ == '__main__':
             mass = SNOMAN_MASS[id]
             # for some reason it's the *second* track which seems to contain the
             # initial energy
-            true_energy = event['mctk'][-1]['energy']
+            true_energy = event['mcgn'][0]['energy']
             # The MCTK bank has the particle's total energy (except for neutrons)
             # so we need to convert it into kinetic energy
             ke = true_energy - mass
             energy = fit[id]['energy']
             dT.append(energy-ke)
-            true_posx = event['mcvx'][0]['posx']
+            true_posx = event['mcgn'][0]['posx']
             posx = fit[id]['posx']
             dx.append(posx-true_posx)
-            true_posy = event['mcvx'][0]['posy']
+            true_posy = event['mcgn'][0]['posy']
             posy = fit[id]['posy']
             dy.append(posy-true_posy)
-            true_posz = event['mcvx'][0]['posz']
+            true_posz = event['mcgn'][0]['posz']
             posz = fit[id]['posz']
             dz.append(posz-true_posz)
-            dirx = event['mctk'][-1]['dirx']
-            diry = event['mctk'][-1]['diry']
-            dirz = event['mctk'][-1]['dirz']
+            dirx = event['mcgn'][0]['dirx']
+            diry = event['mcgn'][0]['diry']
+            dirz = event['mcgn'][0]['dirz']
             true_dir = [dirx,diry,dirz]
             true_dir = np.array(true_dir)/np.linalg.norm(true_dir)
             theta = fit[id]['theta']
diff --git a/utils/plot-fit-results b/utils/plot-fit-results
index 48d922a..c3d3f16 100755
--- a/utils/plot-fit-results
+++ b/utils/plot-fit-results
@@ -84,7 +84,7 @@ if __name__ == '__main__':
 
         for i, event in enumerate(data['data']):
             # get the particle ID
-            id = event['mctk'][-1]['id']
+            id = event['mcgn'][0]['id']
 
             a[i]['id'] = id
 
@@ -95,7 +95,7 @@ if __name__ == '__main__':
             mass = SNOMAN_MASS[id]
             # for some reason it's the *second* track which seems to contain the
             # initial energy
-            true_energy = event['mctk'][-1]['energy']
+            true_energy = event['mcgn'][0]['energy']
             # The MCTK bank has the particle's total energy (except for neutrons)
             # so we need to convert it into kinetic energy
             ke = true_energy - mass
@@ -107,21 +107,21 @@ if __name__ == '__main__':
             a[i]['dT'] = energy-ke
 
             # store the fit position residuals
-            true_posx = event['mcvx'][0]['posx']
+            true_posx = event['mcgn'][0]['posx']
             posx = fit[id]['posx']
             a[i]['dx'] = posx-true_posx
-            true_posy = event['mcvx'][0]['posy']
+            true_posy = event['mcgn'][0]['posy']
             posy = fit[id]['posy']
             a[i]['dy'] = posy-true_posy
-            true_posz = event['mcvx'][0]['posz']
+            true_posz = event['mcgn'][0]['posz']
             posz = fit[id]['posz']
             a[i]['dz'] = posz-true_posz
 
             # compute the angle between the fit direction and the true
             # direction
-            dirx = event['mctk'][-1]['dirx']
-            diry = event['mctk'][-1]['diry']
-            dirz = event['mctk'][-1]['dirz']
+            dirx = event['mcgn'][0]['dirx']
+            diry = event['mcgn'][0]['diry']
+            dirz = event['mcgn'][0]['dirz']
             true_dir = [dirx,diry,dirz]
             true_dir = np.array(true_dir)/np.linalg.norm(true_dir)
             theta = fit[id]['theta']