add the QUAD fitter

4 files changed, 462 insertions, 1 deletions

diff --git a/src/Makefile b/src/Makefile
index 9bae25e..9739428 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -14,7 +14,7 @@ calculate_limits: calculate_limits.c
 
 solid_angle.o: solid_angle.c
 
-test: test.o solid_angle.o optics.o muon.o vector.o quantum_efficiency.o pdg.o scattering.o misc.o mt19937ar.o sno_charge.o path.o random.o pmt_response.o db.o dict.o siphash.o electron.o proton.o likelihood.o pmt.o
+test: test.o solid_angle.o optics.o muon.o vector.o quantum_efficiency.o pdg.o scattering.o misc.o mt19937ar.o sno_charge.o path.o random.o pmt_response.o db.o dict.o siphash.o electron.o proton.o likelihood.o pmt.o quad.o
 
 test-vector: test-vector.o vector.o mt19937ar.o
 
diff --git a/src/quad.c b/src/quad.c
new file mode 100644
index 0000000..c869aed
--- /dev/null
+++ b/src/quad.c
@@ -0,0 +1,237 @@
+#include "quad.h"
+#include "event.h"
+#include <gsl/gsl_randist.h>
+#include <gsl/gsl_rng.h>
+#include <unistd.h> /* for exit() */
+#include <stdlib.h> /* for size_t */
+#include <gsl/gsl_linalg.h>
+#include "pdg.h"
+#include "pmt.h"
+#include <gsl/gsl_statistics_double.h>
+#include "optics.h"
+#include "vector.h"
+#include "likelihood.h"
+#include <gsl/gsl_sort.h>
+#include <gsl/gsl_cblas.h>
+#include "misc.h"
+
+char quad_err[256];
+
+/* Returns an approximate vertex position and time using the QUAD fitter.
+ *
+ * The QUAD fitter was originally developed for SNO and estimates the event
+ * vertex using a sort of Hough Transform[1]. It works by selecting 4 random
+ * PMT hits and computing the position and time of the event which is
+ * consistent with those hits. This procedure is repeated many times and the
+ * median of the resulting "quad cloud" is found as an estimate of the event
+ * position and time.
+ *
+ * Note: In the original formulation I believe instead of taking the median it
+ * used a minimization routine to find the position and time with the highest
+ * density of points, but since we are dealing with mostly high nhit events we
+ * just take the median because it is easier and quicker.
+ *
+ * `ev` should be a pointer to the event.
+ *
+ * `pos` is a pointer to a double array of length 3.
+ *
+ * `t0` is a pointer to a double.
+ *
+ * `npoints` is the number of quad cloud points to compute. This function will
+ * try to compute this many points, but it will stop trying after 10*npoints
+ * times to avoid an infinite loop.
+ *
+ * Returns 0 on success and `pos` and `t0` will be set to the approximate
+ * vertex position and time respectively. On error, returns -1 and sets the
+ * `quad_err` string.
+ *
+ * Example:
+ *
+ *     double pos[3], t0;
+ *     if (quad(&ev, pos, &t0, 1000)) {
+ *         fprintf(stderr, "error running the quad fitter: %s\n", quad_err);
+ *         goto err;
+ *     }
+ *
+ * [1] The only reference to the QUAD fitter that I can find from the SNO days
+ * is Stephen Brice's PHD thesis which cites three SNO technical reports, but I
+ * haven't been able to find these. The first two of these reports are by Bill
+ * Frati who I assume wrote the initial implementation. */
+int quad(event *ev, double *pos, double *t0, size_t npoints)
+{
+    size_t i, j, k;
+    static int index[MAX_PMTS];
+    static int index2[MAX_PMTS];
+    int nhit;
+    const gsl_rng_type *T;
+    gsl_rng *r;
+    int xs[4];
+    size_t nresults;
+    double M[3][3], Minv[3][3];
+    double K[3], g[3], h[3], n[3], tmp[3];
+    double a, b, c;
+    double pos1[3], pos2[3];
+    static double results[QUAD_MAX][4];
+    double c2;
+    double wavelength0;
+    double n_d2o;
+    double t1, t2;
+    int s;
+    double pmt_dir[3];
+    double expected;
+    static double t[MAX_PMTS];
+    static size_t p2[MAX_PMTS];
+    double tmin;
+    size_t max_tries;
+
+    wavelength0 = 400.0;
+    n_d2o = get_index_snoman_d2o(wavelength0);
+
+    c2 = SPEED_OF_LIGHT*SPEED_OF_LIGHT/pow(n_d2o,2);
+
+    if (npoints > QUAD_MAX) {
+        sprintf(quad_err, "npoints must be less than %i", QUAD_MAX);
+        return 1;
+    }
+
+    nhit = 0;
+    for (i = 0; i < MAX_PMTS; i++) {
+        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL) continue;
+
+        if (ev->pmt_hits[i].hit) {
+            index[nhit] = i;
+            t[nhit] = ev->pmt_hits[i].t;
+            nhit++;
+        }
+    }
+
+    /* Only use the first 50% of the hits to avoid counting late/reflected
+     * light. */
+    k = nhit > 10 ? ceil(nhit/2.0) : nhit;
+    gsl_sort_smallest_index(p2,k,t,1,nhit);
+
+    /* Reorder the hits. */
+    for (i = 0; i < k; i++) {
+        index2[i] = index[p2[i]];
+    }
+
+    nhit = k;
+
+    if (nhit < 5) {
+        sprintf(quad_err, "only %i pmt hits. quad needs at least 5 points!", nhit);
+        return 1;
+    }
+
+    T = gsl_rng_default;
+    r = gsl_rng_alloc(T);
+
+    gsl_permutation *p = gsl_permutation_alloc(3);
+
+    i = 0;
+    max_tries = npoints*10;
+    nresults = 0;
+    while (nresults < npoints && i++ < max_tries) {
+        /* Choose 4 random hits. */
+        gsl_ran_choose(r,xs,4,index2,nhit,sizeof(int));
+
+        /* Shuffle them since GSL always returns the random choices in order.
+         *
+         * I'm not actually sure if that affects the quad calculation. */
+        gsl_ran_shuffle(r,xs,4,sizeof(int));
+
+        for (j = 1; j < 4; j++) {
+            for (k = 0; k < 3; k++) {
+                M[j-1][k] = pmts[xs[j]].pos[k] - pmts[xs[0]].pos[k];
+            }
+            n[j-1] = ev->pmt_hits[xs[j]].t - ev->pmt_hits[xs[0]].t;
+
+            K[j-1] = (c2*(pow(ev->pmt_hits[xs[0]].t,2) - pow(ev->pmt_hits[xs[j]].t,2)) - DOT(pmts[xs[0]].pos,pmts[xs[0]].pos) + DOT(pmts[xs[j]].pos,pmts[xs[j]].pos))/2;
+        }
+
+        gsl_matrix_view m = gsl_matrix_view_array(&M[0][0],3,3);
+        gsl_matrix_view minv = gsl_matrix_view_array(&Minv[0][0],3,3);
+        gsl_vector_view k_view = gsl_vector_view_array(K,3);
+        gsl_vector_view g_view = gsl_vector_view_array(g,3);
+        gsl_vector_view h_view = gsl_vector_view_array(h,3);
+        gsl_vector_view n_view = gsl_vector_view_array(n,3);
+
+        gsl_linalg_LU_decomp(&m.matrix, p, &s);
+        gsl_linalg_LU_invert(&m.matrix, p, &minv.matrix);
+
+        gsl_blas_dgemv(CblasNoTrans,1.0,&minv.matrix,&k_view.vector,0.0,&g_view.vector);
+        gsl_blas_dgemv(CblasNoTrans,1.0,&minv.matrix,&n_view.vector,0.0,&h_view.vector);
+
+        SUB(tmp,pmts[xs[0]].pos,g);
+        a = c2*(c2*DOT(h,h) - 1.0);
+        b = -2*c2*(DOT(tmp,h) - ev->pmt_hits[xs[0]].t);
+        c = DOT(tmp,tmp) - c2*pow(ev->pmt_hits[xs[0]].t,2);
+
+        if (b*b - 4*a*c < 0) continue;
+
+        t1 = (-b + sqrt(b*b - 4*a*c))/(2*a);
+        COPY(tmp,h);
+        MUL(tmp,c2*t1);
+        ADD(pos1,g,tmp);
+
+        /* Check the first result. */
+        SUB(pmt_dir,pmts[xs[0]].pos,pos1);
+        expected = t1 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;
+
+        tmin = ev->pmt_hits[xs[0]].t;
+        for (j = 1; j < 4; j++) {
+            if (ev->pmt_hits[xs[j]].t < tmin)
+                tmin = ev->pmt_hits[xs[j]].t;
+        }
+
+        if (t1 < tmin && isclose(ev->pmt_hits[xs[0]].t,expected,0,1e-2)) {
+            results[nresults][0] = pos1[0];
+            results[nresults][1] = pos1[1];
+            results[nresults][2] = pos1[2];
+            results[nresults][3] = t1;
+            nresults++;
+        } else {
+            /* Check the second solution. */
+            t2 = (-b - sqrt(b*b - 4*a*c))/(2*a);
+            COPY(tmp,h);
+            MUL(tmp,c2*t2);
+            ADD(pos2,g,tmp);
+
+            SUB(pmt_dir,pmts[xs[0]].pos,pos2);
+            expected = t2 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;
+
+            if (t2 < tmin && isclose(ev->pmt_hits[xs[0]].t,expected,0,1e-2)) {
+                results[nresults][0] = pos2[0];
+                results[nresults][1] = pos2[1];
+                results[nresults][2] = pos2[2];
+                results[nresults][3] = t2;
+                nresults++;
+            }
+        }
+    }
+
+    if (nresults < 1) {
+        sprintf(quad_err, "no valid quad points found!");
+        goto err;
+    }
+
+    gsl_sort(&results[0][0],4,nresults);
+    gsl_sort(&results[0][1],4,nresults);
+    gsl_sort(&results[0][2],4,nresults);
+    gsl_sort(&results[0][3],4,nresults);
+
+    pos[0] = gsl_stats_median_from_sorted_data(&results[0][0],4,nresults);
+    pos[1] = gsl_stats_median_from_sorted_data(&results[0][1],4,nresults);
+    pos[2] = gsl_stats_median_from_sorted_data(&results[0][2],4,nresults);
+    *t0 = gsl_stats_median_from_sorted_data(&results[0][3],4,nresults);
+
+    gsl_permutation_free(p);
+    gsl_rng_free(r);
+
+    return 0;
+
+err:
+    gsl_permutation_free(p);
+    gsl_rng_free(r);
+
+    return 1;
+}
diff --git a/src/quad.h b/src/quad.h
new file mode 100644
index 0000000..730f45d
--- /dev/null
+++ b/src/quad.h
@@ -0,0 +1,14 @@
+#ifndef QUAD_H
+#define QUAD_H
+
+#include "event.h"
+#include <stdlib.h> /* for size_t */
+
+/* Maximum number of quad cloud points. */
+#define QUAD_MAX 10000
+
+extern char quad_err[256];
+
+int quad(event *ev, double *pos, double *t0, size_t npoints);
+
+#endif
diff --git a/src/test.c b/src/test.c
index 633897d..d7839dc 100644
--- a/src/test.c
+++ b/src/test.c
@@ -22,6 +22,8 @@
 #include <gsl/gsl_errno.h> /* for gsl_strerror() */
 #include <gsl/gsl_randist.h>
 #include <gsl/gsl_cdf.h>
