Skip to main content

torsh_optim/
trust_region.rs

1//! Trust Region optimization methods
2//!
3//! This module provides trust region optimization algorithms including
4//! Trust Region Newton, Trust Region with CG, and general trust region frameworks.
5
6use crate::{
7    Optimizer, OptimizerError, OptimizerResult, OptimizerState, ParamGroup, ParamGroupState,
8};
9use parking_lot::RwLock;
10use std::collections::HashMap;
11use std::ops::Add;
12use std::sync::Arc;
13use torsh_core::error::{Result, TorshError};
14use torsh_tensor::Tensor;
15
16/// Trust region update strategy
17#[derive(Clone, Copy, Debug)]
18pub enum TrustRegionStrategy {
19    /// Standard trust region with reduction ratio
20    Standard,
21    /// Adaptive trust region based on gradient norms
22    Adaptive,
23    /// Conservative trust region with smaller expansion
24    Conservative,
25    /// Aggressive trust region with larger expansion
26    Aggressive,
27}
28
29/// Trust region configuration
30#[derive(Clone)]
31pub struct TrustRegionConfig {
32    /// Initial trust region radius
33    pub initial_radius: f32,
34    /// Maximum trust region radius
35    pub max_radius: f32,
36    /// Minimum trust region radius
37    pub min_radius: f32,
38    /// Acceptance threshold for reduction ratio
39    pub eta1: f32,
40    /// Expansion threshold for reduction ratio
41    pub eta2: f32,
42    /// Radius reduction factor
43    pub gamma1: f32,
44    /// Radius expansion factor
45    pub gamma2: f32,
46    /// Trust region strategy
47    pub strategy: TrustRegionStrategy,
48    /// Maximum number of iterations per step
49    pub max_iter: usize,
50    /// Tolerance for gradient norm
51    pub tolerance_grad: f32,
52    /// Tolerance for step size
53    pub tolerance_step: f32,
54}
55
56impl Default for TrustRegionConfig {
57    fn default() -> Self {
58        Self {
59            initial_radius: 1.0,
60            max_radius: 100.0,
61            min_radius: 1e-6,
62            eta1: 0.25,
63            eta2: 0.75,
64            gamma1: 0.25,
65            gamma2: 2.0,
66            strategy: TrustRegionStrategy::Standard,
67            max_iter: 100,
68            tolerance_grad: 1e-6,
69            tolerance_step: 1e-8,
70        }
71    }
72}
73
74/// Trust region subproblem solver
75#[derive(Clone, Copy, Debug)]
76pub enum SubproblemSolver {
77    /// Cauchy point (steepest descent)
78    CauchyPoint,
79    /// Dogleg method
80    Dogleg,
81    /// Conjugate Gradient (Steihaug-Toint)
82    ConjugateGradient,
83    /// Two-dimensional subspace minimization
84    TwoDSubspace,
85}
86
87/// Objective function evaluated at a flattened parameter vector.
88///
89/// A trust region method accepts or rejects a step by comparing the decrease the
90/// quadratic model predicted with the decrease the objective actually delivered,
91/// so it needs to evaluate the objective at candidate points — exactly like
92/// PyTorch's `LBFGS.step(closure)`.
93pub type ObjectiveFn = Arc<dyn Fn(&Tensor) -> Result<f32> + Send + Sync>;
94
95/// Trust region method implementation
96///
97/// # Objective closure
98///
99/// [`TrustRegionMethod::step`] cannot decide whether to accept a step without
100/// evaluating the objective, so an objective closure must be registered with
101/// [`TrustRegionMethod::set_objective`] before stepping. Without one, `step`
102/// returns an error rather than inventing a reduction ratio.
103pub struct TrustRegionMethod {
104    param_groups: Vec<ParamGroup>,
105    state: HashMap<String, HashMap<String, Tensor>>,
106    step_count: usize,
107    config: TrustRegionConfig,
108    solver: SubproblemSolver,
109    objective: Option<ObjectiveFn>,
110}
111
112impl TrustRegionMethod {
113    pub fn new(
114        params: Vec<Arc<RwLock<Tensor>>>,
115        lr: Option<f32>,
116        config: Option<TrustRegionConfig>,
117        solver: Option<SubproblemSolver>,
118    ) -> Self {
119        let lr = lr.unwrap_or(1.0);
120        let param_group = ParamGroup::new(params, lr);
121
122        Self {
123            param_groups: vec![param_group],
124            state: HashMap::new(),
125            step_count: 0,
126            config: config.unwrap_or_default(),
127            solver: solver.unwrap_or(SubproblemSolver::Dogleg),
128            objective: None,
129        }
130    }
131
132    /// Register the objective function used to measure the actual decrease.
133    ///
134    /// The closure receives the flattened parameter vector (in the same layout
135    /// [`TrustRegionMethod`] uses internally) and returns the objective value at
136    /// that point.
137    pub fn set_objective<F>(&mut self, objective: F)
138    where
139        F: Fn(&Tensor) -> Result<f32> + Send + Sync + 'static,
140    {
141        self.objective = Some(Arc::new(objective));
142    }
143
144    /// Remove a previously registered objective function.
145    pub fn clear_objective(&mut self) {
146        self.objective = None;
147    }
148
149    pub fn builder() -> TrustRegionBuilder {
150        TrustRegionBuilder::new()
151    }
152
153    fn get_param_id(param: &Arc<RwLock<Tensor>>) -> String {
154        format!("{:p}", Arc::as_ptr(param))
155    }
156
157    /// Flatten parameters into a single vector
158    fn flatten_params(&self) -> Result<Tensor> {
159        let mut flattened = Vec::new();
160
161        for group in &self.param_groups {
162            for param in &group.params {
163                let param_read = param.read();
164                let param_flat = param_read.flatten()?;
165                let param_data = param_flat.data()?;
166                flattened.extend_from_slice(&param_data);
167            }
168        }
169
170        let len = flattened.len();
171        Ok(Tensor::from_data(
172            flattened,
173            vec![len],
174            torsh_core::device::DeviceType::Cpu,
175        )?)
176    }
177
178    /// Flatten gradients into a single vector
179    fn flatten_grads(&self) -> Result<Tensor> {
180        let mut flattened = Vec::new();
181
182        for group in &self.param_groups {
183            for param in &group.params {
184                let param_read = param.read();
185                let grad = param_read.grad().ok_or_else(|| {
186                    TorshError::invalid_argument_with_context(
187                        "Parameter has no gradient",
188                        "trust_region_step",
189                    )
190                })?;
191                let grad_flat = grad.flatten()?;
192                let grad_data = grad_flat.data()?;
193                flattened.extend_from_slice(&grad_data);
194            }
195        }
196
197        let len = flattened.len();
198        Ok(Tensor::from_data(
199            flattened,
200            vec![len],
201            torsh_core::device::DeviceType::Cpu,
202        )?)
203    }
204
205    /// Update parameters from flattened vector
206    fn update_params_from_flat(&self, flat_params: &Tensor) -> Result<()> {
207        let flat_data = flat_params.data()?;
208        let mut offset = 0;
209
210        for group in &self.param_groups {
211            for param in &group.params {
212                let mut param_write = param.write();
213                let param_shape = param_write.shape();
214                let param_size = param_shape.numel();
215
216                let param_data = &flat_data[offset..offset + param_size];
217                let new_values = Tensor::from_data(
218                    param_data.to_vec(),
219                    param_shape.dims().to_vec(),
220                    param_write.device(),
221                )?;
222                crate::param_update::assign(&mut param_write, &new_values)?;
223
224                offset += param_size;
225            }
226        }
227
228        Ok(())
229    }
230
231    /// Compute Cauchy point (steepest descent direction projected onto trust region)
232    fn cauchy_point(&self, grad: &Tensor, radius: f32) -> Result<Tensor> {
233        let grad_norm = grad.norm()?.item()?;
234
235        if grad_norm < self.config.tolerance_grad {
236            return Ok(Tensor::zeros(grad.shape().dims(), grad.device())?);
237        }
238
239        // Cauchy point: -τ * Δ / ||g|| * g, where τ is chosen to satisfy ||s|| ≤ Δ
240        let tau = (radius / grad_norm).min(1.0);
241        Ok(grad.mul_scalar(-tau * radius / grad_norm)?)
242    }
243
244    /// Approximate Hessian using gradient differences (simplified)
245    fn approximate_hessian(&self, grad: &Tensor) -> Result<Tensor> {
246        // Simplified Hessian approximation: H ≈ ||g|| * I
247        // In practice, you would use L-BFGS, finite differences, or exact Hessian
248        let grad_norm = grad.norm()?.item()?;
249        let scale = if grad_norm > 1e-8 { grad_norm } else { 1.0 };
250
251        // Return diagonal approximation as a vector (diagonal elements)
252        let n = grad.shape().numel();
253        Ok(Tensor::from_data(vec![scale; n], vec![n], grad.device())?)
254    }
255
256    /// Solve trust region subproblem using dogleg method
257    fn dogleg(&self, grad: &Tensor, hessian_diag: &Tensor, radius: f32) -> Result<Tensor> {
258        let grad_norm = grad.norm()?.item()?;
259
260        if grad_norm < self.config.tolerance_grad {
261            return Ok(Tensor::zeros(grad.shape().dims(), grad.device())?);
262        }
263
264        // Cauchy point
265        let cauchy_step = self.cauchy_point(grad, radius)?;
266        let cauchy_norm = cauchy_step.norm()?.item()?;
267
268        // If Cauchy point is on the boundary, return it
269        if cauchy_norm >= radius - self.config.tolerance_step {
270            return Ok(cauchy_step);
271        }
272
273        // Newton step (simplified): -H^(-1) * g
274        // For diagonal Hessian: -g[i] / H[i][i]
275        let hessian_data = hessian_diag.data()?;
276        let grad_data = grad.data()?;
277        let newton_data: Vec<f32> = grad_data
278            .iter()
279            .zip(hessian_data.iter())
280            .map(|(g, h)| if h.abs() > 1e-12 { -g / h } else { -g })
281            .collect();
282
283        let newton_step =
284            Tensor::from_data(newton_data, grad.shape().dims().to_vec(), grad.device())?;
285        let newton_norm = newton_step.norm()?.item()?;
286
287        // If Newton step is within trust region, return it
288        if newton_norm <= radius {
289            return Ok(newton_step);
290        }
291
292        // Dogleg path: find intersection with trust region boundary
293        // s(τ) = τ * cauchy_step + (1-τ) * newton_step for τ ∈ [0, 1]
294        // Solve ||s(τ)|| = radius
295
296        let diff = newton_step.sub(&cauchy_step)?;
297        let a = diff.dot(&diff)?.item()?;
298        let b = 2.0 * cauchy_step.dot(&diff)?.item()?;
299        let c = cauchy_norm * cauchy_norm - radius * radius;
300
301        let discriminant = b * b - 4.0 * a * c;
302        if discriminant < 0.0 {
303            return Ok(cauchy_step);
304        }
305
306        let tau = (-b + discriminant.sqrt()) / (2.0 * a);
307        let tau = tau.clamp(0.0, 1.0);
308
309        // s = τ * cauchy + (1-τ) * newton
310        let result = cauchy_step
311            .mul_scalar(tau)?
312            .add(&newton_step.mul_scalar(1.0 - tau)?)?;
313        Ok(result)
314    }
315
316    /// Solve trust region subproblem using conjugate gradient (Steihaug-Toint)
317    fn conjugate_gradient_tr(
318        &self,
319        grad: &Tensor,
320        hessian_diag: &Tensor,
321        radius: f32,
322    ) -> Result<Tensor> {
323        let n = grad.shape().numel();
324        let tolerance = self.config.tolerance_grad;
325        let max_iter = n.min(50); // Limit CG iterations
326
327        // Initialize
328        let mut x = Tensor::zeros(&[n], grad.device())?;
329        let mut r = grad.neg()?; // -g
330        let mut p = r.clone();
331        let mut rsold = r.dot(&r)?.item()?;
332
333        for _i in 0..max_iter {
334            if rsold.sqrt() < tolerance {
335                break;
336            }
337
338            // Approximate Hessian-vector product using diagonal approximation
339            let hp_data: Vec<f32> = {
340                let hessian_data = hessian_diag.data()?;
341                let p_data = p.data()?;
342                p_data
343                    .iter()
344                    .zip(hessian_data.iter())
345                    .map(|(p_val, h_val)| p_val * h_val)
346                    .collect()
347            };
348            let hp = Tensor::from_data(hp_data, grad.shape().dims().to_vec(), grad.device())?;
349
350            let pap = p.dot(&hp)?.item()?;
351
352            // Check for negative curvature
353            if pap <= 0.0 {
354                // Find boundary of trust region along direction p
355                let x_norm_sq = x.dot(&x)?.item()?;
356                let xp = x.dot(&p)?.item()?;
357                let p_norm_sq = p.dot(&p)?.item()?;
358
359                let discriminant = xp * xp + p_norm_sq * (radius * radius - x_norm_sq);
360                if discriminant >= 0.0 {
361                    let alpha = (-xp + discriminant.sqrt()) / p_norm_sq;
362                    return Ok(x.add(&p.mul_scalar(alpha)?)?);
363                } else {
364                    return Ok(x);
365                }
366            }
367
368            let alpha = rsold / pap;
369            let x_new = x.add(&p.mul_scalar(alpha)?)?;
370
371            // Check trust region constraint
372            let x_norm = x_new.norm()?.item()?;
373            if x_norm >= radius {
374                // Find intersection with trust region boundary
375                let x_norm_old_sq = x.dot(&x)?.item()?;
376                let xp = x.dot(&p)?.item()?;
377                let p_norm_sq = p.dot(&p)?.item()?;
378
379                let discriminant = xp * xp + p_norm_sq * (radius * radius - x_norm_old_sq);
380                if discriminant >= 0.0 {
381                    let tau = (-xp + discriminant.sqrt()) / p_norm_sq;
382                    return Ok(x.add(&p.mul_scalar(tau)?)?);
383                } else {
384                    return Ok(x);
385                }
386            }
387
388            x = x_new;
389            let r_new = r.sub(&hp.mul_scalar(alpha)?)?;
390            let rsnew = r_new.dot(&r_new)?.item()?;
391
392            if rsnew.sqrt() < tolerance {
393                break;
394            }
395
396            let beta = rsnew / rsold;
397            p = r_new.add(&p.mul_scalar(beta)?)?;
398            r = r_new;
399            rsold = rsnew;
400        }
401
402        Ok(x)
403    }
404
405    /// Solve trust region subproblem
406    fn solve_subproblem(&self, grad: &Tensor, radius: f32) -> Result<Tensor> {
407        match self.solver {
408            SubproblemSolver::CauchyPoint => self.cauchy_point(grad, radius),
409            SubproblemSolver::Dogleg => {
410                let hessian_diag = self.approximate_hessian(grad)?;
411                self.dogleg(grad, &hessian_diag, radius)
412            }
413            SubproblemSolver::ConjugateGradient => {
414                let hessian_diag = self.approximate_hessian(grad)?;
415                self.conjugate_gradient_tr(grad, &hessian_diag, radius)
416            }
417            SubproblemSolver::TwoDSubspace => {
418                // Simplified: fall back to dogleg
419                let hessian_diag = self.approximate_hessian(grad)?;
420                self.dogleg(grad, &hessian_diag, radius)
421            }
422        }
423    }
424
425    /// Decrease predicted by the quadratic model: `m(0) - m(s) = -g^T s - 0.5 s^T H s`.
426    ///
427    /// This is the same model the subproblem solvers minimise, so the ratio
428    /// formed against it is the standard trust-region reduction ratio. `H` is the
429    /// diagonal approximation returned by `approximate_hessian`.
430    fn model_decrease(&self, grad: &Tensor, step: &Tensor, hessian_diag: &Tensor) -> Result<f32> {
431        let linear_term = grad.dot(step)?.item()?;
432        let hessian_step = step.mul_op(hessian_diag)?;
433        let quadratic_term = step.dot(&hessian_step)?.item()?;
434        let decrease = -linear_term - 0.5 * quadratic_term;
435        Ok(decrease.max(0.0))
436    }
437
438    /// Decrease the objective actually delivered: `f(x) - f(x + s)`.
439    ///
440    /// # Errors
441    /// Returns an error if no objective closure has been registered — the actual
442    /// decrease cannot be measured without evaluating the objective, and
443    /// returning a made-up constant would silently defeat the trust-region logic
444    /// that consumes it.
445    fn actual_decrease(&self, old_params: &Tensor, new_params: &Tensor) -> Result<f32> {
446        let objective = self.objective.as_ref().ok_or_else(|| {
447            TorshError::InvalidArgument(
448                "TrustRegionMethod requires an objective function to measure the actual \
449                 decrease; register one with `set_objective` before calling `step`"
450                    .to_string(),
451            )
452        })?;
453        Ok(objective(old_params)? - objective(new_params)?)
454    }
455
456    /// Update trust region radius
457    fn update_radius(&self, current_radius: f32, reduction_ratio: f32) -> f32 {
458        let config = &self.config;
459
460        match config.strategy {
461            TrustRegionStrategy::Standard => {
462                if reduction_ratio < config.eta1 {
463                    (current_radius * config.gamma1).max(config.min_radius)
464                } else if reduction_ratio > config.eta2 {
465                    (current_radius * config.gamma2).min(config.max_radius)
466                } else {
467                    current_radius
468                }
469            }
470            TrustRegionStrategy::Adaptive => {
471                // More aggressive adaptation based on reduction ratio
472                if reduction_ratio < 0.1 {
473                    (current_radius * 0.1).max(config.min_radius)
474                } else if reduction_ratio > 0.9 {
475                    (current_radius * 3.0).min(config.max_radius)
476                } else {
477                    current_radius * (0.5 + reduction_ratio)
478                }
479            }
480            TrustRegionStrategy::Conservative => {
481                // Smaller changes to radius
482                if reduction_ratio < config.eta1 {
483                    (current_radius * 0.5).max(config.min_radius)
484                } else if reduction_ratio > config.eta2 {
485                    (current_radius * 1.5).min(config.max_radius)
486                } else {
487                    current_radius
488                }
489            }
490            TrustRegionStrategy::Aggressive => {
491                // Larger changes to radius
492                if reduction_ratio < config.eta1 {
493                    (current_radius * 0.1).max(config.min_radius)
494                } else if reduction_ratio > 0.5 {
495                    (current_radius * 4.0).min(config.max_radius)
496                } else {
497                    current_radius
498                }
499            }
500        }
501    }
502}
503
504impl Optimizer for TrustRegionMethod {
505    fn step(&mut self) -> OptimizerResult<()> {
506        self.step_count += 1;
507
508        // Get current parameters and gradients
509        let current_params = self.flatten_params()?;
510        let current_grad = self.flatten_grads()?;
511
512        // Check convergence
513        let grad_norm = current_grad.norm()?.item()?;
514        if grad_norm < self.config.tolerance_grad {
515            return Ok(());
516        }
517
518        // Get or initialize trust region radius
519        let state_id = "trust_region_state".to_string();
520        let mut radius = {
521            let param_state = self.state.entry(state_id.clone()).or_default();
522            if let Some(radius_tensor) = param_state.get("radius") {
523                radius_tensor.item()?
524            } else {
525                let initial_radius = self.config.initial_radius;
526                param_state.insert("radius".to_string(), Tensor::scalar(initial_radius)?);
527                initial_radius
528            }
529        };
530
531        // Diagonal Hessian approximation shared by the subproblem solvers and the
532        // predicted-decrease computation, so both use the same model.
533        let hessian_diag = self.approximate_hessian(&current_grad)?;
534
535        // Trust region iteration
536        for _iter in 0..self.config.max_iter {
537            // Solve trust region subproblem
538            let step = self.solve_subproblem(&current_grad, radius)?;
539            let step_norm = step.norm()?.item()?;
540
541            // Check if step is too small
542            if step_norm < self.config.tolerance_step {
543                break;
544            }
545
546            // Compute new parameters
547            let new_params = current_params.add(&step)?;
548
549            // Compute reduction ratio
550            let model_dec = self.model_decrease(&current_grad, &step, &hessian_diag)?;
551            let actual_dec = self.actual_decrease(&current_params, &new_params)?;
552
553            let reduction_ratio = if model_dec > 1e-12 {
554                actual_dec / model_dec
555            } else {
556                0.0
557            };
558
559            // Accept or reject step
560            if reduction_ratio > self.config.eta1 {
561                // Accept step
562                self.update_params_from_flat(&new_params)?;
563            }
564
565            // Update trust region radius
566            radius = self.update_radius(radius, reduction_ratio);
567
568            // Store updated radius
569            let param_state = self
570                .state
571                .get_mut(&state_id)
572                .expect("state should exist for state_id");
573            param_state.insert("radius".to_string(), Tensor::scalar(radius)?);
574
575            // Break if step was accepted and radius is reasonable
576            if reduction_ratio > self.config.eta1 {
577                break;
578            }
579
580            // Break if radius becomes too small
581            if radius < self.config.min_radius {
582                break;
583            }
584        }
585
586        Ok(())
587    }
588
589    fn zero_grad(&mut self) {
590        for group in &self.param_groups {
591            for param in &group.params {
592                param.write().zero_grad();
593            }
594        }
595    }
596
597    fn get_lr(&self) -> Vec<f32> {
598        self.param_groups.iter().map(|g| g.lr).collect()
599    }
600
601    fn set_lr(&mut self, lr: f32) {
602        for group in &mut self.param_groups {
603            group.lr = lr;
604        }
605    }
606
607    fn add_param_group(&mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>) {
608        let lr = options.get("lr").copied().unwrap_or(1.0);
609        let group = ParamGroup::new(params, lr).with_options(options);
610        self.param_groups.push(group);
611    }
612
613    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
614        crate::optimizer::collect_parameters(&self.param_groups)
615    }
616
617    fn state_dict(&self) -> OptimizerResult<OptimizerState> {
618        let param_groups = self
619            .param_groups
620            .iter()
621            .map(|g| ParamGroupState {
622                lr: g.lr,
623                options: g.options.clone(),
624                param_count: g.params.len(),
625            })
626            .collect();
627
628        Ok(OptimizerState {
629            optimizer_type: "TrustRegion".to_string(),
630            version: "0.1.0".to_string(),
631            param_groups,
632            state: self.state.clone(),
633            global_state: HashMap::new(),
634        })
635    }
636
637    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()> {
638        if state.param_groups.len() != self.param_groups.len() {
639            return Err(OptimizerError::TensorError(TorshError::InvalidArgument(
640                "Parameter group count mismatch".to_string(),
641            )));
642        }
643
644        for (i, group_state) in state.param_groups.iter().enumerate() {
645            self.param_groups[i].lr = group_state.lr;
646            self.param_groups[i].options = group_state.options.clone();
647        }
648
649        self.state = state.state;
650        Ok(())
651    }
652}
653
654/// Builder for trust region optimizers
655pub struct TrustRegionBuilder {
656    lr: f32,
657    config: TrustRegionConfig,
658    solver: SubproblemSolver,
659}
660
661impl TrustRegionBuilder {
662    pub fn new() -> Self {
663        Self {
664            lr: 1.0,
665            config: TrustRegionConfig::default(),
666            solver: SubproblemSolver::Dogleg,
667        }
668    }
669
670    pub fn lr(mut self, lr: f32) -> Self {
671        self.lr = lr;
672        self
673    }
674
675    pub fn initial_radius(mut self, radius: f32) -> Self {
676        self.config.initial_radius = radius;
677        self
678    }
679
680    pub fn max_radius(mut self, radius: f32) -> Self {
681        self.config.max_radius = radius;
682        self
683    }
684
685    pub fn min_radius(mut self, radius: f32) -> Self {
686        self.config.min_radius = radius;
687        self
688    }
689
690    pub fn strategy(mut self, strategy: TrustRegionStrategy) -> Self {
691        self.config.strategy = strategy;
692        self
693    }
694
695    pub fn solver(mut self, solver: SubproblemSolver) -> Self {
696        self.solver = solver;
697        self
698    }
699
700    pub fn tolerance_grad(mut self, tol: f32) -> Self {
701        self.config.tolerance_grad = tol;
702        self
703    }
704
705    pub fn tolerance_step(mut self, tol: f32) -> Self {
706        self.config.tolerance_step = tol;
707        self
708    }
709
710    pub fn max_iter(mut self, max_iter: usize) -> Self {
711        self.config.max_iter = max_iter;
712        self
713    }
714
715    pub fn build(self, params: Vec<Arc<RwLock<Tensor>>>) -> TrustRegionMethod {
716        TrustRegionMethod::new(params, Some(self.lr), Some(self.config), Some(self.solver))
717    }
718}
719
720impl Default for TrustRegionBuilder {
721    fn default() -> Self {
722        Self::new()
723    }
724}
725
726#[cfg(test)]
727mod tests {
728    use super::*;
729    use torsh_tensor::creation::randn;
730
731    #[test]
732    fn test_trust_region_creation() -> OptimizerResult<()> {
733        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
734        let optimizer = TrustRegionMethod::new(params, None, None, None);
735        assert_eq!(optimizer.get_lr()[0], 1.0);
736        Ok(())
737    }
738
739    #[test]
740    fn test_trust_region_builder() -> OptimizerResult<()> {
741        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
742        let optimizer = TrustRegionMethod::builder()
743            .lr(0.1)
744            .initial_radius(0.5)
745            .strategy(TrustRegionStrategy::Adaptive)
746            .solver(SubproblemSolver::ConjugateGradient)
747            .build(params);
748
749        assert_eq!(optimizer.get_lr()[0], 0.1);
750        assert_eq!(optimizer.config.initial_radius, 0.5);
751        assert!(matches!(
752            optimizer.config.strategy,
753            TrustRegionStrategy::Adaptive
754        ));
755        assert!(matches!(
756            optimizer.solver,
757            SubproblemSolver::ConjugateGradient
758        ));
759        Ok(())
760    }
761
762    #[test]
763    fn test_cauchy_point() -> OptimizerResult<()> {
764        let params = vec![Arc::new(RwLock::new(randn::<f32>(&[2, 2])?))];
765        let optimizer = TrustRegionMethod::new(params, None, None, None);
766
767        let grad = Tensor::from_data(
768            vec![1.0, 0.0, 0.0, 1.0],
769            vec![4],
770            torsh_core::device::DeviceType::Cpu,
771        )?;
772        let radius = 1.0;
773
774        let cauchy = optimizer.cauchy_point(&grad, radius)?;
775        let cauchy_norm = cauchy.norm()?.item()?;
776
777        // Cauchy point should be within trust region
778        assert!(cauchy_norm <= radius + 1e-6);
779        Ok(())
780    }
781
782    #[test]
783    fn test_trust_region_config() {
784        let config = TrustRegionConfig {
785            initial_radius: 2.0,
786            max_radius: 50.0,
787            min_radius: 1e-5,
788            eta1: 0.1,
789            eta2: 0.9,
790            gamma1: 0.1,
791            gamma2: 3.0,
792            strategy: TrustRegionStrategy::Aggressive,
793            max_iter: 50,
794            tolerance_grad: 1e-8,
795            tolerance_step: 1e-10,
796        };
797
798        assert_eq!(config.initial_radius, 2.0);
799        assert_eq!(config.eta1, 0.1);
800        assert_eq!(config.eta2, 0.9);
801        assert!(matches!(config.strategy, TrustRegionStrategy::Aggressive));
802    }
803}