add interp2d() for fast bilinear 2D interpolation

3 files changed, 110 insertions, 0 deletions

diff --git a/src/misc.c b/src/misc.c
index d20bc26..605f3ea 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -3,6 +3,7 @@
 #include <stdlib.h> /* for size_t */
 #include <gsl/gsl_sf_gamma.h>
 #include "vector.h"
+#include <stdio.h>
 
 static struct {
     int n;
@@ -338,6 +339,52 @@ double kahan_sum(double *x, size_t n)
     return sum;
 }
 
+double interp2d(double x, double y, double *xp, double *yp, double *zp, size_t n1, size_t n2)
+{
+    /* A fast bilinear interpolation routine which assumes that the values in
+     * `xp` and `yp` are evenly spaced.
+     *
+     * `zp` should be a 2D array indexed as zp[i,j] = zp[i*n2 + j].
+     *
+     * If x < xp[0], x > xp[n-1], y < yp[0], or y > yp[n-1] prints an error and
+     * exits.
+     *
+     * See https://en.wikipedia.org/wiki/Bilinear_interpolation. */
+    size_t i, j;
+    double q11, q12, q21, q22;
+
+    if (x <= xp[0]) {
+        fprintf(stderr, "x < xp[0]!\n");
+        exit(1);
+    }
+
+    if (y <= yp[0]) {
+        fprintf(stderr, "y < xp[0]!\n");
+        exit(1);
+    }
+
+    i = (x-xp[0])/(xp[1]-xp[0]);
+
+    if (i > n1-2) {
+        fprintf(stderr, "i > n1 - 2!\n");
+        exit(1);
+    }
+
+    j = (y-yp[0])/(yp[1]-yp[0]);
+
+    if (j > n2-2) {
+        fprintf(stderr, "j > n2 - 2!\n");
+        exit(1);
+    }
+
+    q11 = zp[i*n2+j];
+    q12 = zp[i*n2+j+1];
+    q21 = zp[(i+1)*n2+j];
+    q22 = zp[(i+1)*n2+j+1];
+
+    return (q11*(xp[i+1]-x)*(yp[j+1]-y) + q21*(x-xp[i])*(yp[j+1]-y) + q12*(xp[i+1]-x)*(y-yp[j]) + q22*(x-xp[i])*(y-yp[j]))/((xp[i+1]-xp[i])*(yp[i+1]-yp[i]));
+}
+
 double interp1d(double x, double *xp, double *yp, size_t n)
 {
     /* A fast interpolation routine which assumes that the values in `xp` are
diff --git a/src/misc.h b/src/misc.h
index 80a580f..3a55c2f 100644
--- a/src/misc.h
+++ b/src/misc.h
@@ -17,6 +17,7 @@ double ln(unsigned int n);
 double lnfact(unsigned int n);
 double kahan_sum(double *x, size_t n);
 double interp1d(double x, double *xp, double *yp, size_t n);
+double interp2d(double x, double y, double *xp, double *yp, double *zp, size_t n1, size_t n2);
 int isclose(double a, double b, double rel_tol, double abs_tol);
 int allclose(double *a, double *b, size_t n, double rel_tol, double abs_tol);
 double logsumexp(double *a, size_t n);
diff --git a/src/test.c b/src/test.c
index 242e728..f8cedb9 100644
--- a/src/test.c
+++ b/src/test.c
@@ -18,6 +18,7 @@
 #include "proton.h"
 #include "likelihood.h"
 #include "id_particles.h"
+#include <gsl/gsl_spline2d.h>
 
 typedef int testFunction(char *err);
 
@@ -909,6 +910,66 @@ err:
     return 1;
 }
 
+int test_interp2d(char *err)
+{
+    /* Tests the interp2d() function. */
+    size_t i, j;
+    double xp[100], yp[200], zp[20000], zp2[20000], range, x, y, z, expected;
+    gsl_interp_accel *xacc;
+    gsl_interp_accel *yacc;
+    gsl_spline2d *spline;
+
+    range = 1.0;
+
+    init_genrand(0);
+
+    for (i = 0; i < LEN(xp); i++) {
+        xp[i] = range*i/(LEN(xp)-1);
+    }
+
+    for (i = 0; i < LEN(yp); i++) {
+        yp[i] = range*i/(LEN(yp)-1);
+    }
+
+    spline = gsl_spline2d_alloc(gsl_interp2d_bilinear,LEN(xp),LEN(yp));
+
+    for (i = 0; i < LEN(xp); i++) {
+        for (j = 0; j < LEN(yp); j++) {
+            zp[i*LEN(yp)+j] = genrand_real2();
+            gsl_spline2d_set(spline, zp2, i, j, zp[i*LEN(yp)+j]);
+        }
+    }
+
+    xacc = gsl_interp_accel_alloc();
+    yacc = gsl_interp_accel_alloc();
+
+    gsl_spline2d_init(spline,xp,yp,zp2,LEN(xp),LEN(yp));
+
+    for (i = 0; i < 1000; i++) {
+        x = genrand_real2()*range;
+        y = genrand_real2()*range;
+        expected = gsl_spline2d_eval(spline,x,y,xacc,yacc);
+        z = interp2d(x,y,xp,yp,zp,LEN(xp),LEN(yp));
+        if (!isclose(z, expected, 1e-9, 1e-9)) {
+            sprintf(err, "interp2d() returned %.5g, but expected %.5g", z, expected);
+            goto err;
+        }
+    }
+
+    gsl_interp_accel_free(xacc);
+    gsl_interp_accel_free(yacc);
+    gsl_spline2d_free(spline);
+
+    return 0;
+
+err:
+    gsl_interp_accel_free(xacc);
+    gsl_interp_accel_free(yacc);
+    gsl_spline2d_free(spline);
+
+    return 1;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -937,6 +998,7 @@ struct tests {
     {test_proton_get_energy, "test_proton_get_energy"},
     {test_log_norm, "test_log_norm"},
     {test_trapz, "test_trapz"},
+    {test_interp2d, "test_interp2d"},
 };
 
 int main(int argc, char **argv)