use crate::{
errors::{StatsError, StockTrekError},
statistics,
};
pub fn autocovariance(values: &[f64], lag: usize) -> Result<f64, StockTrekError> {
let n = values.len();
if n == 0 {
return Err(StockTrekError::Stats(StatsError::EmptyInput));
}
if lag >= n {
return Err(StockTrekError::Stats(StatsError::InvalidLag));
}
let mu = statistics::core::mean(values)?;
let sum: f64 = (lag..n)
.map(|t| (values[t] - mu) * (values[t - lag] - mu))
.sum();
let autocovariance = sum / n as f64;
Ok(autocovariance)
}
pub fn autocorrelation(values: &[f64], lag: usize) -> Result<f64, StockTrekError> {
let gamma_0 = autocovariance(values, 0)?;
if gamma_0 == 0.0 {
return Err(StockTrekError::Stats(StatsError::ZeroVariance));
}
let gamma_k = autocovariance(values, lag)?;
let autocorrelation = gamma_k / gamma_0;
Ok(autocorrelation)
}
pub fn partial_autocorrelation(values: &[f64], max_lag: usize) -> Result<Vec<f64>, StockTrekError> {
let n = values.len();
if n == 0 {
return Err(StockTrekError::Stats(StatsError::EmptyInput));
}
if max_lag >= n {
return Err(StockTrekError::Stats(StatsError::InvalidLag));
}
let acf = (0..=max_lag)
.map(|k| autocorrelation(values, k))
.collect::<Result<Vec<f64>, _>>()?;
let mut pacf = vec![0.0; max_lag + 1];
let mut phi = vec![vec![0.0; max_lag + 1]; max_lag + 1];
pacf[0] = 1.0;
if max_lag >= 1 {
phi[1][1] = acf[1];
pacf[1] = acf[1];
}
for k in 2..=max_lag {
let sum: f64 = (1..k).map(|j| phi[k - 1][j] * acf[k - j]).sum();
let denom: f64 = 1.0 - (1..k).map(|j| phi[k - 1][j] * acf[j]).sum::<f64>();
if denom == 0.0 {
return Err(StockTrekError::Stats(StatsError::ZeroVariance));
}
let phi_kk = (acf[k] - sum) / denom;
phi[k][k] = phi_kk;
for j in 1..k {
phi[k][j] = phi[k - 1][j] - phi_kk * phi[k - 1][k - j];
}
pacf[k] = phi_kk;
}
Ok(pacf)
}
pub fn cross_correlation(first: &[f64], second: &[f64], lag: isize) -> Result<f64, StockTrekError> {
let n = first.len();
if n == 0 {
return Err(StockTrekError::Stats(StatsError::EmptyInput));
}
if n != second.len() {
return Err(StockTrekError::Stats(StatsError::MismatchedLengths));
}
let mean_x = statistics::core::mean(first)?;
let mean_y = statistics::core::mean(second)?;
let numerator: f64 = (0..n)
.filter_map(|t| {
let j = t as isize - lag;
if j >= 0 && (j as usize) < n {
Some((first[t] - mean_x) * (second[j as usize] - mean_y))
} else {
None
}
})
.sum();
let var_x = statistics::core::variance(first, 0)?;
let var_y = statistics::core::variance(second, 0)?;
let cross_correlation = numerator / (n as f64 * (var_x.sqrt() * var_y.sqrt()));
Ok(cross_correlation)
}