From df7306bc6566a63867a31fe33971aa2da3c810a4 Mon Sep 17 00:00:00 2001 From: Stephen Waits Date: Mon, 10 Jan 2022 15:06:34 -0700 Subject: [PATCH] refactor main loop for 25% better performance, update README --- README.md | 108 +++++++++++++++++++++++++++++++++++++++++++++++----- src/main.rs | 72 ++++++++++++++++++++--------------- 2 files changed, 140 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index b180a06..7acb67d 100644 --- a/README.md +++ b/README.md @@ -4,17 +4,107 @@ A million permutation hypothesis test utility. ## Installation -1. [install Rust](https://www.rust-lang.org/tools/install) -2. `cargo install --git https://git.sr.ht/~swaits/megapermute` +1. If you don't already have rust installed, [install it](https://www.rust-lang.org/tools/install). +2. Then run: `cargo install --git https://git.sr.ht/~swaits/megapermute` ## Usage -`megapermute` requires two input files, `control.dat` and `treatment.dat`. Each is a text -file containing a list of numbers, treated internally as `f64`. +`megapermute` is a command-line utility which reads two input files, +`control.dat` and `treatment.dat`. Each is a text file containing a list of +numbers, treated internally as 64-bit floating point numbers. -- `control.dat` the control sample -- `treatment.dat` the treatment sample +- `control.dat`: the control sample +- `treatment.dat`: the treatment sample -With these two files in the current directly, running `megapermute` computes -the empircal difference in means and p-value for the null hypothesis (H_0 == -H_a). +With these two files in the current directory, running `megapermute` computes +the empircal difference in means and p-value for the null hypothesis, that H_0 +is equivalent to H_a. You can think of p-value like the probability that the +null hypothesis is true. + +### Example + +Given the following two input files: + +```bash +~/tmp +❯ bat control.dat +───────┬─────────────────────── + │ File: control.dat +───────┼─────────────────────── + 1 │ 52 + 2 │ 104 + 3 │ 146 + 4 │ 10 + 5 │ 51 + 6 │ 30 + 7 │ 40 + 8 │ 27 + 9 │ 46 +───────┴─────────────────────── + +~/tmp +❯ bat treatment.dat +───────┬─────────────────────── + │ File: treatment.dat +───────┼─────────────────────── + 1 │ 94 + 2 │ 197 + 3 │ 16 + 4 │ 38 + 5 │ 99 + 6 │ 141 + 7 │ 23 +───────┴────────────────────── +``` + +Running `megapermute` should yield results like this: + +```bash +~/tmp +❯ megapermute + mu_control = 56.22222222222222 + mu_treatment = 86.85714285714286 +(mu_treatment - mu_control) = 30.63492063492064 + p-value = 0.139556 +``` + +## Performance + +My goal with this was to take advantage of multiple cores by using +[`rayon`](https://crates.io/crates/rayon) iterators to compute the permutation +hypothesis test quickly. As of January 2022, this application can do one +million permutations on the example dataset listed above in 39.4 ms (N=69, +SD=2.6 ms), which is on the order of 25M permutations/second. + +```bash +~/tmp +❯ hyperfine megapermute +Benchmark 1: megapermute + Time (mean ± σ): 39.4 ms ± 2.6 ms [User: 499.3 ms, System: 10.2 ms] + Range (min … max): 35.1 ms … 48.4 ms 69 runs +``` + +Benchmark machine info: + +```bash +~/tmp +❯ neofetch + 'c. stephenwaits@SWaits-Mac-01.local + ,xNMM. -------------------------------- + .OMMMMo OS: macOS 11.6.1 20G224 x86_64 + OMMM0, Host: MacBookPro16,1 + .;loddo:' loolloddol;. Kernel: 20.6.0 + cKMMMMMMMMMMNWMMMMMMMMMM0: Uptime: 31 days, 21 hours, 23 mins + .KMMMMMMMMMMMMMMMMMMMMMMMWd. Packages: 149 (brew) + XMMMMMMMMMMMMMMMMMMMMMMMX. Shell: zsh 5.8 +;MMMMMMMMMMMMMMMMMMMMMMMM: Resolution: 3840x2160, 1920x1080 +:MMMMMMMMMMMMMMMMMMMMMMMM: DE: Aqua +.MMMMMMMMMMMMMMMMMMMMMMMMX. WM: Amethyst + kMMMMMMMMMMMMMMMMMMMMMMMMWd. Terminal: WezTerm + .XMMMMMMMMMMMMMMMMMMMMMMMMMMk CPU: Intel i9-9880H (16) @ 2.30GHz + .XMMMMMMMMMMMMMMMMMMMMMMMMK. GPU: Intel UHD Graphics 630, AMD Radeon Pro 5500M + kMMMMMMMMMMMMMMMMMMMMMMd Memory: 21514MiB / 32768MiB + ;KMMMMMMMWXXWMMMMMMMk. + .cooc,. .,coo:. + +``` diff --git a/src/main.rs b/src/main.rs index dfb0ce5..926d8c5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ const N_THREADS: usize = 1_000; // the number of permutation tests to run per "thread" (above) const N_PERMUTATIONS_PER_THREAD: usize = 1_000; +// load a file of numbers as f64 and panic if we run into any problems fn load_f64s(filename: &str) -> Result> { let f = BufReader::new(File::open(filename)?); Ok(f.lines() @@ -31,13 +32,33 @@ fn mean<'a>(iter: impl Iterator) -> f64 { .fold(0.0, |mu, (i, x)| mu + ((x - mu) / (i + 1) as f64)) } +// this enum is used to denote group memebership during each permutation #[derive(Clone, PartialEq, Eq)] enum Group { Control, Treatment, } -fn permute(control: &[f64], treatment: &[f64], mu_diff: f64) -> f64 { +// Run the permutation tests. +// +// Accepts the control and treatment arrays along with the difference in empircal means (treatment +// minus control). +// +// Runs N_PERMUTATIONS_PER_THREAD on each of N_THREADS using Rayon's `par_iter()`. +// +// The original data (arrays) is never copied. Each of the iterations creates an index array of +// `enum Group` denoting `Control` or `Treatment` at each index. That `index` array is shuffled to +// permute group memebership. Finally, each of the group's means are computed and compared, and +// filtered to only count differences which are larger than the `mu_diff` parameter. +// +// The final p-value is the number of permutations where the diff in means exceeded `mu_diff` +// divided by (N_THREADS * N_PERMUTATIONS_PER_THREAD). +// +// Note: the left tail or right tail is chosen automatically based on whether the empircal mean +// delta is positive or negative. +fn permutation_test(control: &[f64], treatment: &[f64], mu_diff: f64) -> f64 { + // use Rayon to divide this work across N_THREADS, ultimately counting the number of + // permutations where delta(means) > mu_diff let count: f64 = (0..N_THREADS) .into_par_iter() .map(|_| { @@ -57,39 +78,28 @@ fn permute(control: &[f64], treatment: &[f64], mu_diff: f64) -> f64 { // shuffle our group selections index.shuffle(&mut rng); - // comput emean of this "control permutation" - let permuted_mean_control = mean( - control - .iter() - .chain(treatment.iter()) - .enumerate() - .filter_map(|(i, x)| { - if index[i] == Group::Control { - Some(x) - } else { - None - } - }), - ); + // variables for mean computation + let (mut mu_control, mut n_control) = (0.0, 0.0); + let (mut mu_treatment, mut n_treatment) = (0.0, 0.0); - // comput emean of this "treatment permutation" - let permuted_mean_treatment = mean( - control - .iter() - .chain(treatment.iter()) - .enumerate() - .filter_map(|(i, x)| { - if index[i] == Group::Treatment { - Some(x) - } else { - None - } - }), + // walk the combined iterator and add each element to the corresponding mean + // using Welford's online algorithm + control.iter().chain(treatment.iter()).enumerate().for_each( + |(i, x)| match index[i] { + Group::Control => { + n_control += 1.0; + mu_control += (x - mu_control) / n_control + } + Group::Treatment => { + n_treatment += 1.0; + mu_treatment += (x - mu_treatment) / n_treatment + } + }, ); // select this permutation if the diff in these permuted means is more than the // empirical diff of means - (permuted_mean_treatment - permuted_mean_control) > mu_diff + (mu_treatment - mu_control) > mu_diff }) .count() as f64 }) @@ -123,7 +133,7 @@ fn main() -> Result<()> { ); // run permutation test to compute p-value - let pvalue = permute(&control, &treatment, mean_treatment - mean_control); + let pvalue = permutation_test(&control, &treatment, mean_treatment - mean_control); println!(" p-value = {}", pvalue); Ok(()) @@ -160,7 +170,7 @@ mod tests { )); // run permutation test to compute p-value - let pvalue = permute(&control, &treatment, mean_treatment - mean_control); + let pvalue = permutation_test(&control, &treatment, mean_treatment - mean_control); assert!(approx_eq(pvalue, 0.13896357, PVALUE_EPSILON)); } }