1 files changed, 155 insertions, 9 deletions

diff --git a/src/fit.c b/src/fit.c
index 68d66af..9cf1809 100644
--- a/src/fit.c
+++ b/src/fit.c
@@ -50,6 +50,8 @@
 #include "hdf5.h"
 #include "hdf5_utils.h"
 #include <gsl/gsl_sort.h>
+#include <gsl/gsl_cdf.h>
+#include <gsl/gsl_roots.h>
 
 /* Maximum number of fit parameters. Should be at least 4 + 3*MAX_VERTICES. */
 #define MAX_PARS 100
@@ -5210,17 +5212,151 @@ void guess_direction(event *ev, double *pos, double *theta, double *phi)
     *phi = atan2(dir[1],dir[0]);
 }
 
-double guess_energy(event *ev, double *pos, double *dir)
+/* Returns the expected number of detected photons for a particle with type
+ * `id` and kinetic energy `T0`. `distance` should be the distance to the PSUP
+ * in cm. */
+double get_expected_photons(int id, double T0, double distance, double cos_theta)
+{
+    particle *p;
+    double nphotons, E0, p0, beta0;
+
+    p = particle_init(id, T0, 1000);
+
+    /* Calculate total energy */
+    E0 = T0 + p->mass;
+    p0 = sqrt(E0*E0 - p->mass*p->mass);
+    beta0 = p0/E0;
+
+    /* First we calculate the photons from Cerenkov emission of the particle
+     * itself. */
+    switch (id) {
+    case IDP_E_MINUS:
+    case IDP_E_PLUS:
+    case IDP_MU_MINUS:
+    case IDP_MU_PLUS:
+    case IDP_PROTON:
+        nphotons = get_probability2(beta0)*fmin(p->range,distance);
+        break;
+    default:
+        fprintf(stderr, "unknown particle id %i\n", id);
+        exit(1);
+    }
+
+    /* Now, we add the expected number of shower photons. */
+    nphotons += p->shower_photons*gsl_cdf_gamma_P(distance,p->pos_a,p->pos_b)*(1-interp1d(cos_theta,p->cos_theta,p->cdf_shower,p->n));
+
+    /* Now, we add the expected number of delta ray photons. */
+    if (distance > p->range)
+        nphotons += p->delta_ray_photons*(1-interp1d(cos_theta,p->cos_theta,p->cdf_delta,p->n));
+    else
+        nphotons += p->delta_ray_photons*distance/p->range*(1-interp1d(cos_theta,p->cos_theta,p->cdf_delta,p->n));
+
+    /* Finally we multiply by the average quantum efficiency and divide by 2 to
+     * account for the PMT coverage. */
+    return nphotons*get_weighted_quantum_efficiency()/2.0;
+}
+
+typedef struct guessEnergyParams {
+    int id;
+    double distance;
+    double qhs_sum;
+    double cos_theta;
+} guessEnergyParams;
+
+/* Returns the expected number of photons minus the observed number of PE. This
+ * function is used to bisect the most likely energy using Brent's method. */
+static double get_expected_photons_root(double T0, void *params)
+{
+    guessEnergyParams *pars;
+
+    pars = (guessEnergyParams *) params;
+
+    return get_expected_photons(pars->id,T0,pars->distance,pars->cos_theta) - pars->qhs_sum;
+}
+
+/* Returns the most likely kinetic energy for a given particle hypothesis.
+ * `qhs_sum` should be the sum of the observed charge in units of PE.
+ * `distance` is the distance to the PSUP in cm. `id` is the particle ID.
+ *
+ * On success, returns 0 and `T0` is set to the most likely energy. If there
+ * was an error bisecting the most likely energy, returns -1. */
+int bisect_energy(double qhs_sum, double distance, int id, double cos_theta, double *T0)
+{
+    int status;
+    gsl_root_fsolver *s;
+    gsl_function F;
+    int iter = 0, max_iter = 100;
+    double x_lo, x_hi;
+    double mass, Tmin, Tmax;
+    guessEnergyParams pars;
+    double n_d2o;
+
+    n_d2o = get_index_snoman_d2o(400.0);
+
+    pars.id = id;
+    pars.distance = 840.0;
+    pars.qhs_sum = qhs_sum;
+    pars.cos_theta = cos_theta;
+
+    s = gsl_root_fsolver_alloc(gsl_root_fsolver_brent);
+
+    F.function = &get_expected_photons_root;
+    F.params = &pars;
+
+    switch (id) {
+    case IDP_E_MINUS:
+    case IDP_E_PLUS:
+        mass = ELECTRON_MASS;
+        Tmax = electron_get_max_energy();
+        break;
+    case IDP_MU_MINUS:
+    case IDP_MU_PLUS:
+        mass = MUON_MASS;
+        Tmax = muon_get_max_energy();
+        break;
+    case IDP_PROTON:
+        mass = PROTON_MASS;
+        Tmax = proton_get_max_energy();
+        break;
+    default:
+        fprintf(stderr, "unknown particle id %i\n", id);
+        exit(1);
+    }
+
+    /* Calculate the Cerenkov threshold for a muon. */
+    Tmin = mass/sqrt(1.0-1.0/(n_d2o*n_d2o)) - mass;
+
+    gsl_root_fsolver_set(s, &F, Tmin, Tmax);
+
+    do {
+        iter++;
+        status = gsl_root_fsolver_iterate(s);
+        *T0 = gsl_root_fsolver_root(s);
+        x_lo = gsl_root_fsolver_x_lower(s);
+        x_hi = gsl_root_fsolver_x_upper(s);
+
+        /* Find the root to within 1e-10 MeV. */
+        status = gsl_root_test_interval(x_lo, x_hi, 1e-10, 0);
+
+        if (status == GSL_SUCCESS) break;
+    } while (status == GSL_CONTINUE && iter < max_iter);
+
+    gsl_root_fsolver_free(s);
+
+    return status == GSL_SUCCESS ? 0 : -1;
+}
+
+/* Returns the approximate energy for a particle at position `pos` in direction
+ * `dir`. This guess is made by summing up all the charge within a Cerenkov
+ * cone and then bisecting the energy at which that many photons would be
+ * produced. */
+double guess_energy(event *ev, double *pos, double *dir, int id)
 {
-    /* Returns the approximate energy for a particle at position `pos` in
-     * direction `dir`. This guess is made by summing up all the charge within
-     * a Cerenkov cone and making the approximation that the particle creates
-     * approximately 6 photons/MeV.
-     *
-     * FIXME: Should update this to something better. */
     size_t i;
     double qhs_sum, n_d2o;
     double pmt_dir[3];
+    double T0;
+    double distance_to_psup;
 
     n_d2o = get_index_snoman_d2o(400.0);
 
@@ -5233,7 +5369,17 @@ double guess_energy(event *ev, double *pos, double *dir)
             qhs_sum += ev->pmt_hits[i].q;
     }
 
