use approx::assert_abs_diff_eq;
use ndarray::{Array1, Array2, array};
use rustyml::error::Error;
use rustyml::machine_learning::{LinearRegression, RegularizationType};
use crate::common::assert_allclose;
#[test]
fn constructor_zero_learning_rate_is_invalid() {
let result = LinearRegression::new(true, 0.0, 100, 1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_negative_learning_rate_is_invalid() {
let result = LinearRegression::new(true, -0.01, 100, 1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_nan_learning_rate_is_invalid() {
let result = LinearRegression::new(true, f64::NAN, 100, 1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_inf_learning_rate_is_invalid() {
let result = LinearRegression::new(true, f64::INFINITY, 100, 1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_zero_max_iter_is_invalid() {
let result = LinearRegression::new(true, 0.01, 0, 1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_zero_tolerance_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, 0.0);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_negative_tolerance_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, -1e-6);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_nan_tolerance_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, f64::NAN);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_inf_tolerance_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, f64::INFINITY);
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_negative_l2_alpha_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, 1e-6)
.unwrap()
.with_regularization(RegularizationType::L2(-0.1));
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_negative_l1_alpha_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, 1e-6)
.unwrap()
.with_regularization(RegularizationType::L1(-0.5));
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_nan_l2_alpha_is_invalid() {
let result = LinearRegression::new(true, 0.01, 100, 1e-6)
.unwrap()
.with_regularization(RegularizationType::L2(f64::NAN));
assert!(
matches!(result, Err(Error::InvalidParameter { .. })),
"expected InvalidParameter, got {:?}",
result
);
}
#[test]
fn constructor_valid_parameters_succeeds() {
let result = LinearRegression::new(true, 0.01, 1000, 1e-6);
assert!(result.is_ok(), "expected Ok, got {:?}", result);
}
#[test]
fn constructor_getters_round_trip() {
let model = LinearRegression::new(false, 0.05, 500, 1e-4).unwrap();
assert!(!model.get_fit_intercept());
assert_abs_diff_eq!(model.get_learning_rate(), 0.05, epsilon = 1e-15);
assert_eq!(model.get_max_iterations(), 500);
assert_abs_diff_eq!(model.get_tolerance(), 1e-4, epsilon = 1e-20);
assert!(model.get_coefficients().is_none());
assert!(model.get_intercept().is_none());
assert!(model.get_actual_iterations().is_none());
}
#[test]
fn predict_before_fit_returns_not_fitted() {
let model = LinearRegression::new(true, 0.01, 100, 1e-6).unwrap();
let x = array![[1.0, 2.0]];
let result = model.predict(&x);
assert!(
matches!(result, Err(Error::NotFitted(_))),
"expected NotFitted, got {:?}",
result
);
}
#[test]
fn fit_empty_x_returns_empty_input() {
let mut model = LinearRegression::new(true, 0.01, 100, 1e-6).unwrap();
let x: Array2<f64> = Array2::zeros((0, 2));
let y: Array1<f64> = Array1::zeros(0);
let result = model.fit(&x, &y);
assert!(
matches!(result, Err(Error::EmptyInput(_))),
"expected EmptyInput, got {:?}",
result
);
}
#[test]
fn fit_nan_in_x_returns_non_finite() {
let mut model = LinearRegression::new(true, 0.01, 100, 1e-6).unwrap();
let x = array![[1.0, f64::NAN], [2.0, 3.0]];
let y = array![1.0, 2.0];
let result = model.fit(&x, &y);
assert!(
matches!(result, Err(Error::NonFinite(_))),
"expected NonFinite, got {:?}",
result
);
}
#[test]
fn fit_inf_in_x_returns_non_finite() {
let mut model = LinearRegression::new(true, 0.01, 100, 1e-6).unwrap();
let x = array![[1.0, f64::INFINITY], [2.0, 3.0]];
let y = array![1.0, 2.0];
let result = model.fit(&x, &y);
assert!(
matches!(result, Err(Error::NonFinite(_))),
"expected NonFinite, got {:?}",
result
);
}
#[test]
fn fit_y_length_mismatch_returns_dimension_mismatch() {
let mut model = LinearRegression::new(true, 0.01, 100, 1e-6).unwrap();
let x = array![[1.0], [2.0], [3.0]];
let y = array![1.0, 2.0];
let result = model.fit(&x, &y);
assert!(
matches!(result, Err(Error::DimensionMismatch { .. })),
"expected DimensionMismatch, got {:?}",
result
);
}
#[test]
fn predict_empty_matrix_returns_empty_input() {
let mut model = LinearRegression::new(true, 0.01, 5000, 1e-8).unwrap();
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x_train, &y_train).unwrap();
let x_empty: Array2<f64> = Array2::zeros((0, 1));
let result = model.predict(&x_empty);
assert!(
matches!(result, Err(Error::EmptyInput(_))),
"expected EmptyInput, got {:?}",
result
);
}
#[test]
fn predict_wrong_feature_count_returns_dimension_mismatch() {
let mut model = LinearRegression::new(true, 0.01, 5000, 1e-8).unwrap();
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x_train, &y_train).unwrap();
let x_wrong = array![[1.0, 2.0]];
let result = model.predict(&x_wrong);
assert!(
matches!(result, Err(Error::DimensionMismatch { .. })),
"expected DimensionMismatch, got {:?}",
result
);
}
#[test]
fn predict_nan_in_x_returns_non_finite() {
let mut model = LinearRegression::new(true, 0.01, 5000, 1e-8).unwrap();
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x_train, &y_train).unwrap();
let x_nan = array![[f64::NAN]];
let result = model.predict(&x_nan);
assert!(
matches!(result, Err(Error::NonFinite(_))),
"expected NonFinite, got {:?}",
result
);
}
#[test]
fn predict_inf_in_x_returns_non_finite() {
let mut model = LinearRegression::new(true, 0.01, 5000, 1e-8).unwrap();
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x_train, &y_train).unwrap();
let x_inf = array![[f64::INFINITY]];
let result = model.predict(&x_inf);
assert!(
matches!(result, Err(Error::NonFinite(_))),
"expected NonFinite, got {:?}",
result
);
}
#[test]
fn univariate_y_equals_2x_plus_1_coefficient_and_intercept() {
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x, &y).unwrap();
let coeff = model.get_coefficients().unwrap();
assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 3e-3);
let intercept = model.get_intercept().unwrap();
assert_abs_diff_eq!(intercept, 1.0, epsilon = 3e-3);
}
#[test]
fn univariate_y_equals_2x_plus_1_predictions() {
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x, &y).unwrap();
let preds = model.predict(&array![[6.0], [0.0]]).unwrap();
let expected = array![13.0, 1.0];
assert_allclose(&preds, &expected, 1e-3);
}
#[test]
fn fit_sets_n_iter() {
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
model.fit(&x, &y).unwrap();
let n_iter = model.get_actual_iterations();
assert!(n_iter.is_some(), "n_iter should be set after fit");
assert!(n_iter.unwrap() >= 1, "n_iter must be at least 1");
}
#[test]
fn multivariate_y_equals_2x1_plus_3x2_plus_1_coefficients() {
let mut model = LinearRegression::new(true, 0.01, 20_000, 1e-10).unwrap();
let x = Array2::from_shape_vec(
(6, 2),
vec![
1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 3.0, 2.0, 2.0, 3.0, 4.0, 1.0, ],
)
.unwrap();
let y = array![6.0, 8.0, 9.0, 13.0, 14.0, 12.0];
model.fit(&x, &y).unwrap();
let coeff = model.get_coefficients().unwrap();
assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 3e-3);
assert_abs_diff_eq!(coeff[1], 3.0, epsilon = 3e-3);
let intercept = model.get_intercept().unwrap();
assert_abs_diff_eq!(intercept, 1.0, epsilon = 3e-3);
}
#[test]
fn multivariate_predictions_match_closed_form() {
let mut model = LinearRegression::new(true, 0.01, 20_000, 1e-10).unwrap();
let x = Array2::from_shape_vec(
(6, 2),
vec![1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 3.0, 2.0, 2.0, 3.0, 4.0, 1.0],
)
.unwrap();
let y = array![6.0, 8.0, 9.0, 13.0, 14.0, 12.0];
model.fit(&x, &y).unwrap();
let x_new = array![[1.0, 1.0], [2.0, 3.0]];
let preds = model.predict(&x_new).unwrap();
let expected = array![6.0, 14.0];
assert_allclose(&preds, &expected, 5e-3);
}
#[test]
fn no_intercept_stored_intercept_is_zero() {
let mut model = LinearRegression::new(false, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![2.0, 4.0, 6.0, 8.0, 10.0];
model.fit(&x, &y).unwrap();
let intercept = model.get_intercept().unwrap();
assert_abs_diff_eq!(intercept, 0.0, epsilon = 1e-15);
}
#[test]
fn no_intercept_coefficient_converges_to_slope() {
let mut model = LinearRegression::new(false, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![2.0, 4.0, 6.0, 8.0, 10.0];
model.fit(&x, &y).unwrap();
let coeff = model.get_coefficients().unwrap();
assert_abs_diff_eq!(coeff[0], 2.0, epsilon = 1e-4);
}
#[test]
fn no_intercept_getter_returns_false() {
let model = LinearRegression::new(false, 0.01, 1000, 1e-6).unwrap();
assert!(!model.get_fit_intercept());
}
#[test]
fn ols_sanity_y_equals_3x_plus_2_parameters() {
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![5.0, 8.0, 11.0, 14.0, 17.0];
model.fit(&x, &y).unwrap();
let coeff = model.get_coefficients().unwrap();
assert_abs_diff_eq!(coeff[0], 3.0, epsilon = 3e-3);
let intercept = model.get_intercept().unwrap();
assert_abs_diff_eq!(intercept, 2.0, epsilon = 3e-3);
}
#[test]
fn ols_sanity_y_equals_3x_plus_2_predictions() {
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![5.0, 8.0, 11.0, 14.0, 17.0];
model.fit(&x, &y).unwrap();
let preds = model.predict(&array![[6.0], [10.0]]).unwrap();
let expected = array![20.0, 32.0];
assert_allclose(&preds, &expected, 3e-3);
}
#[test]
fn fit_predict_matches_fit_then_predict() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut model_a = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let preds_a = model_a.fit_predict(&x, &y).unwrap();
let mut model_b = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
model_b.fit(&x, &y).unwrap();
let preds_b = model_b.predict(&x).unwrap();
assert_allclose(&preds_a, &preds_b, 1e-12);
}
#[test]
fn fit_predict_values_match_known_true_values() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
let preds = model.fit_predict(&x, &y).unwrap();
let expected = array![3.0, 5.0, 7.0, 9.0, 11.0];
assert_allclose(&preds, &expected, 5e-3);
}
#[test]
fn l2_regularization_shrinks_coefficient_norm() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut unregularized = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
unregularized.fit(&x, &y).unwrap();
let coeff_unreg = unregularized.get_coefficients().unwrap()[0].abs();
let mut ridge = LinearRegression::new(true, 0.01, 10_000, 1e-10)
.unwrap()
.with_regularization(RegularizationType::L2(5.0))
.unwrap();
ridge.fit(&x, &y).unwrap();
let coeff_ridge = ridge.get_coefficients().unwrap()[0].abs();
assert!(
coeff_ridge < coeff_unreg,
"Ridge coefficient {coeff_ridge} should be smaller than unregularized {coeff_unreg}"
);
}
#[test]
fn l1_regularization_shrinks_coefficient() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut unregularized = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
unregularized.fit(&x, &y).unwrap();
let coeff_unreg = unregularized.get_coefficients().unwrap()[0].abs();
let mut lasso = LinearRegression::new(true, 0.01, 10_000, 1e-10)
.unwrap()
.with_regularization(RegularizationType::L1(5.0))
.unwrap();
lasso.fit(&x, &y).unwrap();
let coeff_lasso = lasso.get_coefficients().unwrap()[0].abs();
assert!(
coeff_lasso < coeff_unreg,
"Lasso coefficient {coeff_lasso} should be smaller than unregularized {coeff_unreg}"
);
}
#[test]
fn l2_regularization_intercept_within_reasonable_range() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut ridge = LinearRegression::new(true, 0.01, 10_000, 1e-10)
.unwrap()
.with_regularization(RegularizationType::L2(0.1))
.unwrap();
ridge.fit(&x, &y).unwrap();
let intercept = ridge.get_intercept().unwrap();
assert!(
(intercept - 1.0).abs() < 0.5,
"Intercept {intercept} deviates too far from 1.0 under weak L2 regularization"
);
}
#[test]
fn determinism_same_data_identical_predictions() {
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let x_test = array![[6.0], [7.0], [8.0]];
let mut model_a = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
model_a.fit(&x, &y).unwrap();
let preds_a = model_a.predict(&x_test).unwrap();
let mut model_b = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
model_b.fit(&x, &y).unwrap();
let preds_b = model_b.predict(&x_test).unwrap();
assert_allclose(&preds_a, &preds_b, 0.0);
}
#[test]
fn save_load_round_trip_identical_predictions() {
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
let x_test = array![[6.0], [7.0], [0.5]];
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
model.fit(&x_train, &y_train).unwrap();
let preds_before = model.predict(&x_test).unwrap();
let path = "/tmp/rustyml_linear_regression_test_round_trip.json";
model.save_to_path(path).unwrap();
let loaded = LinearRegression::load_from_path(path).unwrap();
let preds_after = loaded.predict(&x_test).unwrap();
assert_allclose(&preds_before, &preds_after, 0.0);
let _ = std::fs::remove_file(path);
}
#[test]
fn save_load_preserves_model_state() {
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
let mut model = LinearRegression::new(false, 0.005, 8_000, 1e-9).unwrap();
model.fit(&x_train, &y_train).unwrap();
let path = "/tmp/rustyml_linear_regression_test_state.json";
model.save_to_path(path).unwrap();
let loaded = LinearRegression::load_from_path(path).unwrap();
assert_eq!(loaded.get_fit_intercept(), model.get_fit_intercept());
assert_abs_diff_eq!(
loaded.get_learning_rate(),
model.get_learning_rate(),
epsilon = 1e-15
);
assert_eq!(loaded.get_max_iterations(), model.get_max_iterations());
let orig_coeff = model.get_coefficients().unwrap();
let load_coeff = loaded.get_coefficients().unwrap();
assert_allclose(orig_coeff, load_coeff, 0.0);
let _ = std::fs::remove_file(path);
}
#[test]
fn default_constructor_has_expected_hyperparameters() {
let model = LinearRegression::default();
assert!(model.get_fit_intercept());
assert_abs_diff_eq!(model.get_learning_rate(), 0.01, epsilon = 1e-15);
assert_eq!(model.get_max_iterations(), 1000);
assert_abs_diff_eq!(model.get_tolerance(), 1e-5, epsilon = 1e-20);
assert!(model.get_coefficients().is_none());
assert!(model.get_intercept().is_none());
assert!(model.get_actual_iterations().is_none());
}
#[test]
fn default_constructor_can_fit_and_predict() {
let mut model = LinearRegression::default();
let x = array![[1.0], [2.0], [3.0]];
let y = array![3.0, 5.0, 7.0];
model.fit(&x, &y).unwrap();
let preds = model.predict(&array![[4.0]]).unwrap();
assert_abs_diff_eq!(preds[0], 9.0, epsilon = 5e-2);
}
#[test]
fn clone_of_fitted_model_makes_identical_predictions() {
let x_train = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y_train = array![3.0, 5.0, 7.0, 9.0, 11.0];
let x_test = array![[6.0], [0.0]];
let mut model = LinearRegression::new(true, 0.01, 10_000, 1e-10).unwrap();
model.fit(&x_train, &y_train).unwrap();
let preds_orig = model.predict(&x_test).unwrap();
let cloned = model.clone();
let preds_clone = cloned.predict(&x_test).unwrap();
assert_allclose(&preds_orig, &preds_clone, 0.0);
}
#[test]
fn fit_huge_learning_rate_diverges_returns_non_finite() {
let mut model = LinearRegression::new(true, 1e8, 1000, 1e-10).unwrap();
let x = array![[1.0], [2.0], [3.0], [4.0], [5.0]];
let y = array![3.0, 5.0, 7.0, 9.0, 11.0];
let result = model.fit(&x, &y);
assert!(
matches!(result, Err(Error::NonFinite(_))),
"expected NonFinite from in-loop divergence guard, got {:?}",
result
);
}
#[test]
fn l1_regularization_parallel_branch_recovers_informative_feature() {
let n_samples = 12usize;
let n_features = 200usize;
let x = Array2::from_shape_fn((n_samples, n_features), |(i, j)| {
if j == 0 {
(i as f64) - 5.5
} else {
0.01 * (((i * 31 + j * 17) % 7) as f64 - 3.0)
}
});
let y = Array1::from_shape_fn(n_samples, |i| 3.0 * ((i as f64) - 5.5));
let mut model = LinearRegression::new(false, 0.01, 20_000, 1e-12)
.unwrap()
.with_regularization(RegularizationType::L1(1e-3))
.unwrap();
model
.fit(&x, &y)
.expect("fit with >=200 features and L1 should succeed");
let coeffs = model.get_coefficients().unwrap();
assert_eq!(
coeffs.len(),
n_features,
"coefficient vector length must equal feature count"
);
let c0 = coeffs[0];
assert!(
c0 > 1.0,
"informative coefficient[0] = {c0} should be a large positive value (true slope 3.0)"
);
let max_other = coeffs
.iter()
.skip(1)
.fold(0.0_f64, |acc, &w| acc.max(w.abs()));
assert!(
max_other < 0.5,
"uninformative coefficients should stay small; largest |other| = {max_other}"
);
assert!(
c0.abs() > max_other,
"|coefficient[0]| = {} should dominate the largest other |coefficient| = {}",
c0.abs(),
max_other
);
}