K-means++ clustering
K-means
 This data was partitioned into 7 clusters using the K-means algorithm.
The task is to implement the K-means++ algorithm. Produce a function which takes two arguments: the number of clusters K, and the dataset to classify. K is a positive integer and the dataset is a list of points in the Cartesian plane. The output is a list of clusters (related sets of points, according to the algorithm).

For extra credit (in order):

 Provide a function to exercise your code, which generates a list of random points.
 Provide a visualization of your results, including centroids (see example image). 
 Generalize the function to polar coordinates (in radians).
 Generalize the function to points in an arbitrary N space (i.e. instead of x,y pairs, points are an N-tuples representing coordinates in ℝN). 
If this is different or more difficult than the [naive] solution for ℝ2, discuss what had to change to support N dimensions.
Extra credit is only awarded if the examples given demonstrate the feature in question. To earn credit for 1. and 2., visualize 6 clusters of 30,000 points in ℝ2. It is not necessary to provide visualization for spaces higher than ℝ2 but the examples should demonstrate features 3. and 4. if the solution has them.

extern crate csv;
extern crate getopts;
extern crate gnuplot;
extern crate nalgebra;
extern crate num;
extern crate rand;
extern crate rustc_serialize;
extern crate test;
 
use getopts::Options;
use gnuplot::{Axes2D, AxesCommon, Color, Figure, Fix, PointSize, PointSymbol};
use nalgebra::{DVector, Iterable};
use rand::{Rng, SeedableRng, StdRng};
use rand::distributions::{IndependentSample, Range};
use std::f64::consts::PI;
use std::env;
 
type Point = DVector<f64>;
 
struct Cluster<'a> {
    members: Vec<&'a Point>,
    center: Point,
}
 
struct Stats {
    centroids: Vec<Point>,
    mean_d_from_centroid: DVector<f64>,
}
 
/// DVector doesn't implement BaseFloat, so a custom distance function is required.
fn sqdist(p1: &Point, p2: &Point) -> f64 {
    (p1.clone() - p2.clone()).iter().map(|x| x * x).fold(0f64, |a, b| a + b)
}
 
/// Returns (distance^2, index) tuple of winning point.
fn nearest(p: &Point, candidates: &Vec<Point>) -> (f64, usize) {
    let (dsquared, the_index) = candidates.iter()
                                          .enumerate()
                                          .fold((sqdist(p, &candidates[0]), 0),
                                                |(d, index), next| {
                                                    let dprime = sqdist(p, &candidates[next.0]);
                                                    if dprime < d {
                                                        (dprime, next.0)
                                                    } else {
                                                        (d, index)
                                                    }
                                                });
    (dsquared, the_index)
}
 
/// Computes starting centroids and makes initial assignments.
fn kpp(points: &Vec<Point>, k: usize, rng: &mut StdRng) -> Stats {
    let mut centroids: Vec<Point> = Vec::new();
    // Random point for first centroid guess:
    centroids.push(points[rng.gen::<usize>() % points.len()].clone());
    let mut dists: Vec<f64> = vec![0f64; points.len()];
 
    for _ in 1..k {
        let mut sum = 0f64;
        for (j, p) in points.iter().enumerate() {
            let (dsquared, _) = nearest(&p, &centroids);
            dists[j] = dsquared;
            sum += dsquared;
        }
 
        // This part chooses the next cluster center with a probability proportional to d^2
        sum *= rng.next_f64();
        for (j, d) in dists.iter().enumerate() {
            sum -= *d;
            if sum <= 0f64 {
                centroids.push(points[j].clone());
                break;
            }
        }
    }
 
    let clusters = assign_clusters(points, &centroids);
    compute_stats(&clusters)
}
 
fn assign_clusters<'a>(points: &'a Vec<Point>, centroids: &Vec<Point>) -> Vec<Cluster<'a>> {
    let mut clusters: Vec<Cluster> = Vec::new();
 
    for _ in 0..centroids.len() {
        clusters.push(Cluster {
            members: Vec::new(),
            center: DVector::new_zeros(points[0].len()),
        });
    }
 
    for p in points.iter() {
        let (_, nearest_index) = nearest(p, centroids);
        clusters[nearest_index].center = clusters[nearest_index].center.clone() + p.clone();
        clusters[nearest_index].members.push(p);
    }
 
    for i in 0..clusters.len() {
        clusters[i].center = clusters[i].center.clone() / clusters[i].members.len() as f64;
    }
 
    clusters
}
 
/// Computes centroids and mean-distance-from-centroid for each cluster.
fn compute_stats(clusters: &Vec<Cluster>) -> Stats {
    let mut centroids = Vec::new();
    let mut means_vec = Vec::new();
 
    for c in clusters.iter() {
        let pts = &c.members;
        let seed: DVector<f64> = DVector::new_zeros(pts[0].len());
        let centroid = pts.iter().fold(seed, |a, &b| a + b.clone()) / pts.len() as f64;
        means_vec.push(pts.iter().fold(0f64, |acc, pt| acc + sqdist(pt, &centroid).sqrt()) /
                       pts.len() as f64);
        centroids.push(centroid);
    }
 
    Stats {
        centroids: centroids,
        mean_d_from_centroid: DVector::from_slice(means_vec.len(), means_vec.as_slice()),
    }
}
 
