blob: af3dc17857f5ae07af1f1492cda3bfce79def3ac [file] [log] [blame]
Austin Schuh76534f32024-09-02 13:52:45 -07001#!/usr/bin/env python3
2import jax
3
4# casadi uses doubles. jax likes floats. We want to make sure the physics
5# matches really precisely, so force doubles for the tests.
6jax.config.update("jax_enable_x64", True)
Austin Schuh2a1abec2024-07-10 20:31:16 -07007
8import numpy
Austin Schuh76534f32024-09-02 13:52:45 -07009
10numpy.set_printoptions(precision=20)
Austin Schuh2a1abec2024-07-10 20:31:16 -070011import sys, os
12import casadi
Austin Schuhbd75c482024-08-18 00:03:51 -070013import scipy
Austin Schuh2a1abec2024-07-10 20:31:16 -070014from numpy.testing import assert_array_equal, assert_array_almost_equal
15import unittest
16
Austin Schuh27694fa2024-07-20 16:29:49 -070017from frc971.control_loops.swerve import bigcaster_dynamics
Austin Schuh2a1abec2024-07-10 20:31:16 -070018from frc971.control_loops.swerve import dynamics
Austin Schuh27694fa2024-07-20 16:29:49 -070019from frc971.control_loops.swerve import nocaster_dynamics
justinT2107d41ed2024-07-31 21:12:31 -070020from frc971.control_loops.swerve import physics_test_utils as utils
Austin Schuh76534f32024-09-02 13:52:45 -070021from frc971.control_loops.swerve import jax_dynamics
Austin Schuhb8b34be2024-07-14 16:06:19 -070022
23
Austin Schuh6927bc32024-07-14 17:24:56 -070024class TestSwervePhysics(unittest.TestCase):
Austin Schuh2a1abec2024-07-10 20:31:16 -070025 I = numpy.zeros((8, 1))
Austin Schuh76534f32024-09-02 13:52:45 -070026 coefficients = jax_dynamics.Coefficients()
Austin Schuh2a1abec2024-07-10 20:31:16 -070027
Austin Schuh6ea789e2024-07-27 13:45:53 -070028 def to_velocity_state(self, X):
29 return dynamics.to_velocity_state(X)
30
31 def swerve_full_dynamics(self, X, U, skip_compare=False):
32 X_velocity = self.to_velocity_state(X)
33 Xdot = self.position_swerve_full_dynamics(X, U)
Austin Schuh76534f32024-09-02 13:52:45 -070034
Austin Schuh6ea789e2024-07-27 13:45:53 -070035 if not skip_compare:
Austin Schuh76534f32024-09-02 13:52:45 -070036 Xdot_velocity = self.to_velocity_state(Xdot)
Austin Schuh6ea789e2024-07-27 13:45:53 -070037 velocity_physics = self.velocity_swerve_physics(X_velocity, U)
Austin Schuh76534f32024-09-02 13:52:45 -070038
Austin Schuh6ea789e2024-07-27 13:45:53 -070039 self.assertLess(
Austin Schuh76534f32024-09-02 13:52:45 -070040 numpy.linalg.norm(Xdot_velocity - velocity_physics),
Austin Schuh6ea789e2024-07-27 13:45:53 -070041 2e-2,
42 msg=
Austin Schuh76534f32024-09-02 13:52:45 -070043 f'Norm failed, full physics -> {X_velocity.T}, velocity physics -> {velocity_physics}, difference -> {velocity_physics - X_velocity}',
Austin Schuh6ea789e2024-07-27 13:45:53 -070044 )
45
Austin Schuh76534f32024-09-02 13:52:45 -070046 self.validate_dynamics_equality(X, U)
47
Austin Schuh6ea789e2024-07-27 13:45:53 -070048 return Xdot
49
Austin Schuh76534f32024-09-02 13:52:45 -070050 def validate_dynamics_equality(self, X, U):
51 """Tests that both the JAX code and casadi code produce identical answers.
52
53 Note:
54 If the symengine code has been updated, you likely need to update the JAX
55 by hand. We had trouble code generating it with good performance.
56 """
57 X_velocity = self.to_velocity_state(X)
58
59 Xdot = self.position_swerve_full_dynamics(X, U)
60 Xdot_jax = jax_dynamics.full_dynamics(self.coefficients, X[:, 0], U[:,
61 0])
62
63 self.assertLess(
64 numpy.linalg.norm(Xdot[:, 0] - Xdot_jax),
65 2e-8,
66 msg=
67 f'Xdot: {Xdot[:, 0]}, Xdot_jax: {Xdot_jax}, diff: {(Xdot[:, 0] - Xdot_jax)}',
68 )
69
70 velocity_physics = self.velocity_swerve_physics(X_velocity, U)
71 velocity_physics_jax = jax_dynamics.velocity_dynamics(
72 self.coefficients, X_velocity[:, 0], U[:, 0])
73
74 self.assertLess(
75 numpy.linalg.norm(velocity_physics[:, 0] - velocity_physics_jax),
76 2e-8,
77 msg=
78 f'Xdot: {velocity_physics[:, 0]}, Xdot_jax: {velocity_physics_jax}, diff: {(velocity_physics[:, 0] - velocity_physics_jax)}',
79 )
80
81 def wrap_and_validate(self, function, i):
82 """Wraps a function, and validates JAX and casadi agree.
83
84 We want to do it every time we check any intermediate, since the tests
85 are designed to test all the corner cases, but they don't all do it
86 through the main dynamics function above.
87 """
88 wrapped_fn = utils.wrap_module(function, i)
89
90 def do(X, U):
91 self.validate_dynamics_equality(X, U)
92 return wrapped_fn(X, U)
93
94 return do
95
Austin Schuh27694fa2024-07-20 16:29:49 -070096 def wrap(self, python_module):
Austin Schuh76534f32024-09-02 13:52:45 -070097 # Only update on change to avoid re-jiting things.
98 if self.coefficients.caster != python_module.CASTER:
99 self.coefficients = self.coefficients._replace(
100 caster=python_module.CASTER)
101
Austin Schuh6ea789e2024-07-27 13:45:53 -0700102 self.position_swerve_full_dynamics = utils.wrap(
103 python_module.swerve_full_dynamics)
104
105 evaluated_fn = python_module.velocity_swerve_physics(
106 casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
107 casadi.SX.sym("U", 8, 1))
108 self.velocity_swerve_physics = lambda X, U: numpy.array(
109 evaluated_fn(X, U))
110
Austin Schuhb8b34be2024-07-14 16:06:19 -0700111 self.contact_patch_velocity = [
Austin Schuh76534f32024-09-02 13:52:45 -0700112 self.wrap_and_validate(python_module.contact_patch_velocity, i)
Austin Schuh27694fa2024-07-20 16:29:49 -0700113 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700114 ]
115 self.wheel_ground_velocity = [
Austin Schuh76534f32024-09-02 13:52:45 -0700116 self.wrap_and_validate(python_module.wheel_ground_velocity, i)
Austin Schuh27694fa2024-07-20 16:29:49 -0700117 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700118 ]
119 self.wheel_slip_velocity = [
Austin Schuh76534f32024-09-02 13:52:45 -0700120 self.wrap_and_validate(python_module.wheel_slip_velocity, i)
Austin Schuh27694fa2024-07-20 16:29:49 -0700121 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700122 ]
justinT2107d41ed2024-07-31 21:12:31 -0700123 self.wheel_force = [
Austin Schuh76534f32024-09-02 13:52:45 -0700124 self.wrap_and_validate(python_module.wheel_force, i)
justinT2107d41ed2024-07-31 21:12:31 -0700125 for i in range(4)
126 ]
Austin Schuh76534f32024-09-02 13:52:45 -0700127 self.module_angular_accel = [
128 self.wrap_and_validate(python_module.module_angular_accel, i)
129 for i in range(4)
130 ]
131 self.F = [self.wrap_and_validate(python_module.F, i) for i in range(4)]
Austin Schuhb8b34be2024-07-14 16:06:19 -0700132 self.mounting_location = [
Austin Schuh76534f32024-09-02 13:52:45 -0700133 self.wrap_and_validate(python_module.mounting_location, i)
justinT2107d41ed2024-07-31 21:12:31 -0700134 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700135 ]
136
137 self.slip_angle = [
Austin Schuh76534f32024-09-02 13:52:45 -0700138 self.wrap_and_validate(python_module.slip_angle, i)
139 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700140 ]
141 self.slip_ratio = [
Austin Schuh76534f32024-09-02 13:52:45 -0700142 self.wrap_and_validate(python_module.slip_ratio, i)
143 for i in range(4)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700144 ]
Austin Schuh76534f32024-09-02 13:52:45 -0700145 self.Ms = [
146 self.wrap_and_validate(python_module.Ms, i) for i in range(4)
147 ]
Austin Schuh27694fa2024-07-20 16:29:49 -0700148
149 def setUp(self):
150 self.wrap(dynamics)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700151
152 def test_contact_patch_velocity(self):
153 """Tests that the contact patch velocity makes sense."""
154 for i in range(4):
justinT2107d41ed2024-07-31 21:12:31 -0700155 contact_patch_velocity = self.contact_patch_velocity[i]
156 wheel_ground_velocity = self.wheel_ground_velocity[i]
Austin Schuh2a1abec2024-07-10 20:31:16 -0700157
158 # No angular velocity should result in just linear motion.
159 for velocity in [
160 numpy.array([[1.5], [0.0]]),
161 numpy.array([[0.0], [1.0]]),
162 numpy.array([[-1.5], [0.0]]),
163 numpy.array([[0.0], [-1.0]]),
164 numpy.array([[2.0], [-1.7]]),
165 ]:
166 for theta in [0.0, 1.0, -1.0, 100.0]:
Austin Schuhb8b34be2024-07-14 16:06:19 -0700167 patch_velocity = contact_patch_velocity(
justinT2107d41ed2024-07-31 21:12:31 -0700168 utils.state_vector(velocity=velocity, theta=theta),
169 self.I)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700170
171 assert_array_equal(patch_velocity, velocity)
172
173 # Now, test that spinning the robot results in module velocities.
174 # We are assuming that each module is on a square robot.
175 module_center_of_mass_angle = i * numpy.pi / 2.0 + numpy.pi / 4.0
176 for theta in [-module_center_of_mass_angle, 0.0, 1.0, -1.0, 100.0]:
177 for omega in [0.65, -0.1]:
178 # Point each module to the center to make the math easier.
Austin Schuhb8b34be2024-07-14 16:06:19 -0700179 patch_velocity = contact_patch_velocity(
justinT2107d41ed2024-07-31 21:12:31 -0700180 utils.state_vector(
181 velocity=numpy.array([[0.0], [0.0]]),
182 theta=theta,
183 omega=omega,
184 module_angle=module_center_of_mass_angle,
185 ),
186 self.I,
187 )
Austin Schuh2a1abec2024-07-10 20:31:16 -0700188
189 assert_array_almost_equal(
190 patch_velocity,
Austin Schuh6927bc32024-07-14 17:24:56 -0700191 (dynamics.ROBOT_WIDTH / numpy.sqrt(2.0) -
justinT2107d41ed2024-07-31 21:12:31 -0700192 dynamics.CASTER) * omega * numpy.array([
193 [-numpy.sin(theta + module_center_of_mass_angle)],
194 [numpy.cos(theta + module_center_of_mass_angle)],
195 ]),
196 )
Austin Schuh2a1abec2024-07-10 20:31:16 -0700197
198 # Point the wheel along +x, rotate it by theta, then spin it.
199 # Confirm the velocities come out right.
Austin Schuhb8b34be2024-07-14 16:06:19 -0700200 patch_velocity = contact_patch_velocity(
justinT2107d41ed2024-07-31 21:12:31 -0700201 utils.state_vector(
Austin Schuhb8b34be2024-07-14 16:06:19 -0700202 velocity=numpy.array([[0.0], [0.0]]),
203 theta=-module_center_of_mass_angle,
204 module_omega=omega,
justinT2107d41ed2024-07-31 21:12:31 -0700205 module_angle=(theta + module_center_of_mass_angle),
206 ),
207 self.I,
208 )
Austin Schuh2a1abec2024-07-10 20:31:16 -0700209
210 assert_array_almost_equal(
justinT2107d41ed2024-07-31 21:12:31 -0700211 patch_velocity,
212 -dynamics.CASTER * omega *
213 numpy.array([[-numpy.sin(theta)], [numpy.cos(theta)]]),
214 )
Austin Schuh2a1abec2024-07-10 20:31:16 -0700215
216 # Now, test that the rotation back to wheel coordinates works.
217 # The easiest way to do this is to point the wheel in a direction,
218 # move in that direction, and confirm that there is no lateral velocity.
219 for robot_angle in [0.0, 1.0, -5.0]:
220 for module_angle in [0.0, 1.0, -5.0]:
221 wheel_patch_velocity = numpy.array(
222 wheel_ground_velocity(
justinT2107d41ed2024-07-31 21:12:31 -0700223 utils.state_vector(
224 velocity=numpy.array([
225 [numpy.cos(robot_angle + module_angle)],
226 [numpy.sin(robot_angle + module_angle)],
227 ]),
228 theta=robot_angle,
229 module_angle=module_angle,
230 ),
231 self.I,
232 ))
Austin Schuh2a1abec2024-07-10 20:31:16 -0700233
234 assert_array_almost_equal(wheel_patch_velocity,
235 numpy.array([[1], [0]]))
236
237 def test_slip_angle(self):
238 """Tests that the slip_angle calculation works."""
239 velocity = numpy.array([[1.5], [0.0]])
240
241 for i in range(4):
Austin Schuh27694fa2024-07-20 16:29:49 -0700242 for wrap in range(-1, 2):
Austin Schuh2a1abec2024-07-10 20:31:16 -0700243 for theta in [0.0, 0.6, -0.4]:
244 module_angle = numpy.pi * wrap + theta
245
Austin Schuhbd75c482024-08-18 00:03:51 -0700246 # We have redefined the angle to be the softened sin of the angle.
Austin Schuh27694fa2024-07-20 16:29:49 -0700247 # That way, when the module flips directions, the slip angle also flips
248 # directions to keep it stable.
justinT2107d41ed2024-07-31 21:12:31 -0700249 computed_angle = self.slip_angle[i](
250 utils.state_vector(velocity=velocity,
Austin Schuhbd75c482024-08-18 00:03:51 -0700251 module_angle=module_angle),
justinT2107d41ed2024-07-31 21:12:31 -0700252 self.I,
253 )[0, 0]
Austin Schuh2a1abec2024-07-10 20:31:16 -0700254
Austin Schuhbd75c482024-08-18 00:03:51 -0700255 # Compute out the expected value directly to confirm we got it right.
256 loggain = 20.0
257 vy = 1.5 * numpy.sin(-module_angle)
258 vx = 1.5 * numpy.cos(-module_angle)
259 expected = numpy.sin(-numpy.arctan2(
260 vy,
261 scipy.special.logsumexp([1.0, abs(vx) * loggain]) /
262 loggain))
Austin Schuh27694fa2024-07-20 16:29:49 -0700263
Austin Schuh76534f32024-09-02 13:52:45 -0700264 jax_expected = jax.numpy.sin(
265 -jax_dynamics.soft_atan2(vy, vx))
266
267 self.assertAlmostEqual(
268 expected,
269 jax_expected,
270 msg=f"Trying wrap {wrap} theta {theta}",
271 )
272
Austin Schuh27694fa2024-07-20 16:29:49 -0700273 self.assertAlmostEqual(
274 expected,
275 computed_angle,
justinT2107d41ed2024-07-31 21:12:31 -0700276 msg=f"Trying wrap {wrap} theta {theta}",
277 )
Austin Schuh2a1abec2024-07-10 20:31:16 -0700278
Austin Schuhb8b34be2024-07-14 16:06:19 -0700279 def test_wheel_torque(self):
280 """Tests that the per module self aligning forces have the right signs."""
Austin Schuh27694fa2024-07-20 16:29:49 -0700281 # Point all the modules in a little bit.
justinT2107d41ed2024-07-31 21:12:31 -0700282 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700283 velocity=numpy.array([[1.0], [0.0]]),
284 module_angles=[-0.001, -0.001, 0.001, 0.001],
285 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700286 xdot_equal = self.swerve_full_dynamics(X, self.I)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700287
288 self.assertGreater(xdot_equal[2, 0], 0.0)
289 self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
290 self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)
291
292 self.assertGreater(xdot_equal[6, 0], 0.0)
293 self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
294 self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)
295
296 self.assertLess(xdot_equal[10, 0], 0.0)
297 self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
298 self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)
299
300 self.assertLess(xdot_equal[14, 0], 0.0)
301 self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
302 self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)
303
304 # Shouldn't be spinning.
305 self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)
306
Austin Schuh27694fa2024-07-20 16:29:49 -0700307 # Now, make the bot want to go left by going to the other side.
Austin Schuhb8b34be2024-07-14 16:06:19 -0700308 # The wheels will be going too fast based on our calcs, so they should be decelerating.
justinT2107d41ed2024-07-31 21:12:31 -0700309 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700310 velocity=numpy.array([[1.0], [0.0]]),
311 module_angles=[0.01, 0.01, 0.01, 0.01],
312 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700313 xdot_left = self.swerve_full_dynamics(X, self.I)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700314
Austin Schuh27694fa2024-07-20 16:29:49 -0700315 self.assertLess(xdot_left[2, 0], -0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700316 self.assertLess(xdot_left[3, 0], 0.0)
317 self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)
318
Austin Schuh27694fa2024-07-20 16:29:49 -0700319 self.assertLess(xdot_left[6, 0], -0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700320 self.assertLess(xdot_left[7, 0], 0.0)
321 self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)
322
Austin Schuh27694fa2024-07-20 16:29:49 -0700323 self.assertLess(xdot_left[10, 0], -0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700324 self.assertLess(xdot_left[11, 0], 0.0)
325 self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)
326
Austin Schuh27694fa2024-07-20 16:29:49 -0700327 self.assertLess(xdot_left[14, 0], -0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700328 self.assertLess(xdot_left[15, 0], 0.0)
329 self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)
330
331 self.assertGreater(xdot_left[19, 0], 0.0001)
Austin Schuh27694fa2024-07-20 16:29:49 -0700332 self.assertGreater(xdot_left[20, 0], 0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700333 # Shouldn't be spinning.
334 self.assertAlmostEqual(xdot_left[21, 0], 0.0)
335
336 # And now do it to the right too.
justinT2107d41ed2024-07-31 21:12:31 -0700337 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700338 velocity=numpy.array([[1.0], [0.0]]),
339 module_angles=[-0.01, -0.01, -0.01, -0.01],
340 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700341 xdot_right = self.swerve_full_dynamics(X, self.I)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700342
Austin Schuh27694fa2024-07-20 16:29:49 -0700343 self.assertGreater(xdot_right[2, 0], 0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700344 self.assertLess(xdot_right[3, 0], 0.0)
345 self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)
346
Austin Schuh27694fa2024-07-20 16:29:49 -0700347 self.assertGreater(xdot_right[6, 0], 0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700348 self.assertLess(xdot_right[7, 0], 0.0)
349 self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)
350
Austin Schuh27694fa2024-07-20 16:29:49 -0700351 self.assertGreater(xdot_right[10, 0], 0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700352 self.assertLess(xdot_right[11, 0], 0.0)
353 self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)
354
Austin Schuh27694fa2024-07-20 16:29:49 -0700355 self.assertGreater(xdot_right[14, 0], 0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700356 self.assertLess(xdot_right[15, 0], 0.0)
357 self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)
358
359 self.assertGreater(xdot_right[19, 0], 0.0001)
Austin Schuh27694fa2024-07-20 16:29:49 -0700360 self.assertLess(xdot_right[20, 0], -0.05)
361 # Shouldn't be spinning.
362 self.assertAlmostEqual(xdot_right[21, 0], 0.0)
363
364 def test_wheel_torque_backwards_nocaster(self):
365 """Tests that the per module self aligning forces have the right signs when going backwards."""
366 self.wrap(nocaster_dynamics)
367 # Point all the modules in a little bit, going backwards.
justinT2107d41ed2024-07-31 21:12:31 -0700368 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700369 velocity=numpy.array([[1.0], [0.0]]),
370 module_angles=[
371 numpy.pi - 0.001,
372 numpy.pi - 0.001,
373 numpy.pi + 0.001,
374 numpy.pi + 0.001,
375 ],
376 drive_wheel_velocity=-1,
377 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700378 xdot_equal = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700379
380 self.assertGreater(xdot_equal[2, 0], 0.0, msg="Steering backwards")
381 self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
382 self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)
383
384 self.assertGreater(xdot_equal[6, 0], 0.0, msg="Steering backwards")
385 self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
386 self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)
387
388 self.assertLess(xdot_equal[10, 0], 0.0, msg="Steering backwards")
389 self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
390 self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)
391
392 self.assertLess(xdot_equal[14, 0], 0.0, msg="Steering backwards")
393 self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
394 self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)
395
396 # Shouldn't be spinning.
397 self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)
398
399 # Now, make the bot want to go left by going to the other side.
400 # The wheels will be going too fast based on our calcs, so they should be decelerating.
justinT2107d41ed2024-07-31 21:12:31 -0700401 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700402 velocity=numpy.array([[1.0], [0.0]]),
403 module_angles=[numpy.pi + 0.01] * 4,
404 drive_wheel_velocity=-1,
405 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700406 xdot_left = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700407
408 self.assertLess(xdot_left[2, 0], -0.05)
409 self.assertGreater(xdot_left[3, 0], 0.0)
410 self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)
411
412 self.assertLess(xdot_left[6, 0], -0.05)
413 self.assertGreater(xdot_left[7, 0], 0.0)
414 self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)
415
416 self.assertLess(xdot_left[10, 0], -0.05)
417 self.assertGreater(xdot_left[11, 0], 0.0)
418 self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)
419
420 self.assertLess(xdot_left[14, 0], -0.05)
421 self.assertGreater(xdot_left[15, 0], 0.0)
422 self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)
423
424 self.assertGreater(xdot_left[19, 0], 0.0001)
425 self.assertGreater(xdot_left[20, 0], 0.05)
426 # Shouldn't be spinning.
427 self.assertAlmostEqual(xdot_left[21, 0], 0.0)
428
429 # And now do it to the right too.
justinT2107d41ed2024-07-31 21:12:31 -0700430 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700431 velocity=numpy.array([[1.0], [0.0]]),
432 drive_wheel_velocity=-1,
433 module_angles=[-0.01 + numpy.pi] * 4,
434 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700435 xdot_right = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700436
437 self.assertGreater(xdot_right[2, 0], 0.05)
438 self.assertGreater(xdot_right[3, 0], 0.0)
439 self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)
440
441 self.assertGreater(xdot_right[6, 0], 0.05)
442 self.assertGreater(xdot_right[7, 0], 0.0)
443 self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)
444
445 self.assertGreater(xdot_right[10, 0], 0.05)
446 self.assertGreater(xdot_right[11, 0], 0.0)
447 self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)
448
449 self.assertGreater(xdot_right[14, 0], 0.05)
450 self.assertGreater(xdot_right[15, 0], 0.0)
451 self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)
452
453 self.assertGreater(xdot_right[19, 0], 0.0001)
454 self.assertLess(xdot_right[20, 0], -0.05)
455 # Shouldn't be spinning.
456 self.assertAlmostEqual(xdot_right[21, 0], 0.0)
457
458 def test_wheel_torque_backwards_caster(self):
459 """Tests that the per module self aligning forces have the right signs when going backwards with a lot of caster."""
460 self.wrap(bigcaster_dynamics)
461 # Point all the modules in a little bit, going backwards.
justinT2107d41ed2024-07-31 21:12:31 -0700462 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700463 velocity=numpy.array([[1.0], [0.0]]),
464 module_angles=[
465 numpy.pi - 0.001,
466 numpy.pi - 0.001,
467 numpy.pi + 0.001,
468 numpy.pi + 0.001,
469 ],
470 drive_wheel_velocity=-1,
471 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700472 xdot_equal = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700473
474 self.assertLess(xdot_equal[2, 0], 0.0, msg="Steering backwards")
475 self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
476 self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)
477
478 self.assertLess(xdot_equal[6, 0], 0.0, msg="Steering backwards")
479 self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
480 self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)
481
482 self.assertGreater(xdot_equal[10, 0], 0.0, msg="Steering backwards")
483 self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
484 self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)
485
486 self.assertGreater(xdot_equal[14, 0], 0.0, msg="Steering backwards")
487 self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
488 self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)
489
490 # Shouldn't be spinning.
491 self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)
492
493 # Now, make the bot want to go left by going to the other side.
494 # The wheels will be going too fast based on our calcs, so they should be decelerating.
justinT2107d41ed2024-07-31 21:12:31 -0700495 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700496 velocity=numpy.array([[1.0], [0.0]]),
497 module_angles=[numpy.pi + 0.01] * 4,
498 drive_wheel_velocity=-1,
499 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700500 xdot_left = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700501
502 self.assertGreater(xdot_left[2, 0], -0.05)
503 self.assertGreater(xdot_left[3, 0], 0.0)
504 self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)
505
506 self.assertGreater(xdot_left[6, 0], -0.05)
507 self.assertGreater(xdot_left[7, 0], 0.0)
508 self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)
509
510 self.assertGreater(xdot_left[10, 0], -0.05)
511 self.assertGreater(xdot_left[11, 0], 0.0)
512 self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)
513
514 self.assertGreater(xdot_left[14, 0], -0.05)
515 self.assertGreater(xdot_left[15, 0], 0.0)
516 self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)
517
518 self.assertGreater(xdot_left[19, 0], 0.0001)
519 self.assertGreater(xdot_left[20, 0], 0.05)
520 # Shouldn't be spinning.
521 self.assertAlmostEqual(xdot_left[21, 0], 0.0)
522
523 # And now do it to the right too.
justinT2107d41ed2024-07-31 21:12:31 -0700524 X = utils.state_vector(
Austin Schuh27694fa2024-07-20 16:29:49 -0700525 velocity=numpy.array([[1.0], [0.0]]),
526 drive_wheel_velocity=-1,
527 module_angles=[-0.01 + numpy.pi] * 4,
528 )
Austin Schuh6ea789e2024-07-27 13:45:53 -0700529 xdot_right = self.swerve_full_dynamics(X, self.I)
Austin Schuh27694fa2024-07-20 16:29:49 -0700530
531 self.assertLess(xdot_right[2, 0], 0.05)
532 self.assertGreater(xdot_right[3, 0], 0.0)
533 self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)
534
535 self.assertLess(xdot_right[6, 0], 0.05)
536 self.assertGreater(xdot_right[7, 0], 0.0)
537 self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)
538
539 self.assertLess(xdot_right[10, 0], 0.05)
540 self.assertGreater(xdot_right[11, 0], 0.0)
541 self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)
542
543 self.assertLess(xdot_right[14, 0], 0.05)
544 self.assertGreater(xdot_right[15, 0], 0.0)
545 self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)
546
547 self.assertGreater(xdot_right[19, 0], 0.0001)
548 self.assertLess(xdot_right[20, 0], -0.05)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700549 # Shouldn't be spinning.
550 self.assertAlmostEqual(xdot_right[21, 0], 0.0)
551
Austin Schuh2a1abec2024-07-10 20:31:16 -0700552 def test_wheel_forces(self):
553 """Tests that the per module forces have the right signs."""
554 for i in range(4):
justinT2107d41ed2024-07-31 21:12:31 -0700555 wheel_force = self.wheel_force[i]
Austin Schuhb8b34be2024-07-14 16:06:19 -0700556
justinT2107d41ed2024-07-31 21:12:31 -0700557 X = utils.state_vector()
Austin Schuhb8b34be2024-07-14 16:06:19 -0700558 robot_equal = wheel_force(X, self.I)
Austin Schuh6ea789e2024-07-27 13:45:53 -0700559 xdot_equal = self.swerve_full_dynamics(X, self.I)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700560 self.assertEqual(robot_equal[0, 0], 0.0)
561 self.assertEqual(robot_equal[1, 0], 0.0)
562 self.assertEqual(xdot_equal[2 + 4 * i], 0.0)
563 self.assertEqual(xdot_equal[3 + 4 * i], 0.0)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700564
565 # Robot is moving faster than the wheels, it should decelerate.
justinT2107d41ed2024-07-31 21:12:31 -0700566 X = utils.state_vector(dx=0.01)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700567 robot_faster = wheel_force(X, self.I)
Austin Schuh6ea789e2024-07-27 13:45:53 -0700568 xdot_faster = self.swerve_full_dynamics(X,
569 self.I,
570 skip_compare=True)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700571 self.assertLess(robot_faster[0, 0], -0.1)
572 self.assertEqual(robot_faster[1, 0], 0.0)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700573 self.assertGreater(xdot_faster[3 + 4 * i], 0.0)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700574
575 # Robot is now going slower than the wheels. It should accelerate.
justinT2107d41ed2024-07-31 21:12:31 -0700576 X = utils.state_vector(dx=-0.01)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700577 robot_slower = wheel_force(X, self.I)
Austin Schuh6ea789e2024-07-27 13:45:53 -0700578 xdot_slower = self.swerve_full_dynamics(X,
579 self.I,
580 skip_compare=True)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700581 self.assertGreater(robot_slower[0, 0], 0.1)
582 self.assertEqual(robot_slower[1, 0], 0.0)
Austin Schuhb8b34be2024-07-14 16:06:19 -0700583 self.assertLess(xdot_slower[3 + 4 * i], 0.0)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700584
585 # Positive lateral velocity -> negative force.
justinT2107d41ed2024-07-31 21:12:31 -0700586 robot_left = wheel_force(utils.state_vector(dy=0.01), self.I)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700587 self.assertEqual(robot_left[0, 0], 0.0)
588 self.assertLess(robot_left[1, 0], -0.1)
589
590 # Negative lateral velocity -> positive force.
justinT2107d41ed2024-07-31 21:12:31 -0700591 robot_right = wheel_force(utils.state_vector(dy=-0.01), self.I)
Austin Schuh2a1abec2024-07-10 20:31:16 -0700592 self.assertEqual(robot_right[0, 0], 0.0)
593 self.assertGreater(robot_right[1, 0], 0.1)
594
Austin Schuh3233d7f2024-07-27 17:51:52 -0700595 def test_steer_accel(self):
596 """Tests that applying a steer torque accelerates the steer reasonably."""
597
598 for velocity in [
599 numpy.array([[0.0], [0.0]]),
600 numpy.array([[1.0], [0.0]]),
601 numpy.array([[2.0], [0.0]])
602 ]:
603 module_angles = [0.0] * 4
604
605 X = utils.state_vector(
606 velocity=velocity,
607 omega=0.0,
608 module_angles=module_angles,
609 )
610 steer_I = numpy.array([[1.0], [0.0]] * 4)
611 xdot = self.swerve_full_dynamics(X, steer_I)
612
613 self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS0, 0],
614 1.4,
615 places=0)
616 self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS1, 0],
617 1.4,
618 places=0)
619 self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS2, 0],
620 1.4,
621 places=0)
622 self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS3, 0],
623 1.4,
624 places=0)
625
626 def test_drive_accel(self):
627 """Tests that applying a drive torque accelerates the drive reasonably."""
628
629 for velocity in [
630 numpy.array([[0.01], [0.0]]),
631 numpy.array([[1.0], [0.0]]),
632 numpy.array([[2.0], [0.0]])
633 ]:
634 module_angles = [0.0] * 4
635
636 X = utils.state_vector(
637 velocity=velocity,
638 omega=0.0,
639 module_angles=module_angles,
640 )
641 steer_I = numpy.array([[0.0], [100.0]] * 4)
642 # Since the wheels are spinning at the same speed as the ground, there is no force accelerating the robot. Which means there is not steering moment. The two physics diverge pretty heavily.
643 xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)
644
645 self.assertGreater(xdot[dynamics.STATE_OMEGAD0, 0], 100.0)
646 self.assertGreater(xdot[dynamics.STATE_OMEGAD1, 0], 100.0)
647 self.assertGreater(xdot[dynamics.STATE_OMEGAD2, 0], 100.0)
648 self.assertGreater(xdot[dynamics.STATE_OMEGAD3, 0], 100.0)
649
650 X = utils.state_vector(
651 velocity=velocity,
652 omega=0.0,
653 module_angles=module_angles,
654 )
655 steer_I = numpy.array([[0.0], [-100.0]] * 4)
656 xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)
657
658 self.assertLess(xdot[dynamics.STATE_OMEGAD0, 0], -100.0)
659 self.assertLess(xdot[dynamics.STATE_OMEGAD1, 0], -100.0)
660 self.assertLess(xdot[dynamics.STATE_OMEGAD2, 0], -100.0)
661 self.assertLess(xdot[dynamics.STATE_OMEGAD3, 0], -100.0)
662
Austin Schuh98fbbe82024-08-18 23:07:26 -0700663 def test_steer_coupling(self):
664 """Tests that the steer coupling factor cancels out steer coupling torque."""
665 steer_I = numpy.array(
666 [[dynamics.STEER_CURRENT_COUPLING_FACTOR * 10.0], [10.0]] * 4)
667
668 X = utils.state_vector(
669 velocity=numpy.array([[0.0], [0.0]]),
670 omega=0.0,
671 )
672 X_velocity = self.to_velocity_state(X)
673 Xdot = self.velocity_swerve_physics(X_velocity, steer_I)
674
675 self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS0, 0], 0.0)
676 self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS1, 0], 0.0)
677 self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS2, 0], 0.0)
678 self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS3, 0], 0.0)
679
Austin Schuh2a1abec2024-07-10 20:31:16 -0700680
justinT2107d41ed2024-07-31 21:12:31 -0700681if __name__ == "__main__":
Austin Schuh2a1abec2024-07-10 20:31:16 -0700682 unittest.main()