add code to expand the track of a particle using a KL expansion

To fit the path of muons and electrons I use the Karhunen-Loeve expansion of a
random 2D walk in the polar angle in x and y. This allows you to decompose the
path into a sum over sine functions whose coefficients become random variables.
The nice thing about fitting the path in this way is that you can capture
*most* of the variation in the path using a small number of variables by only
summing over the first N terms in the expansion and it is easy to calculate the
probability of the coefficients since they are all uncorrelated.

7 files changed, 470 insertions, 3 deletions

diff --git a/src/Makefile b/src/Makefile
index ae8807b..1ea1685 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -1,7 +1,7 @@
 CFLAGS=-O0 -Wall -g -DSWAP_BYTES
 LDLIBS=-lm -lgsl -lgslcblas -lnlopt_cxx -lstdc++
 
-all: test test-vector test-likelihood fit test-charge
+all: test test-vector test-likelihood fit test-charge test-path
 
 Makefile.dep:
 	-$(CC) -MM *.c > Makefile.dep
@@ -12,12 +12,14 @@ 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
+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
 
 test-vector: test-vector.o vector.o mt19937ar.o
 
 test-likelihood: test-likelihood.o muon.o random.o optics.o quantum_efficiency.o mt19937ar.o pdg.o vector.o solid_angle.o scattering.o
 
+test-path: test-path.o mt19937ar.o vector.o path.o random.o
+
 test-charge: test-charge.o sno_charge.o misc.o
 
 test-zebra: test-zebra.o zebra.o pack2b.o
