Skip to main content

phyz_model/
model.rs

1//! Model definition — static description of a physical system.
2
3use crate::{Body, Joint, State};
4use phyz_math::{GRAVITY, SpatialInertia, SpatialTransform, Vec3};
5
6/// A motor actuator attached to a joint.
7#[derive(Debug, Clone)]
8pub struct Actuator {
9    pub name: String,
10    pub joint_name: String,
11    pub joint_idx: usize,
12    pub gear: f64,
13    pub ctrl_range: Option<[f64; 2]>,
14}
15
16/// Static model describing the topology and parameters of a physical system.
17#[derive(Debug, Clone)]
18pub struct Model {
19    /// Bodies in the kinematic tree (index 0 = first body, no world body).
20    pub bodies: Vec<Body>,
21    /// Joints connecting bodies.
22    pub joints: Vec<Joint>,
23    /// Gravity vector in world frame.
24    pub gravity: Vec3,
25    /// Integration timestep.
26    pub dt: f64,
27    /// Total number of position DOFs.
28    pub nq: usize,
29    /// Total number of velocity DOFs.
30    pub nv: usize,
31    /// Position DOF offset for each joint.
32    pub q_offsets: Vec<usize>,
33    /// Velocity DOF offset for each joint.
34    pub v_offsets: Vec<usize>,
35    /// Actuators (motors) acting on joints.
36    pub actuators: Vec<Actuator>,
37}
38
39impl Model {
40    /// Create a default empty state for this model.
41    pub fn default_state(&self) -> State {
42        State::new(self.nq, self.nv, self.bodies.len())
43    }
44
45    /// Number of bodies.
46    pub fn nbodies(&self) -> usize {
47        self.bodies.len()
48    }
49}
50
51/// Builder for constructing models.
52pub struct ModelBuilder {
53    bodies: Vec<Body>,
54    joints: Vec<Joint>,
55    gravity: Vec3,
56    dt: f64,
57}
58
59impl ModelBuilder {
60    /// Start building a new model.
61    pub fn new() -> Self {
62        Self {
63            bodies: Vec::new(),
64            joints: Vec::new(),
65            gravity: Vec3::new(0.0, 0.0, -GRAVITY),
66            dt: 0.001,
67        }
68    }
69
70    /// Set the gravity vector.
71    pub fn gravity(mut self, g: Vec3) -> Self {
72        self.gravity = g;
73        self
74    }
75
76    /// Set the timestep.
77    pub fn dt(mut self, dt: f64) -> Self {
78        self.dt = dt;
79        self
80    }
81
82    /// Add a body with a revolute joint attached to the given parent.
83    ///
84    /// `parent` is the index of the parent body, or -1 for world.
85    /// `parent_to_joint` is the transform from parent body frame to joint frame.
86    /// `inertia` is the body's spatial inertia in its own frame.
87    pub fn add_revolute_body(
88        mut self,
89        name: &str,
90        parent: i32,
91        parent_to_joint: SpatialTransform,
92        inertia: SpatialInertia,
93    ) -> Self {
94        let joint_idx = self.joints.len();
95        self.joints.push(Joint::revolute(parent_to_joint));
96        self.bodies.push(Body {
97            name: name.to_string(),
98            inertia,
99            parent,
100            joint_idx,
101            geometry: None,
102        });
103        self
104    }
105
106    /// Add a body with a prismatic joint attached to the given parent.
107    pub fn add_prismatic_body(
108        mut self,
109        name: &str,
110        parent: i32,
111        parent_to_joint: SpatialTransform,
112        axis: Vec3,
113        inertia: SpatialInertia,
114    ) -> Self {
115        let joint_idx = self.joints.len();
116        self.joints.push(Joint::prismatic(parent_to_joint, axis));
117        self.bodies.push(Body {
118            name: name.to_string(),
119            inertia,
120            parent,
121            joint_idx,
122            geometry: None,
123        });
124        self
125    }
126
127    /// Add a body with a spherical (ball) joint attached to the given parent.
128    pub fn add_spherical_body(
129        mut self,
130        name: &str,
131        parent: i32,
132        parent_to_joint: SpatialTransform,
133        inertia: SpatialInertia,
134    ) -> Self {
135        let joint_idx = self.joints.len();
136        self.joints.push(Joint::spherical(parent_to_joint));
137        self.bodies.push(Body {
138            name: name.to_string(),
139            inertia,
140            parent,
141            joint_idx,
142            geometry: None,
143        });
144        self
145    }
146
147    /// Add a body with a free joint (6 DOF) attached to the given parent.
148    pub fn add_free_body(
149        mut self,
150        name: &str,
151        parent: i32,
152        parent_to_joint: SpatialTransform,
153        inertia: SpatialInertia,
154    ) -> Self {
155        let joint_idx = self.joints.len();
156        self.joints.push(Joint::free(parent_to_joint));
157        self.bodies.push(Body {
158            name: name.to_string(),
159            inertia,
160            parent,
161            joint_idx,
162            geometry: None,
163        });
164        self
165    }
166
167    /// Add a body with a fixed joint (0 DOF) attached to the given parent.
168    pub fn add_fixed_body(
169        mut self,
170        name: &str,
171        parent: i32,
172        parent_to_joint: SpatialTransform,
173        inertia: SpatialInertia,
174    ) -> Self {
175        let joint_idx = self.joints.len();
176        self.joints.push(Joint::fixed(parent_to_joint));
177        self.bodies.push(Body {
178            name: name.to_string(),
179            inertia,
180            parent,
181            joint_idx,
182            geometry: None,
183        });
184        self
185    }
186
187    /// Add a generic joint and body.
188    pub fn add_body(
189        mut self,
190        name: &str,
191        parent: i32,
192        joint: Joint,
193        inertia: SpatialInertia,
194    ) -> Self {
195        let joint_idx = self.joints.len();
196        self.joints.push(joint);
197        self.bodies.push(Body {
198            name: name.to_string(),
199            inertia,
200            parent,
201            joint_idx,
202            geometry: None,
203        });
204        self
205    }
206
207    /// Add a free body with collision geometry (useful for dropping objects).
208    pub fn add_free_body_with_geometry(
209        mut self,
210        name: &str,
211        parent: i32,
212        parent_to_joint: SpatialTransform,
213        inertia: SpatialInertia,
214        geometry: crate::Body,
215    ) -> Self {
216        let joint_idx = self.joints.len();
217        self.joints.push(Joint::free(parent_to_joint));
218        self.bodies.push(Body {
219            name: name.to_string(),
220            inertia,
221            parent,
222            joint_idx,
223            geometry: geometry.geometry,
224        });
225        self
226    }
227
228    /// Build the model.
229    pub fn build(self) -> Model {
230        let mut nq = 0;
231        let mut nv = 0;
232        let mut q_offsets = Vec::new();
233        let mut v_offsets = Vec::new();
234
235        for joint in &self.joints {
236            q_offsets.push(nq);
237            v_offsets.push(nv);
238            nq += joint.ndof();
239            nv += joint.ndof();
240        }
241
242        Model {
243            bodies: self.bodies,
244            joints: self.joints,
245            gravity: self.gravity,
246            dt: self.dt,
247            nq,
248            nv,
249            q_offsets,
250            v_offsets,
251            actuators: Vec::new(),
252        }
253    }
254}
255
256impl Default for ModelBuilder {
257    fn default() -> Self {
258        Self::new()
259    }
260}