K-means clustering in Rust

Question

I've implemented K-means clustering in Rust. It's my second Rust project (my first one is here: Randomly selecting an adjective and noun, combining them into a message)

I would like advice on whether I am doing stuff idiomatically, and any sensible optimisations I could make. I also appreciate advice on code-style -- I care a lot about having code that is readable and formatted nicely, but I'm still getting the hang of it in Rust.

A brief explanation of the K-means algorithm:

• You have a set of data you wish to partition into a known number of groups, a.k.a. clusters
• The mean of all the data belonging to a cluster is the cluster's centroid. You decide which cluster a datum belongs to by selecting the cluster which is closest, i.e. the one at the smallest Euclidean distance from the data point
• The algorithm works in two iterative steps. First, we initialise the cluster centroids, then:
  • Assign all the data to their nearest centroid
  • Update the cluster centroids so that they are the mean of all the points assigned to them
  • Repeat until a local minimum is found

Here is the contents of lib.rs, now including my unit tests (I realise I could probably have more tests D:):

use std::path::Path;

extern crate csv;
extern crate rustc_serialize;

/// Store one data point's (or one cluster centroid's) x and y co-ordinates
#[derive(Clone, Debug, RustcDecodable)]
pub struct DataPoint {
    pub x: f64,
    pub y: f64,
}

/// Structure for holding data point's assignments to clusters
#[derive(Clone, Debug)]
pub struct Assignment<'a> {
    data_point: &'a DataPoint,
    cluster_ind: usize,
}


pub fn read_data<P>(file_path: P) -> Vec<DataPoint> where P: AsRef<Path> {
    let mut data = vec![];
    let mut reader = csv::Reader::from_file(file_path).unwrap();
    for data_point in reader.decode() {
        let data_point: DataPoint = data_point.unwrap();
        data.push(data_point);
    }
    data
}


pub fn squared_euclidean_distance(point_a: &DataPoint,
                                  point_b: &DataPoint) -> f64 {
   (point_b.x - point_a.x).powi(2) + (point_b.y - point_a.y).powi(2)
}


pub fn get_index_of_min_val(floats: &Vec<f64>) -> usize {

    floats.iter()
          .enumerate()
          .fold(0,
                | min_ind, (ind, &val) |
                if val == f64::min(floats[min_ind], val) { ind }
                else { min_ind })
}

/// Assign points to clusters
fn expectation<'a>(data: &'a Vec<DataPoint>,
                   cluster_centroids: &Vec<DataPoint>) -> Vec<(Assignment<'a>)> {

    let mut assignments: Vec<(Assignment)> = vec![];
    for point in data {
        let mut distance: Vec<f64> = vec![];
        for cluster in cluster_centroids {
            distance.push(squared_euclidean_distance(&point, cluster));
        }
        assignments.push(Assignment{data_point: point,
                                    cluster_ind: get_index_of_min_val(&distance)});
    }
    assignments
}

pub fn count_assignments(assignments: &Vec<Assignment>,
                         cluster_ind: usize) -> usize {
    let points_in_cluster = get_points_in_cluster(assignments, cluster_ind);
    points_in_cluster.len()
}

pub fn get_points_in_cluster<'a>(assignments: &'a Vec<Assignment>,
                                 cluster_ind: usize) -> Vec<Assignment<'a>> {
    let mut points_in_cluster = assignments.clone();
    points_in_cluster.retain(|&Assignment{data_point: _,
                                          cluster_ind: a}| a == cluster_ind);
    points_in_cluster
}

pub fn sum_assigned_values(assignments: &Vec<Assignment>,
                           cluster_ind: usize) -> DataPoint {
    let points_in_cluster = get_points_in_cluster(assignments, cluster_ind);
    let (mut x_tot, mut y_tot) = (0.0_f64, 0.0_f64);
    for point in points_in_cluster {
        x_tot += point.data_point.x;
        y_tot += point.data_point.y;
    }
    DataPoint{x: x_tot, y: y_tot}
}

/// Update cluster centres
fn maximisation(cluster_centroids: &mut Vec<DataPoint>,
                assignments: &Vec<(Assignment)>) {

    for i in 0..cluster_centroids.len() {
        let num_points = count_assignments(&assignments, i);
        let sum_points = sum_assigned_values(&assignments, i);
        cluster_centroids[i] = DataPoint{
            x: sum_points.x/num_points as f64,
            y: sum_points.y/num_points as f64};
    }
}

pub fn get_error_metric(cluster_centroids: &Vec<DataPoint>,
                        assignments: &Vec<Assignment>) -> f64 {
        let mut error = 0.0;
        for i in 0..assignments.len() {
            let cluster_ind = assignments[i].cluster_ind;
            error += squared_euclidean_distance(assignments[i].data_point,
                                                &cluster_centroids[cluster_ind]);
        }
        error
    }

