1use 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#[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 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 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 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 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 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 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#[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 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 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 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 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 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 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}