1use crate::{Body, Joint, State};
4use phyz_math::{GRAVITY, SpatialInertia, SpatialTransform, Vec3};
5
6#[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#[derive(Debug, Clone)]
18pub struct Model {
19 pub bodies: Vec<Body>,
21 pub joints: Vec<Joint>,
23 pub gravity: Vec3,
25 pub dt: f64,
27 pub nq: usize,
29 pub nv: usize,
31 pub q_offsets: Vec<usize>,
33 pub v_offsets: Vec<usize>,
35 pub actuators: Vec<Actuator>,
37}
38
39impl Model {
40 pub fn default_state(&self) -> State {
42 State::new(self.nq, self.nv, self.bodies.len())
43 }
44
45 pub fn nbodies(&self) -> usize {
47 self.bodies.len()
48 }
49}
50
51pub struct ModelBuilder {
53 bodies: Vec<Body>,
54 joints: Vec<Joint>,
55 gravity: Vec3,
56 dt: f64,
57}
58
59impl ModelBuilder {
60 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 pub fn gravity(mut self, g: Vec3) -> Self {
72 self.gravity = g;
73 self
74 }
75
76 pub fn dt(mut self, dt: f64) -> Self {
78 self.dt = dt;
79 self
80 }
81
82 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 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 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 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 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 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 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 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}