kdtree 0.1

data/LICENSE

@@ -0,0 +1,20 @@
+Copyright (c) 2009 Adam Doppelt
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/ext/extconf.rb

@@ -0,0 +1,3 @@
+require 'mkmf'
+
+create_makefile("kdtree")

data/ext/kdtree.c

@@ -0,0 +1,488 @@
+#include "ruby.h"
+#include "rubyio.h"
+#include "version.h"
+
+#ifndef HAVE_RB_IO_T
+#define rb_io_t OpenFile
+#endif
+
+//
+// interface
+//
+
+typedef struct kdtree_data
+{
+    int root;
+    int len;
+    struct kdtree_node *nodes;
+} kdtree_data;
+
+typedef struct kdtree_node
+{
+    float x, y;
+    int id;
+    int left;
+    int right;
+} kdtree_node;
+
+#define KDTREEP \
+    struct kdtree_data *kdtreep; \
+    Data_Get_Struct(kdtree, struct kdtree_data, kdtreep);
+
+static VALUE kdtree_alloc(VALUE klass);
+static void kdtree_free(struct kdtree_data *kdtreep);
+static VALUE kdtree_initialize(VALUE kdtree, VALUE points);
+static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y);
+static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k);
+static VALUE kdtree_persist(VALUE kdtree, VALUE io);
+static VALUE kdtree_to_s(VALUE kdtree);
+
+// helpers
+static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth);
+static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth);
+static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth);
+
+#define KDTREE_MAGIC "KdTr"
+
+//
+// implementation
+//
+
+static VALUE kdtree_alloc(VALUE klass)
+{
+    struct kdtree_data *kdtreep;
+    VALUE obj = Data_Make_Struct(klass, struct kdtree_data, 0, kdtree_free, kdtreep);
+    kdtreep->root = -1;
+    return obj;
+}
+
+static void kdtree_free(struct kdtree_data *kdtreep)
+{
+    if (kdtreep) {
+        free(kdtreep->nodes);
+    }
+}
+
+static void read_all(struct rb_io_t *fptr, char *buf, int len)
+{
+    while (len > 0) {
+        int n = rb_io_fread(buf, len, fptr->f);
+        if (n == 0) {
+            rb_eof_error();
+        }
+        buf += n;
+        len -= n;
+    }
+}
+
+/*
+ * call-seq:
+ *   KDTree.new(points)    => kdtree
+ *   KDTree.new(io)        => kdtree
+ *
+ * Returns a new <code>KDTree</code>. To construct a tree, pass an array of
+ * <i>points</i>. Each point should be an array of the form <code>[x, y,
+ * id]</code>, where <i>x</i> and <i>y</i> are floats and <i>id</i> is an
+ * integer. The <i>id</i> is arbitrary and will be returned to you whenever you
+ * search with nearest or nearestk.
+ *
+ *   # create a new tree
+ *   points = []
+ *   points << [47.6, -122.3, 1] # Seattle
+ *   points << [40.7, -74.0, 2]  # New York
+ *   kd = KDTree.new(points)
+ *
+ * Alternately, you can pass in an <i>IO</i> object containing a persisted
+ * kdtree. This makes it possible to build the tree in advance, persist it, and
+ * start it up quickly later. See persist for more information.
+ */
+static VALUE kdtree_initialize(VALUE kdtree, VALUE arg)
+{
+    KDTREEP;
+
+    if (TYPE(arg) == T_ARRAY) {
+        // init from array of pints
+        VALUE points = arg;
+        kdtreep->len = RARRAY_LEN(points);
+        kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
+
+        int i;
+        for (i = 0; i < RARRAY_LEN(points); ++i) {
+            struct kdtree_node *n = kdtreep->nodes + i;
+            
+            VALUE ptr = RARRAY_PTR(points)[i];
+            VALUE v = rb_check_array_type(ptr);
+            if (NIL_P(v) || RARRAY_LEN(v) != 3) {
+                continue;
+            }
+            VALUE *a = RARRAY_PTR(ptr);
+            n->x = NUM2DBL(a[0]);
+            n->y = NUM2DBL(a[1]);        
+            n->id = NUM2INT(a[2]);
+        }
+
+        // now build the tree
+        kdtreep->root = kdtree_build(kdtreep, 0, kdtreep->len, 0);
+    } else if (rb_respond_to(arg, rb_intern("read"))) {
+        VALUE io = arg;
+        if (rb_respond_to(io, rb_intern("binmode"))) {
+            rb_funcall2(io, rb_intern("binmode"), 0, 0);
+        }
+
+        struct rb_io_t *fptr = RFILE(rb_io_taint_check(io))->fptr;
+        rb_io_check_readable(fptr);
+
+        // check magic
+        char buf[4];
+        read_all(fptr, buf, 4);
+        if (memcmp(KDTREE_MAGIC, buf, 4) != 0) {
+            rb_raise(rb_eRuntimeError, "wrong magic number in kdtree file");
+        }
+        
+        // read start of the struct
+        read_all(fptr, (char *)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
+        // read the nodes
+        kdtreep->nodes = ALLOC_N(struct kdtree_node, kdtreep->len);
+        read_all(fptr, (char *)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
+    } else {
+        rb_raise(rb_eTypeError, "array or IO required to init KDTree");
+    }
+    
+    return kdtree;
+}
+
+static int comparex(const void *pa, const void *pb)
+{
+    float a = ((const struct kdtree_node*)pa)->x;
+    float b = ((const struct kdtree_node*)pb)->x;    
+    return (a < b) ? -1 : ((a > b) ? 1 : 0);
+}
+
+static int comparey(const void *pa, const void *pb)
+{
+    float a = ((const struct kdtree_node*)pa)->y;
+    float b = ((const struct kdtree_node*)pb)->y;    
+    return (a < b) ? -1 : ((a > b) ? 1 : 0);
+}
+
+static int kdtree_build(struct kdtree_data *kdtreep, int min, int max, int depth)
+{
+    if (max <= min) {
+        return -1;
+    }
+
+    // sort nodes from min to max
+    int(*compar)(const void *, const void *) = (depth % 2) ? comparex : comparey;
+    qsort(kdtreep->nodes + min, max - min, sizeof(struct kdtree_node), compar);
+
+    int median = (min + max) / 2;
+    struct kdtree_node *m = kdtreep->nodes + median;
+    m->left = kdtree_build(kdtreep, min, median, depth + 1);
+    m->right = kdtree_build(kdtreep, median + 1, max, depth + 1);
+    return median;
+}
+
+//
+// nearest
+//
+
+static int n_index;
+static float n_dist;
+
+/*
+ * call-seq:
+ *   kd.nearest(x, y)    => id
+ *
+ * Finds the point closest to <i>x</i>, <i>y</i> and returns the id for that
+ * point. Returns -1 if the tree is empty.
+ * 
+ *   points = []
+ *   points << [47.6, -122.3, 1] # Seattle
+ *   points << [40.7, -74.0, 2]  # New York
+ *   kd = KDTree.new(points)
+ *   
+ *   # which city is closest to Portland?
+ *   kd.nearest(45.5, -122.8) #=> 1
+ *   # which city is closest to Boston?
+ *   kd.nearest(42.4, -71.1)  #=> 2
+ */
+static VALUE kdtree_nearest(VALUE kdtree, VALUE x, VALUE y)
+{
+    KDTREEP;
+
+    n_index = -1;
+    n_dist = INT_MAX;
+    kdtree_nearest0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), 0);
+    if (n_index == -1) {
+        return -1;
+    }
+    return INT2NUM((kdtreep->nodes + n_index)->id);
+}
+
+static void kdtree_nearest0(struct kdtree_data *kdtreep, int i, float x, float y, int depth)
+{
+    if (i == -1) {
+        return;
+    }
+    
+    struct kdtree_node *n = kdtreep->nodes + i;
+
+    float ad = (depth % 2) ? (x - n->x) : (y - n->y);
+
+    //
+    // recurse near, and perhaps far as well
+    //
+    
+    int near, far;
+    if (ad <= 0) {
+        near = n->left; far = n->right;
+    } else {
+        near = n->right; far = n->left;
+    }
+    kdtree_nearest0(kdtreep, near,  x, y, depth + 1);    
+    if (ad * ad < n_dist) {
+        kdtree_nearest0(kdtreep, far, x, y, depth + 1);
+    }
+
+    //
+    // do we beat the old distance?
+    //
+    
+    float dx = (x - n->x) * (x - n->x);
+    if (dx < n_dist) {
+        float d = dx + ((y - n->y) * (y - n->y));
+        if (d < n_dist) {
+            n_index = i;
+            n_dist = d;
+        }
+    }
+}
+
+//
+// nearestK
+//
+
+#define MAX_K 255
+
+typedef struct kresult {
+    int index;
+    float distance;
+} kresult;
+// note I leave an extra slot here at the end because of the way our binary insert works
+static struct kresult k_list[MAX_K + 1];
+static int k_len;
+static float k_dist;
+
+/*
+ * call-seq:
+ *   kd.nearestk(x, y, k)    => array
+ *
+ * Finds the <i>k</i> points closest to <i>x</i>, <i>y</i>. Returns an array of
+ * ids, sorted by distance. Returns an empty array if the tree is empty. Note
+ * that <i>k</i> is capped at 255.
+ * 
+ *   points = []
+ *   points << [47.6, -122.3, 1] # Seattle
+ *   points << [45.5, -122.8, 2] # Portland
+ *   points << [40.7, -74.0,  3] # New York
+ *   kd = KDTree.new(points)
+ *   
+ *   # which two cities are closest to San Francisco?
+ *   kd.nearest(34.1, -118.2) #=> [2, 1]
+ */
+static VALUE kdtree_nearestk(VALUE kdtree, VALUE x, VALUE y, VALUE k)
+{
+    KDTREEP;
+
+    k_len = 0;
+    k_dist = INT_MAX;
+
+    int ki = NUM2INT(k);
+    if (ki < 1) {
+        ki = 1;
+    } else if (ki > MAX_K) {
+        ki = MAX_K;
+    }
+    kdtree_nearestk0(kdtreep, kdtreep->root, NUM2DBL(x), NUM2DBL(y), ki, 0);
+
+    // convert result to ruby array
+    VALUE ary = rb_ary_new();
+    int i;
+    for (i = 0; i < k_len; ++i) {
+        rb_ary_push(ary, INT2NUM(kdtreep->nodes[k_list[i].index].id));
+    }
+    return ary;
+}
+
+static void kdtree_nearestk0(struct kdtree_data *kdtreep, int i, float x, float y, int k, int depth)
+{
+    if (i == -1) {
+        return;
+    }
+    
+    struct kdtree_node *n = kdtreep->nodes + i;
+
+    float ad = (depth % 2) ? (x - n->x) : (y - n->y);
+
+    //
+    // recurse near, and then perhaps far as well
+    //
+    
+    int near, far;
+    if (ad <= 0) {
+        near = n->left; far = n->right;
+    } else {
+        near = n->right; far = n->left;
+    }
+    kdtree_nearestk0(kdtreep, near,  x, y, k, depth + 1);
+    if (ad * ad < k_dist) {
+        kdtree_nearestk0(kdtreep, far, x, y, k, depth + 1);
+    }
+
+    //
+    // do we beat the old distance?
+    //
+    
+    float dx = (x - n->x) * (x - n->x);
+    if (dx < k_dist) {
+        float d = dx + ((y - n->y) * (y - n->y));
+        if (d < k_dist) {
+            //
+            // find spot to insert
+            //
+            int lo = 0, hi = k_len;
+            while (lo < hi) {
+                int mid = (lo + hi) / 2;
+                if (k_list[mid].distance < d) {
+                    lo = mid + 1;
+                } else {
+                    hi = mid;
+                }
+            }
+
+            //
+            // insert
+            //
+            
+            memmove(k_list + lo + 1, k_list + lo, (k_len - lo) * sizeof(struct kresult));
+            k_list[lo].index = i;
+            k_list[lo].distance = d;
+
+            //
+            // adjust len/dist if necessary
+            //
+            
+            if (k_len < k) {
+                ++k_len;
+            } else {
+                k_dist = k_list[k - 1].distance;
+            }
+         }
+    }
+}
+
+/*
+ * call-seq:
+ *   kd.persist(io)
+ *
+ * Writes the tree out to <i>io</i> so you can quickly load it later with
+ * KDTree.new. This avoids the startup cost of initializing a tree. Apart from a
+ * small header, the size of the file is proportional to the number of points,
+ * requiring 20 bytes per point.
+ *
+ * This file is <b>NOT PORTABLE</b> across different architectures due to endian
+ * issues.
+ * 
+ *   points = []
+ *   points << [47.6, -122.3, 1] # Seattle
+ *   points << [45.5, -122.8, 2] # Portland
+ *   points << [40.7, -74.0,  3] # New York
+ *   kd = KDTree.new(points)
+ *
+ *   # persist the tree to disk
+ *   File.open("treefile", "w") { |f| kd.persist(f) }
+ *
+ *   ...
+ *   
+ *   # later, read the tree from disk
+ *   kd2 = File.open("treefile") { |f| KDTree.new(f) }
+ */
+static VALUE kdtree_persist(VALUE kdtree, VALUE io)
+{
+    KDTREEP;
+    
+    if (!rb_respond_to(io, rb_intern("write"))) {
+        rb_raise(rb_eTypeError, "instance of IO needed");
+    }
+    if (rb_respond_to(io, rb_intern("binmode"))) {
+        rb_funcall2(io, rb_intern("binmode"), 0, 0);
+    }
+
+    VALUE str = rb_str_buf_new(0);
+    rb_str_buf_cat(str, KDTREE_MAGIC, 4);
+    rb_str_buf_cat(str, (char*)kdtreep, sizeof(struct kdtree_data) - sizeof(struct kdtree_node *));
+    rb_str_buf_cat(str, (char*)kdtreep->nodes, sizeof(struct kdtree_node) * kdtreep->len);
+    rb_io_write(io, str);
+    return io;
+}
+
+/*
+ * call-seq:
+ *   kd.to_s     => string
+ *
+ * A string that tells you a bit about the tree.
+ */
+static VALUE kdtree_to_s(VALUE kdtree)
+{
+    KDTREEP;
+
+    char buf[256];
+    sprintf(buf, "#<%s:%p nodes=%d>", rb_obj_classname(kdtree), (void*)kdtree, kdtreep->len);
+    return rb_str_new(buf, strlen(buf));
+}
+
+//
+// entry point
+//
+
+/*
+ *  KDTree is an insanely fast data structure for finding the nearest
+ *  neighbor(s) to a given point. This implementation only supports 2d
+ *  points. Also, it only supports static points - there is no way to edit the
+ *  tree after it has been initialized. KDTree should scale to millions of
+ *  points, though it's only been tested with around 1 million.
+ *
+ *  Once the tree is constructed, it can be searched with nearest and nearestk.
+ *
+ *  To avoid the startup costs associated with creating a new tree, use persist
+ *  to write the tree to disk. You can then construct the tree later from that
+ *  file.
+ *
+ *   points = []
+ *   points << [47.6, -122.3, 1] # Seattle
+ *   points << [45.5, -122.8, 2] # Portland
+ *   points << [40.7, -74.0,  3] # New York
+ *   kd = KDTree.new(points)
+ *   
+ *   # which city is closest to San Francisco?
+ *   kd.nearest(34.1, -118.2) #=> 2
+ *   # which two cities are closest to San Francisco?
+ *   kd.nearest(34.1, -118.2) #=> [2, 1]
+ *
+ * For more information on kd trees, see:
+ *
+ * http://en.wikipedia.org/wiki/Kd-tree
+ */
+void Init_kdtree()
+{
+    static VALUE clazz;
+
+    clazz = rb_define_class("KDTree", rb_cObject);
+    
+    rb_define_alloc_func(clazz, kdtree_alloc);    
+    rb_define_method(clazz, "initialize", kdtree_initialize, 1);
+    rb_define_method(clazz, "nearest", kdtree_nearest, 2);
+    rb_define_method(clazz, "nearestk", kdtree_nearestk, 3);
+    rb_define_method(clazz, "persist", kdtree_persist, 1);    
+    rb_define_method(clazz, "to_s", kdtree_to_s, 0);        
+}

