1# Copyright 2014 John O. Woods (john.o.woods@gmail.com), West Virginia
2#   University's Applied Space Exploration Lab, and West Virginia Robotic
3#   Technology Center. All rights reserved.
4#
5# THE BSD LICENSE
6#
7# Redistribution and use in source and binary forms, with or without
8# modification, are permitted provided that the following conditions
9# are met:
10#
11# 1. Redistributions of source code must retain the above copyright
12#    notice, this list of conditions and the following disclaimer.
13# 2. Redistributions in binary form must reproduce the above copyright
14#    notice, this list of conditions and the following disclaimer in the
15#    documentation and/or other materials provided with the distribution.
16#
17# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
18# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
19# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
20# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
21# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
22# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
23# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
24# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
26# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27
28class FFI::Pointer
29  class << self
30    def new_from_nmatrix nm
31      raise(StorageError, "dense storage expected") unless nm.dense?
32      c_type = Flann::dtype_to_c(nm.dtype)
33      c_type = :uchar if c_type == :byte
34      ::FFI::Pointer.new(c_type, nm.data_pointer).tap { |p| p.autorelease = false }
35    end
36  end
37end
38
39module Flann
40  class Index
41
42    # Constructor takes a block where we set each of the parameters. We need to be careful to do this since
43    # we're using the C API and not C++; so everything important needs to be initialized or there could be
44    # a segfault. For reasonable default definitions, see:
45    #
46    # * https://github.com/mariusmuja/flann/tree/master/src/cpp/flann/algorithms
47    #
48    def initialize index_dataset = nil, dtype: :float64, parameters: Flann::Parameters::DEFAULT
49      @dataset        = index_dataset
50      #require 'pry'
51      #binding.pry if @dataset.nil?
52      @dtype          = (!index_dataset.nil? && index_dataset.is_a?(NMatrix)) ? index_dataset.dtype : dtype
53      @index_ptr      = nil
54
55      @parameters_ptr, @parameters = Flann::handle_parameters(parameters)
56
57      yield @parameters if block_given?
58    end
59    attr_reader :dtype, :dataset, :parameters, :parameters_ptr, :index_ptr
60
61    # Assign a new dataset. Requires that the old index be freed.
62    def dataset= index_dataset
63      free!
64      @dataset = index_dataset
65    end
66
67    # Build an index
68    def build!
69      raise("no dataset specified") if dataset.nil?
70      c_type   = Flann::dtype_to_c(dtype)
71      c_method = "flann_build_index_#{c_type}".to_sym
72      speedup_float_ptr = FFI::MemoryPointer.new(:float)
73      @index_ptr = Flann.send(c_method, FFI::Pointer.new_from_nmatrix(dataset), dataset.shape[0], dataset.shape[1], speedup_float_ptr, parameters_ptr)
74      if index_ptr.address == 0
75        require 'pry'
76        binding.pry
77        raise("failed to allocate index_ptr")
78      end
79
80
81      # Return the speedup
82      speedup_float_ptr.read_float
83    end
84
85    # Get the nearest neighbors based on this index. Forces a build of the index if one hasn't been done yet.
86    def nearest_neighbors testset, k, parameters = {}
87      parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters))
88
89      self.build! if index_ptr.nil?
90
91      parameters_ptr, parameters = Flann::handle_parameters(parameters)
92      result_size = testset.shape[0] * k
93
94      c_type = Flann::dtype_to_c(dataset.dtype)
95      c_method = "flann_find_nearest_neighbors_index_#{c_type}".to_sym
96      indices_int_ptr, distances_t_ptr = Flann::allocate_results_space(result_size, c_type)
97
98      Flann.send c_method, index_ptr,
99                           FFI::Pointer.new_from_nmatrix(testset),
100                           testset.shape[0],
101                           indices_int_ptr, distances_t_ptr,
102                           k,
103                           parameters_ptr
104
105
106      [indices_int_ptr.read_array_of_int(result_size),
107       c_type == :double ? distances_t_ptr.read_array_of_double(result_size) : distances_t_ptr.read_array_of_float(result_size)]
108    end
109
110    # Perform a radius search on a single query point
111    def radius_search query, radius, max_k=nil, parameters = {}
112      max_k    ||= dataset.shape[0]
113      parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters))
114
115      self.build! if index_ptr.nil?
116      parameters_ptr, parameters = Flann::handle_parameters(parameters)
117
118      c_type = Flann::dtype_to_c(dataset.dtype)
119      c_method = "flann_radius_search_#{c_type}".to_sym
120      indices_int_ptr, distances_t_ptr = Flann::allocate_results_space(max_k, c_type)
121
122      Flann.send(c_method, index_ptr, FFI::Pointer.new_from_nmatrix(query), indices_int_ptr, distances_t_ptr, max_k, radius, parameters_ptr)
123
124      # Return results: two arrays, one of indices and one of distances.
125      indices   = indices_int_ptr.read_array_of_int(max_k)
126      distances = c_type == :double ? distances_t_ptr.read_array_of_double(max_k) : distances_t_ptr.read_array_of_float(max_k)
127
128      # Stop where indices == -1
129      cutoff = indices.find_index(-1)
130      cutoff.nil? ? [indices, distances] : [indices[0...cutoff], distances[0...cutoff]]
131    end
132
133    # Save an index to a file (without the dataset).
134    def save filename
135      raise(IOError, "Cannot write an unbuilt index") if index_ptr.nil?     # FIXME: This should probably have its own exception type.
136      c_method = "flann_save_index_#{Flann::dtype_to_c(dtype)}".to_sym
137      Flann.send(c_method, index_ptr, filename)
138      self
139    end
140
141    # Load an index from a file (with the dataset already known!).
142    #
143    # FIXME: This needs to free the previous dataset first.
144    def load! filename
145      c_method = "flann_load_index_#{Flann::dtype_to_c(dtype)}".to_sym
146
147      @index_ptr = Flann.send(c_method, filename, FFI::Pointer.new_from_nmatrix(dataset), dataset.shape[0], dataset.shape[1])
148      self
149    end
150
151    # Free an index
152    def free! parameters = {}
153      parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters))
154      c_method = "flann_free_index_#{Flann::dtype_to_c(dtype)}".to_sym
155      parameters_ptr, parameters = Flann::handle_parameters(parameters)
156      Flann.send(c_method, index_ptr, parameters_ptr)
157      @index_ptr = nil
158      self
159    end
160
161  end
162end