Skip to main content

sklears_python/metrics/
regression.rs

1//! Python bindings for regression metrics
2
3use super::common::*;
4use sklears_metrics::regression::{
5    mean_absolute_error as skl_mae, mean_squared_error as skl_mse, r2_score as skl_r2,
6};
7
8/// Calculate mean squared error
9#[pyfunction]
10#[pyo3(signature = (y_true, y_pred, sample_weight=None, multioutput="uniform_average", squared=true))]
11pub fn mean_squared_error(
12    y_true: PyReadonlyArray1<f64>,
13    y_pred: PyReadonlyArray1<f64>,
14    sample_weight: Option<PyReadonlyArray1<f64>>,
15    multioutput: &str,
16    squared: bool,
17) -> PyResult<f64> {
18    let _ = (sample_weight, multioutput);
19    let yt = y_true.as_array().to_owned();
20    let yp = y_pred.as_array().to_owned();
21
22    validate_arrays_same_length(&yt, &yp)?;
23
24    match skl_mse(&yt, &yp) {
25        Ok(mse) => Ok(if squared { mse } else { mse.sqrt() }),
26        Err(e) => Err(PyValueError::new_err(format!("mean_squared_error: {}", e))),
27    }
28}
29
30/// Calculate mean absolute error
31#[pyfunction]
32#[pyo3(signature = (y_true, y_pred, sample_weight=None, multioutput="uniform_average"))]
33pub fn mean_absolute_error(
34    y_true: PyReadonlyArray1<f64>,
35    y_pred: PyReadonlyArray1<f64>,
36    sample_weight: Option<PyReadonlyArray1<f64>>,
37    multioutput: &str,
38) -> PyResult<f64> {
39    let _ = (sample_weight, multioutput);
40    let yt = y_true.as_array().to_owned();
41    let yp = y_pred.as_array().to_owned();
42
43    validate_arrays_same_length(&yt, &yp)?;
44
45    match skl_mae(&yt, &yp) {
46        Ok(mae) => Ok(mae),
47        Err(e) => Err(PyValueError::new_err(format!("mean_absolute_error: {}", e))),
48    }
49}
50
51/// Calculate R² (coefficient of determination) score
52#[pyfunction]
53#[pyo3(signature = (y_true, y_pred, sample_weight=None, multioutput="uniform_average"))]
54pub fn r2_score(
55    y_true: PyReadonlyArray1<f64>,
56    y_pred: PyReadonlyArray1<f64>,
57    sample_weight: Option<PyReadonlyArray1<f64>>,
58    multioutput: &str,
59) -> PyResult<f64> {
60    let _ = (sample_weight, multioutput);
61    let yt = y_true.as_array().to_owned();
62    let yp = y_pred.as_array().to_owned();
63
64    validate_arrays_same_length(&yt, &yp)?;
65
66    match skl_r2(&yt, &yp) {
67        Ok(r2) => Ok(r2),
68        Err(e) => Err(PyValueError::new_err(format!("r2_score: {}", e))),
69    }
70}
71
72/// Calculate mean squared logarithmic error
73#[pyfunction]
74#[pyo3(signature = (y_true, y_pred, sample_weight=None, multioutput="uniform_average", squared=true))]
75pub fn mean_squared_log_error(
76    y_true: PyReadonlyArray1<f64>,
77    y_pred: PyReadonlyArray1<f64>,
78    sample_weight: Option<PyReadonlyArray1<f64>>,
79    multioutput: &str,
80    squared: bool,
81) -> PyResult<f64> {
82    let _ = (sample_weight, multioutput);
83    let yt = y_true.as_array().to_owned();
84    let yp = y_pred.as_array().to_owned();
85
86    validate_arrays_same_length(&yt, &yp)?;
87
88    if yt.iter().any(|&x| x < 0.0) || yp.iter().any(|&x| x < 0.0) {
89        return Err(PyValueError::new_err(
90            "Mean Squared Logarithmic Error cannot be used when targets contain negative values.",
91        ));
92    }
93
94    let log_true = yt.mapv(|x| (x + 1.0).ln());
95    let log_pred = yp.mapv(|x| (x + 1.0).ln());
96    let sq_errors = (&log_true - &log_pred).mapv(|x| x * x);
97    let msle = sq_errors.mean().unwrap_or(0.0);
98
99    Ok(if squared { msle } else { msle.sqrt() })
100}
101
102/// Calculate median absolute error
103#[pyfunction]
104#[pyo3(signature = (y_true, y_pred, multioutput="uniform_average", sample_weight=None))]
105pub fn median_absolute_error(
106    y_true: PyReadonlyArray1<f64>,
107    y_pred: PyReadonlyArray1<f64>,
108    multioutput: &str,
109    sample_weight: Option<PyReadonlyArray1<f64>>,
110) -> PyResult<f64> {
111    let _ = multioutput;
112    let yt = y_true.as_array().to_owned();
113    let yp = y_pred.as_array().to_owned();
114
115    validate_arrays_same_length(&yt, &yp)?;
116
117    if sample_weight.is_some() {
118        return Err(PyValueError::new_err(
119            "median_absolute_error does not support sample weights",
120        ));
121    }
122
123    let mut errors: Vec<f64> = (&yt - &yp).mapv(|x| x.abs()).to_vec();
124    errors.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
125
126    let n = errors.len();
127    let median = if n.is_multiple_of(2) {
128        (errors[n / 2 - 1] + errors[n / 2]) / 2.0
129    } else {
130        errors[n / 2]
131    };
132
133    Ok(median)
134}