1use crate::dynamics::Velocity;
2use crate::math::{AngVector, Real, Rot, Vect};
3use crate::utils;
4use bevy::prelude::*;
5use rapier::dynamics::{
6 AxesMask as RapierAxesMask, RigidBody as RapierRigidBody, RigidBodyVelocity,
7};
8
9pub use rapier::control::PdErrors;
10
11#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash, Reflect)]
16#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
17#[reflect(Default, PartialEq, Debug, Hash)]
18pub struct AxesMask(u8);
19
20bitflags::bitflags! {
21 impl AxesMask: u8 {
22 const LIN_X = 1 << 0;
24 const LIN_Y = 1 << 1;
26 #[cfg(feature = "dim3")]
28 const LIN_Z = 1 << 2;
29 #[cfg(feature = "dim3")]
31 const ANG_X = 1 << 3;
32 #[cfg(feature = "dim3")]
34 const ANG_Y = 1 << 4;
35 const ANG_Z = 1 << 5;
37 }
38}
39
40impl AxesMask {
41 #[cfg(feature = "dim2")]
43 pub const LIN_AXES: Self = Self::LIN_X.union(Self::LIN_Y);
44 #[cfg(feature = "dim3")]
46 pub const LIN_AXES: Self = Self::LIN_X.union(Self::LIN_Y).union(Self::LIN_Z);
47 #[cfg(feature = "dim2")]
49 pub const ANG_AXES: Self = Self::ANG_Z;
50 #[cfg(feature = "dim3")]
52 pub const ANG_AXES: Self = Self::ANG_X.union(Self::ANG_Y).union(Self::ANG_Z);
53}
54
55impl From<AxesMask> for RapierAxesMask {
56 fn from(axes: AxesMask) -> RapierAxesMask {
57 RapierAxesMask::from_bits_truncate(axes.bits())
58 }
59}
60
61impl From<RapierAxesMask> for AxesMask {
62 fn from(axes: RapierAxesMask) -> AxesMask {
63 AxesMask::from_bits_truncate(axes.bits())
64 }
65}
66
67#[derive(Copy, Clone, Debug, Default, PartialEq, Reflect)]
73#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
74#[reflect(Default, PartialEq, Debug)]
75pub struct PidTarget {
76 pub translation: Vect,
78 pub rotation: Rot,
80 pub linvel: Vect,
82 pub angvel: AngVector,
84}
85
86impl PidTarget {
87 pub fn new(translation: Vect, rotation: Rot) -> Self {
89 Self {
90 translation,
91 rotation,
92 ..Default::default()
93 }
94 }
95
96 pub fn from_translation(translation: Vect) -> Self {
98 Self {
99 translation,
100 ..Default::default()
101 }
102 }
103
104 pub fn with_velocities(mut self, linvel: Vect, angvel: AngVector) -> Self {
106 self.linvel = linvel;
107 self.angvel = angvel;
108 self
109 }
110
111 fn to_rapier(self) -> (rapier::math::Pose, RigidBodyVelocity<Real>) {
112 (
113 utils::pose_from(self.translation, self.rotation),
114 RigidBodyVelocity {
115 linvel: self.linvel,
116 angvel: self.angvel,
117 },
118 )
119 }
120}
121
122fn to_velocity(vels: RigidBodyVelocity<Real>) -> Velocity {
123 Velocity {
124 linear: vels.linvel,
125 angular: vels.angvel,
126 }
127}
128
129#[cfg(feature = "dim2")]
130fn splat_ang(value: Real) -> AngVector {
131 value
132}
133
134#[cfg(feature = "dim3")]
135fn splat_ang(value: Real) -> AngVector {
136 AngVector::splat(value)
137}
138
139#[derive(Copy, Clone, Debug, PartialEq, Component, Reflect)]
146#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
147#[reflect(Component, Default, PartialEq, Debug)]
148pub struct PdController {
149 pub target: PidTarget,
151 pub lin_kp: Vect,
156 pub lin_kd: Vect,
160 pub ang_kp: AngVector,
162 pub ang_kd: AngVector,
164 pub axes: AxesMask,
166}
167
168impl Default for PdController {
169 fn default() -> Self {
170 Self::from_raw(&rapier::control::PdController::default())
171 }
172}
173
174impl PdController {
175 pub fn new(kp: Real, kd: Real, axes: AxesMask) -> Self {
179 Self::from_raw(&rapier::control::PdController::new(kp, kd, axes.into()))
180 }
181
182 pub fn with_target(mut self, target: PidTarget) -> Self {
184 self.target = target;
185 self
186 }
187
188 pub fn from_raw(raw: &rapier::control::PdController) -> Self {
190 Self {
191 target: PidTarget::default(),
192 lin_kp: raw.lin_kp,
193 lin_kd: raw.lin_kd,
194 ang_kp: raw.ang_kp,
195 ang_kd: raw.ang_kd,
196 axes: raw.axes.into(),
197 }
198 }
199
200 pub fn to_raw(&self) -> rapier::control::PdController {
202 rapier::control::PdController {
203 lin_kp: self.lin_kp,
204 lin_kd: self.lin_kd,
205 ang_kp: self.ang_kp,
206 ang_kd: self.ang_kd,
207 axes: self.axes.into(),
208 }
209 }
210
211 pub fn correction(&self, pose_errors: &PdErrors, vel_errors: &PdErrors) -> Velocity {
213 to_velocity(self.to_raw().correction(pose_errors, vel_errors))
214 }
215
216 pub fn rigid_body_correction(&self, body: &RapierRigidBody) -> Velocity {
218 let (pose, vels) = self.target.to_rapier();
219 to_velocity(self.to_raw().rigid_body_correction(body, pose, vels))
220 }
221}
222
223#[derive(Copy, Clone, Debug, PartialEq, Component, Reflect)]
230#[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
231#[reflect(Component, Default, PartialEq, Debug)]
232pub struct PidController {
233 pub target: PidTarget,
235 pub lin_kp: Vect,
240 pub lin_ki: Vect,
242 pub lin_kd: Vect,
246 pub ang_kp: AngVector,
248 pub ang_ki: AngVector,
250 pub ang_kd: AngVector,
252 pub axes: AxesMask,
254 pub lin_integral: Vect,
256 pub ang_integral: AngVector,
258}
259
260impl Default for PidController {
261 fn default() -> Self {
262 Self::from_raw(&rapier::control::PidController::default())
263 }
264}
265
266impl PidController {
267 pub fn new(kp: Real, ki: Real, kd: Real, axes: AxesMask) -> Self {
271 Self::from_raw(&rapier::control::PidController::new(
272 kp,
273 ki,
274 kd,
275 axes.into(),
276 ))
277 }
278
279 pub fn with_target(mut self, target: PidTarget) -> Self {
281 self.target = target;
282 self
283 }
284
285 pub fn reset_integrals(&mut self) {
287 self.lin_integral = Vect::ZERO;
288 self.ang_integral = splat_ang(0.0);
289 }
290
291 pub fn from_raw(raw: &rapier::control::PidController) -> Self {
293 Self {
294 target: PidTarget::default(),
295 lin_kp: raw.pd.lin_kp,
296 lin_ki: raw.lin_ki,
297 lin_kd: raw.pd.lin_kd,
298 ang_kp: raw.pd.ang_kp,
299 ang_ki: raw.ang_ki,
300 ang_kd: raw.pd.ang_kd,
301 axes: raw.pd.axes.into(),
302 lin_integral: raw.lin_integral,
303 ang_integral: raw.ang_integral,
304 }
305 }
306
307 pub fn to_raw(&self) -> rapier::control::PidController {
309 rapier::control::PidController {
310 pd: rapier::control::PdController {
311 lin_kp: self.lin_kp,
312 lin_kd: self.lin_kd,
313 ang_kp: self.ang_kp,
314 ang_kd: self.ang_kd,
315 axes: self.axes.into(),
316 },
317 lin_integral: self.lin_integral,
318 ang_integral: self.ang_integral,
319 lin_ki: self.lin_ki,
320 ang_ki: self.ang_ki,
321 }
322 }
323
324 fn read_integrals(&mut self, raw: &rapier::control::PidController) {
325 self.lin_integral = raw.lin_integral;
326 self.ang_integral = raw.ang_integral;
327 }
328
329 pub fn correction(
332 &mut self,
333 dt: Real,
334 pose_errors: &PdErrors,
335 vel_errors: &PdErrors,
336 ) -> Velocity {
337 let mut raw = self.to_raw();
338 let result = raw.correction(dt, pose_errors, vel_errors);
339 self.read_integrals(&raw);
340 to_velocity(result)
341 }
342
343 pub fn rigid_body_correction(&mut self, dt: Real, body: &RapierRigidBody) -> Velocity {
346 let (pose, vels) = self.target.to_rapier();
347 let mut raw = self.to_raw();
348 let result = raw.rigid_body_correction(dt, body, pose, vels);
349 self.read_integrals(&raw);
350 to_velocity(result)
351 }
352}