Skip to main content

sklears_python/
neural_network.rs

1//! Python bindings for neural network models
2//!
3//! This module provides PyO3-based Python bindings for sklears neural network algorithms,
4//! including Multi-Layer Perceptron (MLP) classifiers and regressors.
5
6use crate::utils::{numpy_to_ndarray1, numpy_to_ndarray2};
7use numpy::{IntoPyArray, PyArray1, PyArray2};
8use pyo3::exceptions::{PyRuntimeError, PyValueError};
9use pyo3::prelude::*;
10use sklears_core::traits::{Fit, Predict};
11use sklears_neural::solvers::LearningRateSchedule;
12use sklears_neural::{Activation, MLPClassifier, MLPRegressor, Solver};
13
14/// Python wrapper for MLP Classifier
15#[pyclass(name = "MLPClassifier")]
16pub struct PyMLPClassifier {
17    inner: Option<MLPClassifier<sklears_core::traits::Untrained>>,
18    trained: Option<MLPClassifier<sklears_neural::TrainedMLPClassifier>>,
19}
20
21#[pymethods]
22impl PyMLPClassifier {
23    #[new]
24    #[allow(clippy::too_many_arguments)]
25    #[pyo3(signature = (
26        hidden_layer_sizes=None,
27        activation="relu",
28        solver="adam",
29        alpha=0.0001,
30        batch_size=None,
31        learning_rate="constant",
32        learning_rate_init=0.001,
33        power_t=0.5,
34        max_iter=200,
35        shuffle=true,
36        random_state=None,
37        tol=1e-4,
38        verbose=false,
39        warm_start=false,
40        momentum=0.9,
41        nesterovs_momentum=true,
42        early_stopping=false,
43        validation_fraction=0.1,
44        beta_1=0.9,
45        beta_2=0.999,
46        epsilon=1e-8,
47        n_iter_no_change=10,
48        max_fun=15000
49    ))]
50    fn new(
51        hidden_layer_sizes: Option<Vec<usize>>,
52        activation: &str,
53        solver: &str,
54        alpha: f64,
55        batch_size: Option<usize>,
56        learning_rate: &str,
57        learning_rate_init: f64,
58        power_t: f64,
59        max_iter: usize,
60        shuffle: bool,
61        random_state: Option<u64>,
62        tol: f64,
63        verbose: bool,
64        warm_start: bool,
65        momentum: f64,
66        nesterovs_momentum: bool,
67        early_stopping: bool,
68        validation_fraction: f64,
69        beta_1: f64,
70        beta_2: f64,
71        epsilon: f64,
72        n_iter_no_change: usize,
73        max_fun: usize,
74    ) -> PyResult<Self> {
75        let activation = match activation {
76            "identity" => Activation::Identity,
77            "logistic" => Activation::Logistic,
78            "tanh" => Activation::Tanh,
79            "relu" => Activation::Relu,
80            "elu" => Activation::Elu,
81            "swish" => Activation::Swish,
82            "gelu" => Activation::Gelu,
83            "mish" => Activation::Mish,
84            "leaky_relu" => Activation::LeakyRelu,
85            "prelu" => Activation::PRelu,
86            _ => {
87                return Err(PyValueError::new_err(format!(
88                    "Unknown activation: {}",
89                    activation
90                )))
91            }
92        };
93
94        let solver = match solver {
95            "lbfgs" => Solver::Lbfgs,
96            "sgd" => Solver::Sgd,
97            "adam" => Solver::Adam,
98            _ => return Err(PyValueError::new_err(format!("Unknown solver: {}", solver))),
99        };
100
101        let learning_rate_schedule = match learning_rate {
102            "constant" => LearningRateSchedule::Constant,
103            "invscaling" => LearningRateSchedule::InvScaling,
104            "adaptive" => LearningRateSchedule::Adaptive,
105            _ => {
106                return Err(PyValueError::new_err(format!(
107                    "Unknown learning rate schedule: {}",
108                    learning_rate
109                )))
110            }
111        };
112
113        let hidden_sizes = hidden_layer_sizes.unwrap_or_else(|| vec![100]);
114
115        let mut mlp = MLPClassifier::new();
116        mlp.hidden_layer_sizes = hidden_sizes;
117        mlp.activation = activation;
118        mlp.solver = solver;
119        mlp.alpha = alpha;
120        mlp.batch_size = batch_size;
121        mlp.learning_rate = learning_rate_schedule;
122        mlp.learning_rate_init = learning_rate_init;
123        mlp.power_t = power_t;
124        mlp.max_iter = max_iter;
125        mlp.shuffle = shuffle;
126        mlp.random_state = random_state;
127        mlp.tol = tol;
128        mlp.verbose = verbose;
129        mlp.warm_start = warm_start;
130        mlp.momentum = momentum;
131        mlp.nesterovs_momentum = nesterovs_momentum;
132        mlp.early_stopping = early_stopping;
133        mlp.validation_fraction = validation_fraction;
134        mlp.beta_1 = beta_1;
135        mlp.beta_2 = beta_2;
136        mlp.epsilon = epsilon;
137        mlp.n_iter_no_change = n_iter_no_change;
138        mlp.max_fun = max_fun;
139
140        Ok(Self {
141            inner: Some(mlp),
142            trained: None,
143        })
144    }
145
146    /// Fit the MLP classifier
147    fn fit(&mut self, x: &Bound<'_, PyArray2<f64>>, y: &Bound<'_, PyArray1<f64>>) -> PyResult<()> {
148        let x_array = numpy_to_ndarray2(x)?;
149        let y_array = numpy_to_ndarray1(y)?;
150
151        // Convert y to integer vector for classification
152        let y_int: Vec<usize> = y_array.iter().map(|&val| val as usize).collect();
153
154        let model = self.inner.take().ok_or_else(|| {
155            PyRuntimeError::new_err("Model has already been fitted or was not initialized")
156        })?;
157
158        match model.fit(&x_array, &y_int) {
159            Ok(trained_model) => {
160                self.trained = Some(trained_model);
161                Ok(())
162            }
163            Err(e) => Err(PyRuntimeError::new_err(format!(
164                "Failed to fit model: {}",
165                e
166            ))),
167        }
168    }
169
170    /// Make predictions using the fitted model
171    fn predict<'py>(
172        &self,
173        py: Python<'py>,
174        x: &Bound<'py, PyArray2<f64>>,
175    ) -> PyResult<Py<PyArray1<f64>>> {
176        let trained_model = self.trained.as_ref().ok_or_else(|| {
177            PyRuntimeError::new_err("Model must be fitted before making predictions")
178        })?;
179
180        let x_array = numpy_to_ndarray2(x)?;
181
182        match trained_model.predict(&x_array) {
183            Ok(predictions) => {
184                let predictions_f64: Vec<f64> = predictions.iter().map(|&x| x as f64).collect();
185                Ok(PyArray1::from_vec(py, predictions_f64).unbind())
186            }
187            Err(e) => Err(PyRuntimeError::new_err(format!("Prediction failed: {}", e))),
188        }
189    }
190
191    /// Predict class probabilities
192    fn predict_proba<'py>(
193        &self,
194        py: Python<'py>,
195        x: &Bound<'py, PyArray2<f64>>,
196    ) -> PyResult<Py<PyArray2<f64>>> {
197        let trained_model = self.trained.as_ref().ok_or_else(|| {
198            PyRuntimeError::new_err("Model must be fitted before making predictions")
199        })?;
200
201        let x_array = numpy_to_ndarray2(x)?;
202
203        match trained_model.predict_proba(&x_array) {
204            Ok(probabilities) => Ok(probabilities.into_pyarray(py).unbind()),
205            Err(e) => Err(PyRuntimeError::new_err(format!(
206                "Probability prediction failed: {}",
207                e
208            ))),
209        }
210    }
211
212    /// Get the loss after training
213    fn loss_(&self) -> PyResult<f64> {
214        let trained_model = self
215            .trained
216            .as_ref()
217            .ok_or_else(|| PyRuntimeError::new_err("Model must be fitted before accessing loss"))?;
218
219        Ok(trained_model.loss())
220    }
221
222    /// Get number of iterations
223    fn n_iter_(&self) -> PyResult<usize> {
224        let trained_model = self.trained.as_ref().ok_or_else(|| {
225            PyRuntimeError::new_err("Model must be fitted before accessing n_iter")
226        })?;
227
228        Ok(trained_model.n_iter())
229    }
230
231    fn __repr__(&self) -> String {
232        if self.trained.is_some() {
233            "MLPClassifier(fitted=True)".to_string()
234        } else {
235            "MLPClassifier(fitted=False)".to_string()
236        }
237    }
238}
239
240/// Python wrapper for MLP Regressor
241#[pyclass(name = "MLPRegressor")]
242pub struct PyMLPRegressor {
243    inner: Option<MLPRegressor<sklears_core::traits::Untrained>>,
244    trained: Option<MLPRegressor<sklears_neural::TrainedMLPRegressor>>,
245}
246
247#[pymethods]
248impl PyMLPRegressor {
249    #[new]
250    #[allow(clippy::too_many_arguments)]
251    #[pyo3(signature = (
252        hidden_layer_sizes=None,
253        activation="relu",
254        solver="adam",
255        alpha=0.0001,
256        batch_size=None,
257        learning_rate="constant",
258        learning_rate_init=0.001,
259        power_t=0.5,
260        max_iter=200,
261        shuffle=true,
262        random_state=None,
263        tol=1e-4,
264        verbose=false,
265        warm_start=false,
266        momentum=0.9,
267        nesterovs_momentum=true,
268        early_stopping=false,
269        validation_fraction=0.1,
270        beta_1=0.9,
271        beta_2=0.999,
272        epsilon=1e-8,
273        n_iter_no_change=10,
274        max_fun=15000
275    ))]
276    fn new(
277        hidden_layer_sizes: Option<Vec<usize>>,
278        activation: &str,
279        solver: &str,
280        alpha: f64,
281        batch_size: Option<usize>,
282        learning_rate: &str,
283        learning_rate_init: f64,
284        power_t: f64,
285        max_iter: usize,
286        shuffle: bool,
287        random_state: Option<u64>,
288        tol: f64,
289        verbose: bool,
290        warm_start: bool,
291        momentum: f64,
292        nesterovs_momentum: bool,
293        early_stopping: bool,
294        validation_fraction: f64,
295        beta_1: f64,
296        beta_2: f64,
297        epsilon: f64,
298        n_iter_no_change: usize,
299        max_fun: usize,
300    ) -> PyResult<Self> {
301        let activation = match activation {
302            "identity" => Activation::Identity,
303            "logistic" => Activation::Logistic,
304            "tanh" => Activation::Tanh,
305            "relu" => Activation::Relu,
306            "elu" => Activation::Elu,
307            "swish" => Activation::Swish,
308            "gelu" => Activation::Gelu,
309            "mish" => Activation::Mish,
310            "leaky_relu" => Activation::LeakyRelu,
311            "prelu" => Activation::PRelu,
312            _ => {
313                return Err(PyValueError::new_err(format!(
314                    "Unknown activation: {}",
315                    activation
316                )))
317            }
318        };
319
320        let solver = match solver {
321            "lbfgs" => Solver::Lbfgs,
322            "sgd" => Solver::Sgd,
323            "adam" => Solver::Adam,
324            _ => return Err(PyValueError::new_err(format!("Unknown solver: {}", solver))),
325        };
326
327        let learning_rate_schedule = match learning_rate {
328            "constant" => LearningRateSchedule::Constant,
329            "invscaling" => LearningRateSchedule::InvScaling,
330            "adaptive" => LearningRateSchedule::Adaptive,
331            _ => {
332                return Err(PyValueError::new_err(format!(
333                    "Unknown learning rate schedule: {}",
334                    learning_rate
335                )))
336            }
337        };
338
339        let hidden_sizes = hidden_layer_sizes.unwrap_or_else(|| vec![100]);
340
341        let mut mlp = MLPRegressor::new();
342        mlp.hidden_layer_sizes = hidden_sizes;
343        mlp.activation = activation;
344        mlp.solver = solver;
345        mlp.alpha = alpha;
346        mlp.batch_size = batch_size;
347        mlp.learning_rate = learning_rate_schedule;
348        mlp.learning_rate_init = learning_rate_init;
349        mlp.power_t = power_t;
350        mlp.max_iter = max_iter;
351        mlp.shuffle = shuffle;
352        mlp.random_state = random_state;
353        mlp.tol = tol;
354        mlp.verbose = verbose;
355        mlp.warm_start = warm_start;
356        mlp.momentum = momentum;
357        mlp.nesterovs_momentum = nesterovs_momentum;
358        mlp.early_stopping = early_stopping;
359        mlp.validation_fraction = validation_fraction;
360        mlp.beta_1 = beta_1;
361        mlp.beta_2 = beta_2;
362        mlp.epsilon = epsilon;
363        mlp.n_iter_no_change = n_iter_no_change;
364        mlp.max_fun = max_fun;
365
366        Ok(Self {
367            inner: Some(mlp),
368            trained: None,
369        })
370    }
371
372    /// Fit the MLP regressor
373    fn fit(&mut self, x: &Bound<'_, PyArray2<f64>>, y: &Bound<'_, PyArray1<f64>>) -> PyResult<()> {
374        let x_array = numpy_to_ndarray2(x)?;
375        let y_array_1d = numpy_to_ndarray1(y)?;
376
377        // Convert y from 1D to 2D array (n_samples, 1)
378        let y_array = y_array_1d.insert_axis(scirs2_core::ndarray::Axis(1));
379
380        let model = self.inner.take().ok_or_else(|| {
381            PyRuntimeError::new_err("Model has already been fitted or was not initialized")
382        })?;
383
384        match model.fit(&x_array, &y_array) {
385            Ok(trained_model) => {
386                self.trained = Some(trained_model);
387                Ok(())
388            }
389            Err(e) => Err(PyRuntimeError::new_err(format!(
390                "Failed to fit model: {}",
391                e
392            ))),
393        }
394    }
395
396    /// Make predictions using the fitted model
397    fn predict<'py>(
398        &self,
399        py: Python<'py>,
400        x: &Bound<'py, PyArray2<f64>>,
401    ) -> PyResult<Py<PyArray1<f64>>> {
402        let trained_model = self.trained.as_ref().ok_or_else(|| {
403            PyRuntimeError::new_err("Model must be fitted before making predictions")
404        })?;
405
406        let x_array = numpy_to_ndarray2(x)?;
407
408        match trained_model.predict(&x_array) {
409            Ok(predictions_2d) => {
410                // Convert from Array2 (n_samples, 1) to Array1 (n_samples,)
411                let predictions_1d = predictions_2d
412                    .index_axis(scirs2_core::ndarray::Axis(1), 0)
413                    .to_owned();
414                Ok(predictions_1d.into_pyarray(py).unbind())
415            }
416            Err(e) => Err(PyRuntimeError::new_err(format!("Prediction failed: {}", e))),
417        }
418    }
419
420    /// Get the loss after training
421    fn loss_(&self) -> PyResult<f64> {
422        let trained_model = self
423            .trained
424            .as_ref()
425            .ok_or_else(|| PyRuntimeError::new_err("Model must be fitted before accessing loss"))?;
426
427        Ok(trained_model.loss())
428    }
429
430    /// Get number of iterations
431    fn n_iter_(&self) -> PyResult<usize> {
432        let trained_model = self.trained.as_ref().ok_or_else(|| {
433            PyRuntimeError::new_err("Model must be fitted before accessing n_iter")
434        })?;
435
436        Ok(trained_model.n_iter())
437    }
438
439    fn __repr__(&self) -> String {
440        if self.trained.is_some() {
441            "MLPRegressor(fitted=True)".to_string()
442        } else {
443            "MLPRegressor(fitted=False)".to_string()
444        }
445    }
446}