add function to compute unique direction vectors for a multi particle fit

3 files changed, 300 insertions, 0 deletions

diff --git a/src/misc.c b/src/misc.c
index f6e11dd..183f0ce 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -4,6 +4,7 @@
 #include <gsl/gsl_sf_gamma.h>
 #include "vector.h"
 #include <stdio.h>
+#include <stdlib.h>
 
 static struct {
     int n;
@@ -514,3 +515,172 @@ double gamma_pdf(double x, double k, double theta)
      * See https://en.wikipedia.org/wiki/Gamma_distribution. */
     return pow(x,k-1)*exp(-x/theta)/(gsl_sf_gamma(k)*pow(theta,k));
 }
+
+size_t ipow(size_t base, size_t exp)
+{
+    /* Returns base^exp for positive integers `base` and `exp`.
+     *
+     * See https://stackoverflow.com/questions/101439. */
+    size_t result = 1;
+    for (;;) {
+        if (exp & 1)
+            result *= base;
+        exp >>= 1;
+        if (!exp) break;
+        base *= base;
+    }
+
+    return result;
+}
+
+void product(size_t n, size_t r, size_t *result)
+{
+    /* Returns the cartesian product of [1..n] with itself `r` times.
+     *
+     * The resulting array can be indexed as:
+     *
+     *     result[i*r + j]
+     *
+     * where i is the ith product and j is the jth element in the product. For example:
+     *
+     *     size_t result[4];
+     *     product(2,2,result);
+     *     for (i = 0; i < 2; i++)
+     *         print("%zu %zu\n", result[i*2], result[i*2+1]);
+     *
+     * will print
+     *
+     *     0 0
+     *     0 1
+     *     1 0
+     *     1 1
+     *
+     * `result` should be an array with at least `r`*`n`^`r` elements. */
+    size_t i, j;
+
+    for (i = 0; i < r; i++) {
+        for (j = 0; j < ipow(n,r); j++) {
+            result[j*r+i] = (j/ipow(n,r-i-1)) % n;
+        }
+    }
+}
+
+typedef struct direction {
+    int id;
+    size_t index;
+} direction;
+
+static int direction_compare(const void *a, const void *b)
+{
+    const direction *da = (direction *) a;
+    const direction *db = (direction *) b;
+
+    if (da->id > db->id)
+        return 1;
+    else if (da->id < db->id)
+        return -1;
+    else if (da->index > db->index)
+        return 1;
+    else if (da->index < db->index)
+        return -1;
+
+    return 0;
+}
+
+void unique_vertices(int *id, size_t n, size_t npeaks, size_t *result, size_t *nvertices)
+{
+    /* Returns the set of all unique vertices for `n` particles with particle
+     * ids `id` for `npeaks` possible direction vectors. For example, the
+     * unique vertices for two electrons and one muon given 2 possible
+     * directions 1 and 2 would be:
+     *
+     *     1 1 1
+     *     1 1 2
+     *     1 2 1
+     *     1 2 2
+     *     2 2 1
+     *     2 2 2
+     *
+     * where the first column represents the direction for the first electron,
+     * the second for the second electron, and the third for the muon. This is
+     * different from the cartesian product since rows like
+     *
+     *     2 1 1
+     *
+     * and
+     *
+     *     2 1 2
+     *
+     * are duplicates.
+     *
+     * The resulting array can be indexed as:
+     *
+     *     result[i*n + j]
+     *
+     * where i is the ith product and j is the jth direction in the product.
+     * For example:
+     *
+     *     int id[3] = {IDP_E_MINUS, IDP_E_MINUS, IDP_MU_MINUS};
+     *
+     *     size_t result[24];
+     *     size_t nvertices;
+     *     unique_vertices(id, LEN(id), 2, result, &nvertices);
+     *     for (i = 0; i < nvertices; i++)
+     *         print("%zu %zu %zu\n", result[i*3], result[i*3+1], result[i*3+2);
+     *
+     * `result` should be an array with at least `n`*`npeaks`^`n` elements. */
+    size_t i, j, k;
+    direction *ordered_results;
+    size_t *results2 = malloc(sizeof(size_t)*n*ipow(npeaks,n));
+    int unique, equal;
+
+    product(npeaks,n,results2);
+
+    ordered_results = malloc(sizeof(direction)*n*ipow(npeaks,n));
+    direction *unique_results = malloc(sizeof(direction)*n*ipow(npeaks,n));
+
+    for (i = 0; i < ipow(npeaks,n); i++) {
+        for (j = 0; j < n; j++) {
+            ordered_results[i*n+j].id = id[j];
+            ordered_results[i*n+j].index = results2[i*n+j];
+        }
+        /* Sort the vertices in each row so we can compare them. */
+        qsort(ordered_results+i*n,n,sizeof(direction),direction_compare);
+    }
+
+    *nvertices = 0;
+    for (i = 0; i < ipow(npeaks,n); i++) {
+        unique = 1;
+        for (j = 0; j < *nvertices; j++) {
+            equal = 1;
+            for (k = 0; k < n; k++) {
+                if ((unique_results[j*n+k].id != ordered_results[i*n+k].id) || (unique_results[j*n+k].index != ordered_results[i*n+k].index)) {
+                    equal = 0;
+                    break;
+                }
+            }
+            if (equal) {
+                unique = 0;
+                break;
+            }
+        }
+
+        if (unique) {
+            for (j = 0; j < n; j++) {
+                unique_results[(*nvertices)*n+j].id = ordered_results[i*n+j].id;
+                unique_results[(*nvertices)*n+j].index = ordered_results[i*n+j].index;
+            }
+            *nvertices += 1;
+        }
+    }
+
+    for (i = 0; i < *nvertices; i++) {
+        for (j = 0; j < n; j++) {
+            result[i*n+j] = unique_results[i*n+j].index;
+        }
+    }
+
+    free(results2);
+    free(ordered_results);
+    free(unique_results);
+}
diff --git a/src/misc.h b/src/misc.h
index 94ccee7..8450454 100644
--- a/src/misc.h
+++ b/src/misc.h
@@ -27,5 +27,8 @@ double norm_cdf(double x, double mu, double sigma);
 double mean(const double *x, size_t n);
 double std(const double *x, size_t n);
 double gamma_pdf(double x, double k, double theta);
