rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
redclawsystems
2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/specialized/rtx-sklearn-py/src/wrappers/regressors.rs
+713
View File
@@ -0,0 +1,713 @@
/*!
sklearn-compatible regressor wrappers
*/
use pyo3::prelude::*;
use pyo3::types::{PyDict};
use numpy::{PyReadonlyArrayDyn, PyArray1, PyArray2};
use ndarray::{Array1, Array2};
use std::collections::HashMap;
use rtx_ml_classic::linear::{LinearRegressor, ElasticNetRegressor, RidgeRegressor, LassoRegressor};
use rtx_ml_classic::linear::{LinearConfig, ElasticNetConfig, RidgeConfig, LassoConfig};
use crate::error::{SklearnResult, check_is_fitted};
use crate::utils::{ArrayConverter, DeviceConfig, ParamValidator, AsyncHelper};
/// sklearn-compatible LinearRegression
#[pyclass(name = "LinearRegression")]
pub struct LinearRegression {
model: Option<LinearRegressor>,
fit_intercept: bool,
normalize: bool,
device: DeviceConfig,
// Fitted state
is_fitted: bool,
n_features_in: Option<usize>,
coef_: Option<Array1<f64>>,
intercept_: Option<f64>,
}
#[pymethods]
impl LinearRegression {
#[new]
#[pyo3(signature = (fit_intercept=true, normalize=false, device="cpu"))]
fn new(fit_intercept: bool, normalize: bool, device: &str) -> PyResult<Self> {
let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;
Ok(LinearRegression {
model: None,
fit_intercept,
normalize,
device: device_config,
is_fitted: false,
n_features_in: None,
coef_: None,
intercept_: None,
})
}
fn fit(&mut self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
ArrayConverter::validate_fit_input(&x, Some(&y)).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;
self.n_features_in = Some(x_view.shape()[1]);
let config = LinearConfig {
fit_intercept: self.fit_intercept,
normalize: self.normalize,
};
let mut model = LinearRegressor::new(config);
model.fit(x_view.to_owned(), y_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Training failed: {}", e)
))?;
// Extract coefficients
self.coef_ = Some(model.coefficients()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get coefficients: {}", e)
))?);
self.intercept_ = if self.fit_intercept {
Some(model.intercept()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get intercept: {}", e)
))?)
} else {
Some(0.0)
};
self.model = Some(model);
self.is_fitted = true;
Ok(())
}
fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<f64>>> {
check_is_fitted(self.is_fitted).map_err(|e| PyErr::from(e))?;
ArrayConverter::validate_predict_input(&x, self.n_features_in)
.map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Prediction failed: {}", e)
))?;
Python::with_gil(|py| {
Ok(ArrayConverter::from_array1(py, predictions)?.to_owned())
})
}
fn score(&self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<f64> {
let predictions = self.predict(x)?;
Python::with_gil(|py| {
let pred_array = predictions.as_ref(py).readonly();
let pred_view = ArrayConverter::to_array_view1(&pred_array)
.map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y)
.map_err(|e| PyErr::from(e))?;
// Calculate R² score
let y_mean = y_view.iter().sum::<f64>() / y_view.len() as f64;
let ss_res: f64 = pred_view.iter()
.zip(y_view.iter())
.map(|(pred, actual)| (actual - pred).powi(2))
.sum();
let ss_tot: f64 = y_view.iter()
.map(|actual| (actual - y_mean).powi(2))
.sum();
let r2 = 1.0 - ss_res / ss_tot;
Ok(r2)
})
}
fn get_params(&self) -> PyResult<HashMap<String, PyObject>> {
Python::with_gil(|py| {
let mut params = HashMap::new();
params.insert("fit_intercept".to_string(), self.fit_intercept.to_object(py));
params.insert("normalize".to_string(), self.normalize.to_object(py));
params.insert("device".to_string(), self.device.to_string().to_object(py));
Ok(params)
})
}
fn set_params(&mut self, params: &PyDict) -> PyResult<()> {
let valid_params = vec!["fit_intercept", "normalize", "device"];
let validated = ParamValidator::validate_params(params, &valid_params)
.map_err(|e| PyErr::from(e))?;
Python::with_gil(|py| {
if validated.contains_key("fit_intercept") {
self.fit_intercept = ParamValidator::get_param(&validated, "fit_intercept", self.fit_intercept, py)
.map_err(|e| PyErr::from(e))?;
}
if validated.contains_key("normalize") {
self.normalize = ParamValidator::get_param(&validated, "normalize", self.normalize, py)
.map_err(|e| PyErr::from(e))?;
}
if validated.contains_key("device") {
let device_str: String = ParamValidator::get_param(&validated, "device", self.device.to_string(), py)
.map_err(|e| PyErr::from(e))?;
self.device = DeviceConfig::new(&device_str)
.map_err(|e| PyErr::from(e))?;
}
// Reset fitted state if parameters changed
if !validated.is_empty() {
self.is_fitted = false;
self.model = None;
self.coef_ = None;
self.intercept_ = None;
}
Ok(())
})
}
#[getter]
fn coef_(&self) -> PyResult<Option<Py<PyArray1<f64>>>> {
match &self.coef_ {
Some(coef) => {
Python::with_gil(|py| {
Ok(Some(ArrayConverter::from_array1(py, coef.clone())?.to_owned()))
})
}
None => Ok(None),
}
}
#[getter]
fn intercept_(&self) -> Option<f64> {
self.intercept_
}
#[getter]
fn n_features_in_(&self) -> Option<usize> {
self.n_features_in
}
}
/// sklearn-compatible Ridge regression
#[pyclass(name = "Ridge")]
pub struct Ridge {
model: Option<RidgeRegressor>,
alpha: f64,
fit_intercept: bool,
normalize: bool,
max_iter: Option<usize>,
device: DeviceConfig,
// Fitted state
is_fitted: bool,
n_features_in: Option<usize>,
coef_: Option<Array1<f64>>,
intercept_: Option<f64>,
}
#[pymethods]
impl Ridge {
#[new]
#[pyo3(signature = (
alpha=1.0,
fit_intercept=true,
normalize=false,
max_iter=None,
device="cpu"
))]
fn new(
alpha: f64,
fit_intercept: bool,
normalize: bool,
max_iter: Option<usize>,
device: &str,
) -> PyResult<Self> {
if alpha < 0.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
"alpha must be >= 0"
));
}
let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;
Ok(Ridge {
model: None,
alpha,
fit_intercept,
normalize,
max_iter,
device: device_config,
is_fitted: false,
n_features_in: None,
coef_: None,
intercept_: None,
})
}
fn fit(&mut self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
ArrayConverter::validate_fit_input(&x, Some(&y)).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;
self.n_features_in = Some(x_view.shape()[1]);
let config = RidgeConfig {
alpha: self.alpha,
fit_intercept: self.fit_intercept,
normalize: self.normalize,
max_iter: self.max_iter,
};
let mut model = RidgeRegressor::new(config);
model.fit(x_view.to_owned(), y_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Training failed: {}", e)
))?;
self.coef_ = Some(model.coefficients()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get coefficients: {}", e)
))?);
self.intercept_ = if self.fit_intercept {
Some(model.intercept()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get intercept: {}", e)
))?)
} else {
Some(0.0)
};
self.model = Some(model);
self.is_fitted = true;
Ok(())
}
fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<f64>>> {
check_is_fitted(self.is_fitted).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Prediction failed: {}", e)
))?;
Python::with_gil(|py| {
Ok(ArrayConverter::from_array1(py, predictions)?.to_owned())
})
}
fn score(&self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<f64> {
let predictions = self.predict(x)?;
Python::with_gil(|py| {
let pred_array = predictions.as_ref(py).readonly();
let pred_view = ArrayConverter::to_array_view1(&pred_array)
.map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y)
.map_err(|e| PyErr::from(e))?;
let y_mean = y_view.iter().sum::<f64>() / y_view.len() as f64;
let ss_res: f64 = pred_view.iter()
.zip(y_view.iter())
.map(|(pred, actual)| (actual - pred).powi(2))
.sum();
let ss_tot: f64 = y_view.iter()
.map(|actual| (actual - y_mean).powi(2))
.sum();
Ok(1.0 - ss_res / ss_tot)
})
}
#[getter]
fn coef_(&self) -> PyResult<Option<Py<PyArray1<f64>>>> {
match &self.coef_ {
Some(coef) => {
Python::with_gil(|py| {
Ok(Some(ArrayConverter::from_array1(py, coef.clone())?.to_owned()))
})
}
None => Ok(None),
}
}
#[getter]
fn intercept_(&self) -> Option<f64> {
self.intercept_
}
}
/// sklearn-compatible Lasso regression
#[pyclass(name = "Lasso")]
pub struct Lasso {
model: Option<LassoRegressor>,
alpha: f64,
fit_intercept: bool,
normalize: bool,
max_iter: usize,
tol: f64,
device: DeviceConfig,
// Fitted state
is_fitted: bool,
n_features_in: Option<usize>,
coef_: Option<Array1<f64>>,
intercept_: Option<f64>,
n_iter_: Option<usize>,
}
#[pymethods]
impl Lasso {
#[new]
#[pyo3(signature = (
alpha=1.0,
fit_intercept=true,
normalize=false,
max_iter=1000,
tol=1e-4,
device="cpu"
))]
fn new(
alpha: f64,
fit_intercept: bool,
normalize: bool,
max_iter: usize,
tol: f64,
device: &str,
) -> PyResult<Self> {
if alpha < 0.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("alpha must be >= 0"));
}
if max_iter < 1 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("max_iter must be >= 1"));
}
if tol <= 0.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("tol must be > 0"));
}
let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;
Ok(Lasso {
model: None,
alpha,
fit_intercept,
normalize,
max_iter,
tol,
device: device_config,
is_fitted: false,
n_features_in: None,
coef_: None,
intercept_: None,
n_iter_: None,
})
}
fn fit(&mut self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
ArrayConverter::validate_fit_input(&x, Some(&y)).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;
self.n_features_in = Some(x_view.shape()[1]);
let config = LassoConfig {
alpha: self.alpha,
fit_intercept: self.fit_intercept,
normalize: self.normalize,
max_iter: self.max_iter,
tol: self.tol,
};
let mut model = LassoRegressor::new(config);
model.fit(x_view.to_owned(), y_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Training failed: {}", e)
))?;
self.coef_ = Some(model.coefficients()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get coefficients: {}", e)
))?);
self.intercept_ = if self.fit_intercept {
Some(model.intercept()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get intercept: {}", e)
))?)
} else {
Some(0.0)
};
self.n_iter_ = Some(model.n_iterations()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get iteration count: {}", e)
))?);
self.model = Some(model);
self.is_fitted = true;
Ok(())
}
fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<f64>>> {
check_is_fitted(self.is_fitted).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Prediction failed: {}", e)
))?;
Python::with_gil(|py| {
Ok(ArrayConverter::from_array1(py, predictions)?.to_owned())
})
}
fn score(&self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<f64> {
let predictions = self.predict(x)?;
Python::with_gil(|py| {
let pred_array = predictions.as_ref(py).readonly();
let pred_view = ArrayConverter::to_array_view1(&pred_array)
.map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y)
.map_err(|e| PyErr::from(e))?;
let y_mean = y_view.iter().sum::<f64>() / y_view.len() as f64;
let ss_res: f64 = pred_view.iter()
.zip(y_view.iter())
.map(|(pred, actual)| (actual - pred).powi(2))
.sum();
let ss_tot: f64 = y_view.iter()
.map(|actual| (actual - y_mean).powi(2))
.sum();
Ok(1.0 - ss_res / ss_tot)
})
}
#[getter]
fn coef_(&self) -> PyResult<Option<Py<PyArray1<f64>>>> {
match &self.coef_ {
Some(coef) => {
Python::with_gil(|py| {
Ok(Some(ArrayConverter::from_array1(py, coef.clone())?.to_owned()))
})
}
None => Ok(None),
}
}
#[getter]
fn intercept_(&self) -> Option<f64> {
self.intercept_
}
#[getter]
fn n_iter_(&self) -> Option<usize> {
self.n_iter_
}
}
/// sklearn-compatible ElasticNet regression
#[pyclass(name = "ElasticNet")]
pub struct ElasticNet {
model: Option<ElasticNetRegressor>,
alpha: f64,
l1_ratio: f64,
fit_intercept: bool,
normalize: bool,
max_iter: usize,
tol: f64,
device: DeviceConfig,
// Fitted state
is_fitted: bool,
n_features_in: Option<usize>,
coef_: Option<Array1<f64>>,
intercept_: Option<f64>,
n_iter_: Option<usize>,
}
#[pymethods]
impl ElasticNet {
#[new]
#[pyo3(signature = (
alpha=1.0,
l1_ratio=0.5,
fit_intercept=true,
normalize=false,
max_iter=1000,
tol=1e-4,
device="cpu"
))]
fn new(
alpha: f64,
l1_ratio: f64,
fit_intercept: bool,
normalize: bool,
max_iter: usize,
tol: f64,
device: &str,
) -> PyResult<Self> {
if alpha < 0.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("alpha must be >= 0"));
}
if l1_ratio < 0.0 || l1_ratio > 1.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("l1_ratio must be in [0, 1]"));
}
if max_iter < 1 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("max_iter must be >= 1"));
}
if tol <= 0.0 {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>("tol must be > 0"));
}
let device_config = DeviceConfig::new(device).map_err(|e| PyErr::from(e))?;
Ok(ElasticNet {
model: None,
alpha,
l1_ratio,
fit_intercept,
normalize,
max_iter,
tol,
device: device_config,
is_fitted: false,
n_features_in: None,
coef_: None,
intercept_: None,
n_iter_: None,
})
}
fn fit(&mut self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<()> {
ArrayConverter::validate_fit_input(&x, Some(&y)).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y).map_err(|e| PyErr::from(e))?;
self.n_features_in = Some(x_view.shape()[1]);
let config = ElasticNetConfig {
alpha: self.alpha,
l1_ratio: self.l1_ratio,
fit_intercept: self.fit_intercept,
normalize: self.normalize,
max_iter: self.max_iter,
tol: self.tol,
};
let mut model = ElasticNetRegressor::new(config);
model.fit(x_view.to_owned(), y_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Training failed: {}", e)
))?;
self.coef_ = Some(model.coefficients()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get coefficients: {}", e)
))?);
self.intercept_ = if self.fit_intercept {
Some(model.intercept()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get intercept: {}", e)
))?)
} else {
Some(0.0)
};
self.n_iter_ = Some(model.n_iterations()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Failed to get iteration count: {}", e)
))?);
self.model = Some(model);
self.is_fitted = true;
Ok(())
}
fn predict(&self, x: PyReadonlyArrayDyn<f64>) -> PyResult<Py<PyArray1<f64>>> {
check_is_fitted(self.is_fitted).map_err(|e| PyErr::from(e))?;
let x_view = ArrayConverter::to_array_view2(&x).map_err(|e| PyErr::from(e))?;
let predictions = self.model.as_ref().unwrap().predict(x_view.to_owned())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
format!("Prediction failed: {}", e)
))?;
Python::with_gil(|py| {
Ok(ArrayConverter::from_array1(py, predictions)?.to_owned())
})
}
fn score(&self, x: PyReadonlyArrayDyn<f64>, y: PyReadonlyArrayDyn<f64>) -> PyResult<f64> {
let predictions = self.predict(x)?;
Python::with_gil(|py| {
let pred_array = predictions.as_ref(py).readonly();
let pred_view = ArrayConverter::to_array_view1(&pred_array)
.map_err(|e| PyErr::from(e))?;
let y_view = ArrayConverter::to_array_view1(&y)
.map_err(|e| PyErr::from(e))?;
let y_mean = y_view.iter().sum::<f64>() / y_view.len() as f64;
let ss_res: f64 = pred_view.iter()
.zip(y_view.iter())
.map(|(pred, actual)| (actual - pred).powi(2))
.sum();
let ss_tot: f64 = y_view.iter()
.map(|actual| (actual - y_mean).powi(2))
.sum();
Ok(1.0 - ss_res / ss_tot)
})
}
#[getter]
fn coef_(&self) -> PyResult<Option<Py<PyArray1<f64>>>> {
match &self.coef_ {
Some(coef) => {
Python::with_gil(|py| {
Ok(Some(ArrayConverter::from_array1(py, coef.clone())?.to_owned()))
})
}
None => Ok(None),
}
}
#[getter]
fn intercept_(&self) -> Option<f64> {
self.intercept_
}
#[getter]
fn n_iter_(&self) -> Option<usize> {
self.n_iter_
}
}