pub fn kmeans_one_iteration<'a>(cluster_centroids: &mut Vec<DataPoint>,
                                data: &'a Vec<DataPoint>) -> Vec<Assignment<'a>> {
    let assignments = expectation(data, cluster_centroids);
    maximisation(cluster_centroids, &assignments);
    assignments
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_squared_euclidean_distance_simple_case() {
        let origin = DataPoint{x: 0.0, y: 0.0};
        let point = DataPoint{x: 1.0, y: 1.0};
        let expected = 2.0;
        let actual = squared_euclidean_distance(&origin, &point);
        assert_eq!(expected, actual)
    }

    #[test]
    fn test_squared_euclidean_distance_gives_0_for_same_point() {
        let point_a = DataPoint{x: -999.3, y: 10.5};
        let point_b = point_a.clone();
        let expected = 0.0;
        let actual = squared_euclidean_distance(&point_a, &point_b);
        assert_eq!(expected, actual)
    }

    #[test]
    fn test_get_index_of_min_val() {
        let floats = vec![0.0_f64, 1.0_f64, 3.0_f64, -5.5_f64];
        let expected = 3;
        let actual = get_index_of_min_val(&floats);
        assert_eq!(expected, actual)
    }

    #[test]
    fn test_count_assignments_returns_0_when_no_occurences() {
        let dp = DataPoint{x: 0.0, y: 0.0};
        let assignments = vec![Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 1},
                               Assignment{data_point: &dp, cluster_ind: 5},
                               Assignment{data_point: &dp, cluster_ind: 0}];
        let cluster_ind: usize = 4;
        let expected = 0;
        let actual = count_assignments(&assignments, cluster_ind);
        assert_eq!(expected, actual)
    }

    #[test]
    fn test_count_assignments_returns_3_when_3_occurences() {
        let dp = DataPoint{x: 0.0, y: 0.0};
        let assignments = vec![Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 1},
                               Assignment{data_point: &dp, cluster_ind: 5},
                               Assignment{data_point: &dp, cluster_ind: 0}];
        let cluster_ind: usize = 0;
        let expected = 3;
        let actual = count_assignments(&assignments, cluster_ind);
        assert_eq!(expected, actual)
    }

    #[test]
    fn test_sum_assigned_values_returns_0_when_none_assigned() {
        let dp = DataPoint{x: 5.0, y: 5.0};
        let assignments = vec![Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 1},
                               Assignment{data_point: &dp, cluster_ind: 5},
                               Assignment{data_point: &dp, cluster_ind: 0}];
        let cluster_ind: usize = 2;
        let expected = DataPoint{x: 0.0, y: 0.0};
        let actual = sum_assigned_values(&assignments, cluster_ind);
        assert_eq!(expected.x, actual.x);
        assert_eq!(expected.y, actual.y)
    }

    #[test]
    fn test_sum_assigned_values_returns_correctly_when_some_assigned() {
        let dp = DataPoint{x: 1.0, y: 1.0};
        let assignments = vec![Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 0},
                               Assignment{data_point: &dp, cluster_ind: 1},
                               Assignment{data_point: &dp, cluster_ind: 5},
                               Assignment{data_point: &dp, cluster_ind: 0}];
        let cluster_ind: usize = 0;
        let expected = DataPoint{x: 3.0, y: 3.0};
        let actual = sum_assigned_values(&assignments, cluster_ind);
        assert_eq!(expected.x, actual.x);
        assert_eq!(expected.y, actual.y)
    }
}

And here is the contents of main.rs:

extern crate kmeans;

use kmeans::*;

fn main() {
    let data = read_data("../../data/faithful.csv");
    let mut cluster_centroids = vec![DataPoint{x: 2.0, y: 50.0},
                                     DataPoint{x: 7.0, y: 100.0}];
    let (mut error, mut prev_error) = (0.0, -1.0);
    let mut assignments: Vec<Assignment>;
    while error != prev_error {
        prev_error = error;
        assignments = kmeans_one_iteration(&mut cluster_centroids, &data);
        error = get_error_metric(&cluster_centroids, &assignments);
        println!("{}", error);
    }
}

I appreciate any advice! :D

Answer 1

1. DataPoint is a small enough structure, it might as well be Copy.

2. squared_euclidean_distance can be an instance method on DataPoint.

3. where clauses should be placed on a separate line. This makes it easier to see what constraints a function has.

4. It's rarely needed to declare a container (like a Vec) and then fill it up manually. Use iterator adapters instead; map and collect are the big two. Think about iterator adapters any time you use a for loop that's not just for side effects.

