k‑Nearest Neighbors (kNN) (spec model)

This file provides a small kNN classifier/regressor baseline:

• a dataset is an array of (featureTensor, label) pairs,
• prediction is based on the k closest points under a chosen distance.

Equal distances are ordered by the observation's position in the dataset. All classifiers resolve vote ties by the closest observation whose label has the maximum count. Weighted regression averages only the exact matches when any of the selected neighbors has distance zero.

References:

• Cover and Hart (1967), "Nearest Neighbor Pattern Classification": https://ieeexplore.ieee.org/document/1053964

PyTorch / sklearn analogies:

• In the Python ecosystem this is closest to sklearn.neighbors.KNeighborsClassifier / sklearn.neighbors.KNeighborsRegressor.
• kNN is not typically an nn.Module in PyTorch, but it is a common baseline for "classic ML" comparisons and reference checks.

Implementation status

This is a standalone reference specification with targeted executable tests. No API builder or model-correctness theorem is provided here.

Model container

structure Spec.KNN (α : Type) [TorchLean.Storage α] (β : Type) (n : ℕ) : Type

A small kNN model container (parameters + stored dataset).

This is a lazy model: inference consults the stored dataset at query time, rather than learning weights.

• k : ℕ

  Number of neighbors to consult.

• dataset : Array (TorchLean.Tensor α [n] × β)

  Training data: feature vectors paired with labels/targets.

Neighbor selection

The key technical detail here is deterministic tie-breaking: when two points are at exactly the same distance, we prefer the earlier point in the dataset. This makes evaluation stable and keeps formal reasoning about the classifier simpler.

def Spec.findKNearest (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : Array (TorchLean.Tensor α [n] × β)

Find the k nearest neighbors under Euclidean distance.

PyTorch/sklearn analogy: Euclidean L2 distance is the default for many baseline kNN examples.

def Spec.findKNearestWithDistance (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (distanceFn : TorchLean.Tensor α [n] → TorchLean.Tensor α [n] → α) (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : Array (TorchLean.Tensor α [n] × β)

Find the k nearest neighbors under a user-provided distance function.

Notes:

• The distance value is only used for ranking neighbors. It does not need to satisfy metric axioms, but it should be consistent with "smaller means closer".
• We keep deterministic tie-breaking via dataset order.

Classification

def Spec.classify (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [BEq β] [Hashable β] [Inhabited β] (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : β

Majority vote among the neighbors.

Among labels with the same maximum count, choose the one whose closest observation comes first. Equal-distance observations retain dataset order. The nearest observation itself need not belong to a winning class: with ordered labels [A, B, C, C, B], the winner is B. An empty neighborhood, including k = 0, returns default.

def Spec.classifyTreeMap (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [Ord β] [Inhabited β] (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : Option β

Classification using an ordered tree map for label counts.

This uses the same neighbor-order tie rule as classify, with label equality determined by compare. The map supplies counts; its key order does not select the winner. Empty neighborhoods return none.

Regression

def Spec.predict (α : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (knn : KNN α α n) (input : TorchLean.Tensor α [n]) : α

Unweighted kNN regression: average of the neighbor targets.

def Spec.predictWeighted (α : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (knn : KNN α α n) (input : TorchLean.Tensor α [n]) : α

Weighted kNN regression using inverse-distance weights.

Select neighbors using the same (distance, dataset position) order as findKNearest. If one or more selected distances equal zero, return the arithmetic mean of just those targets. Positive distances then have no influence, however small they are. Otherwise use weights w_i = 1 / d_i. If finite positive distances and finite targets overflow the ordinary accumulation, divide all weights by the largest weight, giving d_min / d_i. These weights lie between zero and one. If their weighted sum still overflows, divide each weight by the total before multiplying by its target. No target-dependent scale is introduced. When more than k observations match exactly, dataset order determines which k are averaged. An empty neighborhood or a zero total weight returns 0.

Helpers

def Spec.KNN.fromData (α β : Type) (n k : ℕ) [TorchLean.Storage α] (data : Array (TorchLean.Tensor α [n] × β)) : KNN α β n

Constructor helper (explicit arguments keep elaboration simple in examples).

def Spec.batchPredict (α : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (knn : KNN α α n) (inputs : Array (TorchLean.Tensor α [n])) : Array α

Batch regression: map predict over an array of inputs.

def Spec.batchClassify (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [Hashable β] [Inhabited β] [BEq β] (knn : KNN α β n) (inputs : Array (TorchLean.Tensor α [n])) : Array β

Batch classification: map classify over an array of inputs.

def Spec.classifyWithDistance (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [BEq β] [Hashable β] [Inhabited β] (distanceFn : TorchLean.Tensor α [n] → TorchLean.Tensor α [n] → α) (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : β

Classify with an explicit distance function and the same vote tie rule as classify.

def Spec.classifyWithConfidence (α β : Type) (n : ℕ) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [BEq β] [Hashable β] [Inhabited β] (knn : KNN α β n) (input : TorchLean.Tensor α [n]) : β × α

Classify and return the winning label's fraction of the selected neighbors.

The winner uses the same tie rule as classify. The denominator is the actual number of neighbors, min k dataset.size, so requesting more neighbors than exist does not lower the score. An empty neighborhood returns (default, 0).