frc971/control_loops/swerve/physics_test.py

#!/usr/bin/env python3
import jax

# casadi uses doubles.  jax likes floats.  We want to make sure the physics
# matches really precisely, so force doubles for the tests.
jax.config.update("jax_enable_x64", True)

import numpy

numpy.set_printoptions(precision=20)
import sys, os
import casadi
import scipy
from numpy.testing import assert_array_equal, assert_array_almost_equal
import unittest

from frc971.control_loops.swerve import bigcaster_dynamics
from frc971.control_loops.swerve import dynamics
from frc971.control_loops.swerve import nocaster_dynamics
from frc971.control_loops.swerve import physics_test_utils as utils
from frc971.control_loops.swerve import casadi_velocity_mpc_lib
from frc971.control_loops.swerve import jax_dynamics
from frc971.control_loops.swerve.cpp_dynamics import swerve_dynamics as cpp_dynamics


class TestSwervePhysics(unittest.TestCase):
    I = numpy.zeros((8, 1))
    coefficients = jax_dynamics.Coefficients()

    def to_velocity_state(self, X):
        return dynamics.to_velocity_state(X)

    def swerve_full_dynamics(self, X, U, skip_compare=False):
        X_velocity = self.to_velocity_state(X)
        Xdot = self.position_swerve_full_dynamics(X, U)

        if not skip_compare:
            Xdot_velocity = self.to_velocity_state(Xdot)
            velocity_physics = self.velocity_swerve_physics(X_velocity, U)

            self.assertLess(
                numpy.linalg.norm(Xdot_velocity - velocity_physics),
                2e-2,
                msg=
                f'Norm failed, full physics -> {X_velocity.T}, velocity physics -> {velocity_physics}, difference -> {velocity_physics - X_velocity}',
            )

        self.validate_dynamics_equality(X, U)

        return Xdot

    def validate_dynamics_equality(self, X, U):
        """Tests that both the JAX code and casadi code produce identical answers.

        Note:
          If the symengine code has been updated, you likely need to update the JAX
          by hand.  We had trouble code generating it with good performance.
        """
        X_velocity = self.to_velocity_state(X)

        Xdot = self.position_swerve_full_dynamics(X, U)
        Xdot_jax = jax_dynamics.full_dynamics(self.coefficients, X[:, 0], U[:,
                                                                            0])
        self.assertEqual(Xdot.shape[0], Xdot_jax.shape[0])

        self.assertLess(
            numpy.linalg.norm(Xdot[:, 0] - Xdot_jax),
            2e-8,
            msg=
            f'Xdot: {Xdot[:, 0]}, Xdot_jax: {Xdot_jax}, diff: {(Xdot[:, 0] - Xdot_jax)}',
        )

        velocity_physics = self.velocity_swerve_physics(X_velocity, U)
        velocity_physics_jax = jax_dynamics.velocity_dynamics(
            self.coefficients, X_velocity[:, 0], U[:, 0])

        self.assertEqual(velocity_physics.shape[0],
                         velocity_physics_jax.shape[0])

        self.assertLess(
            numpy.linalg.norm(velocity_physics[:, 0] - velocity_physics_jax),
            2e-8,
            msg=
            f'Xdot: {velocity_physics[:, 0]}, Xdot_jax: {velocity_physics_jax}, diff: {(velocity_physics[:, 0] - velocity_physics_jax)}',
        )

    def wrap_and_validate(self, function, i):
        """Wraps a function, and validates JAX and casadi agree.

        We want to do it every time we check any intermediate, since the tests
        are designed to test all the corner cases, but they don't all do it
        through the main dynamics function above.
        """
        wrapped_fn = utils.wrap_module(function, i)

        def do(X, U):
            self.validate_dynamics_equality(X, U)
            return wrapped_fn(X, U)

        return do

    def wrap(self, python_module):
        # Only update on change to avoid re-jiting things.
        if self.coefficients.caster != python_module.CASTER:
            self.coefficients = self.coefficients._replace(
                caster=python_module.CASTER)

        self.position_swerve_full_dynamics = utils.wrap(
            python_module.swerve_full_dynamics)

        evaluated_fn = python_module.velocity_swerve_physics(
            casadi.SX.sym("X", dynamics.NUM_VELOCITY_STATES, 1),
            casadi.SX.sym("U", 8, 1))
        self.velocity_swerve_physics = lambda X, U: numpy.array(
            evaluated_fn(X, U))

        self.contact_patch_velocity = [
            self.wrap_and_validate(python_module.contact_patch_velocity, i)
            for i in range(4)
        ]
        self.wheel_ground_velocity = [
            self.wrap_and_validate(python_module.wheel_ground_velocity, i)
            for i in range(4)
        ]
        self.wheel_slip_velocity = [
            self.wrap_and_validate(python_module.wheel_slip_velocity, i)
            for i in range(4)
        ]
        self.wheel_force = [
            self.wrap_and_validate(python_module.wheel_force, i)
            for i in range(4)
        ]
        self.module_angular_accel = [
            self.wrap_and_validate(python_module.module_angular_accel, i)
            for i in range(4)
        ]
        self.F = [self.wrap_and_validate(python_module.F, i) for i in range(4)]
        self.mounting_location = [
            self.wrap_and_validate(python_module.mounting_location, i)
            for i in range(4)
        ]

        self.slip_angle = [
            self.wrap_and_validate(python_module.slip_angle, i)
            for i in range(4)
        ]
        self.slip_ratio = [
            self.wrap_and_validate(python_module.slip_ratio, i)
            for i in range(4)
        ]
        self.Ms = [
            self.wrap_and_validate(python_module.Ms, i) for i in range(4)
        ]

    def setUp(self):
        self.wrap(dynamics)

    def test_contact_patch_velocity(self):
        """Tests that the contact patch velocity makes sense."""
        for i in range(4):
            contact_patch_velocity = self.contact_patch_velocity[i]
            wheel_ground_velocity = self.wheel_ground_velocity[i]

            # No angular velocity should result in just linear motion.
            for velocity in [
                    numpy.array([[1.5], [0.0]]),
                    numpy.array([[0.0], [1.0]]),
                    numpy.array([[-1.5], [0.0]]),
                    numpy.array([[0.0], [-1.0]]),
                    numpy.array([[2.0], [-1.7]]),
            ]:
                for theta in [0.0, 1.0, -1.0, 100.0]:
                    patch_velocity = contact_patch_velocity(
                        utils.state_vector(velocity=velocity, theta=theta),
                        self.I)

                    assert_array_equal(patch_velocity, velocity)

            # Now, test that spinning the robot results in module velocities.
            # We are assuming that each module is on a square robot.
            module_center_of_mass_angle = i * numpy.pi / 2.0 + numpy.pi / 4.0
            for theta in [-module_center_of_mass_angle, 0.0, 1.0, -1.0, 100.0]:
                for omega in [0.65, -0.1]:
                    # Point each module to the center to make the math easier.
                    patch_velocity = contact_patch_velocity(
                        utils.state_vector(
                            velocity=numpy.array([[0.0], [0.0]]),
                            theta=theta,
                            omega=omega,
                            module_angle=module_center_of_mass_angle,
                        ),
                        self.I,
                    )

                    assert_array_almost_equal(
                        patch_velocity,
                        (dynamics.ROBOT_WIDTH / numpy.sqrt(2.0) -
                         dynamics.CASTER) * omega * numpy.array([
                             [-numpy.sin(theta + module_center_of_mass_angle)],
                             [numpy.cos(theta + module_center_of_mass_angle)],
                         ]),
                    )

                    # Point the wheel along +x, rotate it by theta, then spin it.
                    # Confirm the velocities come out right.
                    patch_velocity = contact_patch_velocity(
                        utils.state_vector(
                            velocity=numpy.array([[0.0], [0.0]]),
                            theta=-module_center_of_mass_angle,
                            module_omega=omega,
                            module_angle=(theta + module_center_of_mass_angle),
                        ),
                        self.I,
                    )

                    assert_array_almost_equal(
                        patch_velocity,
                        -dynamics.CASTER * omega *
                        numpy.array([[-numpy.sin(theta)], [numpy.cos(theta)]]),
                    )

            # Now, test that the rotation back to wheel coordinates works.
            # The easiest way to do this is to point the wheel in a direction,
            # move in that direction, and confirm that there is no lateral velocity.
            for robot_angle in [0.0, 1.0, -5.0]:
                for module_angle in [0.0, 1.0, -5.0]:
                    wheel_patch_velocity = numpy.array(
                        wheel_ground_velocity(
                            utils.state_vector(
                                velocity=numpy.array([
                                    [numpy.cos(robot_angle + module_angle)],
                                    [numpy.sin(robot_angle + module_angle)],
                                ]),
                                theta=robot_angle,
                                module_angle=module_angle,
                            ),
                            self.I,
                        ))

                    assert_array_almost_equal(wheel_patch_velocity,
                                              numpy.array([[1], [0]]))

    def test_slip_angle(self):
        """Tests that the slip_angle calculation works."""
        velocity = numpy.array([[1.5], [0.0]])

        for i in range(4):
            for wrap in range(-1, 2):
                for theta in [0.0, 0.6, -0.4]:
                    module_angle = numpy.pi * wrap + theta

                    # We have redefined the angle to be the softened sin of the angle.
                    # That way, when the module flips directions, the slip angle also flips
                    # directions to keep it stable.
                    computed_angle = self.slip_angle[i](
                        utils.state_vector(velocity=velocity,
                                           module_angle=module_angle),
                        self.I,
                    )[0, 0]

                    # Compute out the expected value directly to confirm we got it right.
                    loggain = 20.0
                    vy = 1.5 * numpy.sin(-module_angle)
                    vx = 1.5 * numpy.cos(-module_angle)
                    expected = numpy.sin(-numpy.arctan2(
                        vy,
                        scipy.special.logsumexp([1.0, abs(vx) * loggain]) /
                        loggain))

                    jax_expected = jax.numpy.sin(
                        -jax_dynamics.soft_atan2(vy, vx))

                    self.assertAlmostEqual(
                        expected,
                        jax_expected,
                        msg=f"Trying wrap {wrap} theta {theta}",
                    )

                    self.assertAlmostEqual(
                        expected,
                        computed_angle,
                        msg=f"Trying wrap {wrap} theta {theta}",
                    )

    def test_wheel_torque(self):
        """Tests that the per module self aligning forces have the right signs."""
        # Point all the modules in a little bit.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[-0.001, -0.001, 0.001, 0.001],
        )
        xdot_equal = self.swerve_full_dynamics(X, self.I)

        self.assertGreater(xdot_equal[2, 0], 0.0)
        self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
        self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_equal[6, 0], 0.0)
        self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
        self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_equal[10, 0], 0.0)
        self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
        self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_equal[14, 0], 0.0)
        self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
        self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)

        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)

        # Now, make the bot want to go left by going to the other side.
        # The wheels will be going too fast based on our calcs, so they should be decelerating.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[0.01, 0.01, 0.01, 0.01],
        )
        xdot_left = self.swerve_full_dynamics(X, self.I)

        self.assertLess(xdot_left[2, 0], -0.05)
        self.assertLess(xdot_left[3, 0], 0.0)
        self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[6, 0], -0.05)
        self.assertLess(xdot_left[7, 0], 0.0)
        self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[10, 0], -0.05)
        self.assertLess(xdot_left[11, 0], 0.0)
        self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[14, 0], -0.05)
        self.assertLess(xdot_left[15, 0], 0.0)
        self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[19, 0], 0.0001)
        self.assertGreater(xdot_left[20, 0], 0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_left[21, 0], 0.0)

        # And now do it to the right too.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[-0.01, -0.01, -0.01, -0.01],
        )
        xdot_right = self.swerve_full_dynamics(X, self.I)

        self.assertGreater(xdot_right[2, 0], 0.05)
        self.assertLess(xdot_right[3, 0], 0.0)
        self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[6, 0], 0.05)
        self.assertLess(xdot_right[7, 0], 0.0)
        self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[10, 0], 0.05)
        self.assertLess(xdot_right[11, 0], 0.0)
        self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[14, 0], 0.05)
        self.assertLess(xdot_right[15, 0], 0.0)
        self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[19, 0], 0.0001)
        self.assertLess(xdot_right[20, 0], -0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_right[21, 0], 0.0)

    def test_wheel_torque_backwards_nocaster(self):
        """Tests that the per module self aligning forces have the right signs when going backwards."""
        self.wrap(nocaster_dynamics)
        # Point all the modules in a little bit, going backwards.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[
                numpy.pi - 0.001,
                numpy.pi - 0.001,
                numpy.pi + 0.001,
                numpy.pi + 0.001,
            ],
            drive_wheel_velocity=-1,
        )
        xdot_equal = self.swerve_full_dynamics(X, self.I)

        self.assertGreater(xdot_equal[2, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
        self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_equal[6, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
        self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_equal[10, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
        self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_equal[14, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
        self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)

        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)

        # Now, make the bot want to go left by going to the other side.
        # The wheels will be going too fast based on our calcs, so they should be decelerating.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[numpy.pi + 0.01] * 4,
            drive_wheel_velocity=-1,
        )
        xdot_left = self.swerve_full_dynamics(X, self.I)

        self.assertLess(xdot_left[2, 0], -0.05)
        self.assertGreater(xdot_left[3, 0], 0.0)
        self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[6, 0], -0.05)
        self.assertGreater(xdot_left[7, 0], 0.0)
        self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[10, 0], -0.05)
        self.assertGreater(xdot_left[11, 0], 0.0)
        self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_left[14, 0], -0.05)
        self.assertGreater(xdot_left[15, 0], 0.0)
        self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[19, 0], 0.0001)
        self.assertGreater(xdot_left[20, 0], 0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_left[21, 0], 0.0)

        # And now do it to the right too.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            drive_wheel_velocity=-1,
            module_angles=[-0.01 + numpy.pi] * 4,
        )
        xdot_right = self.swerve_full_dynamics(X, self.I)

        self.assertGreater(xdot_right[2, 0], 0.05)
        self.assertGreater(xdot_right[3, 0], 0.0)
        self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[6, 0], 0.05)
        self.assertGreater(xdot_right[7, 0], 0.0)
        self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[10, 0], 0.05)
        self.assertGreater(xdot_right[11, 0], 0.0)
        self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[14, 0], 0.05)
        self.assertGreater(xdot_right[15, 0], 0.0)
        self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[19, 0], 0.0001)
        self.assertLess(xdot_right[20, 0], -0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_right[21, 0], 0.0)

    def test_wheel_torque_backwards_caster(self):
        """Tests that the per module self aligning forces have the right signs when going backwards with a lot of caster."""
        self.wrap(bigcaster_dynamics)
        # Point all the modules in a little bit, going backwards.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[
                numpy.pi - 0.001,
                numpy.pi - 0.001,
                numpy.pi + 0.001,
                numpy.pi + 0.001,
            ],
            drive_wheel_velocity=-1,
        )
        xdot_equal = self.swerve_full_dynamics(X, self.I)

        self.assertLess(xdot_equal[2, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[3, 0], 0.0, places=1)
        self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_equal[6, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[7, 0], 0.0, places=1)
        self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_equal[10, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[11, 0], 0.0, places=1)
        self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_equal[14, 0], 0.0, msg="Steering backwards")
        self.assertAlmostEqual(xdot_equal[15, 0], 0.0, places=1)
        self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)

        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_equal[21, 0], 0.0, places=2)

        # Now, make the bot want to go left by going to the other side.
        # The wheels will be going too fast based on our calcs, so they should be decelerating.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            module_angles=[numpy.pi + 0.01] * 4,
            drive_wheel_velocity=-1,
        )
        xdot_left = self.swerve_full_dynamics(X, self.I)

        self.assertGreater(xdot_left[2, 0], -0.05)
        self.assertGreater(xdot_left[3, 0], 0.0)
        self.assertGreater(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[6, 0], -0.05)
        self.assertGreater(xdot_left[7, 0], 0.0)
        self.assertGreater(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[10, 0], -0.05)
        self.assertGreater(xdot_left[11, 0], 0.0)
        self.assertGreater(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[14, 0], -0.05)
        self.assertGreater(xdot_left[15, 0], 0.0)
        self.assertGreater(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_left[19, 0], 0.0001)
        self.assertGreater(xdot_left[20, 0], 0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_left[21, 0], 0.0)

        # And now do it to the right too.
        X = utils.state_vector(
            velocity=numpy.array([[1.0], [0.0]]),
            drive_wheel_velocity=-1,
            module_angles=[-0.01 + numpy.pi] * 4,
        )
        xdot_right = self.swerve_full_dynamics(X, self.I)

        self.assertLess(xdot_right[2, 0], 0.05)
        self.assertGreater(xdot_right[3, 0], 0.0)
        self.assertLess(self.Ms[0](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_right[6, 0], 0.05)
        self.assertGreater(xdot_right[7, 0], 0.0)
        self.assertLess(self.Ms[1](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_right[10, 0], 0.05)
        self.assertGreater(xdot_right[11, 0], 0.0)
        self.assertLess(self.Ms[2](X, self.I)[0, 0], 0.0)

        self.assertLess(xdot_right[14, 0], 0.05)
        self.assertGreater(xdot_right[15, 0], 0.0)
        self.assertLess(self.Ms[3](X, self.I)[0, 0], 0.0)

        self.assertGreater(xdot_right[19, 0], 0.0001)
        self.assertLess(xdot_right[20, 0], -0.05)
        # Shouldn't be spinning.
        self.assertAlmostEqual(xdot_right[21, 0], 0.0)

    def test_wheel_forces(self):
        """Tests that the per module forces have the right signs."""
        for i in range(4):
            wheel_force = self.wheel_force[i]

            X = utils.state_vector()
            robot_equal = wheel_force(X, self.I)
            xdot_equal = self.swerve_full_dynamics(X, self.I)
            self.assertEqual(robot_equal[0, 0], 0.0)
            self.assertEqual(robot_equal[1, 0], 0.0)
            self.assertEqual(xdot_equal[2 + 4 * i], 0.0)
            self.assertEqual(xdot_equal[3 + 4 * i], 0.0)

            # Robot is moving faster than the wheels, it should decelerate.
            X = utils.state_vector(dx=0.01)
            robot_faster = wheel_force(X, self.I)
            xdot_faster = self.swerve_full_dynamics(X,
                                                    self.I,
                                                    skip_compare=True)
            self.assertLess(robot_faster[0, 0], -0.1)
            self.assertEqual(robot_faster[1, 0], 0.0)
            self.assertGreater(xdot_faster[3 + 4 * i], 0.0)

            # Robot is now going slower than the wheels.  It should accelerate.
            X = utils.state_vector(dx=-0.01)
            robot_slower = wheel_force(X, self.I)
            xdot_slower = self.swerve_full_dynamics(X,
                                                    self.I,
                                                    skip_compare=True)
            self.assertGreater(robot_slower[0, 0], 0.1)
            self.assertEqual(robot_slower[1, 0], 0.0)
            self.assertLess(xdot_slower[3 + 4 * i], 0.0)

            # Positive lateral velocity -> negative force.
            robot_left = wheel_force(utils.state_vector(dy=0.01), self.I)
            self.assertEqual(robot_left[0, 0], 0.0)
            self.assertLess(robot_left[1, 0], -0.1)

            # Negative lateral velocity -> positive force.
            robot_right = wheel_force(utils.state_vector(dy=-0.01), self.I)
            self.assertEqual(robot_right[0, 0], 0.0)
            self.assertGreater(robot_right[1, 0], 0.1)

    def test_steer_accel(self):
        """Tests that applying a steer torque accelerates the steer reasonably."""

        for velocity in [
                numpy.array([[0.0], [0.0]]),
                numpy.array([[1.0], [0.0]]),
                numpy.array([[2.0], [0.0]])
        ]:
            module_angles = [0.0] * 4

            X = utils.state_vector(
                velocity=velocity,
                omega=0.0,
                module_angles=module_angles,
            )
            steer_I = numpy.array([[1.0], [0.0]] * 4)
            xdot = self.swerve_full_dynamics(X, steer_I)

            self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS0, 0],
                                   1.4,
                                   places=0)
            self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS1, 0],
                                   1.4,
                                   places=0)
            self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS2, 0],
                                   1.4,
                                   places=0)
            self.assertAlmostEqual(xdot[dynamics.STATE_OMEGAS3, 0],
                                   1.4,
                                   places=0)

    def test_drive_accel(self):
        """Tests that applying a drive torque accelerates the drive reasonably."""

        for velocity in [
                numpy.array([[0.01], [0.0]]),
                numpy.array([[1.0], [0.0]]),
                numpy.array([[2.0], [0.0]])
        ]:
            module_angles = [0.0] * 4

            X = utils.state_vector(
                velocity=velocity,
                omega=0.0,
                module_angles=module_angles,
            )
            steer_I = numpy.array([[0.0], [100.0]] * 4)
            # 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.
            xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)

            self.assertGreater(xdot[dynamics.STATE_OMEGAD0, 0], 100.0)
            self.assertGreater(xdot[dynamics.STATE_OMEGAD1, 0], 100.0)
            self.assertGreater(xdot[dynamics.STATE_OMEGAD2, 0], 100.0)
            self.assertGreater(xdot[dynamics.STATE_OMEGAD3, 0], 100.0)

            X = utils.state_vector(
                velocity=velocity,
                omega=0.0,
                module_angles=module_angles,
            )
            steer_I = numpy.array([[0.0], [-100.0]] * 4)
            xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)

            self.assertLess(xdot[dynamics.STATE_OMEGAD0, 0], -100.0)
            self.assertLess(xdot[dynamics.STATE_OMEGAD1, 0], -100.0)
            self.assertLess(xdot[dynamics.STATE_OMEGAD2, 0], -100.0)
            self.assertLess(xdot[dynamics.STATE_OMEGAD3, 0], -100.0)

    def test_steer_coupling(self):
        """Tests that the steer coupling factor cancels out steer coupling torque."""
        steer_I = numpy.array(
            [[dynamics.STEER_CURRENT_COUPLING_FACTOR * 10.0], [10.0]] * 4)

        X = utils.state_vector(
            velocity=numpy.array([[0.0], [0.0]]),
            omega=0.0,
        )
        X_velocity = self.to_velocity_state(X)
        Xdot = self.velocity_swerve_physics(X_velocity, steer_I)

        self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS0, 0], 0.0)
        self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS1, 0], 0.0)
        self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS2, 0], 0.0)
        self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS3, 0], 0.0)

    def test_constant_torque(self):
        """Tests that the robot exerts the same amount of torque no matter the orientation"""
        steer_I = numpy.reshape(numpy.array([(i % 2) * 20 for i in range(8)]),
                                (8, 1))

        X = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
                               omega=0.0,
                               module_angles=[
                                   3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
                                   -numpy.pi / 4.0, numpy.pi / 4.0
                               ],
                               drive_wheel_velocity=1.0)
        Xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)

        X_rot = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
                                   omega=0.0,
                                   theta=numpy.pi,
                                   module_angles=[
                                       3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
                                       -numpy.pi / 4.0, numpy.pi / 4.0
                                   ],
                                   drive_wheel_velocity=1.0)
        Xdot_rot = self.swerve_full_dynamics(X_rot, steer_I, skip_compare=True)

        self.assertGreater(Xdot[dynamics.STATE_OMEGA, 0], 0.0)
        self.assertAlmostEqual(Xdot[dynamics.STATE_OMEGA, 0],
                               Xdot_rot[dynamics.STATE_OMEGA, 0])

    def test_cost_equality(self):
        """Tests that the casadi and jax cost functions match."""
        mpc = casadi_velocity_mpc_lib.MPC(jit=False)
        cost = mpc.make_cost()

        for i in range(10):
            X = numpy.random.uniform(size=(dynamics.NUM_VELOCITY_STATES, ))
            U = numpy.random.uniform(low=-10, high=10, size=(8, ))
            R = numpy.random.uniform(low=-1, high=1, size=(3, ))

            J = numpy.array(cost(X, U, R))[0, 0]
            jax_J = jax_dynamics.mpc_cost(self.coefficients, X, U, R)

            self.assertAlmostEqual(J, jax_J)

        R = jax.numpy.array([0.0, 0.0, 1.0])

        # Now try spinning in place and make sure the cost doesn't change.
        # This tells us if we got our rotations right.
        steer_I = numpy.array([(i % 2) * 20 for i in range(8)])

        X = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
                               omega=0.0,
                               module_angles=[
                                   3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
                                   -numpy.pi / 4.0, numpy.pi / 4.0
                               ],
                               drive_wheel_velocity=1.0)

        jax_J_orig = jax_dynamics.mpc_cost(self.coefficients,
                                           self.to_velocity_state(X)[:, 0],
                                           steer_I, R)

        X_rotated = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
                                       omega=0.0,
                                       theta=numpy.pi,
                                       module_angles=[
                                           3 * numpy.pi / 4.0,
                                           -3 * numpy.pi / 4.0,
                                           -numpy.pi / 4.0, numpy.pi / 4.0
                                       ],
                                       drive_wheel_velocity=1.0)
        jax_J_rotated = jax_dynamics.mpc_cost(
            self.coefficients,
            self.to_velocity_state(X_rotated)[:, 0], steer_I, R)

        self.assertAlmostEqual(jax_J_orig, jax_J_rotated)

    def test_cpp_consistency(self):
        """Tests that the C++ physics are consistent with the Python physics."""
        # TODO(james): Currently the physics only match at X = 0 and U = 0.
        # Fix this.
        # Maybe due to different atan2 implementations?
        # TODO(james): Fold this into the general comparisons for JAX versus
        # casadi once the physics actually match.
        for current in [0]:
            print(f"Current: {current}")
            steer_I = numpy.zeros((8, 1)) + current
            for state_values in [0.0]:
                print(f"States all set to: {state_values}")
                X = numpy.zeros((dynamics.NUM_STATES, 1)) + state_values
                Xdot_py = self.swerve_full_dynamics(X,
                                                    steer_I,
                                                    skip_compare=True)
                Xdot_cpp = numpy.array(
                    cpp_dynamics(X.flatten().tolist(),
                                 steer_I.flatten().tolist())).reshape((25, 1))
                for index in range(dynamics.NUM_STATES):
                    self.assertAlmostEqual(Xdot_py[index, 0], Xdot_cpp[index,
                                                                       0])


if __name__ == "__main__":
    unittest.main()