Skip to main content

nabled_sim/
sim.rs

1//! Semi-implicit rigid-body simulation step.
2
3use nabled_core::scalar::NabledReal;
4use nabled_dynamics::fd::forward_dynamics_with_config;
5use nabled_kinematics::fk::end_effector_pose;
6use nabled_linalg::geometry::Transform3;
7use nabled_linalg::lu::LuProviderScalar;
8use ndarray::{Array1, ArrayView1};
9
10use crate::SimError;
11use crate::context::RobotContext;
12
13/// Joint position and velocity state.
14#[derive(Debug, Clone, PartialEq)]
15pub struct SimState<T> {
16    pub q:  Array1<T>,
17    pub qd: Array1<T>,
18}
19
20impl<T: NabledReal> SimState<T> {
21    #[must_use]
22    pub fn new(q: Array1<T>, qd: Array1<T>) -> Self { Self { q, qd } }
23}
24
25/// Integration and logging options.
26#[derive(Debug, Clone, PartialEq)]
27pub struct SimConfig<T> {
28    pub dt:          T,
29    pub log_ee_pose: bool,
30}
31
32impl<T: NabledReal> SimConfig<T> {
33    #[must_use]
34    pub fn new(dt: T) -> Self { Self { dt, log_ee_pose: false } }
35}
36
37/// Output of one semi-implicit Euler step.
38#[derive(Debug, Clone, PartialEq)]
39pub struct SimStepResult<T> {
40    pub state:   SimState<T>,
41    pub qdd:     Array1<T>,
42    pub ee_pose: Option<Transform3<T>>,
43}
44
45/// Advance `(q, qd)` by one semi-implicit Euler step using forward dynamics.
46pub fn semi_implicit_step<T: NabledReal + Default + LuProviderScalar>(
47    ctx: &RobotContext<T>,
48    state: &SimState<T>,
49    tau: &ArrayView1<'_, T>,
50    config: &SimConfig<T>,
51) -> Result<SimStepResult<T>, SimError> {
52    ctx.validate()?;
53    if tau.len() != ctx.chain.num_joints() {
54        return Err(SimError::DimensionMismatch);
55    }
56    if state.q.len() != ctx.chain.num_joints() || state.qd.len() != ctx.chain.num_joints() {
57        return Err(SimError::DimensionMismatch);
58    }
59
60    let qdd = forward_dynamics_with_config(
61        &ctx.model,
62        &ctx.chain,
63        &state.q.view(),
64        &state.qd.view(),
65        tau,
66        &ctx.dynamics_config,
67    )?;
68
69    let mut qd = state.qd.clone();
70    qd += &(&qdd * config.dt);
71    let mut q = state.q.clone();
72    q += &(&qd * config.dt);
73
74    let ee_pose = if config.log_ee_pose { Some(end_effector_pose(&ctx.chain, &q)?) } else { None };
75
76    Ok(SimStepResult { state: SimState { q, qd }, qdd, ee_pose })
77}
78
79#[cfg(test)]
80mod tests {
81    use approx::assert_relative_eq;
82    use nabled_dynamics::DynamicsConfig;
83    use nabled_model::fixture::load_planar2r_json;
84    use ndarray::{Array1, arr1};
85
86    use super::*;
87    use crate::context::RobotContext;
88
89    #[test]
90    fn semi_implicit_step_advances_state() {
91        let fixture = load_planar2r_json().expect("fixture");
92        let ctx = RobotContext::new(
93            fixture.to_robot_model::<f64>().expect("model"),
94            fixture.to_chain_spec::<f64>().expect("chain"),
95            DynamicsConfig {
96                gravity: fixture.gravity.unwrap_or([0.0, -9.81, 0.0]),
97                ..DynamicsConfig::default()
98            },
99        );
100        let state = SimState::new(arr1(&[0.2, 0.4]), Array1::zeros(2));
101        let tau = Array1::zeros(2);
102        let config = SimConfig::new(0.01);
103        let result = semi_implicit_step(&ctx, &state, &tau.view(), &config).expect("step");
104        assert!(result.state.q[0].is_finite());
105        assert!(result.state.q[1].is_finite());
106        assert_relative_eq!(result.state.qd[0], result.qdd[0] * 0.01, epsilon = 1e-12);
107    }
108}