tinympc_rs/
constraint.rs

1use nalgebra::{RealField, SMatrix, convert};
2
3use crate::ProjectMulti;
4
5/// A [`Constraint`] consists of a projection function and a set of associated slack and dual variables.
6pub struct Constraint<T: RealField + Copy, P: ProjectMulti<T, N, H>, const N: usize, const H: usize>
7{
8    pub max_prim_residual: T,
9    pub max_dual_residual: T,
10    pub(crate) slac: SMatrix<T, N, H>,
11    pub(crate) dual: SMatrix<T, N, H>,
12    pub(crate) projector: P,
13}
14
15/// Type alias for a [`Constraint`] that dynamically dispatches its projection function
16pub type DynConstraint<'a, F, const N: usize, const H: usize> =
17    Constraint<F, &'a dyn ProjectMulti<F, N, H>, N, H>;
18
19impl<T: RealField + Copy, const N: usize, const H: usize, P: ProjectMulti<T, N, H>>
20    Constraint<T, P, N, H>
21{
22    /// Construct a new [`Constraint`] from the provided [`Project`] type.
23    pub fn new(projector: P) -> Self {
24        Self {
25            max_prim_residual: convert(1e9),
26            max_dual_residual: convert(1e9),
27            slac: SMatrix::zeros(),
28            dual: SMatrix::zeros(),
29            projector,
30        }
31    }
32
33    /// Get a mutable reference to the projector.
34    /// Allows for modifying or completely swapping the projector at runtime.
35    pub fn projector_mut(&mut self) -> &mut P {
36        &mut self.projector
37    }
38
39    /// Reset all internal state of this constraint
40    pub fn reset(&mut self) {
41        self.max_prim_residual = convert(1e9);
42        self.max_dual_residual = convert(1e9);
43        self.slac = SMatrix::zeros();
44        self.dual = SMatrix::zeros();
45    }
46
47    /// Constrains the set of points, and if `compute_residuals == true`, computes the maximum primal and dual residuals
48    #[profiling::function]
49    pub fn constrain(
50        &mut self,
51        compute_residuals: bool,
52        points: &SMatrix<T, N, H>,
53        reference: Option<&SMatrix<T, N, H>>,
54        scratch: &mut SMatrix<T, N, H>,
55    ) {
56        if compute_residuals {
57            self.constrain_calc_residuals(points, reference, scratch);
58        } else {
59            self.constrain_only(points, reference);
60        }
61    }
62
63    /// Constrains the set of points, and computes the maximum primal and dual residuals
64    #[profiling::function]
65    fn constrain_calc_residuals(
66        &mut self,
67        points: &SMatrix<T, N, H>,
68        reference: Option<&SMatrix<T, N, H>>,
69        scratch: &mut SMatrix<T, N, H>,
70    ) {
71        // Initialize with old slac variables for computing dual residual
72        scratch.copy_from(&self.slac);
73
74        // Offset the slack variables by the reference before projecting
75        points.add_to(&self.dual, &mut self.slac);
76        if let Some(reference) = reference {
77            profiling::scope!("reference offset");
78            self.slac += reference;
79            self.projector.project_multi(&mut self.slac);
80            self.slac -= reference;
81        } else {
82            self.projector.project_multi(&mut self.slac);
83        }
84
85        // Compute dual residual
86        *scratch -= &self.slac;
87        self.max_dual_residual = crate::util::frobenius_norm(scratch);
88
89        // Compute primal residual
90        points.sub_to(&self.slac, scratch);
91
92        // Update dual parameters
93        self.dual += &*scratch;
94
95        // Compute primal residual
96        self.max_prim_residual = crate::util::frobenius_norm(scratch);
97    }
98
99    /// Constrains the set of points
100    #[profiling::function]
101    fn constrain_only(&mut self, points: &SMatrix<T, N, H>, reference: Option<&SMatrix<T, N, H>>) {
102        // Offset the slack variables by the reference before projecting
103        points.add_to(&self.dual, &mut self.slac);
104        if let Some(reference) = reference {
105            profiling::scope!("reference offset");
106            self.slac += reference;
107            self.projector.project_multi(&mut self.slac);
108            self.slac -= reference;
109        } else {
110            self.projector.project_multi(&mut self.slac);
111        }
112
113        // Update dual parameters
114        self.dual += points;
115        self.dual -= &self.slac;
116    }
117
118    /// Add the cost associated with this constraints violation to a cost sum
119    #[profiling::function]
120    pub(crate) fn add_cost(&self, cost: &mut SMatrix<T, N, H>) {
121        *cost += &self.dual;
122        *cost -= &self.slac;
123    }
124
125    /// Add the cost associated with this constraints violation to a cost sum
126    #[profiling::function]
127    pub(crate) fn set_cost(&mut self, cost: &mut SMatrix<T, N, H>) {
128        self.dual.sub_to(&self.slac, cost);
129    }
130
131    /// Re-scale the dual variables for when the value of rho has changed
132    #[profiling::function]
133    pub(crate) fn rescale_dual(&mut self, scalar: T) {
134        self.dual.scale_mut(scalar);
135    }
136}
tinympc_rs/constraint.rs

tinympc_rs/
constraint.rs