fn lloyd<'a>(points: &'a Vec<Point>,
             k: usize,
             stoppage_delta: f64,
             max_iter: u32,
             rng: &mut StdRng)
             -> (Vec<Cluster<'a>>, Stats) {
 
    let mut clusters = Vec::new();
    // Choose starting centroids and make initial assignments
    let mut stats = kpp(points, k, rng);
 
    for i in 1..max_iter {
        let last_means: DVector<f64> = stats.mean_d_from_centroid.clone();
        clusters = assign_clusters(points, &stats.centroids);
        stats = compute_stats(&clusters);
        let err = sqdist(&stats.mean_d_from_centroid, &last_means).sqrt();
        if err < stoppage_delta {
            println!("Stoppage condition reached on iteration {}", i);
            return (clusters, stats);
        }
        // Console output
        print!("Iter {}: ", i);
        for (cen, mu) in stats.centroids.iter().zip(stats.mean_d_from_centroid.iter()) {
            print_dvec(cen);
            print!(" {:1.2} | ", mu);
        }
        print!("{:1.5}\n", err);
    }
 
    println!("Stoppage condition not reached by iteration {}", max_iter);
    (clusters, stats)
}
 
/// Uniform sampling on the unit disk.
fn generate_points(n: u32, rng: &mut StdRng) -> Vec<Point> {
    let r_range = Range::new(0f64, 1f64);
    let theta_range = Range::new(0f64, 2f64 * PI);
    let mut points: Vec<Point> = Vec::new();
 
    for _ in 0..n {
        let root_r = r_range.ind_sample(rng).sqrt();
        let theta = theta_range.ind_sample(rng);
        points.push(DVector::<f64>::from_slice(2, &[root_r * theta.cos(), root_r * theta.sin()]));
    }
 
    points
}
 
// Plot clusters (2d only). Closure idiom allows us to borrow and mutate the Axes2D.
fn viz(clusters: Vec<Cluster>, stats: Stats, k: usize, n: u32, e: f64) {
    let mut fg = Figure::new();
    {
        let prep = |fg: &mut Figure| {
            let axes: &mut Axes2D = fg.axes2d();
            let title: String = format!("k = {}, n = {}, e = {:4}", k, n, e);
            let centroids_x = stats.centroids.iter().map(|c| c[0]);
            let centroids_y = stats.centroids.iter().map(|c| c[1]);
            for cluster in clusters.iter() {
                axes.points(cluster.members.iter().map(|p| p[0]),
                            cluster.members
                                   .iter()
                                   .map(|p| p[1]),
                            &[PointSymbol('O'), PointSize(0.25)]);
            }
            axes.set_aspect_ratio(Fix(1.0))
                .points(centroids_x,
                        centroids_y,
                        &[PointSymbol('o'), PointSize(1.5), Color("black")])
                .set_title(&title[..], &[]);
        };
        prep(&mut fg);
    }
    fg.show();
}
 
fn print_dvec(v: &DVector<f64>) {
    print!("(");
    for elem in v.at.iter().take(v.len() - 1) {
        print!("{:+1.2}, ", elem)
    }
    print!("{:+1.2})", v.at.iter().last().unwrap());
}
 
fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}
 
fn main() {
    let args: Vec<String> = env::args().collect();
    let mut k: usize = 7;
    let mut n: u32 = 30000;
    let mut e: f64 = 1e-3;
    let max_iterations = 100u32;
 
    let mut opts = Options::new();
    opts.optflag("?", "help", "Print this help menu");
    opts.optopt("k",
                "",
                "Number of clusters to assign (default: 7)",
                "<clusters>");
    opts.optopt("n",
                "",
                "Operate on this many points on the unit disk (default: 30000)",
                "<pts>");
    opts.optopt("e",
                "",
                "Min delta in norm of successive cluster centroids to continue (default: 1e-3)",
                "<eps>");
    opts.optopt("f", "", "Read points from file (overrides -n)", "<csv>");
 
    let program = args[0].clone();
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("?") {
        print_usage(&program, opts);
        return;
    }
    match matches.opt_str("k") {
        None => {}
        Some(x) => k = x.parse::<usize>().unwrap(),
    };
    match matches.opt_str("n") {
        None => {}
        Some(x) => n = x.parse::<u32>().unwrap(),
    };
    match matches.opt_str("e") {
        None => {}
        Some(x) => e = x.parse::<f64>().unwrap(),
    };
 
    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);
 
    let mut points: Vec<Point>;
 
    match matches.opt_str("f") {
        None => {
            // Proceed with random 2d data
            points = generate_points(n, &mut rng)
        }
        Some(file) => {
            points = Vec::new();
            let mut rdr = csv::Reader::from_file(file.clone()).unwrap();
            for row in rdr.records().map(|r| r.unwrap()) {
                // row is Vec<String>
                let floats: Vec<f64> = row.iter().map(|s| s.parse::<f64>().unwrap()).collect();
                points.push(DVector::<f64>::from_slice(floats.len(), floats.as_slice()));
            }
            assert!(points.iter().all(|v| v.len() == points[0].len()));
            n = points.len() as u32;
            println!("Read {} points from {}", points.len(), file.clone());
        }
    };
 
    assert!(points.len() >= k);
    let (clusters, stats) = lloyd(&points, k, e, max_iterations, &mut rng);
 
    println!(" k       centroid{}mean dist    pop",
             std::iter::repeat(" ").take((points[0].len() - 2) * 7 + 7).collect::<String>());
    println!("===  {}  ===========  =====",
             std::iter::repeat("=").take(points[0].len() * 7 + 2).collect::<String>());
    for i in 0..clusters.len() {
        print!(" {:>1}    ", i);
        print_dvec(&stats.centroids[i]);
        print!("      {:1.2}       {:>4}\n",
               stats.mean_d_from_centroid[i],
               clusters[i].members.len());
    }
 
    if points[0].len() == 2 {
        viz(clusters, stats, k, n, e)
    }
}
 