interp: add 2D Lagrangian interpolation in C

3 files changed, 144 insertions, 0 deletions

diff --git a/_interp_c.c b/_interp_c.c
new file mode 100644
index 0000000..f2269c8
--- /dev/null
+++ b/_interp_c.c
@@ -0,0 +1,21 @@
+#include <math.h>
+
+#define STENCIL 2
+#include "_interp_c_template.c"
+#undef STENCIL
+
+#define STENCIL 4
+#include "_interp_c_template.c"
+#undef STENCIL
+
+#define STENCIL 6
+#include "_interp_c_template.c"
+#undef STENCIL
+
+#define STENCIL 8
+#include "_interp_c_template.c"
+#undef STENCIL
+
+#define STENCIL 10
+#include "_interp_c_template.c"
+#undef STENCIL
diff --git a/_interp_c_template.c b/_interp_c_template.c
new file mode 100644
index 0000000..e1c5f8a
--- /dev/null
+++ b/_interp_c_template.c
@@ -0,0 +1,59 @@
+
+#define JOIN2(a, b) a ## _ ## b
+#define FUNC2(name, stencil) JOIN2(name, stencil)
+#define FUNC(name) FUNC2(name, STENCIL)
+
+void FUNC(interp2d)(const double *src_start, const double *src_step,
+                    const double *src_val,
+                    const int *src_stride, const int *src_len,
+                    int nb_points, const double *x_vals, const double *y_vals,
+                    double *dst)
+{
+    for (int pt = 0; pt < nb_points; pt++) {
+        const double x = x_vals[pt];
+        const double y = y_vals[pt];
+
+        const int    idx_x = (x - src_start[0]) / src_step[0] - (STENCIL / 2 - 1);
+        const int    idx_y = (y - src_start[1]) / src_step[1] - (STENCIL / 2 - 1);
+
+        const double src_x = idx_x * src_step[0] + src_start[0];
+        const double src_y = idx_y * src_step[1] + src_start[1];
+
+        double fact[2][STENCIL];
+
+        double ret = 0.0;
+
+        if (idx_x < 0 || idx_x > src_len[0] - STENCIL ||
+            idx_y < 0 || idx_y > src_len[1] - STENCIL) {
+            ret = NAN;
+            goto output;
+        }
+
+        for (int i = 0; i < STENCIL; i++) {
+            double fact_x = 1.0;
+            double fact_y = 1.0;
+
+            for (int j = 0; j < STENCIL; j++) {
+                if (i == j)
+                    continue;
+
+                fact_x *= (x - (src_x + j * src_step[0])) / (src_step[0] * (i - j));
+                fact_y *= (y - (src_y + j * src_step[1])) / (src_step[1] * (i - j));
+            }
+
+            fact[0][i] = fact_x;
+            fact[1][i] = fact_y;
+        }
+
+        for (int j = 0; j < STENCIL; j++) {
+            double val = 0.0;
+
+            for (int i = 0; i < STENCIL; i++)
+                val += fact[0][i] * src_val[(idx_y + j) * src_stride[1] + (idx_x + i) * src_stride[0]];
+
+            ret += fact[1][j] * val;
+        }
+output:
+        dst[pt] = ret;
+    }
+}
diff --git a/interp.py b/interp.py
index f0f5693..9e4049e 100644
--- a/interp.py
+++ b/interp.py
@@ -1,3 +1,5 @@
+import ctypes
+
 import numpy as np
 
 def interp1d(src_start, src_step, src_val, dst_coords, stencil):
@@ -50,3 +52,65 @@ def interp2d(src_start, src_step, src_val, dst_coords, stencil):
         ret += fact[0][i] * val
 
     return ret
+
+_doubleptr = ctypes.POINTER(ctypes.c_double)
+
+class Interp2D_C:
+    _lib    = None
+    _interp_func = None
+
+    _src_start_c  = None
+    _src_step_c   = None
+    _src_val_c    = None
+    _src_stride_c = None
+    _src_len_c    = None
+
+    _x     = None
+    _y     = None
+    _ret   = None
+    _x_c   = None
+    _y_c   = None
+    _ret_c = None
+
+    def __init__(self, src_start, src_step, src_val, stencil):
+        self._lib = ctypes.CDLL('lib_interp_c.so')
+
+
+        func = getattr(self._lib, 'interp2d_%d' % stencil)
+        func.argtypes = [_doubleptr, _doubleptr, _doubleptr,
+            ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int),
+            ctypes.c_int, _doubleptr, _doubleptr, _doubleptr]
+
+        self._interp_func = func
+
+        self._src_start_c  = (ctypes.c_double * 2)(src_start[0], src_start[1])
+        self._src_step_c   = (ctypes.c_double * 2)(src_step[0],  src_step[1])
+        self._src_val_c    = ctypes.cast(np.ctypeslib.as_ctypes(src_val), _doubleptr)
+        self._src_stride_c = (ctypes.c_int * 2)(src_val.strides[0] // 8, src_val.strides[1] // 8)
+        self._src_len_c    = (ctypes.c_int * 2)(src_val.shape[0], src_val.shape[1])
+
+        self._setup_arrays((1,))
+
+    def _setup_arrays(self, shape):
+        self._x   = np.empty(shape)
+        self._y   = np.empty(shape)
+        self._ret = np.empty(shape)
+
+        self._x_c   = ctypes.cast(np.ctypeslib.as_ctypes(self._x),   _doubleptr)
+        self._y_c   = ctypes.cast(np.ctypeslib.as_ctypes(self._y),   _doubleptr)
+        self._ret_c = ctypes.cast(np.ctypeslib.as_ctypes(self._ret), _doubleptr)
+
+    def interp(self, x, y):
+        if x.shape != self._x.shape:
+            self._setup_arrays(x.shape)
+
+        np.copyto(self._x,   x)
+        np.copyto(self._y,   y)
+        self._interp_func(self._src_start_c, self._src_step_c, self._src_val_c,
+                          self._src_stride_c, self._src_len_c,
+                          self._x.shape[0],
+                          self._x_c, self._y_c, self._ret_c)
+        return self._ret.copy()
+
+    def __call__(self, x, y):
+        return self.interp(x, y)