trueskill-tt/examples/atp.rs

use plotters::prelude::*;
use time::{Date, Month};
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 from = Date::from_calendar_date(1900, Month::January, 1).unwrap();
    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 date = Date::parse(&row["time_start"], &time_format).unwrap();

        times.push((date - from).whole_days());
    }

    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 = [
        ("aggasi", "a092", 38800),
        ("borg", "b058", 30300),
        ("connors", "c044", 31250),
        ("courier", "c243", 35750),
        ("djokovic", "d643", i64::MAX),
        ("edberg", "e004", 34750),
        ("federer", "f324", i64::MAX),
        ("hewitt", "h432", 40750),
        ("mcenroe", "m047", 33000),
        ("lendl", "l018", 33750),
        ("murray", "mc10", 60750),
        ("nadal", "n409", i64::MAX),
        ("nastase", "n008", 28750),
        ("sampras", "s402", i64::MAX),
        ("wilander", "w023", 32600),
    ];

    let curves = hist.learning_curves();

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

    for (id, cutoff) in players
        .iter()
        .map(|&(_, id, cutoff)| (index_map.get_or_create(id), cutoff))
    {
        for (ts, gs) in &curves[&id] {
            if *ts >= cutoff {
                continue;
            }

            let ts = *ts as f64;

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

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

            let upper = gs.mu + gs.sigma;
            let lower = gs.mu - gs.sigma;

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

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

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

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

    let mut chart = ChartBuilder::on(&root)
        .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, cutoff)) in players
        .iter()
        .map(|&(player, id, cutoff)| (player, index_map.get_or_create(id), cutoff))
        .enumerate()
    {
        let mut data = Vec::new();
        let mut upper = Vec::new();
        let mut lower = Vec::new();

        for (ts, gs) in curves[&id].iter() {
            if *ts >= cutoff {
                continue;
            }

            data.push((*ts as f64, gs.mu));

            upper.push((*ts as f64, gs.mu + gs.sigma));
            lower.push((*ts as f64, gs.mu - gs.sigma));
        }

        let color = Palette99::pick(idx);

        let band = upper
            .into_iter()
            .chain(lower.into_iter().rev())
            .collect::<Vec<_>>();

        chart
            .plotting_area()
            .draw(&Polygon::new(band, &color.mix(0.15)))
            .unwrap();

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