+#include "sno.h"
+#include "quad.h"
 
 typedef int testFunction(char *err);
 
@@ -1358,6 +1360,212 @@ err:
     return 1;
 }
 
+int test_quad(char *err)
+{
+    /* Tests the quad fitter without noise. We draw 1000 random hits, compute
+     * the time each PMT is hit (without noise) and then make sure that the
+     * positions returned by the quad fitter are all within 1 cm and the
+     * initial time is within 1 ns. */
+    size_t i, j;
+    double x0[3];
+    double t0;
+    const gsl_rng_type *T;
+    gsl_rng *r;
+    event ev;
+    int index[MAX_PMTS];
+    int hits[1000];
+    size_t valid_pmts;
+    double fit_pos[3];
+    double pmt_dir[3];
+    double fit_t0;
+    double wavelength0;
+    double n_d2o;
+
+    init_genrand(0);
+
+    T = gsl_rng_default;
+    r = gsl_rng_alloc(T);
+
+    load_pmt_info();
+
+    valid_pmts = 0;
+    for (i = 0; i < MAX_PMTS; i++) {
+        if (pmts[i].pmt_type != PMT_NORMAL) continue;
+
+        index[valid_pmts++] = i;
+    }
+
+    wavelength0 = 400.0;
+    n_d2o = get_index_snoman_d2o(wavelength0);
+
+    for (i = 0; i < 100; i++) {
+        /* Generate a random position within a cube the size of the AV.
+         *
+         * Note: This does produce some points which are outside of the PSUP,
+         * but I think the quad fitter should still be able to fit these. */
+        x0[0] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        x0[1] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        x0[2] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        t0 = genrand_real2()*GTVALID;
+
+        /* Zero out all PMTs. */
+        for (j = 0; j < LEN(ev.pmt_hits); j++) {
+            ev.pmt_hits[j].hit = 0;
+            ev.pmt_hits[j].flags = 0;
+        }
+
+        /* Choose LEN(hits) random PMTs which are hit. */
+        gsl_ran_choose(r,hits,LEN(hits),index,valid_pmts,sizeof(int));
+
+        /* Calculate the time the PMT got hit. */
+        for (j = 0; j < LEN(hits); j++) {
+            SUB(pmt_dir,pmts[hits[j]].pos,x0);
+            ev.pmt_hits[hits[j]].hit = 1;
+            ev.pmt_hits[hits[j]].qhs = genrand_real2();
+            ev.pmt_hits[hits[j]].t = t0 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;
+        }
+
+        if (quad(&ev, fit_pos, &fit_t0, 10000)) {
+            sprintf(err, quad_err);
+            goto err;
+        }
+
+        /* Since there's no noise, these results should be exact. We test to
+         * see that all of the positions are within 1 cm and the time is within
+         * 1 ns. */
+
+        if (!isclose(fit_pos[0], x0[0], 0, 1.0)) {
+            sprintf(err, "quad returned x = %.5f, but expected %.5f", fit_pos[0], x0[0]);
+            goto err;
+        }
+
+        if (!isclose(fit_pos[1], x0[1], 0, 1.0)) {
+            sprintf(err, "quad returned y = %.5f, but expected %.5f", fit_pos[1], x0[1]);
+            goto err;
+        }
+
+        if (!isclose(fit_pos[2], x0[2], 0, 1.0)) {
+            sprintf(err, "quad returned z = %.5f, but expected %.5f", fit_pos[2], x0[2]);
+            goto err;
+        }
+
+        if (!isclose(fit_t0, t0, 0, 1.0)) {
+            sprintf(err, "quad returned t0 = %.5f, but expected %.5f", fit_t0, t0);
+            goto err;
+        }
+    }
+
+    gsl_rng_free(r);
+
+    return 0;
+
+err:
+    gsl_rng_free(r);
+
+    return 1;
+}
+
+int test_quad_noise(char *err)
+{
+    /* Tests the quad fitter with noise. We draw 1000 random hits, compute
+     * the time each PMT is hit, add a gaussian sample with standard deviation
+     * PMT_TTS and then make sure that the positions returned by the quad
+     * fitter are all within 1 m and the initial time is within 10 ns. */
+    size_t i, j;
+    double x0[3];
+    double t0;
+    const gsl_rng_type *T;
+    gsl_rng *r;
+    event ev;
+    int index[MAX_PMTS];
+    int hits[1000];
+    size_t valid_pmts;
+    double fit_pos[3];
+    double pmt_dir[3];
+    double fit_t0;
+    double wavelength0;
+    double n_d2o;
+
+    init_genrand(0);
+
+    T = gsl_rng_default;
+    r = gsl_rng_alloc(T);
+
+    load_pmt_info();
+
+    valid_pmts = 0;
+    for (i = 0; i < MAX_PMTS; i++) {
+        if (pmts[i].pmt_type != PMT_NORMAL) continue;
+
+        index[valid_pmts++] = i;
+    }
+
+    wavelength0 = 400.0;
+    n_d2o = get_index_snoman_d2o(wavelength0);
+
+    for (i = 0; i < 100; i++) {
+        /* Generate a random position within a cube the size of the AV.
+         *
+         * Note: This does produce some points which are outside of the PSUP,
+         * but I think the quad fitter should still be able to fit these. */
+        x0[0] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        x0[1] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        x0[2] = (genrand_real2()*2 - 1)*AV_RADIUS;
+        t0 = genrand_real2()*GTVALID;
+
+        /* Zero out all PMTs. */
+        for (j = 0; j < LEN(ev.pmt_hits); j++) {
+            ev.pmt_hits[j].hit = 0;
+            ev.pmt_hits[j].flags = 0;
+        }
+
+        /* Choose LEN(hits) random PMTs which are hit. */
+        gsl_ran_choose(r,hits,LEN(hits),index,valid_pmts,sizeof(int));
+
+        /* Calculate the time the PMT got hit. */
+        for (j = 0; j < LEN(hits); j++) {
+            SUB(pmt_dir,pmts[hits[j]].pos,x0);
+            ev.pmt_hits[hits[j]].hit = 1;
+            ev.pmt_hits[hits[j]].qhs = genrand_real2();
+            ev.pmt_hits[hits[j]].t = t0 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT + randn()*PMT_TTS;
+        }
+
+        if (quad(&ev, fit_pos, &fit_t0, 10000)) {
+            sprintf(err, quad_err);
+            goto err;
+        }
+
+        if (!isclose(fit_pos[0], x0[0], 0, 100.0)) {
+            sprintf(err, "quad returned x = %.5f, but expected %.5f", fit_pos[0], x0[0]);
+            goto err;
+        }
+
+        if (!isclose(fit_pos[1], x0[1], 0, 100.0)) {
+            sprintf(err, "quad returned y = %.5f, but expected %.5f", fit_pos[1], x0[1]);
+            goto err;
+        }
+
+        if (!isclose(fit_pos[2], x0[2], 0, 100.0)) {
+            sprintf(err, "quad returned z = %.5f, but expected %.5f", fit_pos[2], x0[2]);
+            goto err;
+        }
+
+        if (!isclose(fit_t0, t0, 0, 10.0)) {
+            sprintf(err, "quad returned t0 = %.5f, but expected %.5f", fit_t0, t0);
+            goto err;
+        }
+    }
+
+    gsl_rng_free(r);
+
+    return 0;
+
+err:
+    gsl_rng_free(r);
+
+    return 1;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -1396,6 +1604,8 @@ struct tests {
     {test_norm_cdf, "test_norm_cdf"},
     {test_time_pdf_norm, "test_time_pdf_norm"},
     {test_time_cdf, "test_time_cdf"},
+    {test_quad, "test_quad"},
+    {test_quad_noise, "test_quad_noise"},
 };
 
 int main(int argc, char **argv)