5. Prefer expect over unwrap. For Results, it includes the underlying error message. In either case, it allows you or the user to track down the error easier.

6. Reading from a file should have better error handling than panicking; does tying the library to CSV even make sense? Maybe this really belongs in the executable?

7. Never take &Vec<T> or &String. &[T] or &str is broader and gives you everything you need.

8. There are no spaces inside the closure argument list delimiters (|).

9. Use braces for multi-line closures

10. It doesn't make sense to return 0 for the minimum value of an empty slice (or zero-length iterator). This is what Option is made for.

11. A few places that create a Vec from an iterator just to iterate over it again. It's more efficient to never collect.

12. get is pretty useless in method names. Dropping it loses nothing.

13. When the first thing a function does is convert to an iterator, you might as well just accept anything that can be made into an iterator.

14. Include spaces around a type constructors { and before the }.

15. There's no reason to include parenthesis in Vec<(T)>.

16. When you have a multi-line function argument list, place the { on a new line.

17. Use iterator adapters instead of cloning a slice / Vec and then retaining values.

18. There's no reason to return a vector when you just convert it into an iterator. I used Box<Iterator> for laziness.

19. The code computes points_in_cluster twice. It's probably more efficient to combine the calculations.

20. There's no reason to make a mutable tuple and then convert into a struct, just mutate the struct directly.

21. Instead of a for loop, use fold.

22. Implement addition for DataPoint to make it simpler to understand.

23. Does it make sense to return 0, 0 for the sum of an empty slice?

24. Use iter_mut and enumerate instead of poking at the slices value via array index (iterators are more efficient).

---

use std::path::Path;

extern crate csv;
extern crate rustc_serialize;

/// Store one data point's (or one cluster centroid's) x and y co-ordinates
#[derive(Copy, Clone, Debug, RustcDecodable)]
pub struct DataPoint {
    pub x: f64,
    pub y: f64,
}

impl DataPoint {
    fn zero() -> DataPoint {
        DataPoint {
            x: 0.0,
            y: 0.0,
        }
    }

    pub fn squared_euclidean_distance(&self, other: &DataPoint) -> f64 {
        (other.x - self.x).powi(2) + (other.y - self.y).powi(2)
    }
}

impl std::ops::Add for DataPoint {
    type Output = DataPoint;

    fn add(self, other: DataPoint) -> DataPoint {
        DataPoint {
            x: self.x + other.x,
            y: self.y + other.y,
        }
    }
}

/// Structure for holding data point's assignments to clusters
#[derive(Clone, Debug)]
pub struct Assignment<'a> {
    data_point: &'a DataPoint,
    cluster_ind: usize,
}

pub fn read_data<P>(file_path: P) -> Vec<DataPoint>
    where P: AsRef<Path>
{
    let mut reader = csv::Reader::from_file(file_path).unwrap();
    reader.decode().map(|point| point.unwrap()).collect()
}

pub fn index_of_min_val<I>(floats: I) -> Option<usize>
    where I: IntoIterator<Item = f64>,
{
    let mut iter = floats.into_iter().enumerate();

    iter.next().map(|(i, min)| {
        iter.fold((i, min), |(min_i, min_val), (i, val)| {
            if val < min_val {
                (i, val)
            } else {
                (min_i, min_val)
            }
        }).0
    })
}


/// Assign points to clusters
fn expectation<'a>(data: &'a [DataPoint],
                   cluster_centroids: &[DataPoint]) -> Vec<Assignment<'a>>
{
    data.iter().map(|point| {
        let distances = cluster_centroids.iter().map(|cluster| point.squared_euclidean_distance(cluster));
        let index = index_of_min_val(distances).expect("No minimum value found");
        Assignment { data_point: point, cluster_ind: index }
    }).collect()
}

pub fn count_assignments(assignments: &[Assignment],
                         cluster_ind: usize) -> usize
{
    points_in_cluster(assignments, cluster_ind).count()
}


pub fn points_in_cluster<'a>(assignments: &'a [Assignment],
                                 expected_cluster_ind: usize) -> Box<Iterator<Item = Assignment<'a>> + 'a>
{
    let i = assignments.into_iter()
        .cloned()
        .filter(move |&Assignment { cluster_ind, .. }| expected_cluster_ind == cluster_ind);
    Box::new(i)
}

pub fn sum_assigned_values(assignments: &[Assignment],
                           cluster_ind: usize) -> DataPoint
{
    points_in_cluster(assignments, cluster_ind)
        .into_iter()
        .fold(DataPoint::zero(), |acc, point| acc + *point.data_point)
}