diff --git a/src/path.c b/src/path.c
new file mode 100644
index 0000000..f09ec45
--- /dev/null
+++ b/src/path.c
@@ -0,0 +1,161 @@
+#include "path.h"
+#include <math.h>
+#include <gsl/gsl_spline.h>
+#include <stddef.h> /* for size_t */
+#include <gsl/gsl_integration.h>
+#include "pdg.h"
+#include "vector.h"
+#include <stdlib.h> /* for malloc(), calloc(), etc. */
+
+#define N 100
+
+static double foo(double s, double range, double theta0, int k)
+{
+    return sqrt(2*pow(theta0,2)*range)*sin(M_PI*s*(k-0.5)/range)/(M_PI*(k-0.5));
+}
+
+double path_get_coefficient(unsigned int k, double *s, double *x, double theta0, size_t n)
+{
+    size_t i;
+    double sum, range;
+
+    range = s[n-1];
+
+    sum = 0.0;
+    for (i = 1; i < n; i++) {
+        sum += (foo(s[i],range,theta0,k)*x[i] + foo(s[i-1],range,theta0,k)*x[i-1])*(s[i]-s[i-1])/2.0;
+    }
+
+    return sum*pow(k-0.5,2)*pow(M_PI,2)/pow(theta0*range,2);
+}
+
+path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, double *z1, double *z2, size_t n, double m)
+{
+    size_t i, j;
+    double E, mom, beta, theta, phi;
+
+    path *p = malloc(sizeof(path));
+
+    p->pos[0] = pos[0];
+    p->pos[1] = pos[1];
+    p->pos[2] = pos[2];
+
+    p->dir[0] = dir[0];
+    p->dir[1] = dir[1];
+    p->dir[2] = dir[2];
+
+    double *s = malloc(sizeof(double)*N);
+    double *theta1 = calloc(N,sizeof(double));
+    double *theta2 = calloc(N,sizeof(double));
+    double *x = calloc(N,sizeof(double));
+    double *y = calloc(N,sizeof(double));
+    double *z = calloc(N,sizeof(double));
+    double *T = calloc(N,sizeof(double));
+    double *t = calloc(N,sizeof(double));
+    double *dx = calloc(N,sizeof(double));
+    double *dy = calloc(N,sizeof(double));
+    double *dz = calloc(N,sizeof(double));
+
+    dz[0] = 1.0;
+    for (i = 0; i < N; i++) {
+        s[i] = range*i/(N-1);
+        for (j = 0; j < n; j++) {
+            theta1[i] += z1[j]*foo(s[i],range,theta0,j+1);
+            theta2[i] += z2[j]*foo(s[i],range,theta0,j+1);
+        }
+        T[i] = fun(s[i],T0);
+        if (i > 0) {
+            theta = sqrt(theta1[i]*theta1[i] + theta2[i]*theta2[i]);
+            phi = atan2(theta2[i],theta1[i]);
+
+            dx[i] = (s[i]-s[i-1])*sin(theta)*cos(phi);
+            dy[i] = (s[i]-s[i-1])*sin(theta)*sin(phi);
+            dz[i] = (s[i]-s[i-1])*cos(theta);
+
+            /* Calculate total energy */
+            E = T[i] + m;
+            mom = sqrt(E*E - m*m);
+            beta = mom/E;
+
+            t[i] = t[i-1] + (s[i]-s[i-1])/(beta*SPEED_OF_LIGHT);
+
+            x[i] = x[i-1] + dx[i];
+            y[i] = y[i-1] + dy[i];
+            z[i] = z[i-1] + dz[i];
+        }
+    }
+
+    for (i = 0; i < 8; i++) {
+        p->acc[i] = gsl_interp_accel_alloc();
+        p->spline[i] = gsl_spline_alloc(gsl_interp_linear,N);
+    }
+
+    gsl_spline_init(p->spline[0],s,x,N);
+    gsl_spline_init(p->spline[1],s,y,N);
+    gsl_spline_init(p->spline[2],s,z,N);
+    gsl_spline_init(p->spline[3],s,T,N);
+    gsl_spline_init(p->spline[4],s,t,N);
+    gsl_spline_init(p->spline[5],s,dx,N);
+    gsl_spline_init(p->spline[6],s,dy,N);
+    gsl_spline_init(p->spline[7],s,dz,N);
+
+    free(s);
+    free(theta1);
+    free(theta2);
+    free(x);
+    free(y);
+    free(z);
+    free(T);
+    free(t);
+    free(dx);
+    free(dy);
+    free(dz);
+
+    return p;
+}
+
+void path_eval(path *p, double s, double *pos, double *dir, double *T, double *t)
+{
+    double normal[3], k[3], tmp[3], phi;
+
+    tmp[0] = gsl_spline_eval(p->spline[0],s,p->acc[0]);
+    tmp[1] = gsl_spline_eval(p->spline[1],s,p->acc[1]);
+    tmp[2] = gsl_spline_eval(p->spline[2],s,p->acc[2]);
+
+    *T = gsl_spline_eval(p->spline[3],s,p->acc[3]);
+    *t = gsl_spline_eval(p->spline[4],s,p->acc[4]);
+
+    k[0] = 0.0;
+    k[1] = 0.0;
+    k[2] = 1.0;
+
+    CROSS(normal,k,p->dir);
+
+    normalize(normal);
+
+    phi = acos(DOT(k,p->dir));
+
+    rotate(pos,tmp,normal,phi);
+
+    ADD(pos,pos,p->pos);
+
+    tmp[0] = gsl_spline_eval(p->spline[5],s,p->acc[5]);
+    tmp[1] = gsl_spline_eval(p->spline[6],s,p->acc[6]);
+    tmp[2] = gsl_spline_eval(p->spline[7],s,p->acc[7]);
+
+    normalize(tmp);
+
+    rotate(dir,tmp,normal,phi);
+}
+
+void path_free(path *p)
+{
+    size_t i;
+
+    for (i = 0; i < 5; i++) {
+        gsl_interp_accel_free(p->acc[i]);
+        gsl_spline_free(p->spline[i]);
+    }
+
+    free(p);
+}
diff --git a/src/path.h b/src/path.h
new file mode 100644
index 0000000..8aaa024
--- /dev/null
+++ b/src/path.h
@@ -0,0 +1,21 @@
+#ifndef PATH_H
+#define PATH_H
+
+#include <gsl/gsl_spline.h>
+
+typedef double getKineticEnergyFunc(double x, double T0);
+
+typedef struct path {
+    double pos[3];
+    double dir[3];
+    double theta0;
+    gsl_interp_accel *acc[8];
+    gsl_spline *spline[8];
+} path;
+
+double path_get_coefficient(unsigned int k, double *s, double *x, double theta0, size_t n);
+path *path_init(double *pos, double *dir, double T0, double range, double theta0, getKineticEnergyFunc *fun, double *z1, double *z2, size_t n, double m);
+void path_eval(path *p, double s, double *pos, double *dir, double *T, double *t);
+void path_free(path *p);
+
+#endif
diff --git a/src/random.c b/src/random.c
index f6cf85e..e2e5641 100644
--- a/src/random.c
+++ b/src/random.c
@@ -12,3 +12,19 @@ double randn(void)
 
     return sqrt(-2*log(u1))*cos(2*M_PI*u2);
 }
