blob: 40d25b4a16147a652d3f43f0521125861a1fc172 [file] [log] [blame]
justinT21446e4f62024-06-16 22:36:10 -07001#include <symengine/add.h>
2#include <symengine/matrix.h>
3#include <symengine/number.h>
4#include <symengine/printers.h>
5#include <symengine/real_double.h>
6#include <symengine/simplify.h>
7#include <symengine/solve.h>
8#include <symengine/symbol.h>
9
10#include <array>
11#include <cmath>
12#include <numbers>
13#include <utility>
14
Austin Schuh99f7c6a2024-06-25 22:07:44 -070015#include "absl/flags/flag.h"
16#include "absl/log/check.h"
17#include "absl/log/log.h"
justinT21446e4f62024-06-16 22:36:10 -070018#include "absl/strings/str_format.h"
19#include "absl/strings/str_join.h"
20#include "absl/strings/str_replace.h"
21#include "absl/strings/substitute.h"
justinT21446e4f62024-06-16 22:36:10 -070022
23#include "aos/init.h"
24#include "aos/util/file.h"
25#include "frc971/control_loops/swerve/motors.h"
26
Austin Schuh99f7c6a2024-06-25 22:07:44 -070027ABSL_FLAG(std::string, output_base, "",
28 "Path to strip off the front of the output paths.");
29ABSL_FLAG(std::string, cc_output_path, "",
30 "Path to write generated cc code to");
31ABSL_FLAG(std::string, h_output_path, "",
32 "Path to write generated header code to");
Austin Schuh99f7c6a2024-06-25 22:07:44 -070033ABSL_FLAG(std::string, casadi_py_output_path, "",
34 "Path to write casadi generated py code to");
Austin Schuh27694fa2024-07-20 16:29:49 -070035ABSL_FLAG(double, caster, 0.01, "Caster in meters for the module.");
justinT21446e4f62024-06-16 22:36:10 -070036
Austin Schuh99f7c6a2024-06-25 22:07:44 -070037ABSL_FLAG(bool, symbolic, false, "If true, write everything out symbolically.");
Austin Schuh5371c802024-09-02 13:18:48 -070038ABSL_FLAG(bool, function, true, "If true, make soft_atan2 a function.");
justinT21446e4f62024-06-16 22:36:10 -070039
justinT21942892b2024-07-02 22:33:50 -070040using SymEngine::abs;
justinT21446e4f62024-06-16 22:36:10 -070041using SymEngine::add;
42using SymEngine::atan2;
43using SymEngine::Basic;
44using SymEngine::ccode;
45using SymEngine::cos;
46using SymEngine::DenseMatrix;
47using SymEngine::div;
Austin Schuh78a1b312024-08-18 17:21:34 -070048using SymEngine::exp;
justinT21446e4f62024-06-16 22:36:10 -070049using SymEngine::Inf;
50using SymEngine::integer;
51using SymEngine::map_basic_basic;
52using SymEngine::minus_one;
53using SymEngine::neg;
54using SymEngine::NegInf;
55using SymEngine::pow;
56using SymEngine::RCP;
57using SymEngine::real_double;
58using SymEngine::RealDouble;
59using SymEngine::Set;
60using SymEngine::simplify;
61using SymEngine::sin;
62using SymEngine::solve;
63using SymEngine::symbol;
64using SymEngine::Symbol;
65
66namespace frc971::control_loops::swerve {
67
68// State per module.
69struct Module {
Austin Schuh6ea789e2024-07-27 13:45:53 -070070 DenseMatrix mounting_location;
71
justinT21446e4f62024-06-16 22:36:10 -070072 RCP<const Symbol> Is;
73
74 RCP<const Symbol> Id;
75
76 RCP<const Symbol> thetas;
77 RCP<const Symbol> omegas;
justinT21446e4f62024-06-16 22:36:10 -070078
justinT21446e4f62024-06-16 22:36:10 -070079 RCP<const Symbol> omegad;
justinT21446e4f62024-06-16 22:36:10 -070080
Austin Schuh2a1abec2024-07-10 20:31:16 -070081 DenseMatrix contact_patch_velocity;
Austin Schuhb67a38f2024-07-04 13:48:38 -070082 DenseMatrix wheel_ground_velocity;
Austin Schuhb8b34be2024-07-14 16:06:19 -070083 DenseMatrix wheel_slip_velocity;
Austin Schuhb67a38f2024-07-04 13:48:38 -070084 RCP<const Basic> slip_angle;
85 RCP<const Basic> slip_ratio;
86
Austin Schuhb8b34be2024-07-14 16:06:19 -070087 RCP<const Basic> Ms;
Austin Schuhb67a38f2024-07-04 13:48:38 -070088 RCP<const Basic> Fwy;
89
Austin Schuh6ea789e2024-07-27 13:45:53 -070090 struct Full {
91 RCP<const Basic> Fwx;
92 DenseMatrix F;
93
94 RCP<const Basic> torque;
95
96 RCP<const Basic> alphas_eqn;
97 RCP<const Basic> alphad_eqn;
98 } full;
99
100 struct Direct {
101 RCP<const Basic> Fwx;
102 DenseMatrix F;
103
104 RCP<const Basic> torque;
105
106 RCP<const Basic> alphas_eqn;
107 } direct;
justinT21446e4f62024-06-16 22:36:10 -0700108};
109
Austin Schuh6ea789e2024-07-27 13:45:53 -0700110DenseMatrix SumMatrices(DenseMatrix a) { return a; }
111
112template <typename... Args>
113DenseMatrix SumMatrices(DenseMatrix a, Args... args) {
114 DenseMatrix result = DenseMatrix(2, 1, {integer(0), integer(0)});
115
116 DenseMatrix b = SumMatrices(args...);
117 add_dense_dense(a, b, result);
118 return result;
119}
120
justinT21446e4f62024-06-16 22:36:10 -0700121class SwerveSimulation {
122 public:
123 SwerveSimulation() : drive_motor_(KrakenFOC()), steer_motor_(KrakenFOC()) {
124 auto fx = symbol("fx");
125 auto fy = symbol("fy");
126 auto moment = symbol("moment");
127
Austin Schuh99f7c6a2024-06-25 22:07:44 -0700128 if (absl::GetFlag(FLAGS_symbolic)) {
justinT21446e4f62024-06-16 22:36:10 -0700129 Cx_ = symbol("Cx");
130 Cy_ = symbol("Cy");
131
Austin Schuh2a1abec2024-07-10 20:31:16 -0700132 rw_ = symbol("rw");
justinT21446e4f62024-06-16 22:36:10 -0700133
134 m_ = symbol("m");
135 J_ = symbol("J");
136
137 Gd1_ = symbol("Gd1");
138 rs_ = symbol("rs");
139 rp_ = symbol("rp");
140 Gd2_ = symbol("Gd2");
141
142 rb1_ = symbol("rb1");
143 rb2_ = symbol("rb2");
144
Austin Schuh76534f32024-09-02 13:52:45 -0700145 Gd3_ = symbol("Gd3");
justinT21446e4f62024-06-16 22:36:10 -0700146 Gd_ = symbol("Gd");
147
148 Js_ = symbol("Js");
149
150 Gs_ = symbol("Gs");
151 wb_ = symbol("wb");
152
153 Jdm_ = symbol("Jdm");
154 Jsm_ = symbol("Jsm");
155 Kts_ = symbol("Kts");
156 Ktd_ = symbol("Ktd");
157
158 robot_width_ = symbol("robot_width");
159
160 caster_ = symbol("caster");
161 contact_patch_length_ = symbol("Lcp");
162 } else {
Austin Schuhb8b34be2024-07-14 16:06:19 -0700163 Cx_ = real_double(25.0 * 9.8 / 4.0 / 0.05);
justinT21446e4f62024-06-16 22:36:10 -0700164 Cy_ = real_double(5 * 9.8 / 0.05 / 4.0);
165
Austin Schuh2a1abec2024-07-10 20:31:16 -0700166 rw_ = real_double(2 * 0.0254);
justinT21446e4f62024-06-16 22:36:10 -0700167
168 m_ = real_double(25.0); // base is 20 kg without battery
169 J_ = real_double(6.0);
170
171 Gd1_ = real_double(12.0 / 42.0);
172 rs_ = real_double(28.0 / 20.0 / 2.0);
173 rp_ = real_double(18.0 / 20.0 / 2.0);
174 Gd2_ = div(rs_, rp_);
175
176 // 15 / 45 bevel ratio, calculated using python script ported over to
177 // GetBevelPitchRadius(double
178 // TODO(Justin): Use the function instead of computed constantss
179 rb1_ = real_double(0.3805473);
180 rb2_ = real_double(1.14164);
181
182 Gd3_ = div(rb1_, rb2_);
183 Gd_ = mul(mul(Gd1_, Gd2_), Gd3_);
184
Austin Schuhb8b34be2024-07-14 16:06:19 -0700185 Js_ = real_double(0.001);
justinT21446e4f62024-06-16 22:36:10 -0700186
187 Gs_ = real_double(35.0 / 468.0);
188 wb_ = real_double(0.725);
189
190 Jdm_ = real_double(drive_motor_.motor_inertia);
191 Jsm_ = real_double(steer_motor_.motor_inertia);
192 Kts_ = real_double(steer_motor_.Kt);
193 Ktd_ = real_double(drive_motor_.Kt);
194
195 robot_width_ = real_double(24.75 * 0.0254);
196
Austin Schuh27694fa2024-07-20 16:29:49 -0700197 caster_ = real_double(absl::GetFlag(FLAGS_caster));
justinT21446e4f62024-06-16 22:36:10 -0700198 contact_patch_length_ = real_double(0.02);
199 }
200
201 x_ = symbol("x");
202 y_ = symbol("y");
203 theta_ = symbol("theta");
204
205 vx_ = symbol("vx");
206 vy_ = symbol("vy");
207 omega_ = symbol("omega");
208
209 ax_ = symbol("ax");
210 ay_ = symbol("ay");
211 atheta_ = symbol("atheta");
212
213 // Now, compute the accelerations due to the disturbance forces.
justinT21446e4f62024-06-16 22:36:10 -0700214 DenseMatrix external_accel = DenseMatrix(2, 1, {div(fx, m_), div(fy, m_)});
Austin Schuh6ea789e2024-07-27 13:45:53 -0700215 DenseMatrix external_force = DenseMatrix(2, 1, {fx, fy});
justinT21446e4f62024-06-16 22:36:10 -0700216
217 // And compute the physics contributions from each module.
218 modules_[0] = ModulePhysics(
219 0, DenseMatrix(
220 2, 1,
221 {div(robot_width_, integer(2)), div(robot_width_, integer(2))}));
222 modules_[1] =
223 ModulePhysics(1, DenseMatrix(2, 1,
224 {div(robot_width_, integer(-2)),
225 div(robot_width_, integer(2))}));
226 modules_[2] =
227 ModulePhysics(2, DenseMatrix(2, 1,
228 {div(robot_width_, integer(-2)),
229 div(robot_width_, integer(-2))}));
230 modules_[3] =
231 ModulePhysics(3, DenseMatrix(2, 1,
232 {div(robot_width_, integer(2)),
233 div(robot_width_, integer(-2))}));
234
235 // And convert them into the overall robot contribution.
Austin Schuh6ea789e2024-07-27 13:45:53 -0700236 DenseMatrix net_full_force =
237 SumMatrices(modules_[0].full.F, modules_[1].full.F, modules_[2].full.F,
238 modules_[3].full.F, external_force);
justinT21446e4f62024-06-16 22:36:10 -0700239
Austin Schuh6ea789e2024-07-27 13:45:53 -0700240 DenseMatrix net_direct_force =
241 SumMatrices(modules_[0].direct.F, modules_[1].direct.F,
242 modules_[2].direct.F, modules_[3].direct.F, external_force);
justinT21446e4f62024-06-16 22:36:10 -0700243
Austin Schuh6ea789e2024-07-27 13:45:53 -0700244 full_accel_ = DenseMatrix(2, 1);
245 mul_dense_scalar(net_full_force, pow(m_, minus_one), full_accel_);
justinT21446e4f62024-06-16 22:36:10 -0700246
Austin Schuh6ea789e2024-07-27 13:45:53 -0700247 full_angular_accel_ = div(
248 add(moment, add(add(modules_[0].full.torque, modules_[1].full.torque),
249 add(modules_[2].full.torque, modules_[3].full.torque))),
250 J_);
justinT21446e4f62024-06-16 22:36:10 -0700251
Austin Schuh6ea789e2024-07-27 13:45:53 -0700252 direct_accel_ = DenseMatrix(2, 1);
253 mul_dense_scalar(net_direct_force, pow(m_, minus_one), direct_accel_);
justinT21942892b2024-07-02 22:33:50 -0700254
Austin Schuh6ea789e2024-07-27 13:45:53 -0700255 direct_angular_accel_ =
256 div(add(moment,
257 add(add(modules_[0].direct.torque, modules_[1].direct.torque),
258 add(modules_[2].direct.torque, modules_[3].direct.torque))),
259 J_);
justinT21942892b2024-07-02 22:33:50 -0700260
Austin Schuh6ea789e2024-07-27 13:45:53 -0700261 VLOG(1) << "accel(0, 0) = " << ccode(*full_accel_.get(0, 0));
262 VLOG(1) << "accel(1, 0) = " << ccode(*full_accel_.get(1, 0));
263 VLOG(1) << "angular_accel = " << ccode(*full_angular_accel_);
justinT21942892b2024-07-02 22:33:50 -0700264 }
265
justinT21446e4f62024-06-16 22:36:10 -0700266 // Writes the physics out to the provided .cc and .h path.
267 void Write(std::string_view cc_path, std::string_view h_path) {
268 std::vector<std::string> result_cc;
269 std::vector<std::string> result_h;
270
Austin Schuh0f881092024-06-28 15:36:48 -0700271 std::string_view include_guard_stripped = h_path;
Austin Schuh99f7c6a2024-06-25 22:07:44 -0700272 CHECK(absl::ConsumePrefix(&include_guard_stripped,
273 absl::GetFlag(FLAGS_output_base)));
justinT21446e4f62024-06-16 22:36:10 -0700274 std::string include_guard =
275 absl::StrReplaceAll(absl::AsciiStrToUpper(include_guard_stripped),
276 {{"/", "_"}, {".", "_"}});
277
278 // Write out the header.
279 result_h.emplace_back(absl::Substitute("#ifndef $0_", include_guard));
280 result_h.emplace_back(absl::Substitute("#define $0_", include_guard));
281 result_h.emplace_back("");
282 result_h.emplace_back("#include <Eigen/Dense>");
283 result_h.emplace_back("");
284 result_h.emplace_back("namespace frc971::control_loops::swerve {");
285 result_h.emplace_back("");
286 result_h.emplace_back("// Returns the derivative of our state vector");
287 result_h.emplace_back("// [thetas0, thetad0, omegas0, omegad0,");
288 result_h.emplace_back("// thetas1, thetad1, omegas1, omegad1,");
289 result_h.emplace_back("// thetas2, thetad2, omegas2, omegad2,");
290 result_h.emplace_back("// thetas3, thetad3, omegas3, omegad3,");
291 result_h.emplace_back("// x, y, theta, vx, vy, omega,");
292 result_h.emplace_back("// Fx, Fy, Moment]");
293 result_h.emplace_back("Eigen::Matrix<double, 25, 1> SwervePhysics(");
294 result_h.emplace_back(
295 " Eigen::Map<const Eigen::Matrix<double, 25, 1>> X,");
296 result_h.emplace_back(
297 " Eigen::Map<const Eigen::Matrix<double, 8, 1>> U);");
298 result_h.emplace_back("");
299 result_h.emplace_back("} // namespace frc971::control_loops::swerve");
300 result_h.emplace_back("");
301 result_h.emplace_back(absl::Substitute("#endif // $0_", include_guard));
302
303 // Write out the .cc
304 result_cc.emplace_back(
305 absl::Substitute("#include \"$0\"", include_guard_stripped));
306 result_cc.emplace_back("");
307 result_cc.emplace_back("#include <cmath>");
308 result_cc.emplace_back("");
309 result_cc.emplace_back("namespace frc971::control_loops::swerve {");
310 result_cc.emplace_back("");
311 result_cc.emplace_back("Eigen::Matrix<double, 25, 1> SwervePhysics(");
312 result_cc.emplace_back(
313 " Eigen::Map<const Eigen::Matrix<double, 25, 1>> X,");
314 result_cc.emplace_back(
315 " Eigen::Map<const Eigen::Matrix<double, 8, 1>> U) {");
316 result_cc.emplace_back(" Eigen::Matrix<double, 25, 1> result;");
317
318 // Start by writing out variables matching each of the symbol names we use
319 // so we don't have to modify the computed equations too much.
320 for (size_t m = 0; m < kNumModules; ++m) {
321 result_cc.emplace_back(
322 absl::Substitute(" const double thetas$0 = X($1, 0);", m, m * 4));
323 result_cc.emplace_back(absl::Substitute(
324 " const double omegas$0 = X($1, 0);", m, m * 4 + 2));
325 result_cc.emplace_back(absl::Substitute(
326 " const double omegad$0 = X($1, 0);", m, m * 4 + 3));
327 }
328
329 result_cc.emplace_back(absl::Substitute(" const double theta = X($0, 0);",
330 kNumModules * 4 + 2));
331 result_cc.emplace_back(
332 absl::Substitute(" const double vx = X($0, 0);", kNumModules * 4 + 3));
333 result_cc.emplace_back(
334 absl::Substitute(" const double vy = X($0, 0);", kNumModules * 4 + 4));
335 result_cc.emplace_back(absl::Substitute(" const double omega = X($0, 0);",
336 kNumModules * 4 + 5));
337
338 result_cc.emplace_back(
339 absl::Substitute(" const double fx = X($0, 0);", kNumModules * 4 + 6));
340 result_cc.emplace_back(
341 absl::Substitute(" const double fy = X($0, 0);", kNumModules * 4 + 7));
342 result_cc.emplace_back(absl::Substitute(" const double moment = X($0, 0);",
343 kNumModules * 4 + 8));
344
345 // Now do the same for the inputs.
346 for (size_t m = 0; m < kNumModules; ++m) {
347 result_cc.emplace_back(
348 absl::Substitute(" const double Is$0 = U($1, 0);", m, m * 2));
349 result_cc.emplace_back(
350 absl::Substitute(" const double Id$0 = U($1, 0);", m, m * 2 + 1));
351 }
352
353 result_cc.emplace_back("");
354
355 // And then write out the derivative of each state.
356 for (size_t m = 0; m < kNumModules; ++m) {
357 result_cc.emplace_back(
358 absl::Substitute(" result($0, 0) = omegas$1;", m * 4, m));
359 result_cc.emplace_back(
360 absl::Substitute(" result($0, 0) = omegad$1;", m * 4 + 1, m));
361
Austin Schuh6ea789e2024-07-27 13:45:53 -0700362 result_cc.emplace_back(
363 absl::Substitute(" result($0, 0) = $1;", m * 4 + 2,
364 ccode(*modules_[m].full.alphas_eqn)));
365 result_cc.emplace_back(
366 absl::Substitute(" result($0, 0) = $1;", m * 4 + 3,
367 ccode(*modules_[m].full.alphad_eqn)));
justinT21446e4f62024-06-16 22:36:10 -0700368 }
369
370 result_cc.emplace_back(
Austin Schuh5ddcb472024-07-21 17:55:34 -0700371 absl::Substitute(" result($0, 0) = vx;", kNumModules * 4 + 0));
justinT21446e4f62024-06-16 22:36:10 -0700372 result_cc.emplace_back(
Austin Schuh5ddcb472024-07-21 17:55:34 -0700373 absl::Substitute(" result($0, 0) = vy;", kNumModules * 4 + 1));
justinT21446e4f62024-06-16 22:36:10 -0700374 result_cc.emplace_back(
Austin Schuh5ddcb472024-07-21 17:55:34 -0700375 absl::Substitute(" result($0, 0) = omega;", kNumModules * 4 + 2));
justinT21446e4f62024-06-16 22:36:10 -0700376
justinT21446e4f62024-06-16 22:36:10 -0700377 result_cc.emplace_back(absl::Substitute(" result($0, 0) = $1;",
Austin Schuh5ddcb472024-07-21 17:55:34 -0700378 kNumModules * 4 + 3,
Austin Schuh6ea789e2024-07-27 13:45:53 -0700379 ccode(*full_accel_.get(0, 0))));
justinT21446e4f62024-06-16 22:36:10 -0700380 result_cc.emplace_back(absl::Substitute(" result($0, 0) = $1;",
Austin Schuh5ddcb472024-07-21 17:55:34 -0700381 kNumModules * 4 + 4,
Austin Schuh6ea789e2024-07-27 13:45:53 -0700382 ccode(*full_accel_.get(1, 0))));
383 result_cc.emplace_back(absl::Substitute(" result($0, 0) = $1;",
384 kNumModules * 4 + 5,
385 ccode(*full_angular_accel_)));
justinT21446e4f62024-06-16 22:36:10 -0700386
387 result_cc.emplace_back(
388 absl::Substitute(" result($0, 0) = 0.0;", kNumModules * 4 + 6));
389 result_cc.emplace_back(
390 absl::Substitute(" result($0, 0) = 0.0;", kNumModules * 4 + 7));
391 result_cc.emplace_back(
392 absl::Substitute(" result($0, 0) = 0.0;", kNumModules * 4 + 8));
393
394 result_cc.emplace_back("");
395 result_cc.emplace_back(" return result;");
396 result_cc.emplace_back("}");
397 result_cc.emplace_back("");
398 result_cc.emplace_back("} // namespace frc971::control_loops::swerve");
399
400 aos::util::WriteStringToFileOrDie(cc_path, absl::StrJoin(result_cc, "\n"));
401 aos::util::WriteStringToFileOrDie(h_path, absl::StrJoin(result_h, "\n"));
402 }
403
Austin Schuhb67a38f2024-07-04 13:48:38 -0700404 void WriteCasadiVariables(std::vector<std::string> *result_py) {
405 result_py->emplace_back(" sin = casadi.sin");
406 result_py->emplace_back(" cos = casadi.cos");
Austin Schuh78a1b312024-08-18 17:21:34 -0700407 result_py->emplace_back(" exp = casadi.exp");
Austin Schuh5371c802024-09-02 13:18:48 -0700408 if (absl::GetFlag(FLAGS_function)) {
409 result_py->emplace_back(" atan2 = soft_atan2()");
410 } else {
411 result_py->emplace_back(" atan2 = soft_atan2");
412 }
Austin Schuh2a1abec2024-07-10 20:31:16 -0700413 result_py->emplace_back(" fmax = casadi.fmax");
Austin Schuhb67a38f2024-07-04 13:48:38 -0700414 result_py->emplace_back(" fabs = casadi.fabs");
415
416 // Start by writing out variables matching each of the symbol names we use
417 // so we don't have to modify the computed equations too much.
418 for (size_t m = 0; m < kNumModules; ++m) {
419 result_py->emplace_back(
420 absl::Substitute(" thetas$0 = X[$1, 0]", m, m * 4));
421 result_py->emplace_back(
422 absl::Substitute(" omegas$0 = X[$1, 0]", m, m * 4 + 2));
423 result_py->emplace_back(
424 absl::Substitute(" omegad$0 = X[$1, 0]", m, m * 4 + 3));
425 }
426
427 result_py->emplace_back(
428 absl::Substitute(" theta = X[$0, 0]", kNumModules * 4 + 2));
429 result_py->emplace_back(
430 absl::Substitute(" vx = X[$0, 0]", kNumModules * 4 + 3));
431 result_py->emplace_back(
432 absl::Substitute(" vy = X[$0, 0]", kNumModules * 4 + 4));
433 result_py->emplace_back(
434 absl::Substitute(" omega = X[$0, 0]", kNumModules * 4 + 5));
435
436 result_py->emplace_back(
437 absl::Substitute(" fx = X[$0, 0]", kNumModules * 4 + 6));
438 result_py->emplace_back(
439 absl::Substitute(" fy = X[$0, 0]", kNumModules * 4 + 7));
440 result_py->emplace_back(
441 absl::Substitute(" moment = X[$0, 0]", kNumModules * 4 + 8));
442
443 // Now do the same for the inputs.
444 for (size_t m = 0; m < kNumModules; ++m) {
445 result_py->emplace_back(
446 absl::Substitute(" Is$0 = U[$1, 0]", m, m * 2));
447 result_py->emplace_back(
448 absl::Substitute(" Id$0 = U[$1, 0]", m, m * 2 + 1));
449 }
450 }
451
Austin Schuh6ea789e2024-07-27 13:45:53 -0700452 void WriteCasadiVelocityVariables(std::vector<std::string> *result_py) {
453 result_py->emplace_back(" sin = casadi.sin");
Austin Schuh78a1b312024-08-18 17:21:34 -0700454 result_py->emplace_back(" exp = casadi.exp");
Austin Schuh6ea789e2024-07-27 13:45:53 -0700455 result_py->emplace_back(" cos = casadi.cos");
Austin Schuh5371c802024-09-02 13:18:48 -0700456 if (absl::GetFlag(FLAGS_function)) {
457 result_py->emplace_back(" atan2 = soft_atan2()");
458 } else {
459 result_py->emplace_back(" atan2 = soft_atan2");
460 }
Austin Schuh6ea789e2024-07-27 13:45:53 -0700461 result_py->emplace_back(" fmax = casadi.fmax");
462 result_py->emplace_back(" fabs = casadi.fabs");
463
464 // Start by writing out variables matching each of the symbol names we use
465 // so we don't have to modify the computed equations too much.
466 for (size_t m = 0; m < kNumModules; ++m) {
467 result_py->emplace_back(
468 absl::Substitute(" thetas$0 = X[$1, 0]", m, m * 2 + 0));
469 result_py->emplace_back(
470 absl::Substitute(" omegas$0 = X[$1, 0]", m, m * 2 + 1));
471 }
472
473 result_py->emplace_back(
474 absl::Substitute(" theta = X[$0, 0]", kNumModules * 2 + 0));
475 result_py->emplace_back(
476 absl::Substitute(" vx = X[$0, 0]", kNumModules * 2 + 1));
477 result_py->emplace_back(
478 absl::Substitute(" vy = X[$0, 0]", kNumModules * 2 + 2));
479 result_py->emplace_back(
480 absl::Substitute(" omega = X[$0, 0]", kNumModules * 2 + 3));
481
482 // result_py->emplace_back(
483 // absl::Substitute(" fx = X[$0, 0]", kNumModules * 3 + 4));
484 // result_py->emplace_back(
485 // absl::Substitute(" fy = X[$0, 0]", kNumModules * 3 + 5));
486 // result_py->emplace_back(
487 // absl::Substitute(" moment = X[$0, 0]", kNumModules * 3 + 6));
488 //
489 result_py->emplace_back(" fx = 0");
490 result_py->emplace_back(" fy = 0");
491 result_py->emplace_back(" moment = 0");
492
493 // Now do the same for the inputs.
494 for (size_t m = 0; m < kNumModules; ++m) {
495 result_py->emplace_back(
496 absl::Substitute(" Is$0 = U[$1, 0]", m, m * 2));
497 result_py->emplace_back(
498 absl::Substitute(" Id$0 = U[$1, 0]", m, m * 2 + 1));
499 }
500 }
501
Austin Schuh0f881092024-06-28 15:36:48 -0700502 // Writes the physics out to the provided .cc and .h path.
503 void WriteCasadi(std::string_view py_path) {
504 std::vector<std::string> result_py;
505
506 // Write out the header.
Austin Schuh5371c802024-09-02 13:18:48 -0700507 result_py.emplace_back("#!/usr/bin/env python3");
Austin Schuh0f881092024-06-28 15:36:48 -0700508 result_py.emplace_back("");
Austin Schuh6ea789e2024-07-27 13:45:53 -0700509 result_py.emplace_back("import casadi, numpy");
Austin Schuh0f881092024-06-28 15:36:48 -0700510 result_py.emplace_back("");
Austin Schuh2a1abec2024-07-10 20:31:16 -0700511 result_py.emplace_back(absl::Substitute("WHEEL_RADIUS = $0", ccode(*rw_)));
512 result_py.emplace_back(
513 absl::Substitute("ROBOT_WIDTH = $0", ccode(*robot_width_)));
514 result_py.emplace_back(absl::Substitute("CASTER = $0", ccode(*caster_)));
Austin Schuh6ea789e2024-07-27 13:45:53 -0700515 result_py.emplace_back("STATE_THETAS0 = 0");
516 result_py.emplace_back("STATE_THETAD0 = 1");
517 result_py.emplace_back("STATE_OMEGAS0 = 2");
518 result_py.emplace_back("STATE_OMEGAD0 = 3");
519 result_py.emplace_back("STATE_THETAS1 = 4");
520 result_py.emplace_back("STATE_THETAD1 = 5");
521 result_py.emplace_back("STATE_OMEGAS1 = 6");
522 result_py.emplace_back("STATE_OMEGAD1 = 7");
523 result_py.emplace_back("STATE_THETAS2 = 8");
524 result_py.emplace_back("STATE_THETAD2 = 9");
525 result_py.emplace_back("STATE_OMEGAS2 = 10");
526 result_py.emplace_back("STATE_OMEGAD2 = 11");
527 result_py.emplace_back("STATE_THETAS3 = 12");
528 result_py.emplace_back("STATE_THETAD3 = 13");
529 result_py.emplace_back("STATE_OMEGAS3 = 14");
530 result_py.emplace_back("STATE_OMEGAD3 = 15");
531 result_py.emplace_back("STATE_X = 16");
532 result_py.emplace_back("STATE_Y = 17");
533 result_py.emplace_back("STATE_THETA = 18");
534 result_py.emplace_back("STATE_VX = 19");
535 result_py.emplace_back("STATE_VY = 20");
536 result_py.emplace_back("STATE_OMEGA = 21");
537 result_py.emplace_back("STATE_FX = 22");
538 result_py.emplace_back("STATE_FY = 23");
539 result_py.emplace_back("STATE_MOMENT = 24");
540 result_py.emplace_back("NUM_STATES = 25");
541 result_py.emplace_back("");
542 result_py.emplace_back("VELOCITY_STATE_THETAS0 = 0");
543 result_py.emplace_back("VELOCITY_STATE_OMEGAS0 = 1");
544 result_py.emplace_back("VELOCITY_STATE_THETAS1 = 2");
545 result_py.emplace_back("VELOCITY_STATE_OMEGAS1 = 3");
546 result_py.emplace_back("VELOCITY_STATE_THETAS2 = 4");
547 result_py.emplace_back("VELOCITY_STATE_OMEGAS2 = 5");
548 result_py.emplace_back("VELOCITY_STATE_THETAS3 = 6");
549 result_py.emplace_back("VELOCITY_STATE_OMEGAS3 = 7");
550 result_py.emplace_back("VELOCITY_STATE_THETA = 8");
551 result_py.emplace_back("VELOCITY_STATE_VX = 9");
552 result_py.emplace_back("VELOCITY_STATE_VY = 10");
553 result_py.emplace_back("VELOCITY_STATE_OMEGA = 11");
554 // result_py.emplace_back("VELOCITY_STATE_FX = 16");
555 // result_py.emplace_back("VELOCITY_STATE_FY = 17");
556 // result_py.emplace_back("VELOCITY_STATE_MOMENT = 18");
557 result_py.emplace_back("NUM_VELOCITY_STATES = 12");
558 result_py.emplace_back("");
559 result_py.emplace_back("def to_velocity_state(X):");
560 result_py.emplace_back(" return numpy.array([");
561 result_py.emplace_back(" [X[STATE_THETAS0, 0]],");
562 result_py.emplace_back(" [X[STATE_OMEGAS0, 0]],");
563 result_py.emplace_back(" [X[STATE_THETAS1, 0]],");
564 result_py.emplace_back(" [X[STATE_OMEGAS1, 0]],");
565 result_py.emplace_back(" [X[STATE_THETAS2, 0]],");
566 result_py.emplace_back(" [X[STATE_OMEGAS2, 0]],");
567 result_py.emplace_back(" [X[STATE_THETAS3, 0]],");
568 result_py.emplace_back(" [X[STATE_OMEGAS3, 0]],");
569 result_py.emplace_back(" [X[STATE_THETA, 0]],");
570 result_py.emplace_back(" [X[STATE_VX, 0]],");
571 result_py.emplace_back(" [X[STATE_VY, 0]],");
572 result_py.emplace_back(" [X[STATE_OMEGA, 0]],");
573 // result_py.emplace_back(" [X[STATE_FX, 0]],");
574 // result_py.emplace_back(" [X[STATE_FY, 0]],");
575 // result_py.emplace_back(" [X[STATE_MOMENT, 0]],");
576 result_py.emplace_back(" ])");
Austin Schuh2a1abec2024-07-10 20:31:16 -0700577 result_py.emplace_back("");
Austin Schuhbd75c482024-08-18 00:03:51 -0700578 constexpr double kLogGain = 1.0 / 0.05;
Austin Schuh5371c802024-09-02 13:18:48 -0700579 constexpr double kAbsGain = 1.0 / 0.01;
580 if (absl::GetFlag(FLAGS_function)) {
581 result_py.emplace_back("def soft_atan2():");
582 result_py.emplace_back(" y = casadi.SX.sym('y')");
583 result_py.emplace_back(" x = casadi.SX.sym('x')");
584 result_py.emplace_back(
585 " return casadi.Function('soft_atan2', [y, x], [");
586 result_py.emplace_back(" casadi.arctan2(");
587 result_py.emplace_back(" y,");
588 result_py.emplace_back(" casadi.logsumexp(");
589 result_py.emplace_back(" casadi.SX(");
590 result_py.emplace_back(" numpy.array([");
591 result_py.emplace_back(" 1.0, x * (1.0 - 2.0 /");
592 result_py.emplace_back(
593 absl::Substitute(" (1 + "
594 "casadi.exp($1.0 * x))) * $0.0",
595 kLogGain, kAbsGain));
596 result_py.emplace_back(
597 absl::Substitute(" ]))) / $0.0)", kLogGain));
598 result_py.emplace_back(" ])");
599 } else {
600 result_py.emplace_back("def soft_atan2(y, x):");
601 result_py.emplace_back(" return casadi.arctan2(");
602 result_py.emplace_back(" y,");
603 result_py.emplace_back(" casadi.logsumexp(casadi.SX(numpy.array(");
604 result_py.emplace_back(
605 absl::Substitute(" [1.0, x * (1.0 - 2.0 / (1 + "
606 "casadi.exp($1.0 * x))) * $0.0]))) / $0.0)",
607 kLogGain, kAbsGain));
608 }
609 result_py.emplace_back("");
610 result_py.emplace_back("# Is = STEER_CURRENT_COUPLING_FACTOR * Id");
611 result_py.emplace_back(absl::Substitute(
612 "STEER_CURRENT_COUPLING_FACTOR = $0",
613 ccode(*(neg(
614 mul(div(Gs_, Kts_),
615 mul(div(Ktd_, mul(Gd_, rw_)),
616 neg(mul(add(neg(wb_), mul(add(rs_, rp_),
617 sub(integer(1), div(rb1_, rp_)))),
618 div(rw_, rb2_))))))))));
Austin Schuhbd75c482024-08-18 00:03:51 -0700619 result_py.emplace_back("");
Austin Schuh98fbbe82024-08-18 23:07:26 -0700620 result_py.emplace_back("# Is = STEER_CURRENT_COUPLING_FACTOR * Id");
621 result_py.emplace_back(absl::Substitute(
622 "STEER_CURRENT_COUPLING_FACTOR = $0",
623 ccode(*(neg(
624 mul(div(Gs_, Kts_),
625 mul(div(Ktd_, mul(Gd_, rw_)),
626 neg(mul(add(neg(wb_), mul(add(rs_, rp_),
627 sub(integer(1), div(rb1_, rp_)))),
628 div(rw_, rb2_))))))))));
629 result_py.emplace_back("");
Austin Schuh2a1abec2024-07-10 20:31:16 -0700630
Austin Schuh0f881092024-06-28 15:36:48 -0700631 result_py.emplace_back("# Returns the derivative of our state vector");
632 result_py.emplace_back("# [thetas0, thetad0, omegas0, omegad0,");
633 result_py.emplace_back("# thetas1, thetad1, omegas1, omegad1,");
634 result_py.emplace_back("# thetas2, thetad2, omegas2, omegad2,");
635 result_py.emplace_back("# thetas3, thetad3, omegas3, omegad3,");
636 result_py.emplace_back("# x, y, theta, vx, vy, omega,");
637 result_py.emplace_back("# Fx, Fy, Moment]");
Austin Schuh6ea789e2024-07-27 13:45:53 -0700638 result_py.emplace_back("def swerve_full_dynamics(X, U):");
Austin Schuhb67a38f2024-07-04 13:48:38 -0700639 WriteCasadiVariables(&result_py);
Austin Schuh0f881092024-06-28 15:36:48 -0700640
641 result_py.emplace_back("");
642 result_py.emplace_back(" result = casadi.SX.sym('result', 25, 1)");
643 result_py.emplace_back("");
644
645 // And then write out the derivative of each state.
646 for (size_t m = 0; m < kNumModules; ++m) {
647 result_py.emplace_back(
648 absl::Substitute(" result[$0, 0] = omegas$1", m * 4, m));
649 result_py.emplace_back(
650 absl::Substitute(" result[$0, 0] = omegad$1", m * 4 + 1, m));
651
Austin Schuh6ea789e2024-07-27 13:45:53 -0700652 result_py.emplace_back(
653 absl::Substitute(" result[$0, 0] = $1", m * 4 + 2,
654 ccode(*modules_[m].full.alphas_eqn)));
655 result_py.emplace_back(
656 absl::Substitute(" result[$0, 0] = $1", m * 4 + 3,
657 ccode(*modules_[m].full.alphad_eqn)));
Austin Schuh0f881092024-06-28 15:36:48 -0700658 }
659
660 result_py.emplace_back(
Austin Schuhb8b34be2024-07-14 16:06:19 -0700661 absl::Substitute(" result[$0, 0] = vx", kNumModules * 4 + 0));
Austin Schuh0f881092024-06-28 15:36:48 -0700662 result_py.emplace_back(
Austin Schuhb8b34be2024-07-14 16:06:19 -0700663 absl::Substitute(" result[$0, 0] = vy", kNumModules * 4 + 1));
Austin Schuh0f881092024-06-28 15:36:48 -0700664 result_py.emplace_back(
Austin Schuhb8b34be2024-07-14 16:06:19 -0700665 absl::Substitute(" result[$0, 0] = omega", kNumModules * 4 + 2));
Austin Schuh0f881092024-06-28 15:36:48 -0700666
Austin Schuh0f881092024-06-28 15:36:48 -0700667 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
Austin Schuhb8b34be2024-07-14 16:06:19 -0700668 kNumModules * 4 + 3,
Austin Schuh6ea789e2024-07-27 13:45:53 -0700669 ccode(*full_accel_.get(0, 0))));
Austin Schuh0f881092024-06-28 15:36:48 -0700670 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
Austin Schuhb8b34be2024-07-14 16:06:19 -0700671 kNumModules * 4 + 4,
Austin Schuh6ea789e2024-07-27 13:45:53 -0700672 ccode(*full_accel_.get(1, 0))));
673 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
674 kNumModules * 4 + 5,
675 ccode(*full_angular_accel_)));
Austin Schuh0f881092024-06-28 15:36:48 -0700676
677 result_py.emplace_back(
678 absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 4 + 6));
679 result_py.emplace_back(
680 absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 4 + 7));
681 result_py.emplace_back(
682 absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 4 + 8));
683
684 result_py.emplace_back("");
685 result_py.emplace_back(
686 " return casadi.Function('xdot', [X, U], [result])");
Austin Schuh2a1abec2024-07-10 20:31:16 -0700687
Austin Schuh6ea789e2024-07-27 13:45:53 -0700688 result_py.emplace_back("");
689
690 result_py.emplace_back("# Returns the derivative of our state vector");
691 result_py.emplace_back("# [thetas0, omegas0,");
692 result_py.emplace_back("# thetas1, omegas1,");
693 result_py.emplace_back("# thetas2, omegas2,");
694 result_py.emplace_back("# thetas3, omegas3,");
695 result_py.emplace_back("# theta, vx, vy, omega]");
696 result_py.emplace_back("def velocity_swerve_physics(X, U):");
697 WriteCasadiVelocityVariables(&result_py);
698
699 result_py.emplace_back("");
700 result_py.emplace_back(
701 " result = casadi.SX.sym('result', NUM_VELOCITY_STATES, 1)");
702 result_py.emplace_back("");
703
704 // And then write out the derivative of each state.
705 for (size_t m = 0; m < kNumModules; ++m) {
706 result_py.emplace_back(
707 absl::Substitute(" result[$0, 0] = omegas$1", m * 2 + 0, m));
708 result_py.emplace_back(
709 absl::Substitute(" result[$0, 0] = $1", m * 2 + 1,
710 ccode(*modules_[m].direct.alphas_eqn)));
711 }
712 result_py.emplace_back(
713 absl::Substitute(" result[$0, 0] = omega", kNumModules * 2 + 0));
714
715 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
716 kNumModules * 2 + 1,
717 ccode(*direct_accel_.get(0, 0))));
718 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
719 kNumModules * 2 + 2,
720 ccode(*direct_accel_.get(1, 0))));
721 result_py.emplace_back(absl::Substitute(" result[$0, 0] = $1",
722 kNumModules * 2 + 3,
723 ccode(*direct_angular_accel_)));
724
725 // result_py.emplace_back(
726 // absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 3 + 4));
727 // result_py.emplace_back(
728 // absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 3 + 5));
729 // result_py.emplace_back(
730 // absl::Substitute(" result[$0, 0] = 0.0", kNumModules * 3 + 6));
731
732 result_py.emplace_back("");
733 result_py.emplace_back(
734 " return casadi.Function('xdot', [X, U], [result])");
735
Austin Schuhb8b34be2024-07-14 16:06:19 -0700736 DefineVector2dFunction(
737 "contact_patch_velocity",
738 "# Returns the velocity of the wheel in global coordinates.",
739 [](const Module &m, int dimension) {
740 return ccode(*m.contact_patch_velocity.get(dimension, 0));
741 },
742 &result_py);
743 DefineVector2dFunction(
744 "wheel_ground_velocity",
745 "# Returns the velocity of the wheel in steer module coordinates.",
746 [](const Module &m, int dimension) {
747 return ccode(*m.wheel_ground_velocity.get(dimension, 0));
748 },
749 &result_py);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700750
Austin Schuhb8b34be2024-07-14 16:06:19 -0700751 DefineVector2dFunction(
752 "wheel_slip_velocity",
753 "# Returns the difference in velocities of the wheel surface and the "
754 "ground.",
755 [](const Module &m, int dimension) {
756 return ccode(*m.wheel_slip_velocity.get(dimension, 0));
757 },
758 &result_py);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700759
Austin Schuhb8b34be2024-07-14 16:06:19 -0700760 DefineScalarFunction(
761 "slip_angle", "Returns the slip angle of the ith wheel",
762 [](const Module &m) { return ccode(*m.slip_angle); }, &result_py);
763 DefineScalarFunction(
764 "slip_ratio", "Returns the slip ratio of the ith wheel",
765 [](const Module &m) { return ccode(*m.slip_ratio); }, &result_py);
766 DefineScalarFunction(
767 "module_angular_accel",
768 "Returns the angular acceleration of the robot due to the ith wheel",
Austin Schuh6ea789e2024-07-27 13:45:53 -0700769 [this](const Module &m) { return ccode(*div(m.full.torque, Js_)); },
770 &result_py);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700771
Austin Schuhb8b34be2024-07-14 16:06:19 -0700772 DefineVector2dFunction(
773 "wheel_force",
774 "Returns the force on the wheel in steer module coordinates",
775 [](const Module &m, int dimension) {
Austin Schuh6ea789e2024-07-27 13:45:53 -0700776 return ccode(
777 *std::vector<RCP<const Basic>>{m.full.Fwx, m.Fwy}[dimension]);
Austin Schuhb8b34be2024-07-14 16:06:19 -0700778 },
779 &result_py);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700780
Austin Schuhb8b34be2024-07-14 16:06:19 -0700781 DefineVector2dFunction(
782 "F", "Returns the force on the wheel in absolute coordinates",
783 [](const Module &m, int dimension) {
Austin Schuh6ea789e2024-07-27 13:45:53 -0700784 return ccode(*m.full.F.get(dimension, 0));
Austin Schuhb8b34be2024-07-14 16:06:19 -0700785 },
786 &result_py);
787
justinT21820767f2024-09-11 19:57:55 -0700788 DefineVector2dVelocityFunction(
789 "F_vel",
790 "Returns the force on the wheel in absolute coordinates based on the "
791 "velocity controller",
792 [](const Module &m, int dimension) {
793 return ccode(*m.direct.F.get(dimension, 0));
794 },
795 &result_py);
796
797 DefineVector2dVelocityFunction(
Austin Schuhb8b34be2024-07-14 16:06:19 -0700798 "mounting_location",
799 "Returns the mounting location of wheel in robot coordinates",
800 [](const Module &m, int dimension) {
801 return ccode(*m.mounting_location.get(dimension, 0));
802 },
803 &result_py);
804
805 DefineScalarFunction(
806 "Ms", "Returns the self aligning moment of the ith wheel",
807 [this](const Module &m) {
808 return ccode(*(div(m.Ms, add(Jsm_, div(div(Js_, Gs_), Gs_)))));
809 },
810 &result_py);
Austin Schuh0f881092024-06-28 15:36:48 -0700811
812 aos::util::WriteStringToFileOrDie(py_path, absl::StrJoin(result_py, "\n"));
813 }
814
Austin Schuhb8b34be2024-07-14 16:06:19 -0700815 void DefineScalarFunction(
816 std::string_view name, std::string_view documentation,
817 std::function<std::string(const Module &)> scalar_fn,
818 std::vector<std::string> *result_py) {
819 result_py->emplace_back("");
820 result_py->emplace_back(absl::Substitute("# $0.", documentation));
821 result_py->emplace_back(absl::Substitute("def $0(i, X, U):", name));
822 WriteCasadiVariables(result_py);
823 for (size_t m = 0; m < kNumModules; ++m) {
824 if (m == 0) {
825 result_py->emplace_back(" if i == 0:");
826 } else {
827 result_py->emplace_back(absl::Substitute(" elif i == $0:", m));
828 }
829 result_py->emplace_back(
830 absl::Substitute(" return casadi.Function('$0', [X, U], [$1])",
831 name, scalar_fn(modules_[m])));
832 }
833 result_py->emplace_back(" raise ValueError(\"Invalid module number\")");
834 }
835
836 void DefineVector2dFunction(
837 std::string_view name, std::string_view documentation,
838 std::function<std::string(const Module &, int)> scalar_fn,
839 std::vector<std::string> *result_py) {
840 result_py->emplace_back("");
841 result_py->emplace_back(absl::Substitute("# $0.", documentation));
842 result_py->emplace_back(absl::Substitute("def $0(i, X, U):", name));
843 WriteCasadiVariables(result_py);
844 result_py->emplace_back(
845 absl::Substitute(" result = casadi.SX.sym('$0', 2, 1)", name));
846 for (size_t m = 0; m < kNumModules; ++m) {
847 if (m == 0) {
848 result_py->emplace_back(" if i == 0:");
849 } else {
850 result_py->emplace_back(absl::Substitute(" elif i == $0:", m));
851 }
852 for (int j = 0; j < 2; ++j) {
853 result_py->emplace_back(absl::Substitute(" result[$0, 0] = $1",
854 j, scalar_fn(modules_[m], j)));
855 }
856 }
857 result_py->emplace_back(" else:");
858 result_py->emplace_back(
859 " raise ValueError(\"Invalid module number\")");
860 result_py->emplace_back(absl::Substitute(
861 " return casadi.Function('$0', [X, U], [result])", name));
862 }
863
justinT21820767f2024-09-11 19:57:55 -0700864 void DefineVector2dVelocityFunction(
865 std::string_view name, std::string_view documentation,
866 std::function<std::string(const Module &, int)> scalar_fn,
867 std::vector<std::string> *result_py) {
868 result_py->emplace_back("");
869 result_py->emplace_back(absl::Substitute("# $0.", documentation));
870 result_py->emplace_back(absl::Substitute("def $0(i, X, U):", name));
871 WriteCasadiVelocityVariables(result_py);
872 result_py->emplace_back(
873 absl::Substitute(" result = casadi.SX.sym('$0', 2, 1)", name));
874 for (size_t m = 0; m < kNumModules; ++m) {
875 if (m == 0) {
876 result_py->emplace_back(" if i == 0:");
877 } else {
878 result_py->emplace_back(absl::Substitute(" elif i == $0:", m));
879 }
880 for (int j = 0; j < 2; ++j) {
881 result_py->emplace_back(absl::Substitute(" result[$0, 0] = $1",
882 j, scalar_fn(modules_[m], j)));
883 }
884 }
885 result_py->emplace_back(" else:");
886 result_py->emplace_back(
887 " raise ValueError(\"Invalid module number\")");
888 result_py->emplace_back(absl::Substitute(
889 " return casadi.Function('$0', [X, U], [result])", name));
890 }
891
justinT21446e4f62024-06-16 22:36:10 -0700892 private:
893 static constexpr uint8_t kNumModules = 4;
894
Austin Schuh6ea789e2024-07-27 13:45:53 -0700895 RCP<const Basic> SteerAccel(RCP<const Basic> Fwx, RCP<const Basic> Ms,
896 RCP<const Basic> Is) {
897 RCP<const Basic> lhms =
898 mul(add(neg(wb_), mul(add(rs_, rp_), sub(integer(1), div(rb1_, rp_)))),
899 mul(div(rw_, rb2_), neg(Fwx)));
900 RCP<const Basic> lhs = add(add(Ms, div(mul(Kts_, Is), Gs_)), lhms);
901 RCP<const Basic> rhs = add(Jsm_, div(div(Js_, Gs_), Gs_));
902 return simplify(div(lhs, rhs));
903 }
904
justinT21446e4f62024-06-16 22:36:10 -0700905 Module ModulePhysics(const int m, DenseMatrix mounting_location) {
906 VLOG(1) << "Solving module " << m;
907
908 Module result;
Austin Schuhb8b34be2024-07-14 16:06:19 -0700909 result.mounting_location = mounting_location;
justinT21446e4f62024-06-16 22:36:10 -0700910
911 result.Is = symbol(absl::StrFormat("Is%u", m));
912 result.Id = symbol(absl::StrFormat("Id%u", m));
913
914 RCP<const Symbol> thetamd = symbol(absl::StrFormat("theta_md%u", m));
915 RCP<const Symbol> omegamd = symbol(absl::StrFormat("omega_md%u", m));
916 RCP<const Symbol> alphamd = symbol(absl::StrFormat("alpha_md%u", m));
917
918 result.thetas = symbol(absl::StrFormat("thetas%u", m));
919 result.omegas = symbol(absl::StrFormat("omegas%u", m));
Austin Schuh6ea789e2024-07-27 13:45:53 -0700920 RCP<const Symbol> alphas = symbol(absl::StrFormat("alphas%u", m));
justinT21446e4f62024-06-16 22:36:10 -0700921
justinT21446e4f62024-06-16 22:36:10 -0700922 result.omegad = symbol(absl::StrFormat("omegad%u", m));
Austin Schuh6ea789e2024-07-27 13:45:53 -0700923 RCP<const Symbol> alphad = symbol(absl::StrFormat("alphad%u", m));
justinT21446e4f62024-06-16 22:36:10 -0700924
925 // Velocity of the module in field coordinates
Austin Schuh2a1abec2024-07-10 20:31:16 -0700926 DenseMatrix robot_velocity = DenseMatrix(2, 1, {vx_, vy_});
justinT21446e4f62024-06-16 22:36:10 -0700927 VLOG(1) << "robot velocity: " << robot_velocity.__str__();
928
929 // Velocity of the contact patch in field coordinates
930 DenseMatrix temp_matrix = DenseMatrix(2, 1);
931 DenseMatrix temp_matrix2 = DenseMatrix(2, 1);
Austin Schuhbd75c482024-08-18 00:03:51 -0700932 DenseMatrix temp_matrix3 = DenseMatrix(2, 1);
Austin Schuh2a1abec2024-07-10 20:31:16 -0700933 result.contact_patch_velocity = DenseMatrix(2, 1);
justinT21446e4f62024-06-16 22:36:10 -0700934
Austin Schuhb8b34be2024-07-14 16:06:19 -0700935 mul_dense_dense(R(theta_), result.mounting_location, temp_matrix);
justinT21446e4f62024-06-16 22:36:10 -0700936 add_dense_dense(angle_cross(temp_matrix, omega_), robot_velocity,
937 temp_matrix2);
938 mul_dense_dense(R(add(theta_, result.thetas)),
Austin Schuhbd75c482024-08-18 00:03:51 -0700939 DenseMatrix(2, 1, {neg(caster_), integer(0)}),
940 temp_matrix3);
justinT21446e4f62024-06-16 22:36:10 -0700941 add_dense_dense(temp_matrix2,
Austin Schuhbd75c482024-08-18 00:03:51 -0700942 angle_cross(temp_matrix3, add(omega_, result.omegas)),
Austin Schuh2a1abec2024-07-10 20:31:16 -0700943 result.contact_patch_velocity);
justinT21446e4f62024-06-16 22:36:10 -0700944
945 VLOG(1);
Austin Schuh2a1abec2024-07-10 20:31:16 -0700946 VLOG(1) << "contact patch velocity: "
947 << result.contact_patch_velocity.__str__();
justinT21446e4f62024-06-16 22:36:10 -0700948
949 // Relative velocity of the surface of the wheel to the ground.
Austin Schuhb67a38f2024-07-04 13:48:38 -0700950 result.wheel_ground_velocity = DenseMatrix(2, 1);
Austin Schuh2a1abec2024-07-10 20:31:16 -0700951 mul_dense_dense(R(neg(add(result.thetas, theta_))),
952 result.contact_patch_velocity,
Austin Schuhb67a38f2024-07-04 13:48:38 -0700953 result.wheel_ground_velocity);
justinT21446e4f62024-06-16 22:36:10 -0700954
Austin Schuhb8b34be2024-07-14 16:06:19 -0700955 // Compute the relative velocity between the wheel surface and the ground in
956 // the wheel coordinate system.
957 result.wheel_slip_velocity = DenseMatrix(2, 1);
958 DenseMatrix wheel_velocity =
959 DenseMatrix(2, 1, {mul(rw_, result.omegad), integer(0)});
960 DenseMatrix negative_wheel_ground_velocity =
961 DenseMatrix(2, 1,
962 {neg(result.wheel_ground_velocity.get(0, 0)),
963 neg(result.wheel_ground_velocity.get(1, 0))});
964 add_dense_dense(negative_wheel_ground_velocity, wheel_velocity,
965 result.wheel_slip_velocity);
966
justinT21446e4f62024-06-16 22:36:10 -0700967 VLOG(1);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700968 VLOG(1) << "wheel ground velocity: "
969 << result.wheel_ground_velocity.__str__();
justinT21446e4f62024-06-16 22:36:10 -0700970
Austin Schuh5ddcb472024-07-21 17:55:34 -0700971 result.slip_angle = sin(neg(atan2(result.wheel_ground_velocity.get(1, 0),
972 result.wheel_ground_velocity.get(0, 0))));
justinT21446e4f62024-06-16 22:36:10 -0700973
974 VLOG(1);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700975 VLOG(1) << "slip angle: " << result.slip_angle->__str__();
justinT21446e4f62024-06-16 22:36:10 -0700976
Austin Schuh2a1abec2024-07-10 20:31:16 -0700977 // TODO(austin): Does this handle decel properly?
Austin Schuhb67a38f2024-07-04 13:48:38 -0700978 result.slip_ratio = div(
Austin Schuh2a1abec2024-07-10 20:31:16 -0700979 sub(mul(rw_, result.omegad), result.wheel_ground_velocity.get(0, 0)),
980 SymEngine::max(
981 {real_double(0.02), abs(result.wheel_ground_velocity.get(0, 0))}));
justinT21446e4f62024-06-16 22:36:10 -0700982 VLOG(1);
Austin Schuhb67a38f2024-07-04 13:48:38 -0700983 VLOG(1) << "Slip ratio " << result.slip_ratio->__str__();
justinT21446e4f62024-06-16 22:36:10 -0700984
Austin Schuh6ea789e2024-07-27 13:45:53 -0700985 result.full.Fwx = simplify(mul(Cx_, result.slip_ratio));
Austin Schuhb67a38f2024-07-04 13:48:38 -0700986 result.Fwy = simplify(mul(Cy_, result.slip_angle));
justinT21446e4f62024-06-16 22:36:10 -0700987
Austin Schuh27694fa2024-07-20 16:29:49 -0700988 // The self-aligning moment needs to flip when the module flips direction.
Austin Schuh78a1b312024-08-18 17:21:34 -0700989 RCP<const Basic> softsign_velocity = add(
990 div(integer(-2),
991 add(integer(1), exp(mul(integer(100),
992 result.wheel_ground_velocity.get(0, 0))))),
993 integer(1));
994 result.Ms =
995 mul(neg(result.Fwy),
996 add(div(mul(softsign_velocity, contact_patch_length_), integer(3)),
997 caster_));
justinT21446e4f62024-06-16 22:36:10 -0700998 VLOG(1);
Austin Schuhb8b34be2024-07-14 16:06:19 -0700999 VLOG(1) << "Ms " << result.Ms->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001000 VLOG(1);
Austin Schuh6ea789e2024-07-27 13:45:53 -07001001 VLOG(1) << "full.Fwx " << result.full.Fwx->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001002 VLOG(1);
Austin Schuhb67a38f2024-07-04 13:48:38 -07001003 VLOG(1) << "Fwy " << result.Fwy->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001004
Austin Schuh6ea789e2024-07-27 13:45:53 -07001005 // -K_td * Id / Gd + Fwx * rw = 0
1006 // Fwx = K_td * Id / Gd / rw
1007 result.direct.Fwx = mul(Ktd_, div(result.Id, mul(Gd_, rw_)));
1008
1009 result.direct.alphas_eqn =
1010 SteerAccel(result.direct.Fwx, result.Ms, result.Is);
1011
1012 // d/dt omegas = ...
1013 result.full.alphas_eqn = SteerAccel(result.full.Fwx, result.Ms, result.Is);
justinT21446e4f62024-06-16 22:36:10 -07001014
1015 VLOG(1);
Austin Schuh6ea789e2024-07-27 13:45:53 -07001016 VLOG(1) << alphas->__str__() << " = " << result.full.alphas_eqn->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001017
Austin Schuh6ea789e2024-07-27 13:45:53 -07001018 RCP<const Basic> lhs =
1019 sub(mul(sub(div(add(rp_, rs_), rp_), integer(1)), alphas),
1020 mul(Gd1_, mul(Gd2_, alphamd)));
1021 RCP<const Basic> ddplanitary_eqn = sub(mul(Gd3_, lhs), alphad);
justinT21446e4f62024-06-16 22:36:10 -07001022
Austin Schuh6ea789e2024-07-27 13:45:53 -07001023 RCP<const Basic> full_drive_eqn =
1024 sub(add(mul(neg(Jdm_), div(alphamd, Gd_)),
1025 mul(Ktd_, div(neg(result.Id), Gd_))),
1026 mul(neg(result.full.Fwx), rw_));
justinT21446e4f62024-06-16 22:36:10 -07001027
Austin Schuh6ea789e2024-07-27 13:45:53 -07001028 VLOG(1) << "full_drive_eqn: " << full_drive_eqn->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001029
1030 // Substitute in ddplanitary_eqn so we get rid of alphamd
1031 map_basic_basic map;
1032 RCP<const Set> reals = interval(NegInf, Inf, true, true);
1033 RCP<const Set> solve_solution = solve(ddplanitary_eqn, alphamd, reals);
1034 map[alphamd] = solve_solution->get_args()[1]->get_args()[0];
1035 VLOG(1) << "temp: " << solve_solution->__str__();
Austin Schuh6ea789e2024-07-27 13:45:53 -07001036 RCP<const Basic> drive_eqn_subs = full_drive_eqn->subs(map);
justinT21446e4f62024-06-16 22:36:10 -07001037
1038 map.clear();
Austin Schuh6ea789e2024-07-27 13:45:53 -07001039 map[alphas] = result.full.alphas_eqn;
justinT21446e4f62024-06-16 22:36:10 -07001040 RCP<const Basic> drive_eqn_subs2 = drive_eqn_subs->subs(map);
1041 RCP<const Basic> drive_eqn_subs3 = simplify(drive_eqn_subs2);
Austin Schuh6ea789e2024-07-27 13:45:53 -07001042 VLOG(1) << "full_drive_eqn simplified: " << drive_eqn_subs3->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001043
Austin Schuh6ea789e2024-07-27 13:45:53 -07001044 solve_solution = solve(drive_eqn_subs3, alphad, reals);
justinT21446e4f62024-06-16 22:36:10 -07001045
Austin Schuh6ea789e2024-07-27 13:45:53 -07001046 result.full.alphad_eqn =
justinT21446e4f62024-06-16 22:36:10 -07001047 simplify(solve_solution->get_args()[1]->get_args()[0]);
Austin Schuh6ea789e2024-07-27 13:45:53 -07001048 VLOG(1) << "drive_accel: " << result.full.alphad_eqn->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001049
Austin Schuh2a1abec2024-07-10 20:31:16 -07001050 // Compute the resulting force from the module.
Austin Schuh6ea789e2024-07-27 13:45:53 -07001051 result.full.F = DenseMatrix(2, 1);
Austin Schuhb8b34be2024-07-14 16:06:19 -07001052 mul_dense_dense(R(add(theta_, result.thetas)),
Austin Schuh6ea789e2024-07-27 13:45:53 -07001053 DenseMatrix(2, 1, {result.full.Fwx, result.Fwy}),
1054 result.full.F);
1055 result.full.torque = force_cross(result.mounting_location, result.full.F);
justinT21446e4f62024-06-16 22:36:10 -07001056
Austin Schuh6ea789e2024-07-27 13:45:53 -07001057 result.direct.F = DenseMatrix(2, 1);
1058 mul_dense_dense(R(add(theta_, result.thetas)),
1059 DenseMatrix(2, 1, {result.direct.Fwx, result.Fwy}),
1060 result.direct.F);
1061 result.direct.torque =
1062 force_cross(result.mounting_location, result.direct.F);
justinT21446e4f62024-06-16 22:36:10 -07001063
1064 VLOG(1);
Austin Schuh6ea789e2024-07-27 13:45:53 -07001065 VLOG(1) << "full torque = " << result.full.torque->__str__();
1066 VLOG(1) << "direct torque = " << result.full.torque->__str__();
justinT21446e4f62024-06-16 22:36:10 -07001067
justinT21446e4f62024-06-16 22:36:10 -07001068 return result;
1069 }
1070
1071 DenseMatrix R(const RCP<const Basic> theta) {
1072 return DenseMatrix(2, 2,
1073 {cos(theta), neg(sin(theta)), sin(theta), cos(theta)});
1074 }
1075
1076 DenseMatrix angle_cross(DenseMatrix a, RCP<const Basic> b) {
Austin Schuh2a1abec2024-07-10 20:31:16 -07001077 return DenseMatrix(2, 1, {mul(neg(a.get(1, 0)), b), mul(a.get(0, 0), b)});
justinT21446e4f62024-06-16 22:36:10 -07001078 }
1079
1080 RCP<const Basic> force_cross(DenseMatrix r, DenseMatrix f) {
1081 return sub(mul(r.get(0, 0), f.get(1, 0)), mul(r.get(1, 0), f.get(0, 0)));
1082 }
1083
1084 // z represents the number of teeth per gear, theta is the angle between
1085 // shafts(in degrees), D_02 is the pitch diameter of gear 2 and b_2 is the
1086 // length of the tooth of gear 2
1087 // returns std::pair(r_01, r_02)
1088 std::pair<double, double> GetBevelPitchRadius(double z1, double z2,
1089 double theta, double D_02,
1090 double b_2) {
1091 double gamma_1 = std::atan2(z1, z2);
1092 double gamma_2 = theta / 180.0 * std::numbers::pi - gamma_1;
1093 double R_m = D_02 / 2 / std::sin(gamma_2) - b_2 / 2;
1094 return std::pair(R_m * std::cos(gamma_2), R_m * std::sin(gamma_2));
1095 }
1096
1097 Motor drive_motor_;
1098 Motor steer_motor_;
1099
1100 RCP<const Basic> Cx_;
1101 RCP<const Basic> Cy_;
Austin Schuh2a1abec2024-07-10 20:31:16 -07001102 RCP<const Basic> rw_;
justinT21446e4f62024-06-16 22:36:10 -07001103 RCP<const Basic> m_;
1104 RCP<const Basic> J_;
1105 RCP<const Basic> Gd1_;
1106 RCP<const Basic> rs_;
1107 RCP<const Basic> rp_;
1108 RCP<const Basic> Gd2_;
1109 RCP<const Basic> rb1_;
1110 RCP<const Basic> rb2_;
1111 RCP<const Basic> Gd3_;
1112 RCP<const Basic> Gd_;
1113 RCP<const Basic> Js_;
1114 RCP<const Basic> Gs_;
1115 RCP<const Basic> wb_;
1116 RCP<const Basic> Jdm_;
1117 RCP<const Basic> Jsm_;
1118 RCP<const Basic> Kts_;
1119 RCP<const Basic> Ktd_;
1120 RCP<const Basic> robot_width_;
1121 RCP<const Basic> caster_;
1122 RCP<const Basic> contact_patch_length_;
1123 RCP<const Basic> x_;
1124 RCP<const Basic> y_;
1125 RCP<const Basic> theta_;
1126 RCP<const Basic> vx_;
1127 RCP<const Basic> vy_;
1128 RCP<const Basic> omega_;
1129 RCP<const Basic> ax_;
1130 RCP<const Basic> ay_;
1131 RCP<const Basic> atheta_;
1132
1133 std::array<Module, kNumModules> modules_;
1134
Austin Schuh6ea789e2024-07-27 13:45:53 -07001135 DenseMatrix full_accel_;
1136 RCP<const Basic> full_angular_accel_;
1137 DenseMatrix direct_accel_;
1138 RCP<const Basic> direct_angular_accel_;
justinT21446e4f62024-06-16 22:36:10 -07001139};
1140
1141} // namespace frc971::control_loops::swerve
1142
1143int main(int argc, char **argv) {
1144 aos::InitGoogle(&argc, &argv);
1145
1146 frc971::control_loops::swerve::SwerveSimulation sim;
1147
Austin Schuh99f7c6a2024-06-25 22:07:44 -07001148 if (!absl::GetFlag(FLAGS_cc_output_path).empty() &&
1149 !absl::GetFlag(FLAGS_h_output_path).empty()) {
1150 sim.Write(absl::GetFlag(FLAGS_cc_output_path),
1151 absl::GetFlag(FLAGS_h_output_path));
Austin Schuh0f881092024-06-28 15:36:48 -07001152 }
Austin Schuh99f7c6a2024-06-25 22:07:44 -07001153 if (!absl::GetFlag(FLAGS_casadi_py_output_path).empty()) {
1154 sim.WriteCasadi(absl::GetFlag(FLAGS_casadi_py_output_path));
Austin Schuh0f881092024-06-28 15:36:48 -07001155 }
justinT21446e4f62024-06-16 22:36:10 -07001156
1157 return 0;
1158}