scirs2-neural 0.6.4

//! Advanced Optimizers Example
//!
//! Compares SGD, Adam, AdamW, RAdam, and RMSprop on a synthetic binary
//! classification task, training identically-initialized models with each
//! optimizer and reporting loss/accuracy progression side by side.

use scirs2_core::ndarray::Array2;
use scirs2_core::random::prelude::SliceRandom;
use scirs2_core::random::rngs::SmallRng;
use scirs2_core::random::{Rng, RngExt, SeedableRng};
use scirs2_neural::layers::{Dense, Dropout};
use scirs2_neural::losses::CrossEntropyLoss;
use scirs2_neural::models::{Model, Sequential};
use scirs2_neural::optimizers::{Adam, AdamW, Optimizer, RAdam, RMSprop, SGD};
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Advanced Optimizers Example");
    println!("Comparing SGD, Adam, AdamW, RAdam, and RMSprop optimizers\n");

    // Initialize random number generator
    let mut rng = SmallRng::from_seed([42; 32]);
    // Create a synthetic binary classification dataset
    let num_samples = 1000;
    let num_features = 20;
    let num_classes = 2;
    println!(
        "Generating synthetic dataset with {} samples, {} features...",
        num_samples, num_features
    );
    // Generate random input features
    let mut x_data = Array2::<f32>::zeros((num_samples, num_features));
    for i in 0..num_samples {
        for j in 0..num_features {
            x_data[[i, j]] = rng.random_range(-1.0..1.0);
        }
    }
    // Create true weights and bias for data generation
    let mut true_weights = Array2::<f32>::zeros((num_features, 1));
    for i in 0..num_features {
        true_weights[[i, 0]] = rng.random_range(-1.0..1.0);
    }
    let true_bias = rng.random_range(-1.0..1.0);
    // Generate binary labels (0 or 1) based on linear model with logistic function
    let mut y_data = Array2::<f32>::zeros((num_samples, num_classes));
    for i in 0..num_samples {
        let mut logit = true_bias;
        for j in 0..num_features {
            logit += x_data[[i, j]] * true_weights[[j, 0]];
        }
        // Apply sigmoid to get probability
        let prob = 1.0 / (1.0 + (-logit).exp());
        // Convert to one-hot encoding
        if prob > 0.5 {
            y_data[[i, 1]] = 1.0; // Class 1
        } else {
            y_data[[i, 0]] = 1.0; // Class 0
        }
    }
    // Split into train and test sets (80% train, 20% test)
    let train_size = (num_samples as f32 * 0.8) as usize;
    let test_size = num_samples - train_size;
    let x_train = x_data
        .slice(scirs2_core::ndarray::s![0..train_size, ..])
        .to_owned();
    let y_train = y_data
        .slice(scirs2_core::ndarray::s![0..train_size, ..])
        .to_owned();
    let x_test = x_data
        .slice(scirs2_core::ndarray::s![train_size.., ..])
        .to_owned();
    let y_test = y_data
        .slice(scirs2_core::ndarray::s![train_size.., ..])
        .to_owned();
    println!("Training set: {} samples", train_size);
    println!("Test set: {} samples", test_size);
    // Create a simple neural network model
    let hidden_size = 64;
    let dropout_rate = 0.2;
    let seed_rng = SmallRng::from_seed([42; 32]);
    // Shared function to create identical model architectures for fair comparison
    let create_model = || -> Result<Sequential<f32>, Box<dyn std::error::Error>> {
        let mut model = Sequential::new();
        // Input to hidden layer
        let dense1 = Dense::new(
            num_features,
            hidden_size,
            Some("relu"),
            &mut seed_rng.clone(),
        )?;
        model.add_layer(dense1);
        // Dropout for regularization
        let dropout = Dropout::new(dropout_rate, &mut seed_rng.clone())?;
        model.add_layer(dropout);
        // Hidden to output layer
        let dense2 = Dense::new(
            hidden_size,
            num_classes,
            Some("softmax"),
            &mut seed_rng.clone(),
        )?;
        model.add_layer(dense2);
        Ok(model)
    };
    // Create models for each optimizer
    let mut sgd_model = create_model()?;
    let mut adam_model = create_model()?;
    let mut adamw_model = create_model()?;
    let mut radam_model = create_model()?;
    let mut rmsprop_model = create_model()?;
    // Create the loss function
    let loss_fn = CrossEntropyLoss::new(1e-10);
    // Create optimizers
    let learning_rate = 0.001;
    let batch_size = 32;
    let epochs = 20;
    let mut sgd_optimizer = SGD::new_with_config(learning_rate, 0.9, 0.0);
    let mut adam_optimizer = Adam::new(learning_rate, 0.9, 0.999, 1e-8);
    let mut adamw_optimizer = AdamW::new(learning_rate, 0.9, 0.999, 1e-8, 0.01);
    let mut radam_optimizer = RAdam::new(learning_rate, 0.9, 0.999, 1e-8, 0.0)?;
    let mut rmsprop_optimizer = RMSprop::new_with_config(learning_rate, 0.9, 1e-8, 0.0);
    // Helper function to compute accuracy
    let compute_accuracy = |model: &Sequential<f32>, x: &Array2<f32>, y: &Array2<f32>| -> f32 {
        let predictions = model
            .forward(&x.clone().into_dyn())
            .expect("Operation failed");
        let mut correct = 0;
        for i in 0..x.shape()[0] {
            let mut max_idx = 0;
            let mut max_val = predictions[[i, 0]];
            for j in 1..num_classes {
                if predictions[[i, j]] > max_val {
                    max_val = predictions[[i, j]];
                    max_idx = j;
                }
            }
            let true_idx = if y[[i, 0]] < y[[i, 1]] { 1 } else { 0 };
            if max_idx as i32 == true_idx {
                correct += 1;
            }
        }
        correct as f32 / x.shape()[0] as f32
    };
    // Helper function to train model
    let mut train_model =
        |model: &mut Sequential<f32>, optimizer: &mut dyn Optimizer<f32>, name: &str| -> Vec<f32> {
            println!("\nTraining with {} optimizer...", name);
            let start_time = Instant::now();
            let mut train_losses = Vec::new();
            let num_batches = train_size.div_ceil(batch_size);
            for epoch in 0..epochs {
                let mut epoch_loss = 0.0;
                // Create a permutation for shuffling the data
                let mut indices: Vec<usize> = (0..train_size).collect();
                indices.shuffle(&mut rng);
                for batch_idx in 0..num_batches {
                    let start = batch_idx * batch_size;
                    let end = (start + batch_size).min(train_size);
                    let batch_indices = &indices[start..end];
                    // Create batch data
                    let mut x_batch = Array2::<f32>::zeros((batch_indices.len(), num_features));
                    let mut y_batch = Array2::<f32>::zeros((batch_indices.len(), num_classes));
                    for (i, &idx) in batch_indices.iter().enumerate() {
                        for j in 0..num_features {
                            x_batch[[i, j]] = x_train[[idx, j]];
                        }
                        for j in 0..num_classes {
                            y_batch[[i, j]] = y_train[[idx, j]];
                        }
                    }
                    // Convert to dynamic dimension arrays
                    let x_batch_dyn = x_batch.into_dyn();
                    let y_batch_dyn = y_batch.into_dyn();
                    // Perform a training step
                    let batch_loss = model
                        .train_batch(&x_batch_dyn, &y_batch_dyn, &loss_fn, optimizer)
                        .expect("Operation failed");
                    epoch_loss += batch_loss;
                }
                epoch_loss /= num_batches as f32;
                train_losses.push(epoch_loss);
                // Calculate and print metrics every few epochs
                if epoch % 5 == 0 || epoch == epochs - 1 {
                    let train_accuracy = compute_accuracy(model, &x_train, &y_train);
                    let test_accuracy = compute_accuracy(model, &x_test, &y_test);
                    println!(
                        "Epoch {}/{}: loss = {:.6}, train_acc = {:.2}%, test_acc = {:.2}%",
                        epoch + 1,
                        epochs,
                        epoch_loss,
                        train_accuracy * 100.0,
                        test_accuracy * 100.0
                    );
                }
            }
            let elapsed = start_time.elapsed();
            println!("{name} training completed in {elapsed:.2?}");
            // Final evaluation
            let train_accuracy = compute_accuracy(model, &x_train, &y_train);
            let test_accuracy = compute_accuracy(model, &x_test, &y_test);
            println!("Final metrics for {name}:");
            let train_pct = train_accuracy * 100.0;
            println!("  Train accuracy: {train_pct:.2}%");
            let test_pct = test_accuracy * 100.0;
            println!("  Test accuracy:  {test_pct:.2}%");
            train_losses
        };
    // Train models with different optimizers
    let sgd_losses = train_model(&mut sgd_model, &mut sgd_optimizer, "SGD");
    let adam_losses = train_model(&mut adam_model, &mut adam_optimizer, "Adam");
    let adamw_losses = train_model(&mut adamw_model, &mut adamw_optimizer, "AdamW");
    let radam_losses = train_model(&mut radam_model, &mut radam_optimizer, "RAdam");
    let rmsprop_losses = train_model(&mut rmsprop_model, &mut rmsprop_optimizer, "RMSprop");
    // Print comparison summary
    println!("\nOptimizer Comparison Summary:");
    println!("----------------------------");
    println!("Initial learning rate: {}", learning_rate);
    println!("Batch size: {}", batch_size);
    println!("Epochs: {}", epochs);
    println!();
    println!("Final Loss Values:");
    let sgd_final = sgd_losses.last().expect("Operation failed");
    println!("  SGD:     {sgd_final:.6}");
    let adam_final = adam_losses.last().expect("Operation failed");
    println!("  Adam:    {adam_final:.6}");
    let adamw_final = adamw_losses.last().expect("Operation failed");
    println!("  AdamW:   {adamw_final:.6}");
    let radam_final = radam_losses.last().expect("Operation failed");
    println!("  RAdam:   {radam_final:.6}");
    let rmsprop_final = rmsprop_losses.last().expect("Operation failed");
    println!("  RMSprop: {rmsprop_final:.6}");
    println!("\nLoss progression (first value, middle value, last value):");
    println!(
        "  SGD:     {:.6}, {:.6}, {:.6}",
        sgd_losses.first().expect("Operation failed"),
        sgd_losses[epochs / 2],
        sgd_losses.last().expect("Operation failed")
    );
    println!(
        "  Adam:    {:.6}, {:.6}, {:.6}",
        adam_losses.first().expect("Operation failed"),
        adam_losses[epochs / 2],
        adam_losses.last().expect("Operation failed")
    );
    println!(
        "  AdamW:   {:.6}, {:.6}, {:.6}",
        adamw_losses.first().expect("Operation failed"),
        adamw_losses[epochs / 2],
        adamw_losses.last().expect("Operation failed")
    );
    println!(
        "  RAdam:   {:.6}, {:.6}, {:.6}",
        radam_losses.first().expect("Operation failed"),
        radam_losses[epochs / 2],
        radam_losses.last().expect("Operation failed")
    );
    println!(
        "  RMSprop: {:.6}, {:.6}, {:.6}",
        rmsprop_losses.first().expect("Operation failed"),
        rmsprop_losses[epochs / 2],
        rmsprop_losses.last().expect("Operation failed")
    );
    println!("\nLoss improvement ratio (first loss / last loss):");
    println!(
        "  SGD:     {:.2}x",
        sgd_losses.first().expect("Operation failed")
            / sgd_losses.last().expect("Operation failed")
    );
    println!(
        "  Adam:    {:.2}x",
        adam_losses.first().expect("Operation failed")
            / adam_losses.last().expect("Operation failed")
    );
    println!(
        "  AdamW:   {:.2}x",
        adamw_losses.first().expect("Operation failed")
            / adamw_losses.last().expect("Operation failed")
    );
    println!(
        "  RAdam:   {:.2}x",
        radam_losses.first().expect("Operation failed")
            / radam_losses.last().expect("Operation failed")
    );
    println!(
        "  RMSprop: {:.2}x",
        rmsprop_losses.first().expect("Operation failed")
            / rmsprop_losses.last().expect("Operation failed")
    );
    println!("\nAdvanced optimizers demo completed successfully!");
    Ok(())
}