In Rust for Python: A Match from Heaven

#python #rust #machinelearning

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.

for complete codes, project structure and requirements

complete code and structure code
installation steps

# 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