Skip to content

Commit

Permalink
feat multithreaded acf
Browse files Browse the repository at this point in the history
  • Loading branch information
mrvollger committed Sep 9, 2024
1 parent 2c58af3 commit 5edbfc9
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 3 deletions.
11 changes: 8 additions & 3 deletions src/subcommands/qc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,11 +213,16 @@ impl<'a> QcStats<'a> {
}

log::info!("Calculating m6A auto-correlation.");
let acf =
crate::utils::acf::acf(&self.m6a_acf_starts, Some(self.qc_opts.acf_max_lag), false)?;
let acf = crate::utils::acf::acf_par(
&self.m6a_acf_starts,
Some(self.qc_opts.acf_max_lag),
false,
)?;
log::info!("Done calculating m6A auto-correlation!");
for (i, val) in acf.iter().enumerate() {
out.write_all(format!("m6a_acf\t{}\t{}\n", i, val).as_bytes())?;
out.write_all(
format!("m6a_acf\t{}\t{}\n", i, my_ordered_float(*val as f32)).as_bytes(),
)?;
}
Ok(())
}
Expand Down
65 changes: 65 additions & 0 deletions src/utils/acf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// https://github.com/krfricke/arima
use anyhow::Result;
use num::Float;
use rayon::prelude::*;
use std::cmp;
use std::convert::From;
use std::ops::{Add, AddAssign, Div};
Expand Down Expand Up @@ -73,3 +74,67 @@ pub fn acf<T: Float + From<u32> + From<f64> + Copy + Add + AddAssign + Div>(
}
Ok(y)
}

/// Calculate the auto-correlation function of a time series of length n.
/// but this version is multithreaded over m using rayon.
pub fn acf_par<
T: Float
+ From<u32>
+ From<f64>
+ Copy
+ Add
+ AddAssign
+ Div
+ Send
+ std::iter::Sum
+ std::marker::Sync,
>(
x: &[T],
max_lag: Option<usize>,
covariance: bool,
) -> Result<Vec<T>> {
let max_lag = match max_lag {
// if upper bound for max_lag is n-1
Some(max_lag) => cmp::min(max_lag, x.len() - 1),
None => x.len() - 1,
};
if x.len() <= max_lag {
return Err(anyhow::anyhow!(
"acf-max-lag ({}) must be less than the number of m6A observations ({}).",
max_lag,
x.len()
));
}

let m = max_lag + 1;

let len_x_usize = x.len();
let len_x: T = From::from(len_x_usize as u32);
let sum: T = From::from(0.0);

let sum_x: T = x.iter().fold(sum, |sum, &xi| sum + xi);
let mean_x: T = sum_x / len_x;

let mut y: Vec<T> = vec![From::from(0.0); m];

for t in 0..m {
y[t] = x
.into_par_iter()
.enumerate()
.take(len_x_usize - t)
.map(|(i, xi)| {
let xi = *xi - mean_x;
let xi_t = x[i + t] - mean_x;
(xi * xi_t) / len_x
})
.sum();
// we need y[0] to calculate the correlations, so we set it to 1.0 at the end
if !covariance && t > 0 {
y[t] = y[t] / y[0];
}
}
if !covariance {
y[0] = From::from(1.0);
}
Ok(y)
}

0 comments on commit 5edbfc9

Please sign in to comment.