tinympc_rs/
constraint.rs

1use nalgebra::{RealField, SMatrix, convert};
2
3use crate::project::Project;
4
5/// A [`Constraint`] consists of a projection function and a set of associated slack and dual variables.
6pub struct Constraint<T, P: Project<T, N, H>, const N: usize, const H: usize> {
7    pub max_prim_residual: T,
8    pub max_dual_residual: T,
9    pub(crate) slac: SMatrix<T, N, H>,
10    pub(crate) dual: SMatrix<T, N, H>,
11    pub(crate) projector: P,
12}
13
14/// Type alias for a [`Constraint`] that dynamically dispatches its projection function
15pub type DynConstraint<'a, F, const N: usize, const H: usize> =
16    Constraint<F, &'a dyn Project<F, N, H>, N, H>;
17
18impl<T: RealField + Copy, const N: usize, const H: usize, P: Project<T, N, H>>
19    Constraint<T, P, N, H>
20{
21    /// Construct a new [`Constraint`] from the provided [`Project`] type.
22    pub fn new(projector: P) -> Self {
23        Self {
24            max_prim_residual: convert(1e9),
25            max_dual_residual: convert(1e9),
26            slac: SMatrix::zeros(),
27            dual: SMatrix::zeros(),
28            projector,
29        }
30    }
31
32    /// Get a mutable reference to the projector.
33    /// Allows for modifying or completely swapping the projector at runtime.
34    pub fn projector_mut(&mut self) -> &mut P {
35        &mut self.projector
36    }
37
38    /// Reset all internal state of this constraint
39    pub fn reset(&mut self) {
40        self.max_prim_residual = convert(1e9);
41        self.max_dual_residual = convert(1e9);
42        self.slac = SMatrix::zeros();
43        self.dual = SMatrix::zeros();
44    }
45
46    /// Constrains the set of points, and if `compute_residuals == true`, computes the maximum primal and dual residuals
47    #[inline(always)]
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        match compute_residuals {
57            true => self.constrain_calc_residuals(points, reference, scratch),
58            false => self.constrain_only(points, reference),
59        }
60    }
61
62    /// Constrains the set of points, and computes the maximum primal and dual residuals
63    #[inline(always)]
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        mut 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(&mut self.slac);
80            self.slac -= reference;
81        } else {
82            self.projector.project(&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, &mut 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    #[inline(always)]
101    #[profiling::function]
102    fn constrain_only(&mut self, points: &SMatrix<T, N, H>, reference: Option<&SMatrix<T, N, H>>) {
103        // Offset the slack variables by the reference before projecting
104        points.add_to(&self.dual, &mut self.slac);
105        if let Some(reference) = reference {
106            profiling::scope!("reference offset");
107            self.slac += reference;
108            self.projector.project(&mut self.slac);
109            self.slac -= reference;
110        } else {
111            self.projector.project(&mut self.slac);
112        }
113
114        // Update dual parameters
115        self.dual += points;
116        self.dual -= self.slac;
117    }
118
119    /// Add the cost associated with this constraints violation to a cost sum
120    #[inline(always)]
121    #[profiling::function]
122    pub(crate) fn add_cost(&self, cost: &mut SMatrix<T, N, H>) {
123        *cost += self.dual;
124        *cost -= self.slac;
125    }
126
127    /// Add the cost associated with this constraints violation to a cost sum
128    #[inline(always)]
129    #[profiling::function]
130    pub(crate) fn set_cost(&mut self, cost: &mut SMatrix<T, N, H>) {
131        self.dual.sub_to(&self.slac, cost);
132    }
133
134    /// Re-scale the dual variables for when the value of rho has changed
135    #[inline(always)]
136    #[profiling::function]
137    pub(crate) fn rescale_dual(&mut self, scalar: T) {
138        self.dual.scale_mut(scalar);
139    }
140}
tinympc_rs/constraint.rs

tinympc_rs/
constraint.rs
