Enables vectors to be set and implements interpolation #2

pystreamline/src/streamline.cpp

@@ -1,6 +1,8 @@
 #include <cassert>
-#include <unordered_map>
+#include <cmath>
+#include <cstring>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "kdtree.hpp"
@@ -16,6 +18,10 @@ static double calc_dist_sq( double *a1, double *a2, int dims ) {
     return dist_sq;
 }
 
+static void delete_node_int_data(void* data)
+{
+    delete (int*) data;
+}
 
 template<typename TK, typename TV>
 std::vector<TK> extract_keys_from_unordered_map(std::unordered_map<TK, TV> const& input_map)
@@ -67,6 +73,7 @@ StreamlineIntegrator::StreamlineIntegrator(double *pos, int _n_points, int _dim)
 
 StreamlineIntegrator::~StreamlineIntegrator()
 {
+    tree->destr = delete_node_int_data;
     kd_free(tree);
 }
 
@@ -89,6 +96,44 @@ double* StreamlineIntegrator::get_bounds()
 }
 
 
+int StreamlineIntegrator::set_vec(std::string name, std::string *vec_name_ref,
+  double **vec_ref, bool *vec_set_ref)
+{
+    if ( var_store_double.find(name) != var_store_double.end() ) {
+        *vec_name_ref = name;
+        *vec_ref = var_store_double[name];
+        *vec_set_ref = true;
+    }
+    else {
+        char message[100];
+        sprintf(message, "No array with name: %s", name.c_str());
+        throw std::invalid_argument(message);
+    }
+    return 0;
+}
+
+
+int StreamlineIntegrator::set_vec_x(std::string name)
+{
+    set_vec(name, &vec_x_name, &vec_x, &vec_x_set);
+    return 0;
+}
+
+
+int StreamlineIntegrator::set_vec_y(std::string name)
+{
+    set_vec(name, &vec_y_name, &vec_y, &vec_y_set);
+    return 0;
+}
+
+
+int StreamlineIntegrator::set_vec_z(std::string name)
+{
+    set_vec(name, &vec_z_name, &vec_z, &vec_z_set);
+    return 0;
+}
+
+
 int StreamlineIntegrator::get_points_in_range(double x, double y, double z,
    double range, int* count, int** idx, double** dist_sq)
 {
@@ -113,6 +158,46 @@ int StreamlineIntegrator::get_points_in_range(double x, double y, double z,
 }
 
 
+int StreamlineIntegrator::set_interp_lscale(double _interp_lscale)
+{
+    interp_lscale = _interp_lscale;
+    neg_inv_interp_lscale_sq = -1. / (interp_lscale * interp_lscale);
+    return 0;
+}
+
+
+int StreamlineIntegrator::interpolate_vec_at_point(double *pos, double *vec)
+{
+    int count, *idx;
+    double *dist_sq;
+
+    get_points_in_range(pos[0], pos[1], pos[2], interp_lscale, &count, &idx, &dist_sq);
+
+    memset(vec, 0, 3*sizeof(double));
+
+    double norm = 0.;
+    for (int i=0; i<count; i++)
+    {
+        int tmp_idx = idx[i];
+        double weight = exp(neg_inv_interp_lscale_sq * dist_sq[i]);
+        norm += weight;
+        vec[0] += vec_x[tmp_idx] * weight;
+        vec[1] += vec_y[tmp_idx] * weight;
+        vec[2] += vec_z[tmp_idx] * weight;
+    }
+    norm = count > 0 ? 1. / norm : 1.;
+    for (int i=0; i<3; i++)
+    {
+        vec[i] *= norm;
+    }
+
+    delete [] idx;
+    delete [] dist_sq;
+
+    return 0;
+}
+
+
 int StreamlineIntegrator::add_int_array(std::string name, int *arr)
 {
     if ( var_store_int.find(name) == var_store_int.end() )
@@ -127,7 +212,7 @@ int* StreamlineIntegrator::get_int_array_with_name(std::string name)
 {
     if ( var_store_int.find(name) == var_store_int.end() )
     {
-        char* message;
+        char message[100];
         sprintf(message, "key: %s not in int store", name.c_str());
         throw std::invalid_argument(message);
     }
@@ -158,7 +243,7 @@ double* StreamlineIntegrator::get_double_array_with_name(std::string name)
 {
     if ( var_store_double.find(name) == var_store_double.end() )
     {
-        char* message;
+        char message[100];
         sprintf(message, "key: %s not in double store", name.c_str());
         throw std::invalid_argument(message);
     }

pystreamline/src/streamline.hpp

@@ -12,18 +12,32 @@ class StreamlineIntegrator
     int n_points, dim;
     double bounds[6];
     struct kdtree *tree;
+
     std::unordered_map<std::string, int*> var_store_int;
     std::unordered_map<std::string, double*> var_store_double;
 
+    double interp_lscale = 0., neg_inv_interp_lscale_sq = 0.;
+    double *vec_x, *vec_y, *vec_z;
+    bool vec_x_set = false, vec_y_set = false, vec_z_set = false;
+    std::string vec_x_name, vec_y_name, vec_z_name;
+
+    int set_vec(std::string, std::string*, double**, bool*);
+
   public:
     StreamlineIntegrator(double*, int, int);
     ~StreamlineIntegrator();
-    
+
     int get_n_points();
     int get_dim();
     double* get_bounds();
 
+    int set_vec_x(std::string);
+    int set_vec_y(std::string);
+    int set_vec_z(std::string);
+    int set_interp_lscale(double);
+
     int get_points_in_range(double, double, double, double, int*, int**, double**);
+    int interpolate_vec_at_point(double*, double*);
 
     int add_int_array(std::string, int*);
     std::vector<std::string> get_int_array_names();

pystreamline/streamline_wrapper.pyx

@@ -39,7 +39,12 @@ cdef extern from "src/streamline.hpp" namespace "PyStreamline":
         StreamlineIntegrator(double*, int, int)
         int get_n_points();
         int get_dim();
-        int get_points_in_range(double, double, double, double, int*, int**, double**)
+        int set_vec_x(string);
+        int set_vec_y(string);
+        int set_vec_z(string);
+        int set_interp_lscale(double);
+        int get_points_in_range(double, double, double, double, int*, int**, double**);
+        int interpolate_vec_at_point(double*, double*);
         int add_int_array(string, int*);
         int* get_int_array_with_name(string);
         vector[string] get_int_array_names();
@@ -129,3 +134,41 @@ cdef class _StreamlineIntegrator__wrapper:
         cdef string c_name = <string> name.encode('utf-8')
         cdef double* c_arr_ptr = <double*> self.thisptr.get_double_array_with_name(c_name)
         return data_to_numpy_double_array_with_spec(c_arr_ptr, self.n_points)
+
+    def set_vec_x(self, name):
+        try:
+            assert name in self.double_array_names
+        except AssertionError:
+            raise ValueError('No double array with name: {:s}'.format(name))
+        cdef string c_name = <string> name.encode('utf-8')
+        self.thisptr.set_vec_x(c_name)
+
+    def set_vec_y(self, name):
+        try:
+            assert name in self.double_array_names
+        except AssertionError:
+            raise ValueError('No double array with name: {:s}'.format(name))
+        cdef string c_name = <string> name.encode('utf-8')
+        self.thisptr.set_vec_y(c_name)
+
+    def set_vec_z(self, name):
+        try:
+            assert name in self.double_array_names
+        except AssertionError:
+            raise ValueError('No double array with name: {:s}'.format(name))
+        cdef string c_name = <string> name.encode('utf-8')
+        self.thisptr.set_vec_z(c_name)
+
+    def set_vec_names(self, vec_x_name, vec_y_name, vec_z_name):
+        self.set_vec_x_name(vec_x_name)
+        self.set_vec_y_name(vec_y_name)
+        self.set_vec_z_name(vec_z_name)
+
+    def set_interp_lscale(self, lscale):
+        self.thisptr.set_interp_lscale(<double> lscale)
+
+    def _interp_vec_at_point(self, np.ndarray[np.float64_t, ndim=1, mode='c'] point):
+        point = np.ascontiguousarray(point, dtype=np.double)
+        cdef double c_vec[3]
+        self.thisptr.interpolate_vec_at_point(<double*> point.data, c_vec)
+        return np.array([x for x in c_vec])

pystreamline/tests/test_streamlineintegrator.py

@@ -62,6 +62,9 @@ def test_StreamlineIntegrator_var_store_int(self):
             self.streamline_integrator.get_int_array_with_name('arr4')
 
     def test_StreamlineIntegrator_var_store_double(self):
+        """Tests that double arrays are passed correctly between Python and C++,
+         and stored correctly.
+        """
         # Tests ``arr1`` data can be set correctly
         arr1 = npr.random(self.n_points)
         self.streamline_integrator.add_double_array('arr1', arr1)
@@ -111,3 +114,51 @@ def test_StreamlineIntegrator_get_points_in_range(self):
 
         assert len(idx) == 2
         np.testing.assert_allclose(np.array([0., 0.01]), dist_sq)
+
+    def test_StreamlineIntegrator__interp_vec_at_point(self):
+        """Tests that values are interpolated correctly."""
+
+        def gaussian(r, eps=1.0e-05):
+            return np.exp(-(r/eps)**2)
+
+        def interp(points, values, x, fn=gaussian, eps=1.0e-05):
+            r = np.linalg.norm(points - x, axis=1)
+            w = np.zeros(len(r))
+            w[r <= eps] = fn(r[r <= eps], eps=eps)
+            w /= w.sum()
+            return np.sum(values * w[:, np.newaxis], axis=0)
+
+        pos = np.array([
+            [0.5, 0.6, 0.5],
+            [0.5, 0.5, 0.6],
+            [0.5, 0.6, 0.6],
+            [0.4, 0.4, 0.4],
+            [0., 0., 0.],       # A point outside the interpolation range
+            [0., 1., 0.],       # "
+        ])
+
+        x, radius = np.array([0.5, 0.5, 0.5]), 0.3
+
+        n_points = len(pos)
+        arr1 = npr.random(n_points)
+        arr2 = npr.random(n_points)
+        arr3 = npr.random(n_points)
+
+        # Obtain ground truth interpolated value of vector
+        vec_data = np.column_stack((arr1, arr2, arr3))
+        interp_result_benchmark = interp(pos, vec_data, x, eps=radius)
+
+        # Perform interpolation with StreamlineIntegrator
+        streamline_integrator = StreamlineIntegrator(pos)
+        streamline_integrator.add_double_array('arr1', arr1)
+        streamline_integrator.add_double_array('arr2', arr2)
+        streamline_integrator.add_double_array('arr3', arr3)
+
+        streamline_integrator.set_vec_x('arr1')
+        streamline_integrator.set_vec_y('arr2')
+        streamline_integrator.set_vec_z('arr3')
+
+        streamline_integrator.set_interp_lscale(radius)
+        interp_result_sl = streamline_integrator._interp_vec_at_point(x)
+
+        np.testing.assert_allclose(interp_result_sl, interp_result_benchmark)