Remove yaw frame to reduce the number of operations in the EoMs

src/brim/bicycle/whipple_bicycle.py

@@ -9,6 +9,7 @@
     PinJoint,
     ReferenceFrame,
     System,
+    Vector,
     dynamicsymbols,
 )
 
@@ -107,10 +108,8 @@ def _define_kinematics(self) -> None:
             self.ground.frame,
             self.rear_tire.contact_point.vel(self.ground.frame).xreplace(qd_repl))
         # Define the orientation of the rear frame.
-        yaw_frame = ReferenceFrame("yaw_frame")
         roll_frame = ReferenceFrame("roll_frame")
-        yaw_frame.orient_axis(self.ground.frame, self.ground.frame.z, self.q[2])
-        roll_frame.orient_axis(yaw_frame, yaw_frame.x, self.q[3])
+        roll_frame.orient_body_fixed(self.ground.frame, (*self.q[2:4], 0), "zxy")
         self.rear_frame.wheel_hub.frame.orient_axis(roll_frame, roll_frame.y, self.q[4])
         self.rear_frame.wheel_hub.frame.set_ang_vel(
             self.ground.frame,
@@ -141,8 +140,15 @@ def _define_kinematics(self) -> None:
             normal = self.ground.get_normal(self.rear_tire.contact_point)
             direction = normal.dot(-self.ground.frame.z)
             self.rear_tire.upward_radial_axis = direction * -roll_frame.z
-            self.rear_tire.longitudinal_axis = direction * yaw_frame.x
-            self.rear_tire.lateral_axis = yaw_frame.y
+            self.rear_tire.longitudinal_axis = direction * roll_frame.x
+            # It is for some reason highly inefficient to define a yaw frame between the
+            # ground and roll frame. Instead, we define two disconnected frames to get
+            # the y axis of the yaw frame efficiently expressed in the roll frame.
+            fake_roll = ReferenceFrame("fake_roll")
+            fake_yaw = ReferenceFrame("fake_yaw")
+            fake_roll.orient_axis(fake_yaw, fake_yaw.x, self.q[3])
+            self.rear_tire.lateral_axis = Vector({
+                roll_frame: fake_yaw.y.to_matrix(fake_roll)})
         self.rear_tire.define_kinematics()
         self.front_tire.define_kinematics()
         # Add the coordinates and speeds to the system.

tests/bicycle/test_whipple_bicycle.py

@@ -149,19 +149,19 @@ def test_define_tire_axes(self, _setup_default, normal, upward, longitudinal,
                               lateral) -> None:
         self.bike.ground = FlatGround("ground", normal=normal)
         self.bike.define_all()
-        assert len(self.bike.rear_tire.upward_radial_axis.args) == 1
-        assert self.bike.rear_tire.upward_radial_axis.args[0][0] == upward
-        assert len(self.bike.rear_tire.longitudinal_axis.args) == 1
-        assert self.bike.rear_tire.longitudinal_axis.args[0][0] == longitudinal
-        assert len(self.bike.rear_tire.lateral_axis.args) == 1
-        assert self.bike.rear_tire.lateral_axis.args[0][0] == lateral
+        tire = self.bike.rear_tire
+        assert count_ops(tire.upward_radial_axis.to_matrix(self.bike.ground.frame)) < 30
+        assert count_ops(tire.longitudinal_axis.to_matrix(self.bike.ground.frame)) < 30
+        assert count_ops(tire.lateral_axis.to_matrix(self.bike.ground.frame)) < 30
 
     @pytest.mark.parametrize("normal", ["+x", "-x", "+y", "-y"])
     def test_do_not_define_upward_radial_axis(self, _setup_default, normal) -> None:
         self.bike.ground = FlatGround("ground", normal=normal)
         self.bike.define_all()
-        radial_axis = self.bike.rear_tire.upward_radial_axis
-        assert count_ops(radial_axis.args[0][0]) > 20
+        tire = self.bike.rear_tire
+        assert count_ops(tire.upward_radial_axis.to_matrix(self.bike.ground.frame)) > 30
+        assert count_ops(tire.longitudinal_axis.to_matrix(self.bike.ground.frame)) > 30
+        assert count_ops(tire.lateral_axis.to_matrix(self.bike.ground.frame)) > 30
 
     @pytest.mark.parametrize("compute_rear", [True, False])
     @pytest.mark.parametrize("compute_front", [True, False])