-    return qhs_sum/6.0;
+    if (NORM(pos) > PSUP_RADIUS || !intersect_sphere(pos,dir,PSUP_RADIUS,&distance_to_psup)) {
+        /* If the particle is outside the PSUP or it doesn't intersect the
+         * PSUP we just assume it produces no light. */
+        fprintf(stderr, "guess_energy: particle doesn't intersect PSUP!\n");
+        distance_to_psup = PSUP_RADIUS;
+    }
+
+    if (bisect_energy(qhs_sum, distance_to_psup, id, 1/n_d2o, &T0))
+        return qhs_sum/6.0;
+
+    return T0;
 }
 
 int fit_event2(event *ev, double *xopt, double *fmin, int *id, size_t n, double maxtime)
@@ -5373,7 +5519,7 @@ int fit_event2(event *ev, double *xopt, double *fmin, int *id, size_t n, double
             dir[0] = sin(peak_theta[result[i*n+j]])*cos(peak_phi[result[i*n+j]]);
             dir[1] = sin(peak_theta[result[i*n+j]])*sin(peak_phi[result[i*n+j]]);
             dir[2] = cos(peak_theta[result[i*n+j]]);
-            T0 = guess_energy(ev,x0,dir);
+            T0 = guess_energy(ev,x0,dir,id[j]);
 
             switch (id[j]) {
             case IDP_E_MINUS: