linreg_core/wasm/
regularized.rs

1//! Regularized regression methods for WASM
2//!
3//! Provides WASM bindings for Ridge, Lasso, and Elastic Net regression.
4
5#![cfg(feature = "wasm")]
6
7use wasm_bindgen::prelude::*;
8
9use super::domain::check_domain;
10use crate::error::{error_json, error_to_json};
11use crate::linalg;
12use crate::regularized;
13
14/// Performs Ridge regression via WASM.
15///
16/// Ridge regression adds an L2 penalty to the coefficients, which helps with
17/// multicollinearity and overfitting. The intercept is never penalized.
18///
19/// # Arguments
20///
21/// * `y_json` - JSON array of response variable values
22/// * `x_vars_json` - JSON array of predictor arrays
23/// * `variable_names` - JSON array of variable names
24/// * `lambda` - Regularization strength (>= 0, typical range 0.01 to 100)
25/// * `standardize` - Whether to standardize predictors (recommended: true)
26///
27/// # Returns
28///
29/// JSON string containing:
30/// - `lambda` - Lambda value used
31/// - `intercept` - Intercept coefficient
32/// - `coefficients` - Slope coefficients
33/// - `fitted_values` - Predictions on training data
34/// - `residuals` - Residuals (y - fitted_values)
35/// - `df` - Effective degrees of freedom
36///
37/// # Errors
38///
39/// Returns a JSON error object if parsing fails, lambda is negative,
40/// or domain check fails.
41#[wasm_bindgen]
42pub fn ridge_regression(
43    y_json: &str,
44    x_vars_json: &str,
45    _variable_names: &str,
46    lambda: f64,
47    standardize: bool,
48) -> String {
49    if let Err(e) = check_domain() {
50        return error_to_json(&e);
51    }
52
53    // Parse JSON input
54    let y: Vec<f64> = match serde_json::from_str(y_json) {
55        Ok(v) => v,
56        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
57    };
58
59    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
60        Ok(v) => v,
61        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
62    };
63
64    // Build design matrix with intercept column
65    let (x, n, p) = build_design_matrix(&y, &x_vars);
66
67    if n <= p + 1 {
68        return error_json(&format!(
69            "Insufficient data: need at least {} observations for {} predictors",
70            p + 2,
71            p
72        ));
73    }
74
75    // Configure ridge options
76    let options = regularized::ridge::RidgeFitOptions {
77        lambda,
78        intercept: true,
79        standardize,
80        max_iter: 100000,
81        tol: 1e-7,
82        warm_start: None,
83        weights: None,
84    };
85
86    match regularized::ridge::ridge_fit(&x, &y, &options) {
87        Ok(output) => serde_json::to_string(&output)
88            .unwrap_or_else(|_| error_json("Failed to serialize ridge regression result")),
89        Err(e) => error_json(&e.to_string()),
90    }
91}
92
93/// Performs Lasso regression via WASM.
94///
95/// Lasso regression adds an L1 penalty to the coefficients, which performs
96/// automatic variable selection by shrinking some coefficients to exactly zero.
97/// The intercept is never penalized.
98///
99/// # Arguments
100///
101/// * `y_json` - JSON array of response variable values
102/// * `x_vars_json` - JSON array of predictor arrays
103/// * `variable_names` - JSON array of variable names
104/// * `lambda` - Regularization strength (>= 0, typical range 0.01 to 10)
105/// * `standardize` - Whether to standardize predictors (recommended: true)
106/// * `max_iter` - Maximum coordinate descent iterations (default: 100000)
107/// * `tol` - Convergence tolerance (default: 1e-7)
108///
109/// # Returns
110///
111/// JSON string containing:
112/// - `lambda` - Lambda value used
113/// - `intercept` - Intercept coefficient
114/// - `coefficients` - Slope coefficients (some may be exactly zero)
115/// - `fitted_values` - Predictions on training data
116/// - `residuals` - Residuals (y - fitted_values)
117/// - `n_nonzero` - Number of non-zero coefficients (excluding intercept)
118/// - `iterations` - Number of coordinate descent iterations
119/// - `converged` - Whether the algorithm converged
120///
121/// # Errors
122///
123/// Returns a JSON error object if parsing fails, lambda is negative,
124/// or domain check fails.
125#[wasm_bindgen]
126pub fn lasso_regression(
127    y_json: &str,
128    x_vars_json: &str,
129    _variable_names: &str,
130    lambda: f64,
131    standardize: bool,
132    max_iter: usize,
133    tol: f64,
134) -> String {
135    if let Err(e) = check_domain() {
136        return error_to_json(&e);
137    }
138
139    // Parse JSON input
140    let y: Vec<f64> = match serde_json::from_str(y_json) {
141        Ok(v) => v,
142        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
143    };
144
145    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
146        Ok(v) => v,
147        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
148    };
149
150    // Build design matrix with intercept column
151    let (x, n, p) = build_design_matrix(&y, &x_vars);
152
153    if n <= p + 1 {
154        return error_json(&format!(
155            "Insufficient data: need at least {} observations for {} predictors",
156            p + 2,
157            p
158        ));
159    }
160
161    // Configure lasso options
162    let options = regularized::lasso::LassoFitOptions {
163        lambda,
164        intercept: true,
165        standardize,
166        max_iter,
167        tol,
168        ..Default::default()
169    };
170
171    match regularized::lasso::lasso_fit(&x, &y, &options) {
172        Ok(output) => serde_json::to_string(&output)
173            .unwrap_or_else(|_| error_json("Failed to serialize lasso regression result")),
174        Err(e) => error_json(&e.to_string()),
175    }
176}
177
178/// Performs Elastic Net regression via WASM.
179///
180/// Elastic Net combines L1 (Lasso) and L2 (Ridge) penalties.
181///
182/// # Arguments
183///
184/// * `y_json` - JSON array of response variable values
185/// * `x_vars_json` - JSON array of predictor arrays
186/// * `variable_names` - JSON array of variable names
187/// * `lambda` - Regularization strength (>= 0)
188/// * `alpha` - Elastic net mixing parameter (0 = Ridge, 1 = Lasso)
189/// * `standardize` - Whether to standardize predictors (recommended: true)
190/// * `max_iter` - Maximum coordinate descent iterations
191/// * `tol` - Convergence tolerance
192///
193/// # Returns
194///
195/// JSON string containing regression results (same fields as Lasso).
196///
197/// # Errors
198///
199/// Returns a JSON error object if parsing fails, parameters are invalid,
200/// or domain check fails.
201#[wasm_bindgen]
202#[allow(clippy::too_many_arguments)]
203pub fn elastic_net_regression(
204    y_json: &str,
205    x_vars_json: &str,
206    _variable_names: &str,
207    lambda: f64,
208    alpha: f64,
209    standardize: bool,
210    max_iter: usize,
211    tol: f64,
212) -> String {
213    if let Err(e) = check_domain() {
214        return error_to_json(&e);
215    }
216
217    // Parse JSON input
218    let y: Vec<f64> = match serde_json::from_str(y_json) {
219        Ok(v) => v,
220        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
221    };
222
223    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
224        Ok(v) => v,
225        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
226    };
227
228    // Build design matrix with intercept column
229    let (x, n, p) = build_design_matrix(&y, &x_vars);
230
231    if n <= p + 1 {
232        return error_json(&format!(
233            "Insufficient data: need at least {} observations for {} predictors",
234            p + 2,
235            p
236        ));
237    }
238
239    // Configure elastic net options
240    let options = regularized::elastic_net::ElasticNetOptions {
241        lambda,
242        alpha,
243        intercept: true,
244        standardize,
245        max_iter,
246        tol,
247        ..Default::default()
248    };
249
250    match regularized::elastic_net::elastic_net_fit(&x, &y, &options) {
251        Ok(output) => serde_json::to_string(&output)
252            .unwrap_or_else(|_| error_json("Failed to serialize elastic net regression result")),
253        Err(e) => error_json(&e.to_string()),
254    }
255}
256
257/// Lightweight result for regularization path (excludes residuals/fitted values).
258#[derive(serde::Serialize)]
259struct PathResult {
260    lambdas: Vec<f64>,
261    coefficients: Vec<Vec<f64>>,
262    r_squared: Vec<f64>,
263    aic: Vec<f64>,
264    bic: Vec<f64>,
265    n_nonzero: Vec<usize>,
266}
267
268/// Fits an elastic net regularization path via WASM (Optimized).
269///
270/// Computes the coefficient path for a sequence of lambda values.
271/// Returns a lightweight summary to avoid excessive JSON serialization overhead.
272///
273/// # Arguments
274///
275/// * `y_json` - JSON array of response variable values
276/// * `x_vars_json` - JSON array of predictor arrays
277/// * `n_lambda` - Number of lambda values (default: 100)
278/// * `lambda_min_ratio` - Ratio for smallest lambda
279/// * `alpha` - Mixing parameter (0 = Ridge, 1 = Lasso)
280/// * `standardize` - Whether to standardize predictors
281/// * `max_iter` - Maximum iterations per lambda
282/// * `tol` - Convergence tolerance
283///
284/// # Returns
285///
286/// JSON string containing `PathResult` (lambdas, coefficients, stats).
287#[wasm_bindgen]
288#[allow(clippy::too_many_arguments)]
289pub fn elastic_net_path_wasm(
290    y_json: &str,
291    x_vars_json: &str,
292    n_lambda: usize,
293    lambda_min_ratio: f64,
294    alpha: f64,
295    standardize: bool,
296    max_iter: usize,
297    tol: f64,
298) -> String {
299    if let Err(e) = check_domain() {
300        return error_to_json(&e);
301    }
302
303    // Parse JSON input
304    let y: Vec<f64> = match serde_json::from_str(y_json) {
305        Ok(v) => v,
306        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
307    };
308
309    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
310        Ok(v) => v,
311        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
312    };
313
314    // Build design matrix with intercept column
315    let (x, n, p) = build_design_matrix(&y, &x_vars);
316
317    if n <= p + 1 {
318        return error_json(&format!(
319            "Insufficient data: need at least {} observations for {} predictors",
320            p + 2,
321            p
322        ));
323    }
324
325    // Configure options
326    let path_options = regularized::path::LambdaPathOptions {
327        nlambda: n_lambda.max(1),
328        lambda_min_ratio: if lambda_min_ratio > 0.0 {
329            Some(lambda_min_ratio)
330        } else {
331            None
332        },
333        alpha,
334        ..Default::default()
335    };
336
337    let fit_options = regularized::elastic_net::ElasticNetOptions {
338        lambda: 0.0, // Ignored by path fitting
339        alpha,
340        intercept: true,
341        standardize,
342        max_iter,
343        tol,
344        ..Default::default()
345    };
346
347    match regularized::elastic_net::elastic_net_path(&x, &y, &path_options, &fit_options) {
348        Ok(fits) => {
349            // Convert to lightweight structure
350            let result = PathResult {
351                lambdas: fits.iter().map(|f| f.lambda).collect(),
352                coefficients: fits.iter().map(|f| f.coefficients.clone()).collect(),
353                r_squared: fits.iter().map(|f| f.r_squared).collect(),
354                aic: fits.iter().map(|f| f.aic).collect(),
355                bic: fits.iter().map(|f| f.bic).collect(),
356                n_nonzero: fits.iter().map(|f| f.n_nonzero).collect(),
357            };
358            
359            serde_json::to_string(&result)
360                .unwrap_or_else(|_| error_json("Failed to serialize elastic net path result"))
361        },
362        Err(e) => error_json(&e.to_string()),
363    }
364}
365
366/// Generates a lambda path for regularized regression via WASM.
367///
368/// Creates a logarithmically-spaced sequence of lambda values from lambda_max
369/// (where all penalized coefficients are zero) down to lambda_min. This is
370/// useful for fitting regularization paths and selecting optimal lambda via
371/// cross-validation.
372///
373/// # Arguments
374///
375/// * `y_json` - JSON array of response variable values
376/// * `x_vars_json` - JSON array of predictor arrays
377/// * `n_lambda` - Number of lambda values to generate (default: 100)
378/// * `lambda_min_ratio` - Ratio for smallest lambda (default: 0.0001 if n >= p, else 0.01)
379///
380/// # Returns
381///
382/// JSON string containing:
383/// - `lambda_path` - Array of lambda values in decreasing order
384/// - `lambda_max` - Maximum lambda value
385/// - `lambda_min` - Minimum lambda value
386/// - `n_lambda` - Number of lambda values
387///
388/// # Errors
389///
390/// Returns a JSON error object if parsing fails or domain check fails.
391#[wasm_bindgen]
392pub fn make_lambda_path(
393    y_json: &str,
394    x_vars_json: &str,
395    n_lambda: usize,
396    lambda_min_ratio: f64,
397) -> String {
398    if let Err(e) = check_domain() {
399        return error_to_json(&e);
400    }
401
402    // Parse JSON input
403    let y: Vec<f64> = match serde_json::from_str(y_json) {
404        Ok(v) => v,
405        Err(e) => return error_json(&format!("Failed to parse y: {}", e)),
406    };
407
408    let x_vars: Vec<Vec<f64>> = match serde_json::from_str(x_vars_json) {
409        Ok(v) => v,
410        Err(e) => return error_json(&format!("Failed to parse x_vars: {}", e)),
411    };
412
413    // Build design matrix with intercept column
414    let (x, n, p) = build_design_matrix(&y, &x_vars);
415
416    // Standardize X for lambda path computation
417    let x_mean: Vec<f64> = (0..x.cols)
418        .map(|j| {
419            if j == 0 {
420                1.0 // Intercept column
421            } else {
422                (0..n).map(|i| x.get(i, j)).sum::<f64>() / n as f64
423            }
424        })
425        .collect();
426
427    let x_standardized: Vec<f64> = (0..x.cols)
428        .map(|j| {
429            if j == 0 {
430                0.0 // Intercept column - no centering
431            } else {
432                let mean = x_mean[j];
433                let variance =
434                    (0..n).map(|i| (x.get(i, j) - mean).powi(2)).sum::<f64>() / (n - 1) as f64;
435                variance.sqrt()
436            }
437        })
438        .collect();
439
440    // Build standardized X matrix
441    let mut x_standardized_data = vec![1.0; n * (p + 1)];
442    for j in 0..x.cols {
443        for i in 0..n {
444            if j == 0 {
445                x_standardized_data[i * (p + 1)] = 1.0; // Intercept
446            } else {
447                let std = x_standardized[j];
448                if std > 1e-10 {
449                    x_standardized_data[i * (p + 1) + j] = (x.get(i, j) - x_mean[j]) / std;
450                } else {
451                    x_standardized_data[i * (p + 1) + j] = 0.0;
452                }
453            }
454        }
455    }
456    let x_standardized = linalg::Matrix::new(n, p + 1, x_standardized_data);
457
458    // Center y
459    let y_mean: f64 = y.iter().sum::<f64>() / n as f64;
460    let y_centered: Vec<f64> = y.iter().map(|&yi| yi - y_mean).collect();
461
462    // Configure lambda path options
463    let options = regularized::path::LambdaPathOptions {
464        nlambda: n_lambda.max(1),
465        lambda_min_ratio: if lambda_min_ratio > 0.0 {
466            Some(lambda_min_ratio)
467        } else {
468            None
469        },
470        alpha: 1.0, // Lasso
471        ..Default::default()
472    };
473
474    let lambda_path =
475        regularized::path::make_lambda_path(&x_standardized, &y_centered, &options, None, Some(0));
476
477    let lambda_max = lambda_path.first().copied().unwrap_or(0.0);
478    let lambda_min = lambda_path.last().copied().unwrap_or(0.0);
479
480    // Return as JSON (note: infinity serializes as null in JSON, handled in JS)
481    let result = serde_json::json!({
482        "lambda_path": lambda_path,
483        "lambda_max": lambda_max,
484        "lambda_min": lambda_min,
485        "n_lambda": lambda_path.len()
486    });
487
488    result.to_string()
489}
490
491/// Helper function to build a design matrix from column vectors.
492///
493/// # Arguments
494///
495/// * `y` - Response variable (used to determine n)
496/// * `x_vars` - Predictor column vectors
497///
498/// # Returns
499///
500/// A tuple of (Matrix, n, p) where p is the number of predictors (excluding intercept)
501fn build_design_matrix(y: &[f64], x_vars: &[Vec<f64>]) -> (linalg::Matrix, usize, usize) {
502    let n = y.len();
503    let p = x_vars.len();
504
505    let mut x_data = vec![1.0; n * (p + 1)]; // Intercept column
506    for (j, x_var) in x_vars.iter().enumerate() {
507        for (i, &val) in x_var.iter().enumerate() {
508            x_data[i * (p + 1) + j + 1] = val;
509        }
510    }
511
512    (linalg::Matrix::new(n, p + 1, x_data), n, p)
513}
linreg_core/wasm/regularized.rs

linreg_core/wasm/
regularized.rs
