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

1 files changed, 170 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);
+}