⚡ Bolt: Optimize KNN classification with select_nth_unstable_by and squared distance

.jules/bolt.md

@@ -9,3 +9,7 @@
 ## 2026-06-20 - Fast Spatial Neighborhood Checks
 **Learning:** In hot paths doing spatial scanning (like SparseAttentionGraph), calling `libm::sqrt` for distance thresholds is a massive bottleneck. Moreover, using a squared distance optimization `d^2 < r^2` must be robust to NaN values.
 **Action:** Always optimize spatial neighborhood checks using squared distances (`eps_sq`) inside an inline loop. Explicitly check for invalid thresholds before the loop (e.g. `!(epsilon > 0.0)`), and use the `!(sum < eps_sq)` pattern for early exits to handle NaNs safely while matching standard Euclidean semantics without the `sqrt` overhead.
+
+## 2026-06-30 - O(N) KNN Classification with select_nth_unstable_by
+**Learning:** The KNN classifier previously used a bubble sort $O(k \cdot N)$ to find the $k$ nearest neighbors. In high-performance ML routines, this becomes a bottleneck for large datasets or high $k$ values. Furthermore, it used exact Euclidean distance, invoking `libm::sqrt` on every comparison.
+**Action:** Replace $O(k \cdot N)$ sorting with Rust's `select_nth_unstable_by`, which finds the $k$-th element in $O(N)$ average time. Ensure the zero-based index `k_min - 1` is guarded with `if k_min > 0` to prevent underflow. Optimize distance computations in these hot loops by comparing squared distances to eliminate `libm::sqrt` overhead.

crates/aether-core/src/ml/classification.rs

@@ -13,7 +13,6 @@
 // ═══════════════════════════════════════════════════════════════════════════════
 //
 
-
 #![allow(dead_code)]
 
 // use heapless::Vec as HVec;
@@ -200,21 +199,20 @@ impl<const D: usize> KNNClassifier<D> {
         // Find k nearest neighbors
         let mut distances = [(f64::MAX, 0u32); MAX_POINTS];
         for (i, dist) in distances.iter_mut().enumerate().take(self.n_train) {
-            *dist = (self.distance(x, &self.x_train[i]), self.y_train[i]);
+            *dist = (self.squared_distance(x, &self.x_train[i]), self.y_train[i]);
         }
 
-        // Sort by distance (simple bubble sort for small k)
-        for i in 0..self.k.min(self.n_train) {
-            for j in (i + 1)..self.n_train {
-                if distances[j].0 < distances[i].0 {
-                    distances.swap(i, j);
-                }
-            }
+        let k_min = self.k.min(self.n_train);
+        if k_min > 0 {
+            // Find k nearest neighbors in O(N) using select_nth_unstable_by
+            distances[..self.n_train].select_nth_unstable_by(k_min - 1, |a, b| {
+                a.0.partial_cmp(&b.0).unwrap_or(core::cmp::Ordering::Equal)
+            });
         }
 
         // Vote among k nearest
         let mut votes = [0u32; MAX_CLASSES];
-        for dist in distances.iter().take(self.k.min(self.n_train)) {
+        for dist in distances.iter().take(k_min) {
             let label = dist.1 as usize;
             if label < MAX_CLASSES {
                 votes[label] += 1;
@@ -234,14 +232,14 @@ impl<const D: usize> KNNClassifier<D> {
         best_class
     }
 
-    /// Euclidean distance
-    fn distance(&self, a: &[f64; D], b: &[f64; D]) -> f64 {
+    /// Squared Euclidean distance to avoid sqrt overhead in hot paths
+    fn squared_distance(&self, a: &[f64; D], b: &[f64; D]) -> f64 {
         let mut sum = 0.0;
         for i in 0..D {
             let diff = a[i] - b[i];
             sum += diff * diff;
         }
-        sqrt(sum)
+        sum
     }
 }