path: root/src/fit.c

#include "likelihood.h"
#include <stdio.h>
#include "zebra.h"
#include "Record_Info.h"
#include "event.h"
#include "zdab_utils.h"
#include "scattering.h"
#include "pmt.h"
#include "sno_charge.h"
#include "db.h"
#include "dqxx.h"
#include <nlopt.h>
#include <math.h> /* for sin(), cos(), etc. */
#include <sys/time.h> /* for gettimeofday() */
#include <inttypes.h> /* for PRIu32 macro */
#include <string.h> /* for memcpy() */
#include <errno.h> /* for errno */
#include "pdg.h"
#include "optics.h"
#include "vector.h"
#include "pmt_response.h"

#define EV_RECORD    0x45562020 // 'EV  '  (as written to ZDAB file)

static size_t iter;

typedef struct fitParams {
    event *ev;
    double epsrel;
    int fast;
} fitParams;

/* In order to start the fitter close to the minimum, we first do a series of
 * "quick" minimizations starting at the following points. We keep track of the
 * parameters with the best likelihood value and then start the "real"
 * minimization from those parameters. */

static struct startingParameters {
    double x;
    double y;
    double z;
} startingParameters[] = {
    {   0.0,   0.0,    0.0},
    { 100.0,   0.0,    0.0},
    { 100.0, 100.0,    0.0},
    {   0.0, 100.0,    0.0},
    {-100.0, 100.0,    0.0},
    {-100.0,   0.0,    0.0},
    {-100.0,-100.0,    0.0},
    {   0.0,-100.0,    0.0},
    { 100.0,-100.0,    0.0},
    {   0.0,   0.0,  100.0},
    { 100.0,   0.0,  100.0},
    { 100.0, 100.0,  100.0},
    {   0.0, 100.0,  100.0},
    {-100.0, 100.0,  100.0},
    {-100.0,   0.0,  100.0},
    {-100.0,-100.0,  100.0},
    {   0.0,-100.0,  100.0},
    { 100.0,-100.0,  100.0},
    {   0.0,   0.0, -100.0},
    { 100.0,   0.0, -100.0},
    { 100.0, 100.0, -100.0},
    {   0.0, 100.0, -100.0},
    {-100.0, 100.0, -100.0},
    {-100.0,   0.0, -100.0},
    {-100.0,-100.0, -100.0},
    {   0.0,-100.0, -100.0},
    { 100.0,-100.0, -100.0},
};

double nll(unsigned int n, const double *x, double *grad, void *params)
{
    fitParams *fpars = (fitParams *) params;
    double T, theta, phi, t0;
    double pos[3], dir[3];
    double fval;
    double z1[1], z2[1];
    struct timeval tv_start, tv_stop;

    T = x[0];

    pos[0] = x[1];
    pos[1] = x[2];
    pos[2] = x[3];

    theta = x[4];
    phi = x[5];
    dir[0] = sin(theta)*cos(phi);
    dir[1] = sin(theta)*sin(phi);
    dir[2] = cos(theta);

    t0 = x[6];

    z1[0] = x[7];
    z2[0] = x[8];

    gettimeofday(&tv_start, NULL);
    fval = nll_muon(fpars->ev, T, pos, dir, t0, z1, z2, 1, fpars->epsrel, fpars->fast);
    gettimeofday(&tv_stop, NULL);

    long long elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;

    printf("%5zu %10.2f %7.2f %7.2f %7.2f %5.2f %5.2f %5.2f %5.2f %5.2f f() = %7.3e took %lld ms\n",
           iter++,
           x[0],
           x[1],
           x[2],
           x[3],
           x[4],
           x[5],
           x[6],
           x[7],
           x[8],
           fval,
           elapsed);

    return fval;
}

double guess_t0(event *ev, double *pos)
{
    /* Returns a guess for the t0 of event `ev` given a position `pos`. The t0
     * is calculated by computing the charge weighted average of the time
     * difference between each PMT's hit time and the time it would take a
     * photon to travel from `pos` to the PMT. */
    size_t i;
    double pmt_dir[3], distance, t0, n, qhs_sum;

    /* Compute the index of refraction for heavy water at 400 nm. */
    n = get_index_snoman_d2o(400.0);

    t0 = 0.0;
    qhs_sum = 0.0;
    for (i = 0; i < MAX_PMTS; i++) {
        /* Only look at normal PMTs that were hit and aren't flagged. */
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL || !ev->pmt_hits[i].hit) continue;

        /* Compute the vector between `pos` and the PMT. */
        SUB(pmt_dir,pmts[i].pos,pos);

        /* Compute the distance to the PMT. */
        distance = NORM(pmt_dir);

        /* Add the charge weighted time difference between the PMT hit time and
         * the time it would take a photon to hit the PMT from `pos`. */
        t0 += ev->pmt_hits[i].qhs*(ev->pmt_hits[i].t - distance*n/SPEED_OF_LIGHT);

        qhs_sum += ev->pmt_hits[i].qhs;
    }

    /* Divide by the total QHS sum. */
    t0 /= qhs_sum;

    return t0;
}