data/test/test.rb

@@ -0,0 +1,138 @@
+require "#{File.expand_path(File.dirname(__FILE__))}/../ext/kdtree.o"
+require "test/unit"
+require "tempfile"
+
+#
+# create a tree
+#
+
+class KDTreeTest < Test::Unit::TestCase
+  TMP = "#{Dir.tmpdir}/kdtree_test"
+  
+  def test_nearest
+    setup_tree(1000)
+    100.times do
+      pt = [rand_coord, rand_coord]
+
+      # kdtree search
+      id = @kdtree.nearest(pt[0], pt[1])
+      kdpt = @points[id]
+
+      # slow search
+      sortpt = @points.sort_by { |i| distance(i, pt) }.first
+
+      # assert
+      kdd = distance(kdpt, pt)
+      sortd = distance(sortpt, pt)
+      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
+    end
+  end
+
+  def test_nearestk
+    setup_tree(1000)    
+    100.times do
+      pt = [rand_coord, rand_coord]
+
+      # kdtree search
+      list = @kdtree.nearestk(pt[0], pt[1], 5)
+      kdpt = @points[list.last]
+
+      # slow search
+      sortpt = @points.sort_by { |i| distance(i, pt) }[list.length - 1]
+
+      # assert
+      kdd = distance(kdpt, pt)
+      sortd = distance(sortpt, pt)
+      assert((kdd - sortd).abs < 0.0000001, "kdtree didn't return the closest result")
+    end
+  end
+
+  def test_persist
+    setup_tree(1000)
+
+    begin
+      # write
+      File.open(TMP, "w") { |f| @kdtree.persist(f) }
+      # read
+      kdtree2 = File.open(TMP, "r") { |f| KDTree.new(f) }
+
+      # now test some random points
+      100.times do
+        pt = [rand_coord, rand_coord]
+        id1 = @kdtree.nearest(*pt)
+        id2 = kdtree2.nearest(*pt)
+        assert(id1 == id2, "kdtree2 differed from kdtree")
+      end
+    ensure
+      File.unlink(TMP)
+    end
+
+    # now test magic problems
+    begin
+      File.open(TMP, "w") { |f| f.puts "That ain't right" }
+      assert_raise RuntimeError do
+        File.open(TMP, "r") { |f| KDTree.new(f) }
+      end
+    ensure
+      File.unlink(TMP)
+    end
+  end
+
+  def dont_test_speed
+    printf("\n")
+    sizes = [1, 100, 1000, 10000, 100000, 1000000]
+    ks = [1, 5, 50, 255]
+    sizes.each do |s|
+      points = (0...s).map { |i| [rand_coord, rand_coord, i] }
+
+      # build
+      tm = Time.now
+      kdtree = KDTree.new(points)
+      printf "build %d took %.6fs\n", s, Time.now - tm
+      
+      begin
+        # write
+        tm = Time.now
+        File.open(TMP, "w") { |f| kdtree.persist(f) }
+        printf "write %d took %.6fs\n", s, Time.now - tm
+        # read
+        tm = Time.now
+        File.open(TMP, "r") { |f| KDTree.new(f) }
+        printf "read  %d took %.6fs\n", s, Time.now - tm        
+      ensure
+        File.unlink(TMP)
+      end
+      
+      ks.each do |k|
+        total = count = 0
+        100.times do
+          tm = Time.now
+          if k == 1
+            kdtree.nearest(rand_coord, rand_coord)
+          else
+            kdtree.nearestk(rand_coord, rand_coord, k)
+          end
+          total += Time.now - tm
+          count += 1
+        end
+        printf "avg query time = %.6fs [%d/%d]\n", total / count, s, k
+      end
+    end
+  end
+
+  protected
+  
+  def setup_tree(len)
+    @points = (0...len).map { |i| [rand_coord, rand_coord, i] }
+    @kdtree = KDTree.new(@points)
+  end
+  
+  def distance(a, b)
+    x, y = a[0] - b[0], a[1] - b[1]
+    x * x + y * y
+  end
+
+  def rand_coord
+    rand(0) * 10 - 5
+  end
+end