+size_t ipow(size_t base, size_t exp);
+void product(size_t n, size_t r, size_t *result);
+void unique_vertices(int *id, size_t n, size_t npeaks, size_t *result, size_t *nvertices);
 
 #endif
diff --git a/src/test.c b/src/test.c
index 335d473..6a0bc25 100644
--- a/src/test.c
+++ b/src/test.c
@@ -1633,6 +1633,130 @@ int test_find_peaks_array(char *err)
     return 0;
 }
 
+int test_ipow(char *err)
+{
+    /* Tests the ipow() function. */
+    size_t result;
+
+    result = ipow(2,2);
+
+    if (result != 4) {
+        sprintf(err, "ipow(2,2) returned %zu, but expected %i", result, 4);
+        return 1;
+    }
+
+    result = ipow(2,3);
+
+    if (result != 8) {
+        sprintf(err, "ipow(2,3) returned %zu, but expected %i", result, 8);
+        return 1;
+    }
+
+    return 0;
+}
+
+int test_product(char *err)
+{
+    /* Tests the product() function. */
+    size_t i, j;
+    size_t result[1000];
+
+    size_t expected1[4][2] = {{0,0},{0,1},{1,0},{1,1}};
+
+    product(2,2,result);
+
+    for (i = 0; i < 4; i++) {
+        for (j = 0; j < 2; j++) {
+            if (result[i*2+j] != expected1[i][j]) {
+                sprintf(err, "result[%zu,%zu] = %zu, but expected %zu", i, j, result[i*2+j], expected1[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    size_t expected2[8][3] = {{0,0,0},{0,0,1},{0,1,0},{0,1,1},{1,0,0},{1,0,1},{1,1,0},{1,1,1}};
+
+    product(2,3,result);
+
+    for (i = 0; i < 8; i++) {
+        for (j = 0; j < 3; j++) {
+            if (result[i*3+j] != expected2[i][j]) {
+                sprintf(err, "result[%zu,%zu] = %zu, but expected %zu", i, j, result[i*3+j], expected2[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    return 0;
+}
+
+int test_unique_vertices(char *err)
+{
+    /* Tests the unique_vertices() function. */
+    size_t i, j;
+    size_t result[1000];
+    size_t nvertices;
+
+    int id1[2] = {IDP_E_MINUS, IDP_E_MINUS};
+    size_t expected1[3][2] = {{0,0},{0,1},{1,1}};
+
+    unique_vertices(id1,LEN(id1),2,result,&nvertices);
+
+    if (nvertices != 3) {
+        sprintf(err, "unique vertices returned nvertices = %zu, but expected %i", nvertices, 3);
+        return 1;
+    }
+
+    for (i = 0; i < nvertices; i++) {
+        for (j = 0; j < 2; j++) {
+            if (result[i*2+j] != expected1[i][j]) {
+                sprintf(err, "result[%zu,%zu] = %zu, but expected %zu", i, j, result[i*2+j], expected1[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    int id2[3] = {IDP_E_MINUS, IDP_MU_MINUS, IDP_TAU_MINUS};
+    size_t expected2[8][3] = {{0,0,0},{0,0,1},{0,1,0},{0,1,1},{1,0,0},{1,0,1},{1,1,0},{1,1,1}};
+
+    unique_vertices(id2,LEN(id2),2,result,&nvertices);
+
+    if (nvertices != 8) {
+        sprintf(err, "unique vertices returned nvertices = %zu, but expected %i", nvertices, 8);
+        return 1;
+    }
+
+    for (i = 0; i < 8; i++) {
+        for (j = 0; j < 3; j++) {
+            if (result[i*3+j] != expected2[i][j]) {
+                sprintf(err, "result[%zu,%zu] = %zu, but expected %zu", i, j, result[i*3+j], expected2[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    int id3[3] = {IDP_E_MINUS, IDP_E_MINUS, IDP_MU_MINUS};
+    size_t expected3[6][3] = {{0,0,0},{0,0,1},{0,1,0},{0,1,1},{1,1,0},{1,1,1}};
+
+    unique_vertices(id3,LEN(id3),2,result,&nvertices);
+
+    if (nvertices != 6) {
+        sprintf(err, "unique vertices returned nvertices = %zu, but expected %i", nvertices, 6);
+        return 1;
+    }
+
+    for (i = 0; i < nvertices; i++) {
+        for (j = 0; j < 3; j++) {
+            if (result[i*3+j] != expected3[i][j]) {
+                sprintf(err, "result[%zu,%zu] = %zu, but expected %zu", i, j, result[i*3+j], expected3[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    return 0;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -1674,6 +1798,9 @@ struct tests {
     {test_quad, "test_quad"},
     {test_quad_noise, "test_quad_noise"},
     {test_find_peaks_array, "test_find_peaks_array"},
+    {test_ipow, "test_ipow"},
+    {test_product, "test_product"},
+    {test_unique_vertices, "test_unique_vertices"},
 };
 
 int main(int argc, char **argv)