void guess_direction(event *ev, double *pos, double *theta, double *phi)
{
    /* Returns the approximate direction of the event from a given position.
     * The direction is computed by taking the charge weighted average of the
     * vectors from `pos` to each hit PMT.
     *
     * A possible improvement here is to do something like a Hough transform
     * where we map PMT hits to cones with a 42 degree angle and then search
     * for the highest point in the transformed space. This method might also
     * generalize better when searching for a second ring since we can look for
     * secondary peaks not near the highest peak. */
    size_t i;
    double pmt_dir[3], dir[3];

    dir[0] = 0.0;
    dir[1] = 0.0;
    dir[2] = 0.0;

    for (i = 0; i < MAX_PMTS; i++) {
        /* Only look at normal PMTs that were hit and aren't flagged. */
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL || !ev->pmt_hits[i].hit) continue;

        /* Compute the vector between `pos` and the PMT. */
        SUB(pmt_dir,pmts[i].pos,pos);

        /* Normalize it. */
        normalize(pmt_dir);

        /* Multiply the vector by the QHS charge in the PMT. */
        MUL(pmt_dir,ev->pmt_hits[i].qhs);

        /* Add this to the estimated direction. */
        ADD(dir,dir,pmt_dir);
    }

    /* Normalize the charge weighted sum. */
    normalize(dir);

    /* Compute the spherical coordinates for `dir`. */
    *theta = acos(dir[2]);
    *phi = atan2(dir[1],dir[0]);
}

int fit_event(event *ev, double *xopt, double *fmin)
{
    size_t i;
    fitParams fpars;
    double x[9], ss[9], lb[9], ub[9], fval, n, qhs_sum, x0[9], T0, Tmin;
    int rv;

    nlopt_opt opt = nlopt_create(NLOPT_LN_BOBYQA, 9);
    nlopt_set_min_objective(opt,nll,&fpars);

    /* Make a guess as to the energy. Right now we just use a simple
     * approximation that the muon produces approximately 6 photons/MeV.
     *
     * FIXME: Should update this to something better. */

    qhs_sum = 0.0;
    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL || !ev->pmt_hits[i].hit) continue;
        qhs_sum += ev->pmt_hits[i].qhs;
    }

    T0 = qhs_sum/6.0;

    n = get_index_snoman_d2o(400.0);

    /* Calculate the Cerenkov threshold for a muon. */
    Tmin = MUON_MASS/sqrt(1.0-1.0/(n*n)) - MUON_MASS;

    /* If our guess is below the Cerenkov threshold, start at the Cerenkov
     * threshold. */
    double Tmin2 = sqrt(MUON_MASS*MUON_MASS/(1-pow(0.9,2)))-MUON_MASS;
    if (T0 < Tmin2) T0 = Tmin2;

    x0[0] = T0;
    x0[1] = 0.0;
    x0[2] = 0.0;
    x0[3] = 0.0;
    x0[4] = 1.57;
    x0[5] = 0.0;
    x0[6] = 130.0;
    x0[7] = 0.0;
    x0[8] = 0.0;

    ss[0] = x0[0]*0.02;
    ss[1] = 10.0;
    ss[2] = 10.0;
    ss[3] = 10.0;
    ss[4] = 0.01;
    ss[5] = 0.01;
    ss[6] = 1.0;
    ss[7] = 0.1;
    ss[8] = 0.1;

    lb[0] = Tmin;
    lb[1] = -1000.0;
    lb[2] = -1000.0;
    lb[3] = -1000.0;
    lb[4] = -INFINITY;
    lb[5] = -INFINITY;
    lb[6] = 0.0;
    lb[7] = -10.0;
    lb[8] = -10.0;

    ub[0] = 10000.0;
    ub[1] = 1000.0;
    ub[2] = 1000.0;
    ub[3] = 1000.0;
    ub[4] = INFINITY;
    ub[5] = INFINITY;
    ub[6] = 400.0;
    ub[7] = 10.0;
    ub[8] = 10.0;

    nlopt_set_lower_bounds(opt, lb);
    nlopt_set_upper_bounds(opt, ub);

    nlopt_set_initial_step(opt, ss);

    fpars.ev = ev;

    iter = 0;

    /* First we do a set of "quick" minimizations to try and start the
     * minimizer close to the minimum. To make these function evaluations
     * faster, we set the absolute tolerance on the likelihood to 1.0, the
     * maximum number of function evaluations to 100, and the relative
     * tolerance on the numerical integration to 10%. */
    fpars.epsrel = 1e-1;
    fpars.fast = 1;
    nlopt_set_ftol_abs(opt, 1.0);
    nlopt_set_maxeval(opt, 100);

    for (i = 0; i < sizeof(startingParameters)/sizeof(startingParameters[0]); i++) {
        memcpy(x,x0,sizeof(x));
        x[1] = startingParameters[i].x;
        x[2] = startingParameters[i].y;
        x[3] = startingParameters[i].z;

        guess_direction(ev,x+1,&x[4],&x[5]);

        x[6] = guess_t0(ev,x+1);

        rv = nlopt_optimize(opt,x,&fval);

        if (fval < *fmin || i == 0) {
            *fmin = fval;
            memcpy(xopt,x,sizeof(x));
        }
    }

    memcpy(x,xopt,sizeof(x));

    /* Reset path coefficients. */
    x[7] = 0.0;
    x[8] = 0.0;

    /* Now, we do the "real" minimization. */
    fpars.epsrel = 1e-4;
    fpars.fast = 0;
    nlopt_set_ftol_abs(opt, 1e-5);
    nlopt_set_maxeval(opt, 1000);

    do {
        *fmin = fval;
        memcpy(xopt,x,sizeof(x));
        rv = nlopt_optimize(opt,x,&fval);
    } while (fval < *fmin && fabs(fval-*fmin) > 1e-5);

    if (fval < *fmin) {
        *fmin = fval;
        memcpy(xopt,x,sizeof(x));
    }

    nlopt_destroy(opt);

    return rv;
}

