linreg_core/wasm/
prediction_intervals.rs

1//! Prediction intervals for WASM
2//!
3//! Provides WASM bindings for computing prediction intervals from
4//! OLS, Ridge, Lasso, and Elastic Net regression models.
5
6#![cfg(feature = "wasm")]
7
8use wasm_bindgen::prelude::*;
9
10use super::domain::check_domain;
11use crate::error::{error_json, error_to_json};
12use crate::linalg;
13use crate::regularized;
14
15/// Computes OLS prediction intervals via WASM.
16///
17/// Fits an OLS model to the training data and computes prediction intervals
18/// for the new observations.
19///
20/// # Arguments
21///
22/// * `y_json` - JSON array of response variable values
23/// * `x_vars_json` - JSON array of predictor arrays (training data)
24/// * `new_x_json` - JSON array of predictor arrays (new observations)
25/// * `alpha` - Significance level (e.g., 0.05 for 95% PI)
26///
27/// # Returns
28///
29/// JSON string containing predicted values, lower/upper bounds, SE, leverage.
30#[wasm_bindgen]
31pub fn ols_prediction_intervals(
32    y_json: &str,
33    x_vars_json: &str,
34    new_x_json: &str,
35    alpha: f64,
36) -> String {
37    if let Err(e) = check_domain() {
38        return error_to_json(&e);
39    }
40
41    let y: Vec<f64> = match serde_json::from_str(y_json) {
42        Ok(v) => v,
43        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
44    };
45
46    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
47        Ok(v) => v,
48        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
49    };
50
51    let new_x_vecs: Vec<Vec<f64>> = match serde_json::from_str(new_x_json) {
52        Ok(v) => v,
53        Err(e) => return error_json(&format!("Failed to parse new_x: {}", e)),
54    };
55
56    let new_x_refs: Vec<&[f64]> = new_x_vecs.iter().map(|v| v.as_slice()).collect();
57
58    match crate::prediction_intervals::prediction_intervals(&y, &x_vars, &new_x_refs, alpha) {
59        Ok(output) => serde_json::to_string(&output)
60            .unwrap_or_else(|_| error_json("Failed to serialize prediction intervals")),
61        Err(e) => error_json(&e.to_string()),
62    }
63}
64
65/// Computes approximate Ridge regression prediction intervals via WASM.
66///
67/// Fits a Ridge model and computes conservative prediction intervals using
68/// leverage from unpenalized X'X and MSE from the ridge fit.
69///
70/// # Arguments
71///
72/// * `y_json` - JSON array of response variable values
73/// * `x_vars_json` - JSON array of predictor arrays (training data)
74/// * `new_x_json` - JSON array of predictor arrays (new observations)
75/// * `alpha` - Significance level (e.g., 0.05 for 95% PI)
76/// * `lambda` - Regularization strength
77/// * `standardize` - Whether to standardize predictors
78#[wasm_bindgen]
79pub fn ridge_prediction_intervals(
80    y_json: &str,
81    x_vars_json: &str,
82    new_x_json: &str,
83    alpha: f64,
84    lambda: f64,
85    standardize: bool,
86) -> String {
87    if let Err(e) = check_domain() {
88        return error_to_json(&e);
89    }
90
91    let y: Vec<f64> = match serde_json::from_str(y_json) {
92        Ok(v) => v,
93        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
94    };
95
96    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
97        Ok(v) => v,
98        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
99    };
100
101    let new_x_vecs: Vec<Vec<f64>> = match serde_json::from_str(new_x_json) {
102        Ok(v) => v,
103        Err(e) => return error_json(&format!("Failed to parse new_x: {}", e)),
104    };
105
106    // Build design matrix with intercept column
107    let (x, _n, _p) = build_design_matrix(&y, &x_vars);
108
109    let options = regularized::ridge::RidgeFitOptions {
110        lambda,
111        intercept: true,
112        standardize,
113        ..Default::default()
114    };
115
116    let fit = match regularized::ridge::ridge_fit(&x, &y, &options) {
117        Ok(f) => f,
118        Err(e) => return error_json(&e.to_string()),
119    };
120
121    let new_x_refs: Vec<&[f64]> = new_x_vecs.iter().map(|v| v.as_slice()).collect();
122
123    match crate::prediction_intervals::ridge_prediction_intervals(&fit, &x_vars, &new_x_refs, alpha) {
124        Ok(output) => serde_json::to_string(&output)
125            .unwrap_or_else(|_| error_json("Failed to serialize prediction intervals")),
126        Err(e) => error_json(&e.to_string()),
127    }
128}
129
130/// Computes approximate Lasso regression prediction intervals via WASM.
131///
132/// # Arguments
133///
134/// * `y_json` - JSON array of response variable values
135/// * `x_vars_json` - JSON array of predictor arrays (training data)
136/// * `new_x_json` - JSON array of predictor arrays (new observations)
137/// * `alpha` - Significance level (e.g., 0.05 for 95% PI)
138/// * `lambda` - Regularization strength
139/// * `standardize` - Whether to standardize predictors
140/// * `max_iter` - Maximum coordinate descent iterations
141/// * `tol` - Convergence tolerance
142#[wasm_bindgen]
143#[allow(clippy::too_many_arguments)]
144pub fn lasso_prediction_intervals(
145    y_json: &str,
146    x_vars_json: &str,
147    new_x_json: &str,
148    alpha: f64,
149    lambda: f64,
150    standardize: bool,
151    max_iter: usize,
152    tol: f64,
153) -> String {
154    if let Err(e) = check_domain() {
155        return error_to_json(&e);
156    }
157
158    let y: Vec<f64> = match serde_json::from_str(y_json) {
159        Ok(v) => v,
160        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
161    };
162
163    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
164        Ok(v) => v,
165        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
166    };
167
168    let new_x_vecs: Vec<Vec<f64>> = match serde_json::from_str(new_x_json) {
169        Ok(v) => v,
170        Err(e) => return error_json(&format!("Failed to parse new_x: {}", e)),
171    };
172
173    let (x, _n, _p) = build_design_matrix(&y, &x_vars);
174
175    let options = regularized::lasso::LassoFitOptions {
176        lambda,
177        intercept: true,
178        standardize,
179        max_iter,
180        tol,
181        ..Default::default()
182    };
183
184    let fit = match regularized::lasso::lasso_fit(&x, &y, &options) {
185        Ok(f) => f,
186        Err(e) => return error_json(&e.to_string()),
187    };
188
189    let new_x_refs: Vec<&[f64]> = new_x_vecs.iter().map(|v| v.as_slice()).collect();
190
191    match crate::prediction_intervals::lasso_prediction_intervals(&fit, &x_vars, &new_x_refs, alpha) {
192        Ok(output) => serde_json::to_string(&output)
193            .unwrap_or_else(|_| error_json("Failed to serialize prediction intervals")),
194        Err(e) => error_json(&e.to_string()),
195    }
196}
197
198/// Computes approximate Elastic Net regression prediction intervals via WASM.
199///
200/// # Arguments
201///
202/// * `y_json` - JSON array of response variable values
203/// * `x_vars_json` - JSON array of predictor arrays (training data)
204/// * `new_x_json` - JSON array of predictor arrays (new observations)
205/// * `alpha` - Significance level (e.g., 0.05 for 95% PI)
206/// * `lambda` - Regularization strength
207/// * `enet_alpha` - Elastic net mixing parameter (0 = Ridge, 1 = Lasso)
208/// * `standardize` - Whether to standardize predictors
209/// * `max_iter` - Maximum coordinate descent iterations
210/// * `tol` - Convergence tolerance
211#[wasm_bindgen]
212#[allow(clippy::too_many_arguments)]
213pub fn elastic_net_prediction_intervals(
214    y_json: &str,
215    x_vars_json: &str,
216    new_x_json: &str,
217    alpha: f64,
218    lambda: f64,
219    enet_alpha: f64,
220    standardize: bool,
221    max_iter: usize,
222    tol: f64,
223) -> String {
224    if let Err(e) = check_domain() {
225        return error_to_json(&e);
226    }
227
228    let y: Vec<f64> = match serde_json::from_str(y_json) {
229        Ok(v) => v,
230        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
231    };
232
233    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
234        Ok(v) => v,
235        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
236    };
237
238    let new_x_vecs: Vec<Vec<f64>> = match serde_json::from_str(new_x_json) {
239        Ok(v) => v,
240        Err(e) => return error_json(&format!("Failed to parse new_x: {}", e)),
241    };
242
243    let (x, _n, _p) = build_design_matrix(&y, &x_vars);
244
245    let options = regularized::elastic_net::ElasticNetOptions {
246        lambda,
247        alpha: enet_alpha,
248        intercept: true,
249        standardize,
250        max_iter,
251        tol,
252        ..Default::default()
253    };
254
255    let fit = match regularized::elastic_net::elastic_net_fit(&x, &y, &options) {
256        Ok(f) => f,
257        Err(e) => return error_json(&e.to_string()),
258    };
259
260    let new_x_refs: Vec<&[f64]> = new_x_vecs.iter().map(|v| v.as_slice()).collect();
261
262    match crate::prediction_intervals::elastic_net_prediction_intervals(&fit, &x_vars, &new_x_refs, alpha) {
263        Ok(output) => serde_json::to_string(&output)
264            .unwrap_or_else(|_| error_json("Failed to serialize prediction intervals")),
265        Err(e) => error_json(&e.to_string()),
266    }
267}
268
269/// Helper function to build a design matrix from column vectors.
270fn build_design_matrix(y: &[f64], x_vars: &[Vec<f64>]) -> (linalg::Matrix, usize, usize) {
271    let n = y.len();
272    let p = x_vars.len();
273
274    let mut x_data = vec![1.0; n * (p + 1)];
275    for (j, x_var) in x_vars.iter().enumerate() {
276        for (i, &val) in x_var.iter().enumerate() {
277            x_data[i * (p + 1) + j + 1] = val;
278        }
279    }
280
281    (linalg::Matrix::new(n, p + 1, x_data), n, p)
282}
linreg_core/wasm/prediction_intervals.rs

linreg_core/wasm/
prediction_intervals.rs
