Add tests for wpimath.system

Author: virtuald

subprojects/robotpy-wpimath/tests/test_system.py

@@ -1,6 +1,202 @@
+import math
+
 import wpimath.system
-import wpimath.system.plant
 
+import pytest
+import numpy as np
+
+
+def test_rk4_exponential():
+    """Test that integrating dx/dt = eˣ works"""
+    y0 = np.array([[0.0]])
+
+    y1 = wpimath.system.RK4(lambda x: np.array([[math.exp(x[0, 0])]]), y0, 0.1)
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rk4_exponential_with_u():
+    """Test that integrating dx/dt = eˣ works when we provide a u"""
+    y0 = np.array([[0.0]])
+
+    y1 = wpimath.system.RK4(
+        lambda x, u: np.array([[math.exp(u[0, 0] * x[0, 0])]]),
+        y0,
+        np.array([[1.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rk4_time_varying():
+    """
+    Tests RK4 with a time varying solution. From
+    http://www2.hawaii.edu/~jmcfatri/math407/RungeKuttaTest.html:
+
+        dx/dt = x (2 / (eᵗ + 1) - 1)
+
+    The true (analytical) solution is:
+
+        x(t) = 12eᵗ/(eᵗ + 1)²
+    """
+    y0 = np.array([[12.0 * math.exp(5.0) / math.pow(math.exp(5.0) + 1.0, 2.0)]])
+
+    y1 = wpimath.system.RK4(
+        lambda t, x: np.array([[x[0, 0] * (2.0 / (math.exp(t) + 1.0) - 1.0)]]),
+        5.0,
+        y0,
+        1.0,
+    )
+
+    expected = 12.0 * math.exp(6.0) / math.pow(math.exp(6.0) + 1.0, 2.0)
+    assert math.isclose(y1[0, 0], expected, abs_tol=1e-3)
+
+
+def test_rkdp_zero():
+    """Tests that integrating dx/dt = 0 works with RKDP"""
+    y1 = wpimath.system.RKDP(
+        lambda x, u: np.zeros((1, 1)),
+        np.array([[0.0]]),
+        np.array([[0.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], 0.0, abs_tol=1e-3)
+
+
+def test_rkdp_exponential():
+    """Tests that integrating dx/dt = eˣ works with RKDP"""
+    y0 = np.array([[0.0]])
+
+    y1 = wpimath.system.RKDP(
+        lambda x, u: np.array([[math.exp(x[0, 0])]]),
+        y0,
+        np.array([[0.0]]),
+        0.1,
+    )
+
+    assert math.isclose(y1[0, 0], math.exp(0.1) - math.exp(0.0), abs_tol=1e-3)
+
+
+def test_rkdp_time_varying():
+    """
+    Tests RKDP with a time varying solution. From
+    http://www2.hawaii.edu/~jmcfatri/math407/RungeKuttaTest.html:
+
+        dx/dt = x(2/(eᵗ + 1) - 1)
+
+    The true (analytical) solution is:
+
+        x(t) = 12eᵗ/(eᵗ + 1)²
+    """
+    y0 = np.array([[12.0 * math.exp(5.0) / math.pow(math.exp(5.0) + 1.0, 2.0)]])
+
+    y1 = wpimath.system.RKDP(
+        lambda t, x: np.array([[x[0, 0] * (2.0 / (math.exp(t) + 1.0) - 1.0)]]),
+        5.0,
+        y0,
+        1.0,
+        1e-12,
+    )
+
+    expected = 12.0 * math.exp(6.0) / math.pow(math.exp(6.0) + 1.0, 2.0)
+    assert math.isclose(y1[0, 0], expected, abs_tol=1e-3)
+
+
+def test_numerical_jacobian():
+    """Test that we can recover A from ax_fn() pretty accurately"""
+    a = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+            [1.0, 1.0, 3.0, 7.0],
+        ]
+    )
+
+    def ax_fn(x):
+        return a @ x
+
+    new_a = wpimath.system.numericalJacobian(ax_fn, np.zeros((4, 1)))
+    np.testing.assert_allclose(new_a, a, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_u_square():
+    """Test that we can recover B from axbu_fn() pretty accurately"""
+    a = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+            [1.0, 1.0, 3.0, 7.0],
+        ]
+    )
+    b = np.array([[1.0, 1.0], [2.0, 1.0], [3.0, 2.0], [3.0, 7.0]])
+
+    def axbu_fn(x, u):
+        return a @ x + b @ u
+
+    x0 = np.zeros((4, 1))
+    u0 = np.zeros((2, 1))
+    new_a = wpimath.system.numericalJacobianX(axbu_fn, x0, u0)
+    new_b = wpimath.system.numericalJacobianU(axbu_fn, x0, u0)
+    np.testing.assert_allclose(new_a, a, rtol=1e-6, atol=1e-5)
+    np.testing.assert_allclose(new_b, b, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_u_rectangular():
+    c = np.array(
+        [
+            [1.0, 2.0, 4.0, 1.0],
+            [5.0, 2.0, 3.0, 4.0],
+            [5.0, 1.0, 3.0, 2.0],
+        ]
+    )
+    d = np.array([[1.0, 1.0], [2.0, 1.0], [3.0, 2.0]])
+
+    def cxdu_fn(x, u):
+        return c @ x + d @ u
+
+    x0 = np.zeros((4, 1))
+    u0 = np.zeros((2, 1))
+    new_c = wpimath.system.numericalJacobianX(cxdu_fn, x0, u0)
+    new_d = wpimath.system.numericalJacobianU(cxdu_fn, x0, u0)
+    np.testing.assert_allclose(new_c, c, rtol=1e-6, atol=1e-5)
+    np.testing.assert_allclose(new_d, d, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_x_passes_extra_args():
+    a = np.array([[2.0, -1.0], [0.5, 3.0]])
+    b = np.array([[1.0], [4.0]])
+    x0 = np.zeros((2, 1))
+    u0 = np.zeros((1, 1))
+
+    seen = {}
+
+    def axbu_fn(x, u, scale, bias):
+        seen["args"] = (scale, bias)
+        return scale * (a @ x) + bias * (b @ u)
+
+    new_a = wpimath.system.numericalJacobianX(axbu_fn, x0, u0, 2.5, -3.0)
+
+    assert seen["args"] == (2.5, -3.0)
+    np.testing.assert_allclose(new_a, 2.5 * a, rtol=1e-6, atol=1e-5)
+
+
+def test_numerical_jacobian_u_passes_extra_args():
+    a = np.array([[1.0, 0.0], [0.0, -2.0]])
+    b = np.array([[1.5], [-0.5]])
+    x0 = np.zeros((2, 1))
+    u0 = np.zeros((1, 1))
+
+    seen = {}
+
+    def axbu_fn(x, u, scale, bias):
+        seen["args"] = (scale, bias)
+        return scale * (a @ x) + bias * (b @ u)
+
+    new_b = wpimath.system.numericalJacobianU(axbu_fn, x0, u0, 4.0, 0.25)
 
-def test_todo():
-    pass
+    assert seen["args"] == (4.0, 0.25)
+    np.testing.assert_allclose(new_b, 0.25 * b, rtol=1e-6, atol=1e-5)