use crate::core::{
IntervalType, PredictionResult, RegressionOptions, RegressionOptionsBuilder, RegressionResult,
};
use crate::solvers::traits::{FittedRegressor, RegressionError, Regressor};
use faer::{Col, Mat};
use statrs::distribution::{ContinuousCDF, FisherSnedecor};
#[derive(Debug, Clone)]
pub struct BlsRegressor {
options: RegressionOptions,
lower_bounds: Option<Vec<f64>>,
upper_bounds: Option<Vec<f64>>,
}
impl BlsRegressor {
pub fn new(
options: RegressionOptions,
lower_bounds: Option<Vec<f64>>,
upper_bounds: Option<Vec<f64>>,
) -> Self {
Self {
options,
lower_bounds,
upper_bounds,
}
}
pub fn builder() -> BlsRegressorBuilder {
BlsRegressorBuilder::default()
}
pub fn nnls() -> BlsRegressorBuilder {
BlsRegressorBuilder::default().lower_bound_all(0.0)
}
fn solve_bls(
&self,
x: &Mat<f64>,
y: &Col<f64>,
lower: &[f64],
upper: &[f64],
) -> Result<(Col<f64>, usize), RegressionError> {
let _n_samples = x.nrows();
let n_features = x.ncols();
let mut beta = Col::from_fn(n_features, |j| {
if lower[j].is_finite() {
lower[j]
} else if upper[j].is_finite() {
upper[j]
} else {
0.0
}
});
let mut active = vec![true; n_features];
let mut at_lower = vec![true; n_features];
let xtx = compute_xtx(x);
let xty = compute_xty(x, y);
let max_iter = self.options.max_iterations;
let tol = self.options.tolerance;
for _outer_iter in 0..max_iter {
let gradient = compute_gradient(&xtx, &xty, &beta);
let mut max_violation = 0.0;
let mut best_idx: Option<usize> = None;
for j in 0..n_features {
if active[j] {
if at_lower[j] {
if -gradient[j] > max_violation {
max_violation = -gradient[j];
best_idx = Some(j);
}
} else {
if gradient[j] > max_violation {
max_violation = gradient[j];
best_idx = Some(j);
}
}
}
}
if max_violation < tol {
break;
}
if let Some(idx) = best_idx {
active[idx] = false;
} else {
break;
}
for _inner_iter in 0..max_iter {
let passive_indices: Vec<usize> = (0..n_features).filter(|&j| !active[j]).collect();
if passive_indices.is_empty() {
break;
}
let beta_passive = self.solve_passive_set(x, y, &passive_indices)?;
let mut has_violation = false;
let mut min_alpha = 1.0;
let mut violation_idx = 0;
let mut violation_at_lower = true;
for (i, &j) in passive_indices.iter().enumerate() {
let new_val = beta_passive[i];
let old_val = beta[j];
if new_val < lower[j] - tol {
let alpha = (lower[j] - old_val) / (new_val - old_val);
if alpha < min_alpha {
min_alpha = alpha;
violation_idx = j;
violation_at_lower = true;
has_violation = true;
}
} else if new_val > upper[j] + tol {
let alpha = (upper[j] - old_val) / (new_val - old_val);
if alpha < min_alpha {
min_alpha = alpha;
violation_idx = j;
violation_at_lower = false;
has_violation = true;
}
}
}
if !has_violation {
for (i, &j) in passive_indices.iter().enumerate() {
beta[j] = beta_passive[i];
}
break;
} else {
for (i, &j) in passive_indices.iter().enumerate() {
beta[j] = beta[j] + min_alpha * (beta_passive[i] - beta[j]);
}
active[violation_idx] = true;
at_lower[violation_idx] = violation_at_lower;
beta[violation_idx] = if violation_at_lower {
lower[violation_idx]
} else {
upper[violation_idx]
};
}
}
}
for j in 0..n_features {
beta[j] = beta[j].clamp(lower[j], upper[j]);
}
let rank = beta
.iter()
.enumerate()
.filter(|(j, &b)| {
let at_lower_bound = (b - lower[*j]).abs() < tol;
let at_upper_bound = (b - upper[*j]).abs() < tol;
!(at_lower_bound && lower[*j] == 0.0 || at_upper_bound && upper[*j] == 0.0)
})
.count();
Ok((beta, rank.max(1)))
}
fn solve_passive_set(
&self,
x: &Mat<f64>,
y: &Col<f64>,
passive_indices: &[usize],
) -> Result<Col<f64>, RegressionError> {
let n_samples = x.nrows();
let n_passive = passive_indices.len();
if n_passive == 0 {
return Ok(Col::zeros(0));
}
let x_passive = Mat::from_fn(n_samples, n_passive, |i, j| x[(i, passive_indices[j])]);
let qr = x_passive.col_piv_qr();
let q = qr.compute_Q();
let r = qr.R();
let perm = qr.P();
let qty = q.transpose() * y;
let mut beta_perm = Col::zeros(n_passive);
for i in (0..n_passive).rev() {
let mut sum = qty[i];
for j in (i + 1)..n_passive {
sum -= r[(i, j)] * beta_perm[j];
}
if r[(i, i)].abs() > self.options.rank_tolerance {
beta_perm[i] = sum / r[(i, i)];
} else {
beta_perm[i] = 0.0;
}
}
let mut beta = Col::zeros(n_passive);
for i in 0..n_passive {
beta[perm.inverse().arrays().0[i]] = beta_perm[i];
}
Ok(beta)
}
}
impl Regressor for BlsRegressor {
type Fitted = FittedBls;
fn fit(&self, x: &Mat<f64>, y: &Col<f64>) -> Result<Self::Fitted, RegressionError> {
self.fit_internal(x, y)
}
}
impl BlsRegressor {
#[allow(clippy::too_many_arguments)]
fn build_result(
&self,
x: &Mat<f64>,
y: &Col<f64>,
coefficients: Col<f64>,
intercept: Option<f64>,
rank: usize,
n_params: usize,
_lower: &[f64],
_upper: &[f64],
) -> Result<FittedBls, RegressionError> {
let n_samples = x.nrows();
let n_features = x.ncols();
let mut fitted_values = Col::zeros(n_samples);
let mut residuals = Col::zeros(n_samples);
for i in 0..n_samples {
let mut pred = intercept.unwrap_or(0.0);
for j in 0..n_features {
pred += x[(i, j)] * coefficients[j];
}
fitted_values[i] = pred;
residuals[i] = y[i] - pred;
}
let y_mean: f64 = y.iter().sum::<f64>() / n_samples as f64;
let tss: f64 = y.iter().map(|&yi| (yi - y_mean).powi(2)).sum();
let rss: f64 = residuals.iter().map(|&r| r.powi(2)).sum();
let ess = tss - rss;
let r_squared = if tss > 0.0 {
(1.0 - rss / tss).clamp(0.0, 1.0)
} else if rss < 1e-10 {
1.0
} else {
0.0
};
let df_total = (n_samples - 1) as f64;
let df_resid = (n_samples.saturating_sub(n_params)) as f64;
let adj_r_squared = if df_resid > 0.0 && df_total > 0.0 {
1.0 - (1.0 - r_squared) * df_total / df_resid
} else {
f64::NAN
};
let mse = if df_resid > 0.0 {
rss / df_resid
} else {
f64::NAN
};
let rmse = mse.sqrt();
let df_model = (n_params - if intercept.is_some() { 1 } else { 0 }) as f64;
let f_statistic = if df_model > 0.0 && df_resid > 0.0 && mse > 0.0 {
(ess / df_model) / mse
} else {
f64::NAN
};
let f_pvalue = if f_statistic.is_finite() && df_model > 0.0 && df_resid > 0.0 {
FisherSnedecor::new(df_model, df_resid)
.map(|d| 1.0 - d.cdf(f_statistic))
.unwrap_or(f64::NAN)
} else {
f64::NAN
};
let n = n_samples as f64;
let k = n_params as f64;
let log_likelihood = if mse > 0.0 {
-0.5 * n * (1.0 + (2.0 * std::f64::consts::PI).ln() + mse.ln())
} else {
f64::NAN
};
let aic = if log_likelihood.is_finite() {
2.0 * k - 2.0 * log_likelihood
} else {
f64::NAN
};
let aicc = if log_likelihood.is_finite() && (n - k - 1.0) > 0.0 {
aic + 2.0 * k * (k + 1.0) / (n - k - 1.0)
} else {
f64::NAN
};
let bic = if log_likelihood.is_finite() {
k * n.ln() - 2.0 * log_likelihood
} else {
f64::NAN
};
let mut result = RegressionResult::empty(n_features, n_samples);
result.coefficients = coefficients;
result.intercept = intercept;
result.residuals = residuals;
result.fitted_values = fitted_values;
result.rank = rank;
result.n_parameters = n_params;
result.n_observations = n_samples;
result.aliased = vec![false; n_features]; result.r_squared = r_squared;
result.adj_r_squared = adj_r_squared;
result.mse = mse;
result.rmse = rmse;
result.f_statistic = f_statistic;
result.f_pvalue = f_pvalue;
result.aic = aic;
result.aicc = aicc;
result.bic = bic;
result.log_likelihood = log_likelihood;
result.confidence_level = self.options.confidence_level;
Ok(FittedBls {
result,
options: self.options.clone(),
})
}
}
#[derive(Debug, Clone)]
pub struct FittedBls {
result: RegressionResult,
#[allow(dead_code)]
options: RegressionOptions,
}
impl FittedRegressor for FittedBls {
fn predict(&self, x: &Mat<f64>) -> Col<f64> {
let n_samples = x.nrows();
let n_features = x.ncols();
let intercept = self.result.intercept.unwrap_or(0.0);
Col::from_fn(n_samples, |i| {
let mut pred = intercept;
for j in 0..n_features {
pred += x[(i, j)] * self.result.coefficients[j];
}
pred
})
}
fn result(&self) -> &RegressionResult {
&self.result
}
fn predict_with_interval(
&self,
x: &Mat<f64>,
interval: Option<IntervalType>,
_level: f64,
) -> PredictionResult {
let predictions = self.predict(x);
match interval {
None => PredictionResult::point_only(predictions),
Some(_) => {
let n = x.nrows();
let nan_vec = Col::from_fn(n, |_| f64::NAN);
PredictionResult::with_intervals(
predictions,
nan_vec.clone(),
nan_vec.clone(),
nan_vec,
)
}
}
}
}
fn compute_xtx(x: &Mat<f64>) -> Mat<f64> {
let n_features = x.ncols();
let n_samples = x.nrows();
let mut xtx = Mat::zeros(n_features, n_features);
for i in 0..n_samples {
for j in 0..n_features {
for k in 0..n_features {
xtx[(j, k)] += x[(i, j)] * x[(i, k)];
}
}
}
xtx
}
fn compute_xty(x: &Mat<f64>, y: &Col<f64>) -> Col<f64> {
let n_features = x.ncols();
let n_samples = x.nrows();
let mut xty = Col::zeros(n_features);
for i in 0..n_samples {
for j in 0..n_features {
xty[j] += x[(i, j)] * y[i];
}
}
xty
}
fn compute_gradient(xtx: &Mat<f64>, xty: &Col<f64>, beta: &Col<f64>) -> Col<f64> {
let n = beta.nrows();
let mut grad = Col::zeros(n);
for j in 0..n {
grad[j] = -xty[j];
for k in 0..n {
grad[j] += xtx[(j, k)] * beta[k];
}
}
grad
}
#[derive(Debug, Clone, Default)]
pub struct BlsRegressorBuilder {
options_builder: RegressionOptionsBuilder,
lower_bounds: Option<Vec<f64>>,
upper_bounds: Option<Vec<f64>>,
lower_bound_all: Option<f64>,
upper_bound_all: Option<f64>,
}
impl BlsRegressorBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn with_intercept(mut self, include: bool) -> Self {
self.options_builder = self.options_builder.with_intercept(include);
self
}
pub fn max_iterations(mut self, max_iter: usize) -> Self {
self.options_builder = self.options_builder.max_iterations(max_iter);
self
}
pub fn tolerance(mut self, tol: f64) -> Self {
self.options_builder = self.options_builder.tolerance(tol);
self
}
pub fn lower_bounds(mut self, bounds: Vec<f64>) -> Self {
self.lower_bounds = Some(bounds);
self.lower_bound_all = None;
self
}
pub fn upper_bounds(mut self, bounds: Vec<f64>) -> Self {
self.upper_bounds = Some(bounds);
self.upper_bound_all = None;
self
}
pub fn lower_bound_all(mut self, bound: f64) -> Self {
self.lower_bound_all = Some(bound);
self.lower_bounds = None;
self
}
pub fn upper_bound_all(mut self, bound: f64) -> Self {
self.upper_bound_all = Some(bound);
self.upper_bounds = None;
self
}
pub fn build(self) -> BlsRegressor {
let lower = self.lower_bounds.or_else(|| {
self.lower_bound_all.map(|b| vec![b]) });
let upper = self.upper_bounds.or_else(|| {
self.upper_bound_all.map(|b| vec![b]) });
BlsRegressor {
options: self.options_builder.build_unchecked(),
lower_bounds: if let Some(lb) = self.lower_bound_all {
Some(vec![lb])
} else {
lower
},
upper_bounds: if let Some(ub) = self.upper_bound_all {
Some(vec![ub])
} else {
upper
},
}
}
}
impl BlsRegressor {
fn get_lower_bounds_expanded(&self, n_features: usize) -> Vec<f64> {
match &self.lower_bounds {
Some(bounds) if bounds.len() == 1 => vec![bounds[0]; n_features],
Some(bounds) => bounds.clone(),
None => vec![f64::NEG_INFINITY; n_features],
}
}
fn get_upper_bounds_expanded(&self, n_features: usize) -> Vec<f64> {
match &self.upper_bounds {
Some(bounds) if bounds.len() == 1 => vec![bounds[0]; n_features],
Some(bounds) => bounds.clone(),
None => vec![f64::INFINITY; n_features],
}
}
}
impl BlsRegressor {
fn fit_internal(&self, x: &Mat<f64>, y: &Col<f64>) -> Result<FittedBls, RegressionError> {
let n_samples = x.nrows();
let n_features = x.ncols();
if x.nrows() != y.nrows() {
return Err(RegressionError::DimensionMismatch {
x_rows: x.nrows(),
y_len: y.nrows(),
});
}
if n_samples < 2 {
return Err(RegressionError::InsufficientObservations {
needed: 2,
got: n_samples,
});
}
let lower = self.get_lower_bounds_expanded(n_features);
let upper = self.get_upper_bounds_expanded(n_features);
for j in 0..n_features {
if lower[j] > upper[j] {
return Err(RegressionError::NumericalError(format!(
"lower bound {} exceeds upper bound {} for feature {}",
lower[j], upper[j], j
)));
}
}
if self.options.with_intercept {
let mut x_aug = Mat::zeros(n_samples, n_features + 1);
for i in 0..n_samples {
x_aug[(i, 0)] = 1.0;
for j in 0..n_features {
x_aug[(i, j + 1)] = x[(i, j)];
}
}
let mut lower_aug = vec![f64::NEG_INFINITY];
lower_aug.extend_from_slice(&lower);
let mut upper_aug = vec![f64::INFINITY];
upper_aug.extend_from_slice(&upper);
let (coeffs_aug, rank) = self.solve_bls(&x_aug, y, &lower_aug, &upper_aug)?;
let intercept = Some(coeffs_aug[0]);
let coefficients = Col::from_fn(n_features, |j| coeffs_aug[j + 1]);
self.build_result(
x,
y,
coefficients,
intercept,
rank,
n_features + 1,
&lower,
&upper,
)
} else {
let (coefficients, rank) = self.solve_bls(x, y, &lower, &upper)?;
self.build_result(x, y, coefficients, None, rank, n_features, &lower, &upper)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_nnls_simple() {
let x = Mat::from_fn(5, 2, |i, _| (i + 1) as f64);
let y = Col::from_fn(5, |i| (2 * i + 3) as f64);
let fitted = BlsRegressor::nnls()
.build()
.fit_internal(&x, &y)
.expect("model should fit");
for i in 0..fitted.result.coefficients.nrows() {
assert!(
fitted.result.coefficients[i] >= -1e-10,
"Coefficient {} is negative: {}",
i,
fitted.result.coefficients[i]
);
}
}
#[test]
fn test_bls_box_constraints() {
let x = Mat::from_fn(10, 2, |i, j| (i * 2 + j) as f64);
let y = Col::from_fn(10, |i| (3 * i + 1) as f64);
let fitted = BlsRegressor::builder()
.lower_bounds(vec![0.0, 0.0])
.upper_bounds(vec![1.0, 2.0])
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("model should fit");
assert!(fitted.result.coefficients[0] >= -1e-10);
assert!(fitted.result.coefficients[0] <= 1.0 + 1e-10);
assert!(fitted.result.coefficients[1] >= -1e-10);
assert!(fitted.result.coefficients[1] <= 2.0 + 1e-10);
}
#[test]
fn test_bls_unconstrained_matches_ols() {
let x = Mat::from_fn(20, 2, |i, j| ((i + 1) * (j + 1)) as f64);
let y = Col::from_fn(20, |i| 2.0 * i as f64 + 1.0);
let bls_fitted = BlsRegressor::builder()
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("model should fit");
assert!(bls_fitted.result.r_squared > 0.5);
}
#[test]
fn test_bls_with_intercept() {
let x = Mat::from_fn(10, 1, |i, _| i as f64);
let y = Col::from_fn(10, |i| 5.0 + 2.0 * i as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(true)
.build()
.fit_internal(&x, &y)
.expect("model should fit");
assert!(fitted.result.intercept.is_some());
assert!(fitted.result.coefficients[0] >= -1e-10);
}
#[test]
fn test_bound_all() {
let x = Mat::from_fn(10, 3, |i, j| (i + j) as f64);
let y = Col::from_fn(10, |i| i as f64);
let fitted = BlsRegressor::builder()
.lower_bound_all(0.0)
.upper_bound_all(1.0)
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("model should fit");
for i in 0..3 {
assert!(
fitted.result.coefficients[i] >= -1e-10,
"Coef {} below lower bound",
i
);
assert!(
fitted.result.coefficients[i] <= 1.0 + 1e-10,
"Coef {} above upper bound",
i
);
}
}
#[test]
fn test_bls_new_constructor() {
let options = RegressionOptionsBuilder::default()
.build()
.expect("valid options");
let lower = Some(vec![0.0, 0.0]);
let upper = Some(vec![10.0, 10.0]);
let regressor = BlsRegressor::new(options, lower, upper);
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| i as f64);
let fitted = regressor.fit_internal(&x, &y).expect("should fit");
assert!(fitted.result.coefficients.nrows() == 2);
}
#[test]
fn test_builder_new() {
let builder = BlsRegressorBuilder::new();
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| i as f64);
let fitted = builder
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("should fit");
assert!(fitted.result.coefficients.nrows() == 2);
}
#[test]
fn test_builder_max_iterations() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| i as f64);
let fitted = BlsRegressor::builder()
.max_iterations(100)
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("should fit");
assert!(fitted.result.r_squared >= 0.0);
}
#[test]
fn test_builder_tolerance() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| i as f64);
let fitted = BlsRegressor::builder()
.tolerance(1e-8)
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("should fit");
assert!(fitted.result.r_squared >= 0.0);
}
#[test]
fn test_predict_method() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| 1.0 + 0.5 * i as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(true)
.build()
.fit_internal(&x, &y)
.expect("should fit");
let x_new = Mat::from_fn(5, 2, |i, j| ((i + 20) + j) as f64);
let predictions = fitted.predict(&x_new);
assert_eq!(predictions.nrows(), 5);
for i in 0..5 {
assert!(predictions[i].is_finite());
}
}
#[test]
fn test_result_method() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| 1.0 + 0.5 * i as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(true)
.build()
.fit_internal(&x, &y)
.expect("should fit");
let result = fitted.result();
assert!(result.coefficients.nrows() == 2);
assert!(result.r_squared >= 0.0);
}
#[test]
fn test_predict_with_interval_none() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| 1.0 + 0.5 * i as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(true)
.build()
.fit_internal(&x, &y)
.expect("should fit");
let x_new = Mat::from_fn(5, 2, |i, j| ((i + 20) + j) as f64);
let result = fitted.predict_with_interval(&x_new, None, 0.95);
assert_eq!(result.fit.nrows(), 5);
}
#[test]
fn test_predict_with_interval_confidence() {
let x = Mat::from_fn(20, 2, |i, j| (i + j) as f64);
let y = Col::from_fn(20, |i| 1.0 + 0.5 * i as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(true)
.build()
.fit_internal(&x, &y)
.expect("should fit");
let x_new = Mat::from_fn(5, 2, |i, j| ((i + 20) + j) as f64);
let result = fitted.predict_with_interval(&x_new, Some(IntervalType::Confidence), 0.95);
assert_eq!(result.fit.nrows(), 5);
assert_eq!(result.lower.nrows(), 5);
assert_eq!(result.upper.nrows(), 5);
for i in 0..5 {
assert!(result.lower[i].is_nan());
assert!(result.upper[i].is_nan());
}
}
#[test]
fn test_predict_without_intercept() {
let x = Mat::from_fn(20, 2, |i, j| ((i + 1) * (j + 1)) as f64);
let y = Col::from_fn(20, |i| 0.5 * (i + 1) as f64);
let fitted = BlsRegressor::nnls()
.with_intercept(false)
.build()
.fit_internal(&x, &y)
.expect("should fit");
let x_new = Mat::from_fn(5, 2, |i, j| ((i + 20) * (j + 1)) as f64);
let predictions = fitted.predict(&x_new);
assert_eq!(predictions.nrows(), 5);
}
}