sklears_python/linear/
bayesian_ridge.rs

1//! Python bindings for Bayesian Ridge Regression
2//!
3//! This module provides Python bindings for Bayesian Ridge Regression,
4//! offering scikit-learn compatible interfaces with automatic relevance determination
5//! and uncertainty quantification using the sklears-linear crate.
6
7use super::common::*;
8use pyo3::types::PyDict;
9use pyo3::Bound;
10use sklears_core::traits::{Fit, Predict, Score, Trained};
11use sklears_linear::{BayesianRidge, BayesianRidgeConfig};
12
13/// Python-specific configuration wrapper for BayesianRidge
14#[derive(Debug, Clone)]
15pub struct PyBayesianRidgeConfig {
16    pub max_iter: usize,
17    pub tol: f64,
18    pub alpha_init: Option<f64>,
19    pub lambda_init: Option<f64>,
20    pub fit_intercept: bool,
21    pub compute_score: bool,
22    pub copy_x: bool,
23}
24
25impl Default for PyBayesianRidgeConfig {
26    fn default() -> Self {
27        Self {
28            max_iter: 300,
29            tol: 1e-3,
30            alpha_init: Some(1.0),
31            lambda_init: Some(1.0),
32            fit_intercept: true,
33            compute_score: false,
34            copy_x: true,
35        }
36    }
37}
38
39impl From<PyBayesianRidgeConfig> for BayesianRidgeConfig {
40    fn from(py_config: PyBayesianRidgeConfig) -> Self {
41        BayesianRidgeConfig {
42            max_iter: py_config.max_iter,
43            tol: py_config.tol,
44            alpha_init: py_config
45                .alpha_init
46                .unwrap_or_else(|| BayesianRidgeConfig::default().alpha_init),
47            lambda_init: py_config
48                .lambda_init
49                .unwrap_or_else(|| BayesianRidgeConfig::default().lambda_init),
50            fit_intercept: py_config.fit_intercept,
51            compute_score: py_config.compute_score,
52        }
53    }
54}
55
56/// Bayesian ridge regression.
57///
58/// Fit a Bayesian ridge model. See the Notes section for details on this
59/// implementation and the optimization of the regularization parameters
60/// lambda (precision of the weights) and alpha (precision of the noise).
61///
62/// Parameters
63/// ----------
64/// max_iter : int, default=300
65///     Maximum number of iterations. Should be greater than or equal to 1.
66///
67/// tol : float, default=1e-3
68///     Stop the algorithm if w has converged.
69///
70/// alpha_init : float, default=1.0
71///     Initial value for alpha (precision of the weights).
72///     If not provided, alpha_init is set to 1.0.
73///
74/// lambda_init : float, default=1.0
75///     Initial value for lambda (precision of the noise).
76///     If not provided, lambda_init is set to 1.0.
77///
78/// fit_intercept : bool, default=True
79///     Whether to calculate the intercept for this model.
80///     The intercept is not treated as a probabilistic parameter
81///     and thus has no associated variance. If set
82///     to False, no intercept will be used in calculations
83///     (i.e. data is expected to be centered).
84///
85/// compute_score : bool, default=False
86///     If True, compute the log marginal likelihood at each iteration of the
87///     optimization.
88///
89/// copy_X : bool, default=True
90///     If True, X will be copied; else, it may be overwritten.
91///
92/// Attributes
93/// ----------
94/// coef_ : array-like of shape (n_features,)
95///     Coefficients of the regression model (mean of distribution)
96///
97/// intercept_ : float
98///     Independent term in decision function. Set to 0.0 if
99///     ``fit_intercept = False``.
100///
101/// alpha_ : float
102///     Estimated precision of the weights.
103///
104/// lambda_ : float
105///     Estimated precision of the noise.
106///
107/// sigma_ : array-like of shape (n_features, n_features)
108///     Estimated variance-covariance matrix of the weights
109///
110/// scores_ : array-like of shape (n_iter_ + 1,)
111///     If computed_score is True, value of the log marginal likelihood (to be
112///     maximized) at each iteration of the optimization. The array starts with
113///     the value of the log marginal likelihood obtained for the initial values
114///     of alpha and lambda and ends with the value obtained for the estimated
115///     alpha and lambda.
116///
117/// n_iter_ : int
118///     The actual number of iterations to reach the stopping criterion.
119///
120/// n_features_in_ : int
121///     Number of features seen during fit.
122///
123/// Examples
124/// --------
125/// >>> from sklears_python import BayesianRidge
126/// >>> import numpy as np
127/// >>> X = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
128/// >>> # y = 1 * x_0 + 2 * x_1 + 3
129/// >>> y = np.dot(X, [1, 2]) + 3
130/// >>> reg = BayesianRidge()
131/// >>> reg.fit(X, y)
132/// BayesianRidge()
133/// >>> reg.predict([[1, 0]])
134/// array([4.])
135/// >>> reg.coef_
136/// array([1., 2.])
137///
138/// Notes
139/// -----
140/// There exist several strategies to perform Bayesian ridge regression. This
141/// implementation is based on the algorithm described in Appendix A of
142/// (Tipping, 2001) where updates of the regularization parameters are done as
143/// suggested in (MacKay, 1992). Note that according to A New
144/// View of Automatic Relevance Determination (Wipf and Nagarajan, 2008) these
145/// update rules do not guarantee that the marginal likelihood is increasing
146/// between two consecutive iterations of the optimization.
147///
148/// References
149/// ----------
150/// D. J. C. MacKay, Bayesian Interpolation, Computation and Neural Systems,
151/// Vol. 4, No. 3, 1992.
152///
153/// M. E. Tipping, Sparse Bayesian Learning and the Relevance Vector Machine,
154/// Journal of Machine Learning Research, Vol. 1, 2001.
155#[pyclass(name = "BayesianRidge")]
156pub struct PyBayesianRidge {
157    /// Python-specific configuration
158    py_config: PyBayesianRidgeConfig,
159    /// Trained model instance using the actual sklears-linear implementation
160    fitted_model: Option<BayesianRidge<Trained>>,
161}
162
163#[pymethods]
164impl PyBayesianRidge {
165    #[new]
166    #[pyo3(signature = (max_iter=300, tol=1e-3, alpha_init=1.0, lambda_init=1.0, fit_intercept=true, compute_score=false, copy_x=true))]
167    fn new(
168        max_iter: usize,
169        tol: f64,
170        alpha_init: f64,
171        lambda_init: f64,
172        fit_intercept: bool,
173        compute_score: bool,
174        copy_x: bool,
175    ) -> PyResult<Self> {
176        // Validate parameters
177        if max_iter == 0 {
178            return Err(PyValueError::new_err("max_iter must be greater than 0"));
179        }
180        if tol <= 0.0 {
181            return Err(PyValueError::new_err("tol must be positive"));
182        }
183        if alpha_init <= 0.0 {
184            return Err(PyValueError::new_err("alpha_init must be positive"));
185        }
186        if lambda_init <= 0.0 {
187            return Err(PyValueError::new_err("lambda_init must be positive"));
188        }
189
190        let py_config = PyBayesianRidgeConfig {
191            max_iter,
192            tol,
193            alpha_init: Some(alpha_init),
194            lambda_init: Some(lambda_init),
195            fit_intercept,
196            compute_score,
197            copy_x,
198        };
199
200        Ok(Self {
201            py_config,
202            fitted_model: None,
203        })
204    }
205
206    /// Fit the Bayesian Ridge regression model
207    fn fit(&mut self, x: PyReadonlyArray2<f64>, y: PyReadonlyArray1<f64>) -> PyResult<()> {
208        let x_array = pyarray_to_core_array2(x)?;
209        let y_array = pyarray_to_core_array1(y)?;
210
211        // Validate input arrays
212        validate_fit_arrays(&x_array, &y_array)?;
213
214        // Create sklears-linear model with Bayesian Ridge configuration
215        let model = BayesianRidge::new()
216            .max_iter(self.py_config.max_iter)
217            .tol(self.py_config.tol)
218            .fit_intercept(self.py_config.fit_intercept)
219            .compute_score(self.py_config.compute_score);
220
221        // Fit the model using sklears-linear's implementation
222        match model.fit(&x_array, &y_array) {
223            Ok(fitted_model) => {
224                self.fitted_model = Some(fitted_model);
225                Ok(())
226            }
227            Err(e) => Err(PyValueError::new_err(format!(
228                "Failed to fit Bayesian Ridge model: {:?}",
229                e
230            ))),
231        }
232    }
233
234    /// Predict using the fitted model
235    fn predict(&self, py: Python<'_>, x: PyReadonlyArray2<f64>) -> PyResult<Py<PyArray1<f64>>> {
236        let fitted = self
237            .fitted_model
238            .as_ref()
239            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
240
241        let x_array = pyarray_to_core_array2(x)?;
242        validate_predict_array(&x_array)?;
243
244        match fitted.predict(&x_array) {
245            Ok(predictions) => Ok(core_array1_to_py(py, &predictions)),
246            Err(e) => Err(PyValueError::new_err(format!("Prediction failed: {:?}", e))),
247        }
248    }
249
250    /// Get model coefficients
251    #[getter]
252    fn coef_(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {
253        let fitted = self
254            .fitted_model
255            .as_ref()
256            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
257
258        Ok(core_array1_to_py(py, fitted.coef()))
259    }
260
261    /// Get model intercept
262    #[getter]
263    fn intercept_(&self) -> PyResult<f64> {
264        let fitted = self
265            .fitted_model
266            .as_ref()
267            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
268
269        Ok(fitted.intercept().unwrap_or(0.0))
270    }
271
272    /// Get estimated precision of weights (alpha)
273    #[getter]
274    fn alpha_(&self) -> PyResult<f64> {
275        let fitted = self
276            .fitted_model
277            .as_ref()
278            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
279
280        Ok(fitted.alpha())
281    }
282
283    /// Get estimated precision of noise (lambda)
284    #[getter]
285    fn lambda_(&self) -> PyResult<f64> {
286        let fitted = self
287            .fitted_model
288            .as_ref()
289            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
290
291        Ok(fitted.lambda())
292    }
293
294    /// Calculate R² score
295    fn score(&self, x: PyReadonlyArray2<f64>, y: PyReadonlyArray1<f64>) -> PyResult<f64> {
296        let fitted = self
297            .fitted_model
298            .as_ref()
299            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
300
301        let x_array = pyarray_to_core_array2(x)?;
302        let y_array = pyarray_to_core_array1(y)?;
303
304        match fitted.score(&x_array, &y_array) {
305            Ok(score) => Ok(score),
306            Err(e) => Err(PyValueError::new_err(format!(
307                "Score calculation failed: {:?}",
308                e
309            ))),
310        }
311    }
312
313    /// Get number of features
314    #[getter]
315    fn n_features_in_(&self) -> PyResult<usize> {
316        let fitted = self
317            .fitted_model
318            .as_ref()
319            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
320
321        // Infer number of features from coefficient array length
322        Ok(fitted.coef().len())
323    }
324
325    /// Return parameters for this estimator (sklearn compatibility)
326    fn get_params(&self, py: Python<'_>, deep: Option<bool>) -> PyResult<Py<PyDict>> {
327        let _deep = deep.unwrap_or(true);
328
329        let dict = PyDict::new(py);
330
331        dict.set_item("max_iter", self.py_config.max_iter)?;
332        dict.set_item("tol", self.py_config.tol)?;
333        dict.set_item("alpha_init", self.py_config.alpha_init)?;
334        dict.set_item("lambda_init", self.py_config.lambda_init)?;
335        dict.set_item("fit_intercept", self.py_config.fit_intercept)?;
336        dict.set_item("compute_score", self.py_config.compute_score)?;
337        dict.set_item("copy_X", self.py_config.copy_x)?;
338
339        Ok(dict.into())
340    }
341
342    /// Set parameters for this estimator (sklearn compatibility)
343    fn set_params(&mut self, kwargs: &Bound<'_, PyDict>) -> PyResult<()> {
344        // Update configuration parameters
345        if let Some(max_iter) = kwargs.get_item("max_iter")? {
346            let max_iter_val: usize = max_iter.extract()?;
347            if max_iter_val == 0 {
348                return Err(PyValueError::new_err("max_iter must be greater than 0"));
349            }
350            self.py_config.max_iter = max_iter_val;
351        }
352        if let Some(tol) = kwargs.get_item("tol")? {
353            let tol_val: f64 = tol.extract()?;
354            if tol_val <= 0.0 {
355                return Err(PyValueError::new_err("tol must be positive"));
356            }
357            self.py_config.tol = tol_val;
358        }
359        if let Some(alpha_init) = kwargs.get_item("alpha_init")? {
360            let alpha_init_val: f64 = alpha_init.extract()?;
361            if alpha_init_val <= 0.0 {
362                return Err(PyValueError::new_err("alpha_init must be positive"));
363            }
364            self.py_config.alpha_init = Some(alpha_init_val);
365        }
366        if let Some(lambda_init) = kwargs.get_item("lambda_init")? {
367            let lambda_init_val: f64 = lambda_init.extract()?;
368            if lambda_init_val <= 0.0 {
369                return Err(PyValueError::new_err("lambda_init must be positive"));
370            }
371            self.py_config.lambda_init = Some(lambda_init_val);
372        }
373        if let Some(fit_intercept) = kwargs.get_item("fit_intercept")? {
374            self.py_config.fit_intercept = fit_intercept.extract()?;
375        }
376        if let Some(compute_score) = kwargs.get_item("compute_score")? {
377            self.py_config.compute_score = compute_score.extract()?;
378        }
379        if let Some(copy_x) = kwargs.get_item("copy_X")? {
380            self.py_config.copy_x = copy_x.extract()?;
381        }
382
383        // Clear fitted model since config changed
384        self.fitted_model = None;
385
386        Ok(())
387    }
388
389    /// String representation
390    fn __repr__(&self) -> String {
391        format!(
392            "BayesianRidge(max_iter={}, tol={}, alpha_init={:?}, lambda_init={:?}, fit_intercept={}, compute_score={}, copy_X={})",
393            self.py_config.max_iter,
394            self.py_config.tol,
395            self.py_config.alpha_init,
396            self.py_config.lambda_init,
397            self.py_config.fit_intercept,
398            self.py_config.compute_score,
399            self.py_config.copy_x
400        )
401    }
402}