K-NN Classifier

lib/scholar/neighbors/kd_tree.ex

@@ -1,6 +1,6 @@
 defmodule Scholar.Neighbors.KDTree do
   @moduledoc """
-  Implements a kd-tree, a space-partitioning data structure for organizing points
+  Implements a k-d tree, a space-partitioning data structure for organizing points
   in a k-dimensional space.
 
   It can be used to predict the K-Nearest Neighbors of a given input.
@@ -19,14 +19,13 @@ defmodule Scholar.Neighbors.KDTree do
 
   import Nx.Defn
   import Scholar.Shared
-  alias Scholar.Metrics.Distance
 
-  @derive {Nx.Container, keep: [:levels], containers: [:indices, :data]}
-  @enforce_keys [:levels, :indices, :data]
-  defstruct [:levels, :indices, :data]
+  @derive {Nx.Container, keep: [:levels, :num_neighbors, :metric], containers: [:indices, :data]}
+  @enforce_keys [:levels, :indices, :data, :num_neighbors, :metric]
+  defstruct [:levels, :indices, :data, :num_neighbors, :metric]
 
   opts = [
-    k: [
+    num_neighbors: [
       type: :pos_integer,
       default: 3,
       doc: "The number of neighbors to use by default for `k_neighbors` queries"
@@ -45,22 +44,48 @@ defmodule Scholar.Neighbors.KDTree do
     ]
   ]
 
-  @predict_schema NimbleOptions.new!(opts)
+  @opts_schema NimbleOptions.new!(opts)
 
   @doc """
   Builds a KDTree.
 
   ## Examples
 
-      iex> Scholar.Neighbors.KDTree.fit(Nx.iota({5, 2}))
-      %Scholar.Neighbors.KDTree{
-        data: Nx.iota({5, 2}),
-        levels: 3,
-        indices: Nx.u32([3, 1, 4, 0, 2])
-      }
+      iex> tree = Scholar.Neighbors.KDTree.fit(Nx.iota({5, 2}))
+      iex> tree.data
+      Nx.tensor(
+        [
+          [0, 1],
+          [2, 3],
+          [4, 5],
+          [6, 7],
+          [8, 9]
+        ]
+      )
+      iex> tree.levels
+      3
+      iex> tree.indices
+      Nx.u32([3, 1, 4, 0, 2])
   """
-  deftransform fit(tensor, _opts \\ []) do
-    %__MODULE__{levels: levels(tensor), indices: fit_n(tensor), data: tensor}
+  deftransform fit(tensor, opts \\ []) do
+    opts = NimbleOptions.validate!(opts, @opts_schema)
+
+    metric =
+      case opts[:metric] do
+        {:minkowski, p} ->
+          &Scholar.Metrics.Distance.minkowski(&1, &2, p: p)
+
+        :cosine ->
+          &Scholar.Metrics.Distance.pairwise_cosine/2
+      end
+
+    %__MODULE__{
+      levels: levels(tensor),
+      indices: fit_n(tensor),
+      data: tensor,
+      num_neighbors: opts[:num_neighbors],
+      metric: metric
+    }
   end
 
   defnp fit_n(tensor) do
@@ -247,8 +272,8 @@ defmodule Scholar.Neighbors.KDTree do
 
       iex> x = Nx.iota({10, 2})
       iex> x_predict = Nx.tensor([[2, 5], [1, 9], [6, 4]])
-      iex> kdtree = Scholar.Neighbors.KDTree.fit(x)
-      iex> Scholar.Neighbors.KDTree.predict(kdtree, x_predict, k: 3)
+      iex> kdtree = Scholar.Neighbors.KDTree.fit(x, num_neighbors: 3)
+      iex> Scholar.Neighbors.KDTree.predict(kdtree, x_predict)
       #Nx.Tensor<
         s64[3][3]
         [
@@ -257,7 +282,11 @@ defmodule Scholar.Neighbors.KDTree do
           [2, 3, 1]
         ]
       >
-      iex> Scholar.Neighbors.KDTree.predict(kdtree, x_predict, k: 3, metric: {:minkowski, 1})
+
+      iex> x = Nx.iota({10, 2})
+      iex> x_predict = Nx.tensor([[2, 5], [1, 9], [6, 4]])
+      iex> kdtree = Scholar.Neighbors.KDTree.fit(x, num_neighbors: 3, metric: {:minkowski, 1})
+      iex> Scholar.Neighbors.KDTree.predict(kdtree, x_predict)
       #Nx.Tensor<
         s64[3][3]
         [
@@ -267,25 +296,18 @@ defmodule Scholar.Neighbors.KDTree do
         ]
       >
   """
-  deftransform predict(tree, data, opts \\ []) do
-    predict_n(tree, data, NimbleOptions.validate!(opts, @predict_schema))
+  deftransform predict(tree, data) do
+    predict_n(tree, data)
   end
 
   defnp sort_by_distances(distances, point_indices) do
     indices = Nx.argsort(distances)
     {Nx.take(distances, indices), Nx.take(point_indices, indices)}
   end
 
-  defnp compute_distance(x1, x2, opts) do
-    case opts[:metric] do
-      {:minkowski, 2} -> Distance.squared_euclidean(x1, x2)
-      {:minkowski, p} -> Distance.minkowski(x1, x2, p: p)
-      :cosine -> Distance.cosine(x1, x2)
-    end
-  end
-
   defnp update_knn(nearest_neighbors, distances, data, indices, curr_node, point, k, opts) do
-    curr_dist = compute_distance(data[[indices[curr_node]]], point, opts)
+    metric = opts[:metric]
+    curr_dist = metric.(data[[indices[curr_node]]], point)
 
     if curr_dist < distances[[-1]] do
       nearest_neighbors =
@@ -311,8 +333,8 @@ defmodule Scholar.Neighbors.KDTree do
     end
   end
 
-  defnp predict_n(tree, point, opts) do
-    k = opts[:k]
+  defnp predict_n(tree, point) do
+    k = tree.num_neighbors
     node = Nx.as_type(root(), :s64)
 
     input_vectorized_axes = point.vectorized_axes
@@ -330,6 +352,7 @@ defmodule Scholar.Neighbors.KDTree do
 
     indices = tree.indices |> Nx.as_type(:s64)
     data = tree.data
+    metric = tree.metric
 
     mode = down()
     i = Nx.s64(0)
@@ -384,7 +407,7 @@ defmodule Scholar.Neighbors.KDTree do
                       indices,
                       point,
                       k,
-                      opts
+                      metric: metric
                     )
 
                   {parent(node), i - 1, visited, nearest_neighbors, distances, up()}
@@ -402,14 +425,13 @@ defmodule Scholar.Neighbors.KDTree do
                       indices,
                       point,
                       k,
-                      opts
+                      metric: metric
                     )
 
                   if Nx.any(
-                       compute_distance(
+                       metric.(
                          point[[coord_indicator]],
-                         data[[indices[right_child(node)], coord_indicator]],
-                         opts
+                         data[[indices[right_child(node)], coord_indicator]]
                        ) <
                          distances
                      ) do
@@ -431,14 +453,13 @@ defmodule Scholar.Neighbors.KDTree do
                       indices,
                       point,
                       k,
-                      opts
+                      metric: metric
                     )
 
                   if Nx.any(
-                       compute_distance(
+                       metric.(
                          point[[coord_indicator]],
-                         data[[indices[left_child(node)], coord_indicator]],
-                         opts
+                         data[[indices[left_child(node)], coord_indicator]]
                        ) <
                          distances
                      ) do