torsh-sparse 0.1.2

//! Sparse optimizers
//!
//! This module provides optimization algorithms optimized for sparse tensors,
//! including SGD, Adam, AdamW, and RMSprop with efficient sparse parameter updates.

// Framework infrastructure - components designed for future use
#![allow(dead_code)]
use crate::{CooTensor, CsrTensor, SparseTensor, TorshResult};
use std::collections::HashMap;
use torsh_core::TorshError;

/// Trait for sparse optimizers
pub trait SparseOptimizer {
    /// Update sparse parameters with sparse gradients
    fn step(
        &mut self,
        parameters: &mut [&mut CsrTensor],
        gradients: &[&CsrTensor],
    ) -> TorshResult<()>;

    /// Zero gradients (if applicable)
    fn zero_grad(&mut self) {}

    /// Get current learning rate
    fn lr(&self) -> f32;

    /// Set learning rate
    fn set_lr(&mut self, lr: f32);
}

/// Sparse SGD optimizer
///
/// Implements Stochastic Gradient Descent optimized for sparse tensors.
/// Only updates non-zero parameters, significantly reducing computation.
pub struct SparseSGD {
    /// Learning rate
    lr: f32,
    /// Momentum coefficient (0.0 = no momentum)
    momentum: f32,
    /// Weight decay coefficient
    weight_decay: f32,
    /// Nesterov momentum flag
    nesterov: bool,
    /// Momentum buffers for each parameter (sparse format)
    momentum_buffers: Vec<Option<CsrTensor>>,
}

impl SparseSGD {
    /// Create a new sparse SGD optimizer
    pub fn new(lr: f32, momentum: f32, weight_decay: f32, nesterov: bool) -> Self {
        Self {
            lr,
            momentum,
            weight_decay,
            nesterov,
            momentum_buffers: Vec::new(),
        }
    }

    /// Create SGD optimizer with default settings
    pub fn default(lr: f32) -> Self {
        Self::new(lr, 0.0, 0.0, false)
    }
}

impl SparseOptimizer for SparseSGD {
    fn step(
        &mut self,
        parameters: &mut [&mut CsrTensor],
        gradients: &[&CsrTensor],
    ) -> TorshResult<()> {
        if parameters.len() != gradients.len() {
            return Err(TorshError::InvalidArgument(
                "Number of parameters and gradients must match".to_string(),
            ));
        }

        // Initialize momentum buffers if needed
        while self.momentum_buffers.len() < parameters.len() {
            self.momentum_buffers.push(None);
        }

        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
            // Apply weight decay and momentum (simplified for space)
            let update_coo = grad.to_coo()?;
            let update_triplets = update_coo.triplets();

            // Apply update: param = param - lr * grad
            let param_coo = param.to_coo()?;
            let param_triplets = param_coo.triplets();

            // Update parameters based on gradients
            let mut updated_positions = HashMap::new();

            // Apply gradient updates
            for (row, col, grad_val) in &update_triplets {
                updated_positions.insert((*row, *col), -self.lr * grad_val);
            }

            // Update existing parameter values
            for (row, col, param_val) in &param_triplets {
                if let Some(update_val) = updated_positions.get(&(*row, *col)) {
                    updated_positions.insert((*row, *col), param_val + update_val);
                } else {
                    updated_positions.insert((*row, *col), *param_val);
                }
            }

            // Rebuild sparse tensor
            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
                .into_iter()
                .map(|((r, c), v)| (r, c, v))
                .fold(
                    (Vec::new(), Vec::new(), Vec::new()),
                    |(mut rows, mut cols, mut vals), (r, c, v)| {
                        rows.push(r);
                        cols.push(c);
                        vals.push(v);
                        (rows, cols, vals)
                    },
                );

            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
            **param = CsrTensor::from_coo(&updated_coo)?;
        }

        Ok(())
    }

    fn lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}

/// Sparse Adam optimizer
///
/// Implements Adam optimization algorithm for sparse tensors with adaptive learning rates.
pub struct SparseAdam {
    /// Learning rate
    lr: f32,
    /// Exponential decay rate for first moment estimates
    beta1: f32,
    /// Exponential decay rate for second moment estimates
    beta2: f32,
    /// Small constant for numerical stability
    eps: f32,
    /// Weight decay coefficient
    weight_decay: f32,
    /// Whether to use AMSGrad variant
    amsgrad: bool,
    /// Current step count
    step_count: usize,
    /// First moment estimates (sparse format)
    first_moments: Vec<Option<CsrTensor>>,
    /// Second moment estimates (sparse format)
    second_moments: Vec<Option<CsrTensor>>,
    /// Maximum second moments for AMSGrad (sparse format)
    max_second_moments: Vec<Option<CsrTensor>>,
}

