feat: initial implementation

This commit is contained in:
swaits committed 2023-12-04 19:59:30 -07:00
commit 6c565ff0a3
9 files changed
+624

No files matched your search

+90
View File
@@ -0,0 +1,90 @@
use std::str::FromStr;
use indoc::indoc;
use lazy_static::lazy_static;
use regex::Regex;
use crate::elo::Rating;
// a chess player - the minimum data we care about, just name and rating
#[derive(Clone, Debug)]
pub struct Player {
pub name: String,
pub rating: Rating,
}
// a single chess game, white/black players and the outcome (actual or simulated)
#[derive(Clone, Debug)]
pub struct Game {
pub white: Player,
pub black: Player,
pub outcome: Outcome,
}
// possible outocmes of a chess game
#[derive(Clone, Debug)]
pub enum Outcome {
WhiteWin,
BlackWin,
Draw,
}
// this regex captures all the possible fields in a PGN as exported from chess.com's API
lazy_static! {
static ref RE: Regex = Regex::new(indoc! {r#"
\[Event\s+"(?P<Event>.+?)"\]\s*
\[Site\s+"(?P<Site>.+?)"\]\s*
\[Date\s+"(?P<Date>.+?)"\]\s*
\[Round\s+"(?P<Round>.+?)"\]\s*
\[White\s+"(?P<White>.+?)"\]\s*
\[Black\s+"(?P<Black>.+?)"\]\s*
\[Result\s+"(?P<Result>.+?)"\]\s*
\[CurrentPosition\s+"(?P<CurrentPosition>.+?)"\]\s*
\[Timezone\s+"(?P<Timezone>.+?)"\]\s*
\[ECO\s+"(?P<ECO>.+?)"\]\s*
\[ECOUrl\s+"(?P<ECOUrl>.+?)"\]\s*
\[UTCDate\s+"(?P<UTCDate>.+?)"\]\s*
\[UTCTime\s+"(?P<UTCTime>.+?)"\]\s*
\[WhiteElo\s+"(?P<WhiteElo>.+?)"\]\s*
\[BlackElo\s+"(?P<BlackElo>.+?)"\]\s*
\[TimeControl\s+"(?P<TimeControl>.+?)"\]\s*
\[Termination\s+"(?P<Termination>.+?)"\]\s*
\[StartTime\s+"(?P<StartTime>.+?)"\]\s*
\[EndDate\s+"(?P<EndDate>.+?)"\]\s*
\[EndTime\s+"(?P<EndTime>.+?)"\]\s*
\[Link\s+"(?P<Link>.+?)"\]\s*
(?P<PGN>.+?)\n
"#})
.unwrap();
}
// given a sting (expected to be a dump of chess.com's API), parse and return a list of games
pub fn parse_games(input: &str) -> Vec<Game> {
let mut games = Vec::new();
for caps in RE.captures_iter(input) {
let white_player = Player {
name: caps["White"].to_string(),
rating: Rating::from_str(caps["WhiteElo"].to_string().as_str()).unwrap(),
};
let black_player = Player {
name: caps["Black"].to_string(),
// rating: caps["BlackElo"].to_string().as_str().into(),
rating: Rating::from_str(caps["BlackElo"].to_string().as_str()).unwrap(),
};
let outcome = match &caps["Result"] {
"1-0" => Outcome::WhiteWin,
"0-1" => Outcome::BlackWin,
"1/2-1/2" => Outcome::Draw,
_ => continue, // or handle unknown result
};
games.push(Game {
white: white_player,
black: black_player,
outcome,
});
}
games
}
+73
View File
@@ -0,0 +1,73 @@
// standard imports
use std::{num::ParseIntError, str::FromStr};
// crate imports
use assert_approx_eq::assert_approx_eq;
// probabilities of white or black winning or drawing
#[derive(Debug)]
pub struct ExpectedOutcome {
pub white: f64,
pub black: f64,
pub draw: f64,
}
// a new type for an Elo rating
#[derive(Clone, Debug)]
pub struct Rating(pub u32);
// create a `Rating` from a string (used during loading/parsing)
impl FromStr for Rating {
type Err = ParseIntError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let num = s.parse::<u32>()?; // Or whatever the primitive type should be
Ok(Self(num)) // Replace with the correct method to construct a Rating
}
}
// estimate probability of draw given two Elo ratings
fn probability_of_draw(white_elo: &Rating, black_elo: &Rating) -> f64 {
// Constants for draw probability estimation
let max_draw_probability = 0.18; // Upper limit of draw probability
let rating_difference_sensitivity = 0.05; // Adjusts sensitivity to rating difference
let rating_diff = (black_elo.0 as f64 - white_elo.0 as f64).abs();
max_draw_probability / (1.0 + rating_difference_sensitivity * rating_diff)
}
// estimate probability of white winning, given two Elo ratings
fn probability_of_white_win(white_elo: &Rating, black_elo: &Rating) -> f64 {
// Expected outcomes for players A and B
// E_a = 1 / (1 + 10^((b - a) / 400)
// E_b = 1 - E_a
1.0 / (1.0 + 10.0f64.powf((black_elo.0 as f64 - white_elo.0 as f64) / 400.0))
}
// estimate probability of black winning given white's probability of winning
fn probability_of_black_win(probability_of_white_win: f64) -> f64 {
1.0 - probability_of_white_win
}
// given two ratings, compute probabilities of win/loss/draw and return in `ExpectedOutcome`
pub fn calculate_expected_outcome(white_elo: &Rating, black_elo: &Rating) -> ExpectedOutcome {
// Estimate win probabilities
let mut p_white = probability_of_white_win(white_elo, black_elo);
let mut p_black = probability_of_black_win(p_white);
// Estimate draw probability based on rating difference
let p_draw = probability_of_draw(white_elo, black_elo);
// Adjust win probabilities to account for draw probability
p_white *= 1.0 - p_draw;
p_black *= 1.0 - p_draw;
// Sanity check
assert_approx_eq!((p_white + p_black + p_draw), 1.0);
ExpectedOutcome {
// Estimate win probabilities
white: p_white,
black: p_black,
draw: p_draw,
}
}
+124
View File
@@ -0,0 +1,124 @@
// project modules
mod chess;
mod elo;
mod stochastic;
// standard imports
use std::{collections::HashMap, fs, sync::Arc};
// crate imports
use rayon::prelude::*;
// project imports
use chess::{parse_games, Game, Outcome};
use stochastic::simulate_games;
// load games into `Vec<Game>` from the chess.com API output
fn load_games() -> Vec<Game> {
let filename = "games.pgn";
// Read the contents of the file into a string
let contents = fs::read_to_string(filename).unwrap();
// Parse the contents into a Vec<Game>
parse_games(&contents)
}
// find all the win streaks in a `Vec<Game>` for players matching `player_name`
fn generate_win_streak_histogram(player_name: &str, games: &[Game]) -> HashMap<usize, usize> {
let mut histogram = HashMap::new();
let mut current_streak = 0;
for game in games {
let is_player_white = game.white.name == player_name;
let is_player_black = game.black.name == player_name;
assert!(is_player_white || is_player_black);
if !(is_player_white || is_player_black) {
continue;
}
let won_as_white = is_player_white && matches!(game.outcome, Outcome::WhiteWin);
let won_as_black = is_player_black && matches!(game.outcome, Outcome::BlackWin);
if won_as_white || won_as_black {
// our player won
current_streak += 1;
} else if current_streak > 0 {
// our playe lost or drew, record streak and reset it
*histogram.entry(current_streak).or_insert(0) += 1;
current_streak = 0;
}
}
// Handle the last streak if the games list ends on a winning streak
if current_streak > 0 {
*histogram.entry(current_streak).or_insert(0) += 1;
}
histogram
}
// dump the histogram to stdout
fn print_histogram(histogram: &HashMap<usize, usize>, divisor: Option<u32>) {
let max_streak = histogram.keys().max().cloned().unwrap_or(0);
println!("Win Streak Histogram:");
for streak_length in 1..=max_streak {
let count = histogram.get(&streak_length).cloned().unwrap_or(0);
match divisor {
Some(div) => {
let expected_count = (count as f64) / (div as f64);
let probability = expected_count.min(1.0) * 100.0;
println!(
"Streak Length: {:>3}, Expected Count: {:>5.0}, Probability: {:>6.2}%",
streak_length, expected_count, probability
);
}
None => {
println!("Streak Length: {:>3}, Count: {}", streak_length, count);
}
}
}
}
fn main() {
// load actual games
println!("Loading...");
let games = Arc::new(load_games()); // Arc because of rayon use below
println!("Loaded {} games", games.len());
println!();
println!("Actual Results Histogram:");
let histogram = generate_win_streak_histogram("Hikaru", &games);
print_histogram(&histogram, None);
println!();
// Run the simulations in parallel and collect the histograms
println!("Simulating...");
println!();
let iterations = 10_000;
let histograms: Vec<HashMap<usize, usize>> = (0..iterations)
.into_par_iter()
.map(|_| {
let simulated_games = simulate_games(&games);
generate_win_streak_histogram("Hikaru", &simulated_games)
})
.collect();
// Reduce the histograms into a single histogram
let merged_histogram = histograms
.into_iter()
.reduce(|mut acc, h| {
for (k, v) in h {
*acc.entry(k).or_insert(0) += v;
}
acc
})
.unwrap();
// And print it!
println!("Simulated Results ({} games):", iterations);
print_histogram(&merged_histogram, Some(iterations));
}
+37
View File
@@ -0,0 +1,37 @@
use rand::Rng;
use crate::{
chess::{Game, Outcome},
elo::{calculate_expected_outcome, Rating},
};
// simulates a single game
fn simulate_game(white_elo: &Rating, black_elo: &Rating) -> Outcome {
// Get expected outcomes
let e = calculate_expected_outcome(white_elo, black_elo);
// Roll the dice
let mut rng = rand::thread_rng();
let random_value = rng.gen::<f64>();
// Determine and return outcome
if random_value <= e.white {
Outcome::WhiteWin
} else if random_value <= (e.white + e.black) {
Outcome::BlackWin
} else {
Outcome::Draw
}
}
// simulate a series of games and return a new simulated `Vec<Game>` series
pub fn simulate_games(actual: &[Game]) -> Vec<Game> {
actual
.iter()
.map(|g| Game {
white: g.white.to_owned(),
black: g.black.to_owned(),
outcome: simulate_game(&g.white.rating, &g.black.rating),
})
.collect()
}