image-registration/src/optimize.rs

use algos::{ObjectiveFunction, OptimizationConfig, OptimizationResult};
use num::Float;
use std::collections::VecDeque;
use std::fmt::Debug;

/// L-BFGS optimizer with proper line search and initial Hessian scaling
pub fn lbfgs_minimize<T, F>(
    f: &F,
    initial_point: &[T],
    config: &OptimizationConfig<T>,
) -> OptimizationResult<T>
where
    T: Float + Debug,
    F: ObjectiveFunction<T>,
{
    const M: usize = 10;
    let n = initial_point.len();
    let mut current_point = initial_point.to_vec();
    let mut iterations = 0;
    let mut converged = false;

    let mut s_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
    let mut y_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
    let mut rho_list: VecDeque<T> = VecDeque::with_capacity(M);

    let mut gradient = match f.gradient(&current_point) {
        Some(g) => g,
        None => {
            return OptimizationResult {
                optimal_point: current_point.clone(),
                optimal_value: f.evaluate(&current_point),
                iterations: 0,
                converged: false,
            };
        }
    };

    while iterations < config.max_iterations {
        let gradient_norm = gradient
            .iter()
            .fold(T::zero(), |acc, &x| acc + x * x)
            .sqrt();
        if gradient_norm < config.tolerance {
            converged = true;
            break;
        }

        let mut q = gradient.clone();
        let mut alpha_list = Vec::with_capacity(s_list.len());

        for i in (0..s_list.len()).rev() {
            let alpha = rho_list[i]
                * s_list[i]
                    .iter()
                    .zip(q.iter())
                    .fold(T::zero(), |acc, (&s, &q)| acc + s * q);
            alpha_list.push(alpha);
            for (q_j, y_j) in q.iter_mut().zip(y_list[i].iter()) {
                *q_j = *q_j - alpha * *y_j;
            }
        }

        let mut r = if !s_list.is_empty() {
            let i = s_list.len() - 1;
            let yy = y_list[i].iter().fold(T::zero(), |acc, &y| acc + y * y);
            let ys = y_list[i]
                .iter()
                .zip(s_list[i].iter())
                .fold(T::zero(), |acc, (&y, &s)| acc + y * s);
            q.iter_mut().for_each(|r_j| *r_j = *r_j * (ys / yy));
            q
        } else {
            q.iter_mut()
                .for_each(|r_j| *r_j = *r_j * config.learning_rate);
            q
        };

        for i in 0..s_list.len() {
            let beta = rho_list[i]
                * y_list[i]
                    .iter()
                    .zip(r.iter())
                    .fold(T::zero(), |acc, (&y, &r)| acc + y * r);
            let alpha = alpha_list[s_list.len() - 1 - i];
            for (r_j, s_j) in r.iter_mut().zip(s_list[i].iter()) {
                *r_j = *r_j + (alpha - beta) * *s_j;
            }
        }

        let direction: Vec<T> = r.iter().map(|&x| -x).collect();

        let mut alpha = T::one();
        let mut new_point = vec![T::zero(); n];
        let current_value = f.evaluate(&current_point);
        let mut improved = false;
        let g_dot_d: T = gradient
            .iter()
            .zip(direction.iter())
            .fold(T::zero(), |acc, (&g, &d)| acc + g * d);
        let c1 = T::from(1e-4).unwrap();

        for _ in 0..30 {
            for i in 0..n {
                new_point[i] = current_point[i] + alpha * direction[i];
            }
            let new_value = f.evaluate(&new_point);
            if new_value <= current_value + c1 * alpha * g_dot_d {
                improved = true;
                break;
            }
            alpha = alpha * T::from(0.5).unwrap();
        }

        if !improved {
            break;
        }

        let new_gradient = match f.gradient(&new_point) {
            Some(g) => g,
            None => break,
        };

        let s = new_point
            .iter()
            .zip(current_point.iter())
            .map(|(&x_new, &x_old)| x_new - x_old)
            .collect::<Vec<T>>();
        let y = new_gradient
            .iter()
            .zip(gradient.iter())
            .map(|(&g_new, &g_old)| g_new - g_old)
            .collect::<Vec<T>>();

        let ys = y
            .iter()
            .zip(s.iter())
            .fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i);

        if ys <= T::zero() {
            break;
        }

        let rho = T::one() / ys;

        if s_list.len() == M {
            s_list.pop_front();
            y_list.pop_front();
            rho_list.pop_front();
        }
        s_list.push_back(s);
        y_list.push_back(y);
        rho_list.push_back(rho);

        current_point = new_point;
        gradient = new_gradient;
        iterations += 1;
    }

    OptimizationResult {
        optimal_point: current_point.clone(),
        optimal_value: f.evaluate(&current_point),
        iterations,
        converged,
    }
}