+
+void rand_sphere(double *dir)
+{
+    /* Generates a random point on the unit sphere. */
+    double u, v, theta, phi;
+
+    u = genrand_real1();
+    v = genrand_real1();
+
+    phi = 2*M_PI*u;
+    theta = acos(2*v-1);
+
+    dir[0] = sin(theta)*cos(phi);
+    dir[1] = sin(theta)*sin(phi);
+    dir[2] = cos(theta);
+}
diff --git a/src/random.h b/src/random.h
index 309cd71..733279c 100644
--- a/src/random.h
+++ b/src/random.h
@@ -4,5 +4,6 @@
 #include "mt19937ar.h"
 
 double randn(void);
+void rand_sphere(double *dir);
 
 #endif
diff --git a/src/test-path.c b/src/test-path.c
new file mode 100644
index 0000000..fc46908
--- /dev/null
+++ b/src/test-path.c
@@ -0,0 +1,185 @@
+#include <stdio.h>
+#include "path.h"
+#include <stddef.h> /* for size_t */
+#include "mt19937ar.h"
+#include <math.h> /* for sin(), cos(), etc. */
+#include "vector.h"
+#include <errno.h> /* for errno */
+#include <string.h> /* for strerror() */
+#include "random.h"
+
+typedef struct sim {
+    double *pos;
+    path *p;
+} sim;
+
+static double getKineticEnergy(double x, double T0)
+{
+    return 1.0;
+}
+
+sim *simulate_path(double *pos0, double *dir0, double range, double theta0, size_t N, size_t n)
+{
+    size_t i;
+    double *s, *theta1, *theta2, *x, *y, *z, *z1, *z2, tmp1[3], tmp2[3], k[3], normal[3];
+    double dx, dy, dz, theta, phi;
+
+    sim *simvalue = malloc(sizeof(sim));
+
+    s = calloc(N,sizeof(double));
+    theta1 = calloc(N,sizeof(double));
+    theta2 = calloc(N,sizeof(double));
+    x = calloc(N,sizeof(double));
+    y = calloc(N,sizeof(double));
+    z = calloc(N,sizeof(double));
+    simvalue->pos = calloc(N*3,sizeof(double));
+    z1 = calloc(n,sizeof(double));
+    z2 = calloc(n,sizeof(double));
+
+    init_genrand(0);
+
+    simvalue->pos[0] = pos0[0];
+    simvalue->pos[1] = pos0[1];
+    simvalue->pos[2] = pos0[2];
+
+    for (i = 1; i < N; i++) {
+        s[i] = range*i/(N-1);
+
+        theta1[i] = theta1[i-1] + randn()*theta0*sqrt(s[i]-s[i-1]);
+        theta2[i] = theta2[i-1] + randn()*theta0*sqrt(s[i]-s[i-1]);
+
+        theta = sqrt(theta1[i]*theta1[i] + theta2[i]*theta2[i]);
+        phi = atan2(theta2[i],theta1[i]);
+
+        dx = (s[i]-s[i-1])*sin(theta)*cos(phi);
+        dy = (s[i]-s[i-1])*sin(theta)*sin(phi);
+        dz = (s[i]-s[i-1])*cos(theta);
+
+        x[i] = x[i-1] + dx;
+        y[i] = y[i-1] + dy;
+        z[i] = z[i-1] + dz;
+
+        tmp1[0] = x[i];
+        tmp1[1] = y[i];
+        tmp1[2] = z[i];
+
+        k[0] = 0.0;
+        k[1] = 0.0;
+        k[2] = 1.0;
+
+        CROSS(normal,k,dir0);
+
+        normalize(normal);
+
+        phi = acos(DOT(k,dir0));
+
+        rotate(tmp2,tmp1,normal,phi);
+
+        simvalue->pos[i*3] = pos0[0] + tmp2[0];
+        simvalue->pos[i*3+1] = pos0[1] + tmp2[1];
+        simvalue->pos[i*3+2] = pos0[2] + tmp2[2];
+    }
+
+    for (i = 0; i < n; i++) {
+        z1[i] = path_get_coefficient(i+1,s,theta1,theta0,N);
+        z2[i] = path_get_coefficient(i+1,s,theta2,theta0,N);
+    }
+
+    simvalue->p = path_init(pos0,dir0,1.0,range,theta0,getKineticEnergy,z1,z2,n,1.0);
+
+    free(s);
+    free(theta1);
+    free(theta2);
+    free(x);
+    free(y);
+    free(z);
+    free(z1);
+    free(z2);
+
+    return simvalue;
+}
+
+void usage(void)
+{
+    fprintf(stderr,"Usage: ./test-path [options]\n");
+    fprintf(stderr,"  -N      number of points along track\n");
+    fprintf(stderr,"  -n      number of terms in KL expansion\n");
+    fprintf(stderr,"  -t      standard deviation of angular distribution\n");
+    fprintf(stderr,"  --range range of track\n");
+    fprintf(stderr,"  -h      display this help message\n");
+    exit(1);
+}
+
+int main(int argc, char **argv)
+{
+    size_t i, n, N;
+    double pos0[3], dir0[3], theta0, range, pos2[3], dir[3], T, t;
+
+    n = 2;
+    N = 1000;
+    theta0 = 0.1;
+    range = 100.0;
+
+    for (i = 1; i < argc; i++) {
+        if (!strncmp(argv[i], "--", 2)) {
+            if (!strcmp(argv[i]+2,"range")) {
+                range = strtod(argv[++i],NULL);
+                continue;
+            }
+        } else if (argv[i][0] == '-') {
+            switch (argv[i][1]) {
+            case 'N':
+                N = atoi(argv[++i]);
+                break;
+            case 'n':
+                n = atoi(argv[++i]);
+                break;
+            case 't':
+                theta0 = strtod(argv[++i],NULL);
+                break;
+            case 'h':
+                usage();
+            default:
+                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
+                exit(1);
+            }
+        }
+    }
+
+    pos0[0] = 0.0;
+    pos0[1] = 0.0;
+    pos0[2] = 0.0;
+
+    dir0[0] = 1.0;
+    dir0[1] = 0.0;
+    dir0[2] = 0.0;
+
+    sim *simvalue = simulate_path(pos0,dir0,range,theta0,N,n);
+
+    FILE *pipe = popen("graph -T X --bitmap-size 2000x2000 -X X -Y Y", "w");
+
+    if (!pipe) {
+        fprintf(stderr, "error running graph command: %s\n", strerror(errno));
+        exit(1);
+    }
+
+    for (i = 0; i < N; i++) {
+        fprintf(pipe,"%.10g %.10g\n", simvalue->pos[i*3], simvalue->pos[i*3+1]);
+        path_eval(simvalue->p,i*range/9999.0,pos2,dir,&T,&t);
+    }
+    fprintf(pipe,"\n\n");
+
+    for (i = 0; i < N; i++) {
+        path_eval(simvalue->p,i*range/(N-1),pos2,dir,&T,&t);
+        fprintf(pipe,"%.10g %.10g\n", pos2[0], pos2[1]);
+    }
+    fprintf(pipe,"\n\n");
+
+    if (pclose(pipe)) {
+        fprintf(stderr, "error closing graph command: %s\n", strerror(errno));
+        exit(1);
+    }
+
+    return 0;
+}
+
diff --git a/src/test.c b/src/test.c
index 747e664..b2faedc 100644
--- a/src/test.c
+++ b/src/test.c
@@ -8,6 +8,9 @@
 #include "misc.h"
 #include "sno_charge.h"
 #include <gsl/gsl_integration.h>
