trueskill-tt/examples/atp.rs

use plotters::prelude::*;
use time::Date;
use trueskill_tt::{History, IndexMap};

fn main() {
    let mut csv = csv::Reader::open("examples/atp.csv").unwrap();

    let mut composition = Vec::new();
    let mut results = Vec::new();
    let mut times = Vec::new();

    let time_format = time::format_description::parse("[year]-[month]-[day]").unwrap();

    let mut index_map = IndexMap::new();

    for row in csv.records() {
        if &row["double"] == "t" {
            let w1_id = index_map.get_or_create(&row["w1_id"]);
            let w2_id = index_map.get_or_create(&row["w2_id"]);

            let l1_id = index_map.get_or_create(&row["l1_id"]);
            let l2_id = index_map.get_or_create(&row["l2_id"]);

            composition.push(vec![vec![w1_id, w2_id], vec![l1_id, l2_id]]);
        } else {
            let w1_id = index_map.get_or_create(&row["w1_id"]);

            let l1_id = index_map.get_or_create(&row["l1_id"]);

            composition.push(vec![vec![w1_id], vec![l1_id]]);
        }

        results.push(vec![1.0, 0.0]);

        let time = Date::parse(&row["time_start"], &time_format)
            .unwrap()
            .midnight()
            .assume_utc()
            .unix_timestamp();

        times.push(time / (60 * 60 * 24));
    }

    let mut hist = History::builder().sigma(1.6).gamma(0.036).build();

    hist.add_events(composition, results, times, vec![]);
    hist.convergence(10, 0.01, true);

    let players = [
        ("djokovic", "d643"),
        ("federer", "f324"),
        ("sampras", "s402"),
        ("lendl", "l018"),
        ("connors", "c044"),
        ("nadal", "n409"),
        ("john_mcenroe", "m047"),
        ("bjorn_borg", "b058"),
        ("aggasi", "a092"),
        ("hewitt", "h432"),
        ("edberg", "e004"),
        ("vilas", "v028"),
        ("nastase", "n008"),
        ("courier", "c243"),
        ("kuerten", "k293"),
        ("murray", "mc10"),
        ("wilander", "w023"),
        ("roddick", "r485"),
    ];

    let curves = hist.learning_curves();

    let mut x_spec = (f64::MAX, f64::MIN);
    let mut y_spec = (f64::MAX, f64::MIN);

    for id in players.iter().map(|&(_, id)| index_map.get_or_create(id)) {
        for (ts, gs) in &curves[&id] {
            let ts = *ts as f64;

            if ts < x_spec.0 {
                x_spec.0 = ts;
            }

            if ts > x_spec.1 {
                x_spec.1 = ts;
            }

            let mu = gs.mu as f64;

            if mu < y_spec.0 {
                y_spec.0 = mu;
            }

            if mu > y_spec.1 {
                y_spec.1 = mu;
            }
        }
    }

    let root = SVGBackend::new("plot.svg", (1024, 1024)).into_drawing_area();

    root.fill(&WHITE).unwrap();

    let mut chart = ChartBuilder::on(&root)
        .caption("Hello world", ("sans-serif", 50).into_font())
        .margin(5)
        .x_label_area_size(30)
        .y_label_area_size(30)
        .build_cartesian_2d(x_spec.0..x_spec.1, y_spec.0..y_spec.1)
        .unwrap();

    chart.configure_mesh().draw().unwrap();

    for (idx, (player, id)) in players
        .iter()
        .map(|&(player, id)| (player, index_map.get_or_create(id)))
        .enumerate()
    {
        let mut data = Vec::new();

        for (ts, gs) in &curves[&id] {
            data.push((*ts as f64, gs.mu));
        }

        let color = Palette99::pick(idx);

        chart
            .draw_series(LineSeries::new(data, &color))
            .unwrap()
            .label(player)
            .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &color));
    }

    chart
        .configure_series_labels()
        .background_style(&WHITE.mix(0.8))
        .border_style(&BLACK)
        .draw()
        .unwrap();
}

mod csv {
    use std::fs::File;
    use std::io::{self, BufRead, BufReader, Lines};
    use std::ops;
    use std::path::Path;

    pub struct Reader {
        header_map: Vec<String>,
        lines: Lines<BufReader<File>>,
    }

    impl Reader {
        pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, io::Error> {
            let mut lines = File::open(path).map(BufReader::new)?.lines();

            let header_map = if let Some(header) = lines.next() {
                let header = header?;

                header.split(',').map(Into::into).collect::<Vec<_>>()
            } else {
                Vec::new()
            };

            Ok(Self { header_map, lines })
        }

        pub fn records(&mut self) -> Records<'_> {
            Records {
                header_map: &self.header_map,
                lines: &mut self.lines,
            }
        }
    }

    pub struct Records<'a> {
        header_map: &'a Vec<String>,
        lines: &'a mut Lines<BufReader<File>>,
    }

    impl<'a> Iterator for Records<'a> {
        type Item = Record<'a>;

        fn next(&mut self) -> Option<Self::Item> {
            let line = self.lines.next()?;

            Some(Record {
                header_map: self.header_map,
                columns: line.unwrap().split(',').map(Into::into).collect::<Vec<_>>(),
            })
        }
    }

    pub struct Record<'a> {
        header_map: &'a Vec<String>,
        columns: Vec<String>,
    }

    impl<'a> ops::Index<&str> for Record<'a> {
        type Output = str;

        fn index(&self, index: &str) -> &Self::Output {
            &self.columns[self
                .header_map
                .iter()
                .position(|header| header == index)
                .unwrap()]
        }
    }
}