#[derive(Clone, Debug)]
pub struct AsgdConfig {
    pub max_iterations: usize,
    pub tolerance: f64,
    pub maximum_step_length: f64,
    pub sp_a: f64,
    pub sp_alpha: f64,
    pub sigmoid_max: f64,
    pub sigmoid_min: f64,
    pub sigmoid_scale: f64,
    pub scales: Option<Vec<f64>>,
}

impl Default for AsgdConfig {
    fn default() -> Self {
        Self {
            max_iterations: 250,
            tolerance: 1e-6,
            maximum_step_length: 1.0,
            sp_a: 20.0,
            sp_alpha: 0.602,
            sigmoid_max: 1.0,
            sigmoid_min: -0.01,
            sigmoid_scale: 1e-8,
            scales: None,
        }
    }
}

pub fn asgd_minimize<T, F>(f: &F, initial_point: &[T], config: &AsgdConfig) -> OptimizationResult<T>
where
    T: Float + Debug,
    F: ObjectiveFunction<T>,
{
    let n = initial_point.len();
    let mut current_point = initial_point.to_vec();
    let mut iterations = 0;
    let mut converged = false;

    let mut gradient = match f.gradient(&current_point) {
        Some(g) => g,
        None => {
            return OptimizationResult {
                optimal_point: current_point.clone(),
                optimal_value: f.evaluate(&current_point),
                iterations: 0,
                converged: false,
            };
        }
    };

    // Estimate ASGD parameters from gradient statistics
    let alpha = T::from(config.sp_alpha).unwrap();
    let a_param = T::from(config.sp_a).unwrap();
    let delta = T::from(config.maximum_step_length).unwrap();
    let fmax = T::from(config.sigmoid_max).unwrap();
    let fmin = T::from(config.sigmoid_min).unwrap();
    let omega = T::from(config.sigmoid_scale).unwrap();

    // Compute max_j = max(|g_i / scales_i|)
    let max_j = if let Some(ref scales) = config.scales {
        let mut mj = T::zero();
        for (&g, &s) in gradient.iter().zip(scales.iter()) {
            let s_t = T::from(s.max(1e-10)).unwrap();
            let gs = (g / s_t).abs();
            if gs > mj {
                mj = gs;
            }
        }
        mj
    } else {
        let mut mj = T::zero();
        for &g in gradient.iter() {
            let ga = g.abs();
            if ga > mj {
                mj = ga;
            }
        }
        mj
    };

    // a = delta * (A+1)^alpha / max_j  (Elastix AutomaticParameterEstimation)
    let a = if max_j > T::from(1e-14).unwrap() {
        delta * (a_param + T::one()).powf(alpha) / max_j
    } else {
        delta * (a_param + T::one()).powf(alpha)
    };

    let mut current_time = T::zero();
    let mut previous_gradient: Option<Vec<T>> = None;

    while iterations < config.max_iterations {
        let gradient_norm = gradient
            .iter()
            .fold(T::zero(), |acc, &x| acc + x * x)
            .sqrt();
        let tol: T = num::cast(config.tolerance).unwrap();
        if gradient_norm < tol {
            converged = true;
            break;
        }

        // a(k) = min(a / (A + t_k + 1)^alpha, delta)
        let t_k = current_time + T::from(iterations).unwrap();
        let a_t = T::min(a / (a_param + t_k + T::one()).powf(alpha), delta);

        // Update time: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1}))
        if let Some(ref prev_g) = previous_gradient {
            let dot_product: T = if let Some(ref scales) = config.scales {
                gradient.iter().zip(prev_g.iter()).zip(scales.iter()).fold(
                    T::zero(),
                    |acc, ((&g, &pg), &s)| {
                        let s2: T = num::cast(s.max(1e-10)).unwrap();
                        let s4 = s2 * s2;
                        acc + g * pg / s4
                    },
                )
            } else {
                gradient
                    .iter()
                    .zip(prev_g.iter())
                    .fold(T::zero(), |acc, (&g, &pg)| acc + g * pg)
            };
            let sigmoid_val = sigmoid(-dot_product, fmax, fmin, omega);
            current_time = T::max(T::zero(), current_time + sigmoid_val);
        }
        previous_gradient = Some(gradient.clone());

        // direction = -g / scales
        let direction: Vec<T> = match &config.scales {
            Some(scales) => gradient
                .iter()
                .zip(scales.iter())
                .map(|(&g, &s)| {
                    let s2: T = num::cast(s.max(1e-10)).unwrap();
                    -g / s2
                })
                .collect(),
            None => gradient.iter().map(|&g| -g).collect(),
        };

        let mut new_point = current_point.clone();
        for i in 0..n {
            new_point[i] = current_point[i] + direction[i] * a_t;
        }

        let new_gradient = match f.gradient(&new_point) {
            Some(g) => g,
            None => break,
        };

        current_point = new_point;
        gradient = new_gradient;
        iterations += 1;
    }

    OptimizationResult {
        optimal_point: current_point.clone(),
        optimal_value: f.evaluate(&current_point),
        iterations,
        converged,
    }
}

fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
    let beta = omega * (-fmax / fmin).ln();
    let z = (x - beta) / omega;
    let sigmoid_raw = if z > T::from(20.0).unwrap() {
        T::one()
    } else if z < T::from(-20.0).unwrap() {
        T::zero()
    } else {
        T::one() / (T::one() + (-z).exp())
    };
    (fmax - fmin) * sigmoid_raw + fmin
}

#[cfg(test)]
mod tests {
    use super::*;
    use algos::OptimizationConfig;

    struct Quadratic;

    impl ObjectiveFunction<f64> for Quadratic {
        fn evaluate(&self, point: &[f64]) -> f64 {
            point.iter().map(|x| x * x).sum()
        }

        fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
            Some(point.iter().map(|x| 2.0 * x).collect())
        }
    }

    #[test]
    fn test_lbfgs_quadratic() {
        let f = Quadratic;
        let initial_point = vec![1.0, 1.0];
        let config = OptimizationConfig {
            max_iterations: 100,
            tolerance: 1e-6,
            learning_rate: 1.0,
        };

        let result = lbfgs_minimize(&f, &initial_point, &config);

        assert!(result.converged);
        assert!(result.optimal_value < 1e-10);
        for x in result.optimal_point {
            assert!(x.abs() < 1e-5);
        }
    }

    #[test]
    fn test_lbfgs_quadratic_with_minimum() {
        let f = QuadraticWithMinimum;
        let initial_point = vec![0.0];
        let config = OptimizationConfig {
            max_iterations: 100,
            tolerance: 1e-6,
            learning_rate: 1.0,
        };

        let result = lbfgs_minimize(&f, &initial_point, &config);

        assert!(result.converged);
        assert!((result.optimal_point[0] - 2.0).abs() < 1e-5);
    }

    #[test]
    fn test_asgd_quadratic() {
        let f = Quadratic;
        let initial_point = vec![1.0, 1.0];
        let config = AsgdConfig {
            max_iterations: 1000,
            tolerance: 1e-4,
            sigmoid_min: -0.8,
            ..Default::default()
        };

        let result = asgd_minimize(&f, &initial_point, &config);

        assert!(result.optimal_value < 0.1);
        for x in result.optimal_point {
            assert!(x.abs() < 0.5);
        }
    }
}

struct QuadraticWithMinimum;

impl ObjectiveFunction<f64> for QuadraticWithMinimum {
    fn evaluate(&self, point: &[f64]) -> f64 {
        let x = point[0];
        (x - 2.0).powi(2)
    }

    fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
        let x = point[0];
        Some(vec![2.0 * (x - 2.0)])
    }
}