DEV Community

Cover image for In Rust for Python: A Match from Heaven
Prayson Wilfred Daniel
Prayson Wilfred Daniel

Posted on

In Rust for Python: A Match from Heaven

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.


crab snake

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
# 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
Enter fullscreen mode Exit fullscreen mode

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()
    }
}
Enter fullscreen mode Exit fullscreen mode

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]);
    }
}
Enter fullscreen mode Exit fullscreen mode

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(())
}

Enter fullscreen mode Exit fullscreen mode

src/lib.rs
mod snake;
pub mod crab; 
Enter fullscreen mode Exit fullscreen mode


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
Enter fullscreen mode Exit fullscreen mode

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))
Enter fullscreen mode Exit fullscreen mode

final

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
Enter fullscreen mode Exit fullscreen mode

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 (6)

Collapse
 
rustmeup profile image
Ben

Great article! How does Rust access nets like ResNet-50?

Collapse
 
proteusiq profile image
Prayson Wilfred Daniel • Edited

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

Collapse
 
rustmeup profile image
Ben

Specific on reading ResNet onto Rust. Thanks for the links!

Thread Thread
 
proteusiq profile image
Prayson Wilfred Daniel

You can actually do both :) PyO3 also allows calling Python in Rust.

Collapse
 
proteusiq profile image
Prayson Wilfred Daniel • Edited

I should also add Candle -> GitHub, and burn πŸ”₯ which have PyTorchβ€˜s look-and-feel

Collapse
 
erickwdaniel profile image
Erick W Daniel

This is usefull