manopt_rs/optimizers/
multiple.rs

1//! Multi-manifold optimizer for handling modules with different manifold constraints.
2//!
3//! This module provides optimizers that can handle multiple manifold types within
4//! a single optimization step, allowing complex models with heterogeneous constraints.
5//!
6//! Users can implement their own manifolds and use them with the multi-manifold optimizer.
7
8use std::collections::HashMap;
9use std::fmt::Debug;
10use std::marker::PhantomData;
11
12use burn::module::Module;
13
14use crate::constrained_module::Constrained;
15use crate::prelude::*;
16
17/// Multi-manifold optimizer configuration
18#[derive(Debug, Clone)]
19pub struct MultiManifoldOptimizerConfig {
20    pub learning_rate: f64,
21    pub beta1: f64,
22    pub beta2: f64,
23    pub eps: f64,
24    pub weight_decay: f64,
25}
26
27impl Default for MultiManifoldOptimizerConfig {
28    fn default() -> Self {
29        Self {
30            learning_rate: 1e-3,
31            beta1: 0.9,
32            beta2: 0.999,
33            eps: 1e-8,
34            weight_decay: 0.0,
35        }
36    }
37}
38
39/// Multi-manifold optimizer that can handle different manifold types
40#[derive(Debug)]
41pub struct MultiManifoldOptimizer<B: Backend> {
42    #[allow(unused)]
43    config: MultiManifoldOptimizerConfig,
44    _backend: PhantomData<B>,
45}
46
47impl<B: Backend> MultiManifoldOptimizer<B> {
48    #[must_use]
49    pub fn new(config: MultiManifoldOptimizerConfig) -> Self {
50        Self {
51            config,
52            _backend: PhantomData,
53        }
54    }
55
56    /// Collect manifold constraints from a module - simplified version
57    pub fn collect_manifolds<M: Module<B>>(&mut self, _module: &M) {
58        // In a simple version, we don't need to do anything here
59        // The manifold information is already encoded in the types
60    }
61
62    /// Register a manifold for a specific parameter path - simplified version
63    pub fn register_manifold<M: Manifold<B> + Send + Sync + 'static>(&mut self, _path: String) {
64        // In a simple version, this is handled at the type level
65    }
66
67    /// Apply manifold constraints to a module
68    pub fn apply_constraints<M: Module<B>>(self, module: M) -> M {
69        // For now, just return the module as-is
70        // The constraints are already applied at the type level
71        module
72    }
73}
74
75/// Extension trait for modules with manifold constraints
76pub trait ManifoldOptimizable<B: Backend>: Module<B> {
77    /// Apply manifold constraints to the module
78    #[must_use]
79    fn apply_manifold_constraints(self) -> Self;
80
81    /// Get information about manifold constraints
82    fn get_manifold_info(&self) -> HashMap<String, String>;
83}
84
85// Blanket implementation for constrained modules
86impl<B, M, Man> ManifoldOptimizable<B> for Constrained<M, Man>
87where
88    M: Module<B>,
89    B: Backend,
90    Man: Manifold<B> + Clone + Debug + Send,
91{
92    fn apply_manifold_constraints(self) -> Self {
93        // Apply constraints to the inner module and wrap it back
94        self
95    }
96
97    fn get_manifold_info(&self) -> HashMap<String, String> {
98        let mut info = HashMap::new();
99        info.insert("manifold_type".to_string(), Man::name().to_string());
100        info
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use burn::backend::NdArray;
108    use burn::nn::LinearConfig;
109
110    type TestBackend = NdArray;
111
112    #[test]
113    fn test_multi_manifold_optimizer() {
114        let config = MultiManifoldOptimizerConfig::default();
115        let optimizer = MultiManifoldOptimizer::<TestBackend>::new(config);
116
117        // Test basic construction
118        assert_eq!(optimizer.config.learning_rate, 1e-3);
119    }
120
121    #[test]
122    fn test_constrained_module_trait() {
123        let device = Default::default();
124        let linear = LinearConfig::new(2, 2).init::<TestBackend>(&device);
125        let constrained_linear = Constrained::<_, Euclidean>::new(linear);
126
127        let info = constrained_linear.get_manifold_info();
128        assert_eq!(info.get("manifold_type"), Some(&"Euclidean".to_string()));
129    }
130
131    #[test]
132    fn test_apply_constraints() {
133        let config = MultiManifoldOptimizerConfig::default();
134        let optimizer = MultiManifoldOptimizer::<TestBackend>::new(config);
135
136        let device = Default::default();
137        let linear = LinearConfig::new(2, 2).init::<TestBackend>(&device);
138        let constrained_linear = Constrained::<_, Euclidean>::new(linear);
139
140        // Test applying constraints
141        let result = optimizer.apply_constraints(constrained_linear);
142
143        // Should return the same module since we have a simplified implementation
144        assert_eq!(
145            result.get_manifold_info().get("manifold_type"),
146            Some(&"Euclidean".to_string())
147        );
148    }
149}
manopt_rs/optimizers/multiple.rs

manopt_rs/optimizers/
multiple.rs