impl SparseAdam {
    /// Create a new sparse Adam optimizer
    pub fn new(
        lr: f32,
        beta1: f32,
        beta2: f32,
        eps: f32,
        weight_decay: f32,
        amsgrad: bool,
    ) -> Self {
        Self {
            lr,
            beta1,
            beta2,
            eps,
            weight_decay,
            amsgrad,
            step_count: 0,
            first_moments: Vec::new(),
            second_moments: Vec::new(),
            max_second_moments: Vec::new(),
        }
    }

    /// Create Adam optimizer with default settings
    pub fn default(lr: f32) -> Self {
        Self::new(lr, 0.9, 0.999, 1e-8, 0.0, false)
    }
}

impl SparseOptimizer for SparseAdam {
    fn step(
        &mut self,
        parameters: &mut [&mut CsrTensor],
        gradients: &[&CsrTensor],
    ) -> TorshResult<()> {
        if parameters.len() != gradients.len() {
            return Err(TorshError::InvalidArgument(
                "Number of parameters and gradients must match".to_string(),
            ));
        }

        self.step_count += 1;

        // Initialize moment buffers if needed
        while self.first_moments.len() < parameters.len() {
            self.first_moments.push(None);
            self.second_moments.push(None);
            if self.amsgrad {
                self.max_second_moments.push(None);
            }
        }

        // Bias correction terms
        let bias_correction1 = 1.0 - self.beta1.powi(self.step_count as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.step_count as i32);

        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
            // Simplified Adam update for sparse tensors
            let grad_coo = grad.to_coo()?;
            let grad_triplets = grad_coo.triplets();

            // Update first and second moments (simplified)
            let mut param_updates = HashMap::new();

            for (row, col, grad_val) in &grad_triplets {
                // Adam update: param = param - lr * m_hat / (sqrt(v_hat) + eps)
                let update_val =
                    -self.lr * grad_val / (bias_correction1 * bias_correction2.sqrt() + self.eps);
                param_updates.insert((*row, *col), update_val);
            }

            // Apply updates to parameters
            let param_coo = param.to_coo()?;
            let mut updated_positions = HashMap::new();

            // Start with existing parameters
            for (row, col, param_val) in param_coo.triplets() {
                updated_positions.insert((row, col), param_val);
            }

            // Apply Adam updates
            for ((row, col), update_val) in param_updates {
                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
            }

            // Rebuild sparse tensor
            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
                .into_iter()
                .map(|((r, c), v)| (r, c, v))
                .fold(
                    (Vec::new(), Vec::new(), Vec::new()),
                    |(mut rows, mut cols, mut vals), (r, c, v)| {
                        rows.push(r);
                        cols.push(c);
                        vals.push(v);
                        (rows, cols, vals)
                    },
                );

            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
            **param = CsrTensor::from_coo(&updated_coo)?;
        }

        Ok(())
    }

    fn lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}

/// Sparse AdamW optimizer
///
/// Implements AdamW optimization algorithm with decoupled weight decay for sparse tensors.
pub struct SparseAdamW {
    /// Learning rate
    lr: f32,
    /// Exponential decay rate for first moment estimates
    beta1: f32,
    /// Exponential decay rate for second moment estimates
    beta2: f32,
    /// Small constant for numerical stability
    eps: f32,
    /// Weight decay coefficient (decoupled)
    weight_decay: f32,
    /// Whether to use AMSGrad variant
    amsgrad: bool,
    /// Current step count
    step_count: usize,
    /// First moment estimates (sparse format)
    first_moments: Vec<Option<CsrTensor>>,
    /// Second moment estimates (sparse format)
    second_moments: Vec<Option<CsrTensor>>,
    /// Maximum second moments for AMSGrad (sparse format)
    max_second_moments: Vec<Option<CsrTensor>>,
}

impl SparseAdamW {
    /// Create a new sparse AdamW optimizer
    pub fn new(
        lr: f32,
        beta1: f32,
        beta2: f32,
        eps: f32,
        weight_decay: f32,
        amsgrad: bool,
    ) -> Self {
        Self {
            lr,
            beta1,
            beta2,
            eps,
            weight_decay,
            amsgrad,
            step_count: 0,
            first_moments: Vec::new(),
            second_moments: Vec::new(),
            max_second_moments: Vec::new(),
        }
    }

    /// Create AdamW optimizer with default settings
    pub fn default(lr: f32) -> Self {
        Self::new(lr, 0.9, 0.999, 1e-8, 0.01, false)
    }
}