void usage(void)
{
    fprintf(stderr,"Usage: ./fit [options] FILENAME\n");
    fprintf(stderr,"  -o      output file\n");
    fprintf(stderr,"  -h      display this help message\n");
    exit(1);
}

int main(int argc, char **argv)
{
    int i;
    zebraFile *f;
    bank b;
    int rv;
    PMTBank bpmt;
    EVBank bev;
    event ev = {0};
    int crate, card, channel;
    int id;
    double fmin;
    double xopt[9];
    char *filename = NULL;
    char *output = NULL;
    FILE *fout = NULL;

    for (i = 1; i < argc; i++) {
        if (argv[i][0] == '-') {
            switch (argv[i][1]) {
            case 'o':
                output = argv[++i];
                break;
            case 'h':
                usage();
            default:
                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
                exit(1);
            }
        } else {
            filename = argv[i];
        }
    }

    if (!filename)
        usage();

    f = zebra_open(filename);

    if (!f) {
        fprintf(stderr, "%s\n", zebra_err);
        return 1;
    }

    if (output) {
        fout = fopen(output, "w");

        if (!fout) {
            fprintf(stderr, "failed to open '%s': %s\n", output, strerror(errno));
            return 1;
        }
    }

    load_pmt_info();

    for (i = 0; i < MAX_PMTS; i++) {
        ev.pmt_hits[i].hit = 0;
        ev.pmt_hits[i].flags = 0;
    }

    ev.run = -1;

    init_interpolation();
    init_charge();
    dict *db = db_init();

    if (load_file(db, "DQXX_0000010000.dat")) {
        fprintf(stderr, "failed to load DQXX_0000010000.dat: %s\n", db_err);
        exit(1);
    }

    if (dqxx_init(db, &ev)) {
        fprintf(stderr, "failed to initialize DQXX bank: %s\n", dqxx_err);
        exit(1);
    }

    if (load_file(db, "pmt_response.dat")) {
        fprintf(stderr, "failed to load pmt_response.dat: %s\n", db_err);
        exit(1);
    }

    if (pmt_response_init(db)) {
        fprintf(stderr, "failed to initialize PMTR bank: %s\n", pmtr_err);
        exit(1);
    }

    while (1) {
        rv = next_bank(f, &b);

        if (rv == -1) {
            fprintf(stderr, "error getting bank: %s\n", zebra_err);
            goto err;
        } else if (rv == 1) {
            /* EOF */
            break;
        }

        if (b.name == 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.pn & KPF_DIS;
        } else if (b.name == EV_RECORD) {
            unpack_ev(b.data, &bev);
            if (ev.run != -1) {
                if (ev.run != bev.run || ev.gtid != bev.gtr_id) {
                    /* New event, so we need to fit the old event. */
                    fit_event(&ev,xopt,&fmin);

                    if (fout) {
                        fprintf(fout, "%" PRIu32 " %10.2f %7.2f %7.2f %7.2f %5.2f %5.2f %5.2f %5.2f %5.2f f() = %7.3e\n",
                                ev.gtid,
                                xopt[0],
                                xopt[1],
                                xopt[2],
                                xopt[3],
                                xopt[4],
                                xopt[5],
                                xopt[6],
                                xopt[7],
                                xopt[8],
                                fmin);
                        fflush(fout);
                    }
                }
            }
            ev.run = bev.run;
            ev.gtid = bev.gtr_id;

            for (i = 0; i < MAX_PMTS; i++) {
                ev.pmt_hits[i].hit = 0;
            }
        }
    }

    free_interpolation();
    pmt_response_free();

    db_free(db);

    if (fout) fclose(fout);

    zebra_close(f);

    return 0;

err:

    zebra_close(f);
    return 1;
}