wls_alloc/
wls.rs

1//! Encapsulated weighted least-squares control allocator.
2
3use nalgebra::{allocator::Allocator, Const, DefaultAllocator, DimMin, DimName, OMatrix, OVector};
4
5use crate::setup;
6use crate::solver;
7use crate::types::SolverStats;
8
9/// Stateful WLS control allocator: owns the static problem data and the
10/// warm-start solver state across solves.
11///
12/// Build once via [`ControlAllocator::new`] when the effectiveness matrix or
13/// weights change, then call [`solve`](Self::solve) on every control tick.
14/// The previous solution is automatically reused as the warm-start for the
15/// next solve.
16///
17/// Const generics:
18/// - `NU`: number of actuators
19/// - `NV`: number of pseudo-controls
20/// - `NC`: must equal `NU + NV` (compile-time checked)
21///
22/// # Example
23///
24/// ```
25/// use wls_alloc::wls::ControlAllocator;
26/// use wls_alloc::ExitCode;
27/// use nalgebra::{SMatrix, SVector, Vector4, Vector3};
28///
29/// // 3 pseudo-controls × 4 motors (e.g. roll/pitch/yaw mixer)
30/// #[rustfmt::skip]
31/// let g = SMatrix::<f32, 3, 4>::new(
32///     -0.5,  0.5,  1.0,   // motor 1
33///      0.5,  0.5, -1.0,   // motor 2
34///     -0.5, -0.5, -1.0,   // motor 3
35///      0.5, -0.5,  1.0,   // motor 4
36/// );
37/// let wv = Vector3::new(1.0_f32, 1.0_f32, 0.5_f32);
38/// let wu = Vector4::from_element(1.0_f32);
39///
40/// let mut alloc = ControlAllocator::<4, 3, 7>::new(&g, &wv, wu, 2e-9, 4e5);
41///
42/// let v = Vector3::new(0.1_f32, -0.2_f32, 0.05_f32);
43/// let ud = Vector4::from_element(0.5_f32);
44/// let umin = Vector4::from_element(0.0_f32);
45/// let umax = Vector4::from_element(1.0_f32);
46///
47/// let stats = alloc.solve(&v, &ud, &umin, &umax, 100);
48/// assert_eq!(stats.exit_code, ExitCode::Success);
49/// let u = alloc.solution();
50/// ```
51pub struct ControlAllocator<const NU: usize, const NV: usize, const NC: usize>
52where
53    Const<NC>: DimName,
54    Const<NU>: DimName,
55    Const<NV>: DimName,
56    DefaultAllocator: Allocator<Const<NC>, Const<NU>> + Allocator<Const<NU>> + Allocator<Const<NV>>,
57{
58    a: OMatrix<f32, Const<NC>, Const<NU>>,
59    wv: OVector<f32, Const<NV>>,
60    wu_norm: OVector<f32, Const<NU>>,
61    gamma: f32,
62    us: OVector<f32, Const<NU>>,
63    ws: [i8; NU],
64}
65
66impl<const NU: usize, const NV: usize, const NC: usize> ControlAllocator<NU, NV, NC>
67where
68    Const<NC>: DimName + DimMin<Const<NU>, Output = Const<NU>>,
69    Const<NU>: DimName,
70    Const<NV>: DimName,
71    DefaultAllocator: Allocator<Const<NC>, Const<NU>>
72        + Allocator<Const<NC>, Const<NC>>
73        + Allocator<Const<NU>, Const<NU>>
74        + Allocator<Const<NC>>
75        + Allocator<Const<NU>>
76        + Allocator<Const<NV>>,
77{
78    /// Build the allocator: factor the augmented matrix `A`, compute the
79    /// regularisation scalar `γ`, and normalise the actuator weights.
80    ///
81    /// `wu` is consumed so the in-place normalisation is fully internal — the
82    /// caller's data is never mutated through aliasing. The warm-start is
83    /// initialised to zero; use [`set_warmstart`](Self::set_warmstart) to
84    /// seed a non-zero initial guess before the first [`solve`](Self::solve).
85    pub fn new(
86        g: &OMatrix<f32, Const<NV>, Const<NU>>,
87        wv: &OVector<f32, Const<NV>>,
88        mut wu: OVector<f32, Const<NU>>,
89        theta: f32,
90        cond_bound: f32,
91    ) -> Self {
92        const { assert!(NC == NU + NV, "ControlAllocator requires NC == NU + NV") };
93        let (a, gamma) = setup::setup_a::<NU, NV, NC>(g, wv, &mut wu, theta, cond_bound);
94        Self {
95            a,
96            wv: wv.clone_owned(),
97            wu_norm: wu,
98            gamma,
99            us: OVector::zeros(),
100            ws: [0i8; NU],
101        }
102    }
103
104    /// Run one constrained least-squares solve.
105    ///
106    /// Builds the right-hand side `b` from `v` and `ud` using the stored
107    /// weights and `γ`, then runs the active-set solver continuing from the
108    /// current warm-start. The solution is left in the allocator and can be
109    /// read via [`solution`](Self::solution).
110    pub fn solve(
111        &mut self,
112        v: &OVector<f32, Const<NV>>,
113        ud: &OVector<f32, Const<NU>>,
114        umin: &OVector<f32, Const<NU>>,
115        umax: &OVector<f32, Const<NU>>,
116        imax: usize,
117    ) -> SolverStats {
118        let b = setup::setup_b::<NU, NV, NC>(v, ud, &self.wv, &self.wu_norm, self.gamma);
119        solver::solve::<NU, NV, NC>(&self.a, &b, umin, umax, &mut self.us, &mut self.ws, imax)
120    }
121
122    /// The current actuator solution (also the warm-start for the next solve).
123    pub fn solution(&self) -> &OVector<f32, Const<NU>> {
124        &self.us
125    }
126
127    /// The regularisation scalar `γ` chosen at construction time.
128    pub fn gamma(&self) -> f32 {
129        self.gamma
130    }
131
132    /// Seed the warm-start with an explicit initial guess and clear the
133    /// active set. Call before [`solve`](Self::solve) when starting a new
134    /// trajectory; otherwise the previous solution is reused automatically.
135    pub fn set_warmstart(&mut self, us: &OVector<f32, Const<NU>>) {
136        self.us = us.clone_owned();
137        self.ws = [0i8; NU];
138    }
139
140    /// Zero the warm-start solution and clear the active set.
141    pub fn reset_warmstart(&mut self) {
142        self.us = OVector::zeros();
143        self.ws = [0i8; NU];
144    }
145}
wls_alloc/wls.rs

wls_alloc/
wls.rs