metadata

@@ -0,0 +1,61 @@
+--- !ruby/object:Gem::Specification 
+name: kdtree
+version: !ruby/object:Gem::Version 
+  version: "0.1"
+platform: ruby
+authors: 
+- Adam Doppelt
+autorequire: 
+bindir: bin
+cert_chain: []
+
+date: 2010-01-21 00:00:00 -08:00
+default_executable: 
+dependencies: []
+
+description: 
+email: amd@gurge.com
+executables: []
+
+extensions: 
+- ext/extconf.rb
+extra_rdoc_files: []
+
+files: 
+- ext/extconf.rb
+- ext/kdtree.c
+- LICENSE
+- test/test.rb
+has_rdoc: true
+homepage: 
+licenses: []
+
+post_install_message: 
+rdoc_options: 
+- --exclude
+- test
+- --exclude
+- extconf
+require_paths: 
+- .
+required_ruby_version: !ruby/object:Gem::Requirement 
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      version: 1.8.5
+  version: 
+required_rubygems_version: !ruby/object:Gem::Requirement 
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      version: "0"
+  version: 
+requirements: []
+
+rubyforge_project: 
+rubygems_version: 1.3.5
+signing_key: 
+specification_version: 3
+summary: Blazingly fast 2d kdtree.
+test_files: 
+- test/test.rb