Python, much like the cheerful Flounder, is known for its friendliness and accessibility. It darts through development waters with ease and grace, beloved by many for its simplicity. But every now and then, Python finds the currents of performance and efficiency challenging to navigate.
Enter Rust, the Sebastian of the coding sea. Wise and performance-oriented, Rust offers unparalleled safety and speed. However, Rust's strictness can sometimes feel confining, contrasting with Python's free-flowing nature.
A match made in heaven, or perhaps more fittingly, under the ocean, occurs when Rust guides Python through the trickier, more performance-intensive waters
This story unfolds as a captivating journey where the agile Flounder, representing the Python programming language, navigates the vast seas of coding under the wise guidance of Sebastian, symbolizing Rust. Central to their adventure are three powerful tridents: cargo, PyO3, and maturin.
Chapter 1: The Call for Optimization
In our tale, Python seeks to enhance its capabilities, longing for the speed and efficiency needed to process a complex task: implementing a k-Nearest Neighbors (kNN) algorithm. The kNN, simple yet effective, is the perfect test for combining Python's simplicity with Rust's performance.
installation stepsfor complete codes, project structure and requirements
# rust, pixi and python
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
curl -fsSL https://pixi.sh/install.sh | bash
# download the repository knn and install
pixi install
pixi run maturin build && pixi run python -m use.ml
Rust kNN
// src/crab.rc
use ndarray::Array1;
use std::collections::HashMap;
// calculates the Euclidean distance between two points
pub fn euclidean_distance(x1: &Array1<f64>, x2: &Array1<f64>) -> f64 {
(x1 - x2).mapv(|a| a.powi(2)).sum().sqrt()
}
pub struct KNN {
k: usize, // number of nearest neighbors
x_train: Option<Vec<Array1<f64>>>,
y_train: Option<Vec<i32>>,
}
impl KNN {
pub fn new(k: usize) -> KNN {
KNN {
k,
x_train: None,
y_train: None,
}
}
pub fn fit(&mut self, x: Vec<Array1<f64>>, y: Vec<i32>) {
self.x_train = Some(x);
self.y_train = Some(y);
}
pub fn predict(&self, x: Vec<Array1<f64>>) -> Vec<i32> {
x.iter().map(|xi| self.predict_one(xi)).collect()
}
fn predict_one(&self, x: &Array1<f64>) -> i32 {
let x_train = self.x_train.as_ref().expect("Model not fitted");
let y_train = self.y_train.as_ref().expect("Model not fitted");
// compute the distance of the input point from all training points
let distances = x_train
.iter()
.map(|x_train| euclidean_distance(x, x_train))
.collect::<Vec<_>>();
// sort the indices of training points by their distance to the input point
let mut indices: Vec<usize> = (0..distances.len()).collect();
indices.sort_by(|&i, &j| distances[i].partial_cmp(&distances[j]).unwrap());
// count the labels of the k nearest neighbors
let mut counter = HashMap::new();
for &i in &indices[..self.k] {
*counter.entry(y_train[i]).or_insert(0) += 1;
}
// return the most common label among the k nearest neighbors
counter
.into_iter()
.max_by_key(|&(_, count)| count)
.map(|(label, _)| label)
.unwrap()
}
}
crab.rs tests
// above code ...
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_euclidean_distance() {
let x1 = Array1::from(vec![1.0, 2.0]);
let x2 = Array1::from(vec![4.0, 6.0]);
let distance = euclidean_distance(&x1, &x2);
let expected_distance = 5.0;
assert!((distance - expected_distance).abs() < 1e-5);
}
#[test]
fn test_knn_predict() {
let mut knn = KNN::new(1);
// Train the model with iris first and last data
let x_train = vec![
Array1::from(vec![5.1, 3.5, 1.4, 0.2]),
Array1::from(vec![5.9, 3., 5.1, 1.8]),
];
let y_train = vec![0, 2];
knn.fit(x_train, y_train);
// Predict a known value iris second point
let x_test = vec![Array1::from(vec![4.7, 3.2, 1.3, 0.2])];
let predictions = knn.predict(x_test);
assert_eq!(predictions, vec![0]);
}
}
Chapter 2: Enchanted Tools - maturin and PyO3
To bridge their worlds, Python and Rust turn to two magical tools: maturin
and PyO3
. Like Triton's trident, these tools hold the power to unite different realms. maturin
, a build system that compiles Rust libraries into Python packages, works seamlessly alongside PyO3, which allows Rust functions to be called from Python.
These tools promise to simplify the journey, making the integration process feel like part of an enchanting undersea adventure.
// src/snake.rs
use crate::crab::KNN as rKNN;
use ndarray::Array1;
use pyo3::prelude::*;
#[pyclass]
struct KNN {
knn: rKNN,
}
#[pymethods]
impl KNN {
#[new]
fn new(k: usize) -> Self {
KNN { knn: rKNN::new(k) }
}
fn fit(&mut self, x_train: Vec<Vec<f64>>, y_train: Vec<i32>) {
let x_train: Vec<_> = x_train.into_iter().map(Array1::from).collect();
self.knn.fit(x_train, y_train);
}
fn predict(&self, x_test: Vec<Vec<f64>>) -> Vec<i32> {
let x_test: Vec<_> = x_test.into_iter().map(Array1::from).collect();
self.knn.predict(x_test)
}
}
#[pymodule]
fn knn(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<KNN>()?;
Ok(())
}
src/lib.rs
mod snake;
pub mod crab;
Chapter 3: Overcoming the Ursula of Complexity
But no adventure is without its challenges. Complexity and integration issues loom like Ursula, threatening to hinder the progress of our heroes. Rust's strict type system and ownership rules can be daunting, and Python's dynamic nature poses its own set of integration challenges.
Yet, with the power of maturin and PyO3, these obstacles are skillfully navigated. The tools weave Rust's functions into Python's script like a spell, allowing Python to call upon the kNN algorithm as if it were its own.
# develop | build depending on the stage knn package
pixi run maturin build
As the final step, our heroes test their creation. Python, with its extensive libraries and tools, puts the Rust-implemented kNN to the test, feeding it data as it would in any Pythonic environment. The algorithm performs flawlessly, classifying data points with the speed and accuracy that only Rust can provide, all within the comfortable and familiar embrace of Python.
# use/ml.py
from numpy.typing import NDArray
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from knn import KNN
def get_iris(
train_size: float = 0.8, random_state: int = 42
) -> tuple[NDArray, NDArray, NDArray, NDArray]:
iris = load_iris()
X, y = iris.data, iris.target
return train_test_split(
X, y, train_size=train_size, stratify=y, random_state=random_state
)
if __name__ == "__main__":
# get data
X_train, X_test, y_train, y_test = get_iris(train_size=0.8)
# get number of classes
k = len(set(y_train))
# create a kNN predictor and train
predictor = KNN(k)
predictor.fit(X_train, y_train)
# predict on test set
y_pred = predictor.predict(X_test)
# evaluate performance
print(classification_report(y_true=y_test, y_pred=y_pred))
The Final Showdown: A New Dawn in the Programming Kingdom
As our story concludes, Python and Rust, along with their magical allies maturin
and PyO3
, have shown that together, they are more than the sum of their parts. They prove that combining the strengths of different languages can lead to powerful and efficient solutions, pushing the boundaries of what's possible in the programming world.
pixi run python -m use.ml
In the end, the sea of code is vast and deep, filled with opportunities for those brave enough to explore its depths. Python and Rust, once separate, now swim together in harmony, ready for the next great adventure.
⚠️ Rust code is not, and was not meant, to be the effient way of performing kNN. If you do have a more effient way, do share on the comments.
"Now, [until then
] let's eat before this crab wanders off my plate." -Grimsby
Top comments (5)
Great article! How does Rust access nets like ResNet-50?
I am not sure what you mean? Do you mean a Rust Torch Vision or specific on reading ResNet onto Rust as Loading and Running a PyTorch Model in Rust
Specific on reading ResNet onto Rust. Thanks for the links!
You can actually do both :) PyO3 also allows calling Python in Rust.
I should also add Candle -> GitHub, and burn 🔥 which have PyTorch‘s look-and-feel