manopt_rs/
constrained_module.rs

1//! Riemannian manifolds for constrained optimization.
2//!
3//! This module defines manifolds and their operations for Riemannian optimization.
4
5use std::{fmt::Debug, marker::PhantomData};
6
7use crate::prelude::*;
8
9use burn::{
10    module::{AutodiffModule, ModuleDisplay},
11    tensor::backend::AutodiffBackend,
12};
13
14#[derive(Clone, Debug)]
15pub struct Constrained<M, Man> {
16    module: M,
17    _manifold: PhantomData<Man>,
18}
19
20impl<B, M, Man> Module<B> for Constrained<M, Man>
21where
22    M: Module<B>,
23    B: Backend,
24    Man: Clone + Debug + Send,
25{
26    type Record = M::Record;
27
28    fn collect_devices(&self, devices: burn::module::Devices<B>) -> burn::module::Devices<B> {
29        self.module.collect_devices(devices)
30    }
31
32    fn fork(self, device: &B::Device) -> Self {
33        let module = self.module.fork(device);
34        Self {
35            module,
36            _manifold: PhantomData,
37        }
38    }
39
40    fn to_device(self, device: &B::Device) -> Self {
41        let module = self.module.to_device(device);
42        Self {
43            module,
44            _manifold: PhantomData,
45        }
46    }
47
48    fn visit<Visitor: burn::module::ModuleVisitor<B>>(&self, visitor: &mut Visitor) {
49        self.module.visit(visitor);
50    }
51
52    fn map<Mapper: burn::module::ModuleMapper<B>>(self, mapper: &mut Mapper) -> Self {
53        let module = self.module.map(mapper);
54        Self {
55            module,
56            _manifold: PhantomData,
57        }
58    }
59
60    fn load_record(self, record: Self::Record) -> Self {
61        let module = self.module.load_record(record);
62        Self {
63            module,
64            _manifold: PhantomData,
65        }
66    }
67
68    fn into_record(self) -> Self::Record {
69        self.module.into_record()
70    }
71}
72
73impl<B, M, Man> AutodiffModule<B> for Constrained<M, Man>
74where
75    M: AutodiffModule<B>,
76    B: AutodiffBackend,
77    Man: Clone + Debug + Send,
78{
79    type InnerModule = M::InnerModule;
80
81    fn valid(&self) -> Self::InnerModule {
82        self.module.valid()
83    }
84}
85
86impl<M, Man> burn::module::ModuleDisplayDefault for Constrained<M, Man>
87where
88    M: burn::module::ModuleDisplayDefault,
89    Man: Clone + Debug + Send,
90{
91    fn content(&self, content: burn::module::Content) -> Option<burn::module::Content> {
92        self.module.content(content)
93    }
94}
95
96impl<M, Man> ModuleDisplay for Constrained<M, Man>
97where
98    M: ModuleDisplay,
99    Man: Clone + Debug + Send,
100{
101    fn format(&self, passed_settings: burn::module::DisplaySettings) -> String {
102        format!("Constrained<{}>", self.module.format(passed_settings))
103    }
104}
105
106impl<M, Man> Constrained<M, Man> {
107    pub fn new(module: M) -> Self {
108        Self {
109            module,
110            _manifold: PhantomData,
111        }
112    }
113
114    /// Get a reference to the inner module
115    pub fn inner(&self) -> &M {
116        &self.module
117    }
118
119    /// Get a mutable reference to the inner module
120    pub fn inner_mut(&mut self) -> &mut M {
121        &mut self.module
122    }
123
124    /// Apply manifold projection to a tensor - requires explicit Backend type
125    pub fn project_tensor<B, const D: usize>(
126        &self,
127        point: Tensor<B, D>,
128        vector: Tensor<B, D>,
129    ) -> Tensor<B, D>
130    where
131        B: Backend,
132        M: Module<B>,
133        Man: Manifold<B> + Clone + Debug + Send,
134    {
135        Man::project(point, vector)
136    }
137
138    /// Apply manifold retraction to a tensor - requires explicit Backend type
139    pub fn retract_tensor<B, const D: usize>(
140        &self,
141        point: Tensor<B, D>,
142        direction: Tensor<B, D>,
143    ) -> Tensor<B, D>
144    where
145        B: Backend,
146        M: Module<B>,
147        Man: Manifold<B> + Clone + Debug + Send,
148    {
149        Man::retract(point, direction)
150    }
151
152    /// Convert Euclidean gradient to Riemannian gradient - requires explicit Backend type
153    pub fn euclidean_to_riemannian<B, const D: usize>(
154        &self,
155        point: Tensor<B, D>,
156        grad: Tensor<B, D>,
157    ) -> Tensor<B, D>
158    where
159        B: Backend,
160        M: Module<B>,
161        Man: Manifold<B> + Clone + Debug + Send,
162    {
163        Man::egrad2rgrad(point, grad)
164    }
165
166    /// Project point onto manifold - requires explicit Backend type
167    pub fn project_to_manifold<B, const D: usize>(&self, point: Tensor<B, D>) -> Tensor<B, D>
168    where
169        B: Backend,
170        M: Module<B>,
171        Man: Manifold<B> + Clone + Debug + Send,
172    {
173        Man::proj(point)
174    }
175
176    /// Get the manifold name
177    pub fn manifold_name<B>(&self) -> &'static str
178    where
179        B: Backend,
180        Man: Manifold<B>,
181    {
182        Man::name()
183    }
184}
185
186/// Trait for modules that have manifold constraints
187pub trait ConstrainedModule<B: Backend> {
188    /// Apply manifold constraints to all parameters in the module
189    #[must_use]
190    fn apply_manifold_constraints(self) -> Self;
191
192    /// Get information about the manifold constraints
193    fn get_manifold_info(&self) -> std::collections::HashMap<String, String>;
194
195    /// Check if this module has manifold constraints
196    fn has_manifold_constraints(&self) -> bool {
197        true
198    }
199}
200
201/// Blanket implementation for Constrained wrapper
202impl<B, M, Man> ConstrainedModule<B> for Constrained<M, Man>
203where
204    M: Module<B>,
205    B: Backend,
206    Man: Manifold<B> + Clone + Debug + Send,
207{
208    fn apply_manifold_constraints(self) -> Self {
209        self
210    }
211
212    fn get_manifold_info(&self) -> std::collections::HashMap<String, String> {
213        let mut info = std::collections::HashMap::new();
214        info.insert("manifold_type".to_string(), Man::name().to_string());
215        info
216    }
217}
manopt_rs/constrained_module.rs

manopt_rs/
constrained_module.rs
