From 6fd09802773b1f6551290da46440ec7d2d048098 Mon Sep 17 00:00:00 2001
From: RunasSudo <runassudo@yingtongli.me>
Date: Fri, 27 Oct 2023 23:28:01 +1100
Subject: [PATCH] turnbull: Parallelise compute_pi

37% speedup
---
 src/turnbull.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/turnbull.rs b/src/turnbull.rs
index 84a6e05..caac5e9 100644
--- a/src/turnbull.rs
+++ b/src/turnbull.rs
@@ -328,13 +328,35 @@ fn get_sum_fail_prob(data: &TurnbullData, s: &Vec<f64>) -> Vec<f64> {
 }
 
 fn compute_pi(data: &TurnbullData, s: &Vec<f64>, sum_fail_prob: Vec<f64>) -> Vec<f64> {
+	/*
 	let mut pi: Vec<f64> = vec![0.0; data.num_intervals()];
-	
 	for ((idx_left, idx_right), sum_fail_prob_i) in data.data_time_interval_indexes.iter().zip(sum_fail_prob.iter()) {
 		for j in *idx_left..(*idx_right + 1) {
 			pi[j] += s[j] / sum_fail_prob_i / data.num_obs() as f64;
 		}
 	}
+	*/
+	
+	let pi = data.data_time_interval_indexes.par_iter().zip(sum_fail_prob.par_iter())
+		.fold_with(
+			// Compute the contributions to pi[j] for each observation and sum them in parallel using fold_with
+			vec![0.0; data.num_intervals()],
+			|mut acc, ((idx_left, idx_right), sum_fail_prob_i)| {
+				// Contributions to pi[j] for the i-th observation
+				for j in *idx_left..(*idx_right + 1) {
+					acc[j] += s[j] / sum_fail_prob_i / data.num_obs() as f64;
+				}
+				acc
+			}
+		)
+		.reduce(
+			// Reduce all the sub-sums from fold_with into the total sum
+			|| vec![0.0; data.num_intervals()],
+			|mut acc, subsum| {
+				acc.iter_mut().zip(subsum.iter()).for_each(|(x, y)| *x += y);
+				acc
+			}
+		);
 	
 	return pi;
 }