Skip to main content

linreg_core/polynomial/
features.rs

1use crate::error::{Error, Result};
2
3/// Expand a single predictor into higher-order polynomial features.
4///
5/// Returns feature vectors for degrees **2 through `degree`** only. The linear
6/// x term (degree 1) is intentionally excluded; the caller adds it separately.
7///
8/// # Arguments
9///
10/// * `x` - Predictor variable (n observations, already centered if desired)
11/// * `degree` - Polynomial degree (must be ≥ 1)
12/// * `center` - Whether to center `x` before raising to each power
13/// * `x_mean` - Pre-computed mean of `x` (used only when `center = true`)
14///
15/// # Returns
16///
17/// A `Vec` of length `degree - 1` where element `i` corresponds to `x^(i+2)`.
18///
19/// # Example
20///
21/// ```
22/// use linreg_core::polynomial::features::polynomial_features;
23///
24/// let x = vec![1.0, 2.0, 3.0];
25/// let features = polynomial_features(&x, 3, false, 0.0).unwrap();
26/// // features[0] = [1.0, 4.0, 9.0]    (x²)
27/// // features[1] = [1.0, 8.0, 27.0]   (x³)
28/// ```
29pub fn polynomial_features(
30    x: &[f64],
31    degree: usize,
32    center: bool,
33    x_mean: f64,
34) -> Result<Vec<Vec<f64>>> {
35    if degree < 1 {
36        return Err(Error::InvalidInput(
37            "Polynomial degree must be at least 1".into(),
38        ));
39    }
40
41    let n = x.len();
42    if n < 2 {
43        return Err(Error::InsufficientData {
44            required: 2,
45            available: n,
46        });
47    }
48
49    for (i, &xi) in x.iter().enumerate() {
50        if !xi.is_finite() {
51            return Err(Error::InvalidInput(format!(
52                "x[{}] is not finite: {}",
53                i, xi
54            )));
55        }
56    }
57
58    let x_centered: Vec<f64> = if center {
59        x.iter().map(|&xi| xi - x_mean).collect()
60    } else {
61        x.to_vec()
62    };
63
64    // Create polynomial features: x², x³, …, x^degree
65    let mut features = Vec::with_capacity(degree.saturating_sub(1));
66    for d in 2..=degree {
67        let power = d as i32;
68        let poly_feature: Vec<f64> = x_centered.iter().map(|&xi| xi.powi(power)).collect();
69
70        for (i, &val) in poly_feature.iter().enumerate() {
71            if !val.is_finite() {
72                return Err(Error::InvalidInput(format!(
73                    "x^{} at index {} is not finite: {}",
74                    d, i, val
75                )));
76            }
77        }
78
79        features.push(poly_feature);
80    }
81
82    Ok(features)
83}
84
85/// Generate feature names for polynomial regression terms.
86///
87/// Returns names for the linear term and all higher-order terms:
88/// `["x", "x^2", "x^3", …]` (or centered variants).
89pub fn polynomial_feature_names(degree: usize, centered: bool) -> Vec<String> {
90    let mut names = Vec::with_capacity(degree);
91
92    if centered {
93        names.push("x_centered".to_string());
94        for d in 2..=degree {
95            names.push(format!("x^{}_centered", d));
96        }
97    } else {
98        names.push("x".to_string());
99        for d in 2..=degree {
100            names.push(format!("x^{}", d));
101        }
102    }
103
104    names
105}
106
107/// Compute the mean of a slice.
108pub fn compute_mean(x: &[f64]) -> f64 {
109    x.iter().sum::<f64>() / x.len() as f64
110}
111
112/// Compute the sample standard deviation of a slice given a pre-computed mean.
113pub fn compute_std(x: &[f64], x_mean: f64) -> f64 {
114    let variance = x
115        .iter()
116        .map(|&xi| (xi - x_mean).powi(2))
117        .sum::<f64>()
118        / (x.len() - 1) as f64;
119    variance.sqrt()
120}