use approx::assert_relative_eq;
use fleximl_models::ml::mlp::{Activation, MLP};
use fleximl_models::utils::tasks::Task;
use ndarray::{arr1, arr2, Array1};
#[test]
fn test_mlp_creation() {
let mlp = MLP::new(
&[2, 3, 1],
&[Activation::ReLU, Activation::Sigmoid],
0.01,
Task::BinaryClassification,
42,
);
assert_eq!(mlp.layers.len(), 2);
assert_eq!(mlp.layers[0].weights.shape(), &[3, 2]);
assert_eq!(mlp.layers[1].weights.shape(), &[1, 3]);
assert_eq!(mlp.learning_rate, 0.01);
assert_eq!(mlp.task, Task::BinaryClassification);
}
#[test]
fn test_mlp_prediction_binary_classification() {
let mut mlp = MLP::new(
&[2, 2, 1],
&[Activation::ReLU, Activation::Sigmoid],
0.01,
Task::BinaryClassification,
42,
);
mlp.layers[0].weights = arr2(&[[1.0, -1.0], [-1.0, 1.0]]);
mlp.layers[0].biases = arr1(&[0.0, 0.0]);
mlp.layers[1].weights = arr2(&[[1.0, 1.0]]);
mlp.layers[1].biases = arr1(&[0.0]);
let input = arr1(&[1.0, 0.0]);
let prediction = mlp.predict(&input);
assert_relative_eq!(prediction[0], 0.7310585786300049, epsilon = 1e-8);
}
#[test]
fn test_mlp_prediction_regression() {
let mut mlp = MLP::new(
&[2, 2, 1],
&[Activation::ReLU, Activation::Linear],
0.01,
Task::Regression,
42,
);
mlp.layers[0].weights = arr2(&[[1.0, -1.0], [-1.0, 1.0]]);
mlp.layers[0].biases = arr1(&[0.0, 0.0]);
mlp.layers[1].weights = arr2(&[[1.0, 1.0]]);
mlp.layers[1].biases = arr1(&[0.0]);
let input = arr1(&[1.0, 0.0]);
let prediction = mlp.predict(&input);
assert_relative_eq!(prediction[0], 1.0, epsilon = 1e-8);
}
#[test]
fn test_mlp_fit_binary_classification() {
let mut mlp = MLP::new(
&[2, 8, 4, 1],
&[Activation::ReLU, Activation::ReLU, Activation::Sigmoid],
0.005,
Task::BinaryClassification,
42,
);
let x = arr2(&[[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]);
let y = arr2(&[[0.0], [1.0], [1.0], [0.0]]);
mlp.fit(&x, &y, 20000);
let predictions: Vec<f64> = x
.outer_iter()
.map(|input| mlp.predict(&Array1::from(input.to_vec()))[0])
.collect();
assert!(predictions[0] < 0.1);
assert!(predictions[1] > 0.9);
assert!(predictions[2] > 0.9);
assert!(predictions[3] < 0.1);
}
#[test]
fn test_mlp_fit_regression() {
let mut mlp = MLP::new(
&[1, 16, 16, 1],
&[Activation::ReLU, Activation::ReLU, Activation::Linear],
0.001,
Task::Regression,
42,
);
let x = arr2(&[[0.0], [0.25], [0.5], [0.75], [1.0]]);
let y = arr2(&[[0.0], [0.5], [1.0], [1.5], [2.0]]);
mlp.fit(&x, &y, 10000);
let predictions: Vec<f64> = x
.outer_iter()
.map(|input| mlp.predict(&Array1::from(input.to_vec()))[0])
.collect();
for (pred, target) in predictions.iter().zip(y.column(0).iter()) {
println!("Prediction: {:.4}, Target: {:.4}", pred, target);
assert_relative_eq!(pred, target, epsilon = 0.2, max_relative = 0.1);
}
}
#[test]
fn test_mlp_multi_class_classification() {
let mut mlp = MLP::new(
&[2, 4, 3],
&[Activation::ReLU, Activation::Linear],
0.01,
Task::MultiClassClassification,
42,
);
let x = arr2(&[[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]);
let y = arr2(&[
[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[1.0, 0.0, 0.0],
]);
mlp.fit(&x, &y, 5000);
for (input, target) in x.outer_iter().zip(y.outer_iter()) {
let prediction = mlp.predict(&Array1::from(input.to_vec()));
let predicted_class = prediction
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
.unwrap()
.0;
let target_class = target
.iter()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
.unwrap()
.0;
assert_eq!(predicted_class, target_class);
}
}