3 files changed, 53 insertions, 0 deletions

diff --git a/src/misc.c b/src/misc.c
index f705c74..003b170 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -881,3 +881,25 @@ double fast_acos(double x)
 
     return interp1d(x,xs,ys,LEN(xs));
 }
+
+/* Fast version of sqrt() for `x` between 0 and 1 which uses a lookup table
+ * computed on the first call. */
+double fast_sqrt(double x)
+{
+    size_t i;
+    static int initialized = 0;
+    static double xs[N_SQRT];
+    static double ys[N_SQRT];
+
+    if (!initialized) {
+        for (i = 0; i < LEN(xs); i++) {
+            xs[i] = 1.0*i/(LEN(xs)-1);
+            ys[i] = sqrt(xs[i]);
+        }
+        initialized = 1;
+    }
+
+    if (x > 1.0) return sqrt(x);
+
+    return interp1d(x,xs,ys,LEN(xs));
+}
diff --git a/src/misc.h b/src/misc.h
index 60e54ea..e40c1f5 100644
--- a/src/misc.h
+++ b/src/misc.h
@@ -30,6 +30,9 @@
 /* Number of points to create a lookup table for the fast_acos() function. */
 #define N_ACOS 10000
 
+/* Number of points to create a lookup table for the fast_sqrt() function. */
+#define N_SQRT 10000
+
 double trapz(const double *y, double dx, size_t n);
 int intersect_sphere(double *pos, double *dir, double R, double *l);
 void get_path_length(double *pos1, double *pos2, double R, double *l1, double *l2);
@@ -55,5 +58,6 @@ size_t argmax(double *a, size_t n);
 size_t argmin(double *a, size_t n);
 void get_dir(double *dir, double theta, double phi);
 double fast_acos(double x);
+double fast_sqrt(double x);
 
 #endif
diff --git a/src/test.c b/src/test.c
index b1e0426..c2b7466 100644
--- a/src/test.c
+++ b/src/test.c
@@ -2100,6 +2100,32 @@ err:
     return 1;
 }
 
+int test_fast_sqrt(char *err)
+{
+    /* Tests that the fast_sqrt() function returns values within 0.1% of sqrt(). */
+    size_t i;
+    double x, result, expected;
+
+    init_genrand(0);
+
+    for (i = 0; i < 100; i++) {
+        x = genrand_real2()*2;
+
+        result = fast_sqrt(x);
+        expected = sqrt(x);
+
+        if (!isclose(result, expected, 0, 1e-3)) {
+            sprintf(err, "fast_sqrt() returned %.5g, but expected %.5g", result, expected);
+            goto err;
+        }
+    }
+
+    return 0;
+
+err:
+    return 1;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -2152,6 +2178,7 @@ struct tests {
     {test_argmin, "test_argmin"},
     {test_electron_get_angular_pdf_norm, "test_electron_get_angular_pdf_norm"},
     {test_fast_acos, "test_fast_acos"},
+    {test_fast_sqrt, "test_fast_sqrt"},
 };
 
 int main(int argc, char **argv)