1use 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#[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#[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#[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
45pub 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}