KDTree: Supported 'cosine', 'manhattan' and 'hamming' distance (#216)

Author: gyrdym

.gitignore

@@ -11,3 +11,6 @@ doc/api/
 pubspec.lock
 
 test/.test_coverage.dart
+
+.DS_Store
+*/**/.DS_Store

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 16.9.0
+- KDTree:
+    - Supported `cosine`, `manhattan` and `hamming` distance
+
 ## 16.8.0
 - DecisionTreeClassifier:
     - Added Gini index assessor type

benchmark/kd_tree/kd_tree_building.dart

@@ -1,4 +1,4 @@
-// 0.5 sec (MacBook Air mid 2017)
+// 0.14 sec (MacBook Air mid 2017)
 import 'package:benchmark_harness/benchmark_harness.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree.dart';
 import 'package:ml_dataframe/ml_dataframe.dart';

benchmark/kd_tree/kd_tree_querying.dart

@@ -30,7 +30,7 @@ Future main() async {
   final points = Matrix.random(20000, 10, seed: 1, min: -5000, max: 5000);
 
   trainData = DataFrame.fromMatrix(points);
-  tree = KDTree(trainData, leafSize: 1);
+  tree = KDTree(trainData);
   point = Vector.randomFilled(trainData.rows.first.length,
       seed: 10, min: -5000, max: 5000);
 

lib/src/retrieval/kd_tree/kd_tree.dart

@@ -5,6 +5,7 @@ import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_impl.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_neighbour.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_split_strategy.dart';
 import 'package:ml_dataframe/ml_dataframe.dart';
+import 'package:ml_linalg/distance.dart';
 import 'package:ml_linalg/dtype.dart';
 import 'package:ml_linalg/matrix.dart';
 import 'package:ml_linalg/vector.dart';
@@ -24,28 +25,28 @@ abstract class KDTree implements Serializable {
   /// The bigger the number, the less effective search is. If [leafSize] is
   /// equal to the number of [points], a regular KNN-search will take place.
   ///
-  /// Extremely small [leafSize] leads to ineffective memory usage since in
-  /// this case a lot of kd-tree nodes will be allocated
+  /// Small [leafSize] leads to an increasing amount of time for tree building,
+  /// but querying will be very fast
   ///
   /// [dtype] A data type which will be used to convert raw data from [points]
   /// into internal numerical representation
   ///
   /// [splitStrategy] Describes how to choose a split dimension. Default value
-  /// is [KDTreeSplitStrategy.largestVariance]
+  /// is [KDTreeSplitStrategy.inOrder]
+  ///
+  /// if [splitStrategy] is [KDTreeSplitStrategy.inOrder], dimension for data
+  /// splits will be chosen one by one in order, in this case, tree building is
+  /// very fast
   ///
   /// if [splitStrategy] is [KDTreeSplitStrategy.largestVariance], dimension with
   /// the widest column (in terms of variance) will be chosen to split the data
   ///
-  /// if [splitStrategy] is [KDTreeSplitStrategy.inOrder], dimension for data
-  /// splits will be chosen one by one in order
-  ///
-  /// [KDTreeSplitStrategy.largestVariance] provides more accurate KNN-search,
+  /// [KDTreeSplitStrategy.largestVariance] results in more balanced tree,
   /// but this strategy takes much more time to build the tree than [KDTreeSplitStrategy.inOrder]
   factory KDTree(DataFrame points,
           {int leafSize = 1,
           DType dtype = DType.float32,
-          KDTreeSplitStrategy splitStrategy =
-              KDTreeSplitStrategy.largestVariance}) =>
+          KDTreeSplitStrategy splitStrategy = KDTreeSplitStrategy.inOrder}) =>
       createKDTree(points, leafSize, dtype, splitStrategy);
 
   /// [pointsSrc] Data points which will be used to build the tree.
@@ -115,7 +116,8 @@ abstract class KDTree implements Serializable {
   /// final kdTree = KDTree(data);
   /// final neighbours = kdTree.query(Vector.fromList([1, 2, 3, 4]), 2);
   ///
-  /// print(neighbours.index); // let's say, it outputs `3` which means that the nearest neighbour is kdTree.points[3]
+  /// print(neighbours[0].index); // let's say, it outputs `3` which means that the nearest neighbour is kdTree.points[3]
   /// ```
-  Iterable<KDTreeNeighbour> query(Vector point, int k);
+  Iterable<KDTreeNeighbour> query(Vector point, int k,
+      [Distance distance = Distance.euclidean]);
 }

lib/src/retrieval/kd_tree/kd_tree_impl.dart

@@ -6,6 +6,7 @@ import 'package:ml_algo/src/retrieval/kd_tree/kd_tree.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_json_keys.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_neighbour.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_node.dart';
+import 'package:ml_linalg/distance.dart';
 import 'package:ml_linalg/dtype.dart';
 import 'package:ml_linalg/matrix.dart';
 import 'package:ml_linalg/vector.dart';
@@ -45,61 +46,62 @@ class KDTreeImpl with SerializableMixin implements KDTree {
   int searchIterationCount = 0;
 
   @override
-  Iterable<KDTreeNeighbour> query(Vector point, int k) {
+  Iterable<KDTreeNeighbour> query(Vector point, int k,
+      [Distance distanceType = Distance.euclidean]) {
     if (point.length != points.columnsNum) {
       throw InvalidQueryPointLength(point.length, points.columnsNum);
     }
 
     searchIterationCount = 0;
 
-    final neighbours = HeapPriorityQueue<KDTreeNeighbour>((a, b) =>
-        (point.distanceTo(points[b.index]) - point.distanceTo(points[a.index]))
-            .toInt());
+    final neighbours = _createQueue(point, distanceType);
 
-    _findKNNRecursively(root, point, k, neighbours);
+    _findKNNRecursively(root, point, k, neighbours, distanceType);
 
     return neighbours.toList().reversed;
   }
 
   void _findKNNRecursively(KDTreeNode? node, Vector point, int k,
-      HeapPriorityQueue<KDTreeNeighbour> neighbours) {
+      HeapPriorityQueue<KDTreeNeighbour> neighbours, Distance distanceType) {
     if (node == null) {
       return;
     }
 
     if (node.isLeaf) {
-      _knnSearch(point, node.pointIndices, neighbours, k);
+      _knnSearch(point, node.pointIndices, neighbours, k, distanceType);
 
       return;
     }
 
     final nodePoint = points[node.pointIndices[0]];
-    final isNodeTooFar = neighbours.length > 0 &&
-        (point[node.splitIndex] - nodePoint[node.splitIndex]).abs() >
-            neighbours.first.distance;
     final isQueueFilled = neighbours.length == k;
 
-    if (isQueueFilled && isNodeTooFar) {
+    if (isQueueFilled && _isNodeToFar(node, point, neighbours, distanceType)) {
       return;
     }
 
-    _knnSearch(point, node.pointIndices, neighbours, k);
+    _knnSearch(point, node.pointIndices, neighbours, k, distanceType);
 
     if (point[node.splitIndex] < nodePoint[node.splitIndex]) {
-      _findKNNRecursively(node.left, point, k, neighbours);
-      _findKNNRecursively(node.right, point, k, neighbours);
+      _findKNNRecursively(node.left, point, k, neighbours, distanceType);
+      _findKNNRecursively(node.right, point, k, neighbours, distanceType);
     } else {
-      _findKNNRecursively(node.right, point, k, neighbours);
-      _findKNNRecursively(node.left, point, k, neighbours);
+      _findKNNRecursively(node.right, point, k, neighbours, distanceType);
+      _findKNNRecursively(node.left, point, k, neighbours, distanceType);
     }
   }
 
-  void _knnSearch(Vector point, List<int> pointIndices,
-      HeapPriorityQueue<KDTreeNeighbour> neighbours, int k) {
+  void _knnSearch(
+      Vector point,
+      List<int> pointIndices,
+      HeapPriorityQueue<KDTreeNeighbour> neighbours,
+      int k,
+      Distance distanceType) {
     pointIndices.forEach((candidateIdx) {
       searchIterationCount++;
       final candidate = points[candidateIdx];
-      final candidateDistance = candidate.distanceTo(point);
+      final candidateDistance =
+          candidate.distanceTo(point, distance: distanceType);
       final lastNeighbourDistance =
           neighbours.length > 0 ? neighbours.first.distance : candidateDistance;
       final isGoodCandidate = candidateDistance < lastNeighbourDistance;
@@ -114,4 +116,51 @@ class KDTreeImpl with SerializableMixin implements KDTree {
       }
     });
   }
+
+  bool _isNodeToFar(KDTreeNode node, Vector point,
+      HeapPriorityQueue<KDTreeNeighbour> neighbours, Distance distanceType) {
+    if (neighbours.length == 0) {
+      return false;
+    }
+
+    final nodePoint = points[node.pointIndices[0]];
+
+    switch (distanceType) {
+      case Distance.euclidean:
+      case Distance.manhattan:
+        return (point[node.splitIndex] - nodePoint[node.splitIndex]).abs() >
+            neighbours.first.distance;
+
+      case Distance.hamming:
+      case Distance.cosine:
+        final other = nodePoint.set(node.splitIndex, point[node.splitIndex]);
+
+        return other.distanceTo(nodePoint, distance: distanceType) >
+            neighbours.first.distance;
+
+      default:
+        throw UnsupportedError(
+            'Distance type $distanceType is not supported yet');
+    }
+  }
+
+  HeapPriorityQueue<KDTreeNeighbour> _createQueue(
+      Vector point, Distance distanceType) {
+    return HeapPriorityQueue<KDTreeNeighbour>((a, b) {
+      final distanceA =
+          point.distanceTo(points[a.index], distance: distanceType);
+      final distanceB =
+          point.distanceTo(points[b.index], distance: distanceType);
+
+      if (distanceA < distanceB) {
+        return 1;
+      }
+
+      if (distanceA > distanceB) {
+        return -1;
+      }
+
+      return 0;
+    });
+  }
 }

pubspec.yaml

@@ -1,6 +1,6 @@
 name: ml_algo
 description: Machine learning algorithms, Machine learning models performance evaluation functionality
-version: 16.8.0
+version: 16.9.0
 homepage: https://github.com/gyrdym/ml_algo
 
 environment:

test/retrieval/kd_tree/kd_tree_test.dart

@@ -3,6 +3,7 @@ import 'package:ml_algo/src/retrieval/kd_tree/kd_tree.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_impl.dart';
 import 'package:ml_algo/src/retrieval/kd_tree/kd_tree_split_strategy.dart';
 import 'package:ml_dataframe/ml_dataframe.dart';
+import 'package:ml_linalg/distance.dart';
 import 'package:ml_linalg/dtype.dart';
 import 'package:ml_linalg/vector.dart';
 import 'package:test/test.dart';
@@ -61,10 +62,10 @@ void main() {
     });
 
     test(
-        'should find the closest neighbours for [2.79, -9.15, 6.56, -18.59, 13.53], leafSize=3, splitStrategy=KDTreeSplitStrategy.inOrder',
+        'should find the closest neighbours for [2.79, -9.15, 6.56, -18.59, 13.53], leafSize=3, splitStrategy=KDTreeSplitStrategy.largestVariance',
         () {
       final kdTree = KDTree(DataFrame(data, headerExists: false),
-          leafSize: 3, splitStrategy: KDTreeSplitStrategy.inOrder);
+          leafSize: 3, splitStrategy: KDTreeSplitStrategy.largestVariance);
       final sample = Vector.fromList([2.79, -9.15, 6.56, -18.59, 13.53]);
       final result = kdTree.query(sample, 3).toList();
 
@@ -84,7 +85,7 @@ void main() {
       expect(result[0].index, 4);
       expect(result[1].index, 12);
       expect(result[2].index, 3);
-      expect(result[3].index, 1);
+      expect(result[3].index, 18);
       expect(result, hasLength(4));
     });
 
@@ -98,7 +99,7 @@ void main() {
       expect(result[0].index, 4);
       expect(result[1].index, 12);
       expect(result[2].index, 3);
-      expect(result[3].index, 1);
+      expect(result[3].index, 18);
       expect(result, hasLength(4));
     }, skip: false);
 
@@ -112,8 +113,24 @@ void main() {
       expect(result[0].index, 19);
       expect(result[1].index, 11);
       expect(result[2].index, 6);
-      expect(result[4].index, 9);
-      expect(result[3].index, 18);
+      expect(result[3].index, 9);
+      expect(result[4].index, 18);
+      expect(result[5].index, 2);
+      expect(result[6].index, 14);
+      expect(result[7].index, 10);
+      expect(result[8].index, 15);
+      expect(result[9].index, 5);
+      expect(result[10].index, 7);
+      expect(result[10].index, 7);
+      expect(result[11].index, 13);
+      expect(result[12].index, 3);
+      expect(result[13].index, 12);
+      expect(result[14].index, 17);
+      expect(result[15].index, 1);
+      expect(result[16].index, 0);
+      expect(result[17].index, 16);
+      expect(result[18].index, 4);
+      expect(result[19].index, 8);
       expect(result, hasLength(20));
     });
 
@@ -143,7 +160,7 @@ void main() {
 
       kdTree.query(sample, 1).toList();
 
-      expect((kdTree as KDTreeImpl).searchIterationCount, 14);
+      expect((kdTree as KDTreeImpl).searchIterationCount, 4);
     });
 
     test(
@@ -154,7 +171,68 @@ void main() {
 
       kdTree.query(sample, 1).toList();
 
-      expect((kdTree as KDTreeImpl).searchIterationCount, 7);
+      expect((kdTree as KDTreeImpl).searchIterationCount, 6);
+    });
+
+    test(
+        'should find the closest neighbours for [12, 23, 22, 11, -20], k=1, leafSize=1, cosine distance',
+        () {
+      final kdTree = KDTree(DataFrame(data, headerExists: false), leafSize: 1);
+      final sample = Vector.fromList([12, 23, 22, 11, -20]);
+      final result = kdTree.query(sample, 1, Distance.cosine).toList();
+
+      expect(result, hasLength(1));
+      expect(result[0].index, 17);
+    });
+
+    test(
+        'should find the closest neighbours for [12, 23, 22, 11, -20], k=2, leafSize=1, cosine distance',
+        () {
+      final kdTree = KDTree(DataFrame(data, headerExists: false), leafSize: 1);
+      final sample = Vector.fromList([12, 23, 22, 11, -20]);
+      final result = kdTree.query(sample, 2, Distance.cosine).toList();
+
+      expect(result, hasLength(2));
+      expect(result[0].index, 17);
+      expect(result[1].index, 8);
+    });
+
+    test(
+        'should find the closest neighbours for [12, 23, 22, 11, -20], k=3, leafSize=1, cosine distance',
+        () {
+      final kdTree = KDTree(DataFrame(data, headerExists: false), leafSize: 1);
+      final sample = Vector.fromList([12, 23, 22, 11, -20]);
+      final result = kdTree.query(sample, 3, Distance.cosine).toList();
+
+      expect(result, hasLength(3));
+      expect(result[0].index, 17);
+      expect(result[1].index, 8);
+      expect(result[2].index, 0);
+    });
+
+    test(
+        'should find the closest neighbours for [12, 23, 22, 11, -20], k=3, leafSize=3, cosine distance',
+        () {
+      final kdTree = KDTree(DataFrame(data, headerExists: false), leafSize: 3);
+      final sample = Vector.fromList([12, 23, 22, 11, -20]);
+      final result = kdTree.query(sample, 3, Distance.cosine).toList();
+
+      expect(result, hasLength(3));
+      expect(result[0].index, 17);
+      expect(result[1].index, 8);
+      expect(result[2].index, 0);
+    });
+
+    test(
+        'should find the closest neighbours for [12, 23, 22, 11, -20], k=3, leafSize=3, cosine distance for conceivable amount of iterations',
+        () {
+      final kdTree = KDTree(DataFrame(data, headerExists: false),
+          splitStrategy: KDTreeSplitStrategy.largestVariance, leafSize: 1);
+      final sample = Vector.fromList([12, 23, 22, 11, -20]);
+
+      kdTree.query(sample, 3, Distance.cosine).toList();
+
+      expect((kdTree as KDTreeImpl).searchIterationCount, 13);
     });
 
     test('should throw an exception if the query point is of invalid length',