use crate::error::FdarError;
use crate::iter_maybe_parallel;
use crate::matrix::FdMatrix;
use crate::smoothing;
#[cfg(feature = "parallel")]
use rayon::iter::ParallelIterator;
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ConcurrentRegrResult {
pub beta_curve: FdMatrix,
pub intercept: Vec<f64>,
pub fitted: FdMatrix,
pub residuals: FdMatrix,
pub argvals: Vec<f64>,
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn concurrent_regression(
response: &FdMatrix,
predictors: &[FdMatrix],
argvals: Option<&[f64]>,
bandwidth: f64,
kernel: &str,
) -> Result<ConcurrentRegrResult, FdarError> {
if predictors.is_empty() {
return Err(FdarError::InvalidDimension {
parameter: "predictors",
expected: "at least 1".to_string(),
actual: "0".to_string(),
});
}
let (n, m) = response.shape();
if n < 2 {
return Err(FdarError::InvalidDimension {
parameter: "response",
expected: "at least 2 rows (observations)".to_string(),
actual: format!("{n}"),
});
}
if m == 0 {
return Err(FdarError::InvalidDimension {
parameter: "response",
expected: "non-zero columns (grid points)".to_string(),
actual: "0".to_string(),
});
}
for (k, pred) in predictors.iter().enumerate() {
if pred.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "predictors[k]",
expected: format!("{n} rows (matching response)"),
actual: format!("{} (predictor index {k})", pred.nrows()),
});
}
if pred.ncols() != m {
return Err(FdarError::InvalidDimension {
parameter: "predictors[k]",
expected: format!("{m} columns (matching response)"),
actual: format!("{} (predictor index {k})", pred.ncols()),
});
}
}
if !bandwidth.is_finite() || bandwidth <= 0.0 {
return Err(FdarError::InvalidParameter {
parameter: "bandwidth",
message: format!("must be finite and positive, got {bandwidth}"),
});
}
if let Some(av) = argvals {
if av.len() != m {
return Err(FdarError::InvalidDimension {
parameter: "argvals",
expected: format!("{m} elements (matching response columns)"),
actual: format!("{}", av.len()),
});
}
}
let argvals_owned: Vec<f64> = argvals
.map(|v| v.to_vec())
.unwrap_or_else(|| (0..m).map(|j| j as f64 / (m - 1).max(1) as f64).collect());
let p = predictors.len();
let q = p + 1;
if n <= p {
return Err(FdarError::InvalidDimension {
parameter: "response",
expected: format!(
"at least {} rows (more observations than predictors p={})",
p + 1,
p
),
actual: format!("{n}"),
});
}
let raw_cols: Vec<(f64, Vec<f64>)> = iter_maybe_parallel!(0..m)
.map(|j| {
let mut xtx = vec![0.0_f64; q * q];
let mut xty = vec![0.0_f64; q];
for i in 0..n {
let mut row = vec![1.0_f64; q];
for (k, pred) in predictors.iter().enumerate() {
row[k + 1] = pred[(i, j)];
}
for a in 0..q {
for b in 0..q {
xtx[a * q + b] += row[a] * row[b];
}
xty[a] += row[a] * response[(i, j)];
}
}
let eps = 1e-10 * (xtx[0] + 1.0);
for d in 0..q {
xtx[d * q + d] += eps;
}
let coef = smoothing::solve_gaussian_pub(&mut xtx, &mut xty, q);
let raw_intercept_j = coef[0];
let raw_beta_j: Vec<f64> = coef[1..q].to_vec();
(raw_intercept_j, raw_beta_j)
})
.collect();
let mut raw_intercept = vec![0.0_f64; m];
let mut raw_beta = vec![0.0_f64; p * m];
for (j, (ic, bk)) in raw_cols.into_iter().enumerate() {
raw_intercept[j] = ic;
for k in 0..p {
raw_beta[k * m + j] = bk[k];
}
}
let intercept = smoothing::local_linear(
&argvals_owned,
&raw_intercept,
&argvals_owned,
bandwidth,
kernel,
)?;
let mut beta_curve = FdMatrix::zeros(p, m);
for k in 0..p {
let raw_k: Vec<f64> = (0..m).map(|j| raw_beta[k * m + j]).collect();
let smooth_k =
smoothing::local_linear(&argvals_owned, &raw_k, &argvals_owned, bandwidth, kernel)?;
for j in 0..m {
beta_curve[(k, j)] = smooth_k[j];
}
}
let mut fitted = FdMatrix::zeros(n, m);
let mut residuals = FdMatrix::zeros(n, m);
for j in 0..m {
for i in 0..n {
let mut val = intercept[j];
for (k, pred) in predictors.iter().enumerate() {
val += beta_curve[(k, j)] * pred[(i, j)];
}
fitted[(i, j)] = val;
residuals[(i, j)] = response[(i, j)] - val;
}
}
Ok(ConcurrentRegrResult {
beta_curve,
intercept,
fitted,
residuals,
argvals: argvals_owned,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::uniform_grid;
#[test]
fn test_shape_smoke() {
let n = 6;
let m = 8;
let argvals = uniform_grid(m);
let mut response = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * argvals[j];
}
}
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
predictor[(i, j)] = (i as f64 + 1.0) + argvals[j];
}
}
let result = concurrent_regression(&response, &[predictor], None, 0.2, "gaussian");
assert!(result.is_ok(), "expected Ok, got {result:?}");
let r = result.unwrap();
assert_eq!(r.beta_curve.shape(), (1, m), "beta_curve shape");
assert_eq!(r.intercept.len(), m, "intercept length");
assert_eq!(r.fitted.shape(), (n, m), "fitted shape");
assert_eq!(r.residuals.shape(), (n, m), "residuals shape");
assert_eq!(r.argvals.len(), m, "argvals length");
}
#[test]
fn test_multi_predictor_shape() {
let n = 6;
let m = 8;
let p = 3;
let argvals = uniform_grid(m);
let mut response = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * argvals[j];
}
}
let predictors: Vec<FdMatrix> = (0..p)
.map(|k| {
let mut pred = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
pred[(i, j)] = (i as f64 + k as f64 + 1.0) * argvals[j] + 0.1;
}
}
pred
})
.collect();
let result = concurrent_regression(&response, &predictors, None, 0.2, "gaussian");
assert!(result.is_ok(), "expected Ok for p=3, got {result:?}");
let r = result.unwrap();
assert_eq!(r.beta_curve.shape(), (p, m), "beta_curve shape for p=3");
}
#[test]
fn test_parallel_sequential_equivalence() {
let n = 6;
let m = 8;
let argvals = uniform_grid(m);
let mut response = FdMatrix::zeros(n, m);
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * argvals[j] + 0.3;
predictor[(i, j)] = argvals[j] * (i as f64 + 0.5);
}
}
let r1 = concurrent_regression(&response, &[predictor.clone()], None, 0.2, "gaussian")
.expect("first call should succeed");
let r2 = concurrent_regression(&response, &[predictor], None, 0.2, "gaussian")
.expect("second call should succeed");
assert_eq!(
r1, r2,
"two calls on identical inputs must produce equal results"
);
}
fn lcg_noise(seed: u64, n: usize) -> Vec<f64> {
let mut state = seed;
(0..n)
.map(|_| {
state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
(state >> 33) as f64 / (u32::MAX as f64) - 0.5
})
.collect()
}
#[test]
fn test_recovery_known_beta() {
let n = 50usize;
let m = 50usize;
let argvals = uniform_grid(m);
let true_beta0 = 0.5_f64;
let true_beta: Vec<f64> = argvals
.iter()
.map(|&t| (std::f64::consts::PI * t).sin())
.collect();
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
predictor[(i, j)] =
(2.0 * std::f64::consts::PI * (i as f64 / n as f64) + argvals[j]).sin();
}
}
let mut x_max = 0.0_f64;
for i in 0..n {
for j in 0..m {
let v = predictor[(i, j)].abs();
if v > x_max {
x_max = v;
}
}
}
let noise_scale = 0.05 * x_max;
let noise = lcg_noise(42, n * m);
let mut response = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] =
true_beta0 + true_beta[j] * predictor[(i, j)] + noise_scale * noise[i * m + j];
}
}
let result =
concurrent_regression(&response, &[predictor], Some(&argvals), 0.15, "gaussian")
.expect("recovery test should succeed");
for j in 5..45 {
let diff = (result.beta_curve[(0, j)] - true_beta[j]).abs();
assert!(
diff < 0.15,
"recovery failed at j={j}: got {}, expected {}, diff={diff}",
result.beta_curve[(0, j)],
true_beta[j]
);
}
}
fn roughness(row: &[f64]) -> f64 {
let m = row.len();
if m < 3 {
return 0.0;
}
(1..m - 1)
.map(|j| {
let d = row[j + 1] - 2.0 * row[j] + row[j - 1];
d * d
})
.sum()
}
#[test]
fn test_monotone_roughness() {
let n = 50usize;
let m = 50usize;
let argvals = uniform_grid(m);
let true_beta0 = 0.5_f64;
let true_beta: Vec<f64> = argvals
.iter()
.map(|&t| (2.0 * std::f64::consts::PI * t).sin())
.collect();
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
predictor[(i, j)] =
(2.0 * std::f64::consts::PI * (i as f64 / n as f64) + argvals[j]).sin();
}
}
let mut x_max = 0.0_f64;
for i in 0..n {
for j in 0..m {
let v = predictor[(i, j)].abs();
if v > x_max {
x_max = v;
}
}
}
let noise_scale = 0.05 * x_max;
let noise = lcg_noise(42, n * m);
let mut response = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] =
true_beta0 + true_beta[j] * predictor[(i, j)] + noise_scale * noise[i * m + j];
}
}
let bandwidths = [0.05, 0.15, 0.35];
let roughnesses: Vec<f64> = bandwidths
.iter()
.map(|&bw| {
let r = concurrent_regression(
&response,
&[predictor.clone()],
Some(&argvals),
bw,
"gaussian",
)
.expect("monotone roughness call should succeed");
let beta_row: Vec<f64> = (0..m).map(|j| r.beta_curve[(0, j)]).collect();
roughness(&beta_row)
})
.collect();
assert!(
roughnesses[0] > roughnesses[1],
"roughness(bw=0.05)={} should be > roughness(bw=0.15)={}",
roughnesses[0],
roughnesses[1]
);
assert!(
roughnesses[1] > roughnesses[2],
"roughness(bw=0.15)={} should be > roughness(bw=0.35)={}",
roughnesses[1],
roughnesses[2]
);
}
#[test]
fn test_residuals_consistency() {
let n = 6;
let m = 8;
let argvals = uniform_grid(m);
let mut response = FdMatrix::zeros(n, m);
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * argvals[j] + 0.1;
predictor[(i, j)] = argvals[j] + i as f64 * 0.1;
}
}
let result = concurrent_regression(&response, &[predictor], None, 0.2, "gaussian")
.expect("residuals consistency call should succeed");
for i in 0..n {
for j in 0..m {
let expected_resid = response[(i, j)] - result.fitted[(i, j)];
let diff = (result.residuals[(i, j)] - expected_resid).abs();
assert!(
diff < 1e-10,
"residual inconsistency at ({i},{j}): stored={}, computed={expected_resid}, diff={diff}",
result.residuals[(i, j)]
);
}
}
}
#[test]
fn test_invalid_inputs() {
let n = 6;
let m = 8;
let argvals = uniform_grid(m);
let mut response = FdMatrix::zeros(n, m);
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * argvals[j];
predictor[(i, j)] = argvals[j];
}
}
let err = concurrent_regression(&response, &[], None, 0.2, "gaussian")
.expect_err("empty predictors should return Err");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "predictors",
..
}
),
"expected InvalidDimension for empty predictors, got {err:?}"
);
let tiny_response = FdMatrix::zeros(1, m);
let tiny_pred = FdMatrix::zeros(1, m);
let err = concurrent_regression(&tiny_response, &[tiny_pred], None, 0.2, "gaussian")
.expect_err("n<2 response should return Err");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "response",
..
}
),
"expected InvalidDimension for response n<2, got {err:?}"
);
let wrong_n_pred = FdMatrix::zeros(n + 1, m);
let err = concurrent_regression(&response, &[wrong_n_pred], None, 0.2, "gaussian")
.expect_err("wrong nrows predictor should return Err");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "predictors[k]",
..
}
),
"expected InvalidDimension for predictor nrows mismatch, got {err:?}"
);
let wrong_m_pred = FdMatrix::zeros(n, m + 1);
let err = concurrent_regression(&response, &[wrong_m_pred], None, 0.2, "gaussian")
.expect_err("wrong ncols predictor should return Err");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "predictors[k]",
..
}
),
"expected InvalidDimension for predictor ncols mismatch, got {err:?}"
);
let err = concurrent_regression(&response, &[predictor.clone()], None, 0.0, "gaussian")
.expect_err("bandwidth=0.0 should return Err");
assert!(
matches!(
err,
FdarError::InvalidParameter {
parameter: "bandwidth",
..
}
),
"expected InvalidParameter for bandwidth=0.0, got {err:?}"
);
let err = concurrent_regression(&response, &[predictor.clone()], None, -0.1, "gaussian")
.expect_err("negative bandwidth should return Err");
assert!(
matches!(
err,
FdarError::InvalidParameter {
parameter: "bandwidth",
..
}
),
"expected InvalidParameter for negative bandwidth, got {err:?}"
);
let wrong_argvals: Vec<f64> = uniform_grid(m + 3);
let err = concurrent_regression(
&response,
&[predictor],
Some(&wrong_argvals),
0.2,
"gaussian",
)
.expect_err("wrong argvals length should return Err");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "argvals",
..
}
),
"expected InvalidDimension for argvals length mismatch, got {err:?}"
);
}
#[test]
fn test_nan_inf_bandwidth_returns_error() {
let n = 4;
let m = 6;
let mut response = FdMatrix::zeros(n, m);
let mut predictor = FdMatrix::zeros(n, m);
for i in 0..n {
for j in 0..m {
response[(i, j)] = (i as f64 + 1.0) * (j as f64 + 1.0);
predictor[(i, j)] = (i as f64 + 1.0) + (j as f64 * 0.1);
}
}
let err =
concurrent_regression(&response, &[predictor.clone()], None, f64::NAN, "gaussian")
.expect_err("NaN bandwidth must return Err");
assert!(
matches!(
err,
FdarError::InvalidParameter {
parameter: "bandwidth",
..
}
),
"expected InvalidParameter for NaN bandwidth, got {err:?}"
);
let err = concurrent_regression(
&response,
&[predictor.clone()],
None,
f64::INFINITY,
"gaussian",
)
.expect_err("+Inf bandwidth must return Err");
assert!(
matches!(
err,
FdarError::InvalidParameter {
parameter: "bandwidth",
..
}
),
"expected InvalidParameter for +Inf bandwidth, got {err:?}"
);
concurrent_regression(&response, &[predictor], None, 0.2, "gaussian")
.expect("positive finite bandwidth must succeed");
}
#[test]
fn test_underdetermined_system_returns_error() {
let m = 8;
let n = 3usize;
let p = 3usize;
let response = FdMatrix::zeros(n, m);
let predictors: Vec<FdMatrix> = (0..p).map(|_| FdMatrix::zeros(n, m)).collect();
let err = concurrent_regression(&response, &predictors, None, 0.2, "gaussian")
.expect_err("n == p should return Err (underdetermined)");
assert!(
matches!(
err,
FdarError::InvalidDimension {
parameter: "response",
..
}
),
"expected InvalidDimension for n==p, got {err:?}"
);
let n2 = 2usize;
let p2 = 4usize;
let response2 = FdMatrix::zeros(n2, m);
let predictors2: Vec<FdMatrix> = (0..p2).map(|_| FdMatrix::zeros(n2, m)).collect();
let err2 = concurrent_regression(&response2, &predictors2, None, 0.2, "gaussian")
.expect_err("n < p should return Err (underdetermined)");
assert!(
matches!(
err2,
FdarError::InvalidDimension {
parameter: "response",
..
}
),
"expected InvalidDimension for n<p, got {err2:?}"
);
let n3 = 6usize;
let p3 = 2usize;
let mut resp3 = FdMatrix::zeros(n3, m);
let argvals = uniform_grid(m);
for i in 0..n3 {
for j in 0..m {
resp3[(i, j)] = (i as f64 + 1.0) * argvals[j];
}
}
let preds3: Vec<FdMatrix> = (0..p3)
.map(|k| {
let mut pred = FdMatrix::zeros(n3, m);
for i in 0..n3 {
for j in 0..m {
pred[(i, j)] = (i as f64 + k as f64 + 1.0) * argvals[j] + 0.1;
}
}
pred
})
.collect();
concurrent_regression(&resp3, &preds3, None, 0.2, "gaussian")
.expect("n > p should succeed");
}
}