+#include "path.h"
+#include "random.h"
+#include "pdg.h"
 
 typedef int testFunction(char *err);
 
@@ -324,6 +327,83 @@ int test_logsumexp(char *err)
     return 0;
 }
 
+double getKineticEnergy(double x, double T0)
+{
+    return 1.0;
+}
+
+int test_path(char *err)
+{
+    /* Tests the logsumexp() function. */
+    int i, j, k;
+    double pos0[3], dir0[3], T0, range, m;
+    double T, t;
+    double pos_expected[3], t_expected;
+    double pos[3], dir[3];
+    double E, mom, beta;
+    double s;
+
+    T0 = 1.0;
+    range = 1.0;
+    m = 1.0;
+
+    init_genrand(0);
+
+    for (i = 0; i < 100; i++) {
+        pos0[0] = genrand_real2();
+        pos0[1] = genrand_real2();
+        pos0[2] = genrand_real2();
+
+        rand_sphere(dir0);
+
+        path *p = path_init(pos0,dir0,T0,range,0.1,getKineticEnergy,NULL,NULL,0,m);
+
+        for (j = 0; j < 100; j++) {
+            s = range*j/99;
+            pos_expected[0] = pos0[0] + dir0[0]*s;
+            pos_expected[1] = pos0[1] + dir0[1]*s;
+            pos_expected[2] = pos0[2] + dir0[2]*s;
+
+            /* Calculate total energy */
+            E = T0 + m;
+            mom = sqrt(E*E - m*m);
+            beta = mom/E;
+
+            t_expected = s/(beta*SPEED_OF_LIGHT);
+
+            path_eval(p,s,pos,dir,&T,&t);
+
+            for (k = 0; k < 3; k++) {
+                if (!isclose(pos[k], pos_expected[k], 1e-9, 1e-9)) {
+                    sprintf(err, "path_eval(%.2g) returned pos[%i] = %.5g, but expected %.5g", s, k, pos[k], pos_expected[k]);
+                    return 1;
+                }
+            }
+
+            for (k = 0; k < 3; k++) {
+                if (!isclose(dir[k], dir0[k], 1e-9, 1e-9)) {
+                    sprintf(err, "path_eval(%.2g) returned dir[%i] = %.5g, but expected %.5g", s, k, dir[k], dir0[k]);
+                    return 1;
+                }
+            }
+
+            if (!isclose(T, 1.0, 1e-9, 1e-9)) {
+                sprintf(err, "path_eval(%.2g) returned T = %.5g, but expected %.5g", s, T, 1.0);
+                return 1;
+            }
+
+            if (!isclose(t, t_expected, 1e-9, 1e-9)) {
+                sprintf(err, "path_eval(%.2g) returned t = %.5g, but expected %.5g", s, t, t_expected);
+                return 1;
+            }
+        }
+
+        path_free(p);
+    }
+
+    return 0;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -335,7 +415,8 @@ struct tests {
     {test_muon_get_range, "test_muon_get_range"},
     {test_muon_get_T, "test_muon_get_T"},
     {test_logsumexp, "test_logsumexp"},
-    {test_sno_charge_integral, "test_sno_charge_integral"}
+    {test_sno_charge_integral, "test_sno_charge_integral"},
+    {test_path, "test_path"}
 };
 
 int main(int argc, char **argv)