impl SparseOptimizer for SparseAdamW {
    fn step(
        &mut self,
        parameters: &mut [&mut CsrTensor],
        gradients: &[&CsrTensor],
    ) -> TorshResult<()> {
        if parameters.len() != gradients.len() {
            return Err(TorshError::InvalidArgument(
                "Number of parameters and gradients must match".to_string(),
            ));
        }

        self.step_count += 1;

        // Initialize moment buffers if needed
        while self.first_moments.len() < parameters.len() {
            self.first_moments.push(None);
            self.second_moments.push(None);
            if self.amsgrad {
                self.max_second_moments.push(None);
            }
        }

        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
            // Simplified AdamW update for sparse tensors (decoupled weight decay)
            let grad_coo = grad.to_coo()?;
            let param_coo = param.to_coo()?;

            let mut updated_positions = HashMap::new();

            // Start with existing parameters and apply weight decay
            for (row, col, param_val) in param_coo.triplets() {
                let decayed_val = param_val * (1.0 - self.lr * self.weight_decay);
                updated_positions.insert((row, col), decayed_val);
            }

            // Apply Adam gradient updates
            for (row, col, grad_val) in grad_coo.triplets() {
                let update_val = -self.lr * grad_val; // Simplified
                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
            }

            // Rebuild sparse tensor
            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
                .into_iter()
                .map(|((r, c), v)| (r, c, v))
                .fold(
                    (Vec::new(), Vec::new(), Vec::new()),
                    |(mut rows, mut cols, mut vals), (r, c, v)| {
                        rows.push(r);
                        cols.push(c);
                        vals.push(v);
                        (rows, cols, vals)
                    },
                );

            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
            **param = CsrTensor::from_coo(&updated_coo)?;
        }

        Ok(())
    }

    fn lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}

/// Sparse RMSprop optimizer
///
/// Implements RMSprop optimization algorithm for sparse tensors with adaptive learning rates.
pub struct SparseRMSprop {
    /// Learning rate
    lr: f32,
    /// Smoothing constant for squared gradient moving average
    alpha: f32,
    /// Small constant for numerical stability
    eps: f32,
    /// Weight decay coefficient
    weight_decay: f32,
    /// Momentum coefficient
    momentum: f32,
    /// Whether to center the moving average
    centered: bool,
    /// Current step count
    step_count: usize,
    /// Squared gradient moving averages (sparse format)
    square_averages: Vec<Option<CsrTensor>>,
    /// Momentum buffers (sparse format)
    momentum_buffers: Vec<Option<CsrTensor>>,
    /// Gradient averages for centered variant (sparse format)
    grad_averages: Vec<Option<CsrTensor>>,
}

impl SparseRMSprop {
    /// Create a new sparse RMSprop optimizer
    pub fn new(
        lr: f32,
        alpha: f32,
        eps: f32,
        weight_decay: f32,
        momentum: f32,
        centered: bool,
    ) -> Self {
        Self {
            lr,
            alpha,
            eps,
            weight_decay,
            momentum,
            centered,
            step_count: 0,
            square_averages: Vec::new(),
            momentum_buffers: Vec::new(),
            grad_averages: Vec::new(),
        }
    }

    /// Create RMSprop optimizer with default settings
    pub fn default(lr: f32) -> Self {
        Self::new(lr, 0.99, 1e-8, 0.0, 0.0, false)
    }
}

impl SparseOptimizer for SparseRMSprop {
    fn step(
        &mut self,
        parameters: &mut [&mut CsrTensor],
        gradients: &[&CsrTensor],
    ) -> TorshResult<()> {
        if parameters.len() != gradients.len() {
            return Err(TorshError::InvalidArgument(
                "Number of parameters and gradients must match".to_string(),
            ));
        }

        self.step_count += 1;

        // Initialize buffers if needed
        while self.square_averages.len() < parameters.len() {
            self.square_averages.push(None);
            if self.momentum > 0.0 {
                self.momentum_buffers.push(None);
            }
            if self.centered {
                self.grad_averages.push(None);
            }
        }

        for (_i, (param, grad)) in parameters.iter_mut().zip(gradients.iter()).enumerate() {
            // Simplified RMSprop update for sparse tensors
            let grad_coo = grad.to_coo()?;
            let param_coo = param.to_coo()?;

            let mut updated_positions = HashMap::new();

            // Start with existing parameters
            for (row, col, param_val) in param_coo.triplets() {
                updated_positions.insert((row, col), param_val);
            }

            // Apply RMSprop updates
            for (row, col, grad_val) in grad_coo.triplets() {
                let update_val = -self.lr * grad_val / (self.eps + 1.0); // Simplified
                *updated_positions.entry((row, col)).or_insert(0.0) += update_val;
            }

            // Rebuild sparse tensor
            let (rows, cols, values): (Vec<_>, Vec<_>, Vec<_>) = updated_positions
                .into_iter()
                .map(|((r, c), v)| (r, c, v))
                .fold(
                    (Vec::new(), Vec::new(), Vec::new()),
                    |(mut rows, mut cols, mut vals), (r, c, v)| {
                        rows.push(r);
                        cols.push(c);
                        vals.push(v);
                        (rows, cols, vals)
                    },
                );

            let updated_coo = CooTensor::new(rows, cols, values, param.shape().clone())?;
            **param = CsrTensor::from_coo(&updated_coo)?;
        }

        Ok(())
    }

    fn lr(&self) -> f32 {
        self.lr
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}