/// Update cluster centres
fn maximisation(cluster_centroids: &mut [DataPoint],
                assignments: &[Assignment])
{
    for (i, centroid) in cluster_centroids.iter_mut().enumerate() {
        let num_points = count_assignments(&assignments, i);
        let sum_points = sum_assigned_values(&assignments, i);
        *centroid = DataPoint {
            x: sum_points.x / num_points as f64,
            y: sum_points.y / num_points as f64
        };
    }
}

pub fn get_error_metric(cluster_centroids: &[DataPoint],
                        assignments: &[Assignment]) -> f64
{
    assignments.iter().fold(0.0, |error, assignment| {
        let centroid = &cluster_centroids[assignment.cluster_ind];
        error + assignment.data_point.squared_euclidean_distance(centroid)
    })
}

pub fn kmeans_one_iteration<'a>(cluster_centroids: &mut [DataPoint],
                                data: &'a [DataPoint]) -> Vec<Assignment<'a>>
{
    let assignments = expectation(data, cluster_centroids);
    maximisation(cluster_centroids, &assignments);
    assignments
}

---

The tests are straight-forward, but there are some small changes I would make to tighten them up:

1. Because we now accept slices instead of &Vec, we don't need to create as many vectors. This is a small efficiency gain.
2. Implementing PartialEq (via derive) for DataPoint means you don't have to compare the properties individually.
3. I didn't find many of the variables introduced to be useful, so I inlined most of them while keeping line length shorter. I wish that Rust's testing framework made more of a distinction between expected and actual values, but it doesn't...

---

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_squared_euclidean_distance_simple_case() {
        let origin = DataPoint { x: 0.0, y: 0.0 };
        let point = DataPoint { x: 1.0, y: 1.0 };
        assert_eq!(2.0, origin.squared_euclidean_distance(&point))
    }

    #[test]
    fn test_squared_euclidean_distance_gives_0_for_same_point() {
        let point_a = DataPoint { x: -999.3, y: 10.5 };
        assert_eq!(0.0, point_a.squared_euclidean_distance(&point_a));
    }

    #[test]
    fn test_index_of_min_val() {
        let floats = vec![0.0_f64, 1.0_f64, 3.0_f64, -5.5_f64];
        assert_eq!(Some(3), index_of_min_val(floats))
    }

    #[test]
    fn test_count_assignments_returns_0_when_no_occurences() {
        let dp = DataPoint { x: 0.0, y: 0.0 };
        let assignments = [Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 1 },
                           Assignment { data_point: &dp, cluster_ind: 5 },
                           Assignment { data_point: &dp, cluster_ind: 0 }];
        assert_eq!(0, count_assignments(&assignments, 4))
    }

    #[test]
    fn test_count_assignments_returns_3_when_3_occurences() {
        let dp = DataPoint { x: 0.0, y: 0.0 };
        let assignments = [Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 1 },
                           Assignment { data_point: &dp, cluster_ind: 5 },
                           Assignment { data_point: &dp, cluster_ind: 0 }];
        assert_eq!(3, count_assignments(&assignments, 0));
    }

    #[test]
    fn test_sum_assigned_values_returns_0_when_none_assigned() {
        let dp = DataPoint { x: 5.0, y: 5.0 };
        let assignments = [Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 1 },
                           Assignment { data_point: &dp, cluster_ind: 5 },
                           Assignment { data_point: &dp, cluster_ind: 0 }];
        assert_eq!(DataPoint { x: 0.0, y: 0.0 }, sum_assigned_values(&assignments, 2))
    }

    #[test]
    fn test_sum_assigned_values_returns_correctly_when_some_assigned() {
        let dp = DataPoint { x: 1.0, y: 1.0 };
        let assignments = [Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 0 },
                           Assignment { data_point: &dp, cluster_ind: 1 },
                           Assignment { data_point: &dp, cluster_ind: 5 },
                           Assignment { data_point: &dp, cluster_ind: 0 }];
        assert_eq!(DataPoint { x: 3.0, y: 3.0 }, sum_assigned_values(&assignments, 0));
    }
}

---

Can we not just return i from points_in_cluster? Could you explain why you need to create a Box?

This is probably best answered by Correct way to return an Iterator?.

My understanding: returning a trait causes problems as it doesn't have a defined size. Allocating on the heap makes sense when you don't statically know the size needed. Iterator is a trait, so this applies here.

That sounds about right. I said I was being lazy because if I cared about maximal performance, I could have defined my own iterator that wouldn't need any heap allocation. A future version of Rust is likely to allow returning a statically-known concrete type that is guaranteed to implement a trait but without explicitly listing the type.

In the same function you appear to use a move keyword. I can't seem to find up to date docs for this.

The The Rust Programming Language section on closures discusses the move keyword. In this case, we are transferring ownership of expected_cluster_ind to the closure.