feat: add threaded monte carlo simulation example (#1)

Author: eegli

Cargo.lock

@@ -22,8 +22,10 @@ pyrust-cli = { path = "crates/pyrust-cli" }
 
 
 num_cpus = "1.16.0"
+rand = "0.9.0"
 anyhow = "1.0"
 clap = { version = "4.5.28", features = ["derive"] }
+
 # "extension-module" tells pyo3 we want to build an extension module (skips linking against libpython.so)
 # "abi3-py39" tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.9
 pyo3 = { version = "0.23.4", features = ["extension-module", "abi3-py39"] }

Cargo.toml

@@ -10,10 +10,25 @@ pub struct Cli {
 
 #[derive(clap::Subcommand)]
 pub enum Subcommand {
-    #[clap(name = "greet")]
-    Greet { name: Option<String> },
+    /// Greet someone.
+    Greet {
+        #[clap(long, short, default_value = "world")]
+        name: String,
+    },
+
+    /// Print the number of CPUs available on the system.
     #[clap(name = "num-cpus")]
     NumCPUs,
+
+    /// Estimate the value of pi using a Monte Carlo method.
+    EstimatePi {
+        /// Number of samples to use.
+        #[clap(long, short)]
+        num_samples: usize,
+        #[clap(long, short = 't')]
+        /// Number of threads to use. If not provided, the number of CPUs available on the system will be used.
+        num_threads: Option<usize>,
+    },
 }
 
 pub fn run(args: Option<&Vec<String>>) {
@@ -26,10 +41,19 @@ pub fn run(args: Option<&Vec<String>>) {
 
     match cli.command {
         Subcommand::Greet { name } => {
-            println!("Hello, {}!", name.unwrap_or_else(|| "world".to_string()));
+            println!("Hello, {name}!",);
         }
         Subcommand::NumCPUs => {
             println!("{}", num_cpus_available());
         }
+        Subcommand::EstimatePi {
+            num_samples,
+            num_threads,
+        } => {
+            let pi_estimated = pyrust_internal::pi::monte_carlo_pi(num_samples, num_threads);
+            let pi_actual = std::f32::consts::PI;
+            let error = (pi_estimated - pi_actual).abs();
+            println!("🍰 Estimated Pi to be {pi_estimated}. That's an error of {error}.");
+        }
     };
 }

crates/pyrust-cli/src/cli.rs

@@ -12,3 +12,4 @@ repository = { workspace = true }
 
 [dependencies]
 num_cpus = { workspace = true }
+rand = { workspace = true }

crates/pyrust-internal/Cargo.toml

@@ -1,3 +1,5 @@
+pub mod pi;
+
 pub fn num_cpus_available() -> usize {
     num_cpus::get()
 }

crates/pyrust-internal/src/lib.rs

@@ -0,0 +1,73 @@
+use rand::distr::StandardUniform;
+use rand::Rng;
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
+
+#[derive(Debug, Default)]
+struct MonteCarloPiEstimate {
+    points_in_circle: usize,
+    points_total: usize,
+}
+
+fn get_random_points(num_samples: usize) -> Vec<(f32, f32)> {
+    let mut rng = rand::rng();
+    let x: Vec<f32> = (&mut rng)
+        .sample_iter(StandardUniform)
+        .take(num_samples)
+        .collect();
+    let y: Vec<f32> = (&mut rng)
+        .sample_iter(StandardUniform)
+        .take(num_samples)
+        .collect();
+
+    x.into_iter().zip(y).collect()
+}
+
+fn is_in_circle(point: &(f32, f32)) -> bool {
+    let (x, y) = point;
+    x * x + y * y <= 1.0
+}
+
+pub fn monte_carlo_pi(num_samples: usize, n_threads: Option<usize>) -> f32 {
+    let n_threads = n_threads.unwrap_or_else(num_cpus::get);
+    let samples_per_thread = num_samples.div_ceil(n_threads);
+
+    println!("Using {} threads", n_threads);
+    println!(
+        "Using {} samples per thread, {} total samples",
+        samples_per_thread, num_samples
+    );
+
+    let mut thread_handles = Vec::with_capacity(n_threads);
+    // Sure enough, we could simply compute an estimate for each thread and average them at the end. To show that mutexes allow sharing data between threads, we'll use a single estimate struct instead
+    let estimates = Arc::new(Mutex::new(MonteCarloPiEstimate::default()));
+
+    let start = Instant::now();
+
+    for thread_id in 0..n_threads {
+        let estimates = estimates.clone();
+        let thread = std::thread::spawn(move || {
+            println!("Thread {thread_id} started");
+
+            let points = get_random_points(samples_per_thread);
+            let points_in_circle = points.iter().filter(|&p| is_in_circle(p)).count();
+            let mut estimates = estimates.lock().unwrap();
+            estimates.points_in_circle += points_in_circle;
+            estimates.points_total += samples_per_thread;
+
+            println!("Thread {thread_id} finished");
+        });
+
+        thread_handles.push(thread);
+    }
+
+    for thread_handle in thread_handles {
+        thread_handle.join().unwrap();
+    }
+
+    println!("Estimating pi took {:?}", start.elapsed());
+    let estimates = estimates.lock().unwrap();
+    println!("Simulation summary: {estimates:?}");
+
+    4.0 * estimates.points_in_circle as f32 / estimates.points_total as f32
+}

crates/pyrust-internal/src/pi.rs

@@ -0,0 +1 @@
+max_width = 100

rustfmt.toml

@@ -20,6 +20,7 @@ test_command() {
 
 # entry points
 test_command "uv run pyrust num-cpus"
+test_command "uv run pyrust estimate-pi -n 100 -t 2"
 test_command "uv run python -m pyrust greet"
 
 echo -e "${GREEN}All tests passed!${NC}"