IQ.Pilot Release Commit @ bec7652

artifacts/package_sources/rednose/examples/test_compare.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+import pytest
+import os
+import sys
+import sympy as sp
+import numpy as np
+
+if __name__ == '__main__':  # generating sympy code
+  from rednose.helpers.ekf_sym import gen_code
+else:
+  from rednose.helpers.ekf_sym_pyx import EKF_sym_pyx # pylint: disable=no-name-in-module
+  from rednose.helpers.ekf_sym import EKF_sym as EKF_sym2
+
+
+GENERATED_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), 'generated'))
+
+
+class ObservationKind:
+  UNKNOWN = 0
+  NO_OBSERVATION = 1
+  POSITION = 1
+
+  names = [
+    'Unknown',
+    'No observation',
+    'Position'
+  ]
+
+  @classmethod
+  def to_string(cls, kind):
+    return cls.names[kind]
+
+
+class States:
+  POSITION = slice(0, 1)
+  VELOCITY = slice(1, 2)
+
+
+class CompareFilter:
+  name = "compare"
+
+  initial_x = np.array([0.5, 0.0])
+  initial_P_diag = np.array([1.0**2, 1.0**2])
+  Q = np.diag([0.1**2, 2.0**2])
+  obs_noise = {ObservationKind.POSITION: np.atleast_2d(0.1**2)}
+
+  @staticmethod
+  def generate_code(generated_dir):
+    name = CompareFilter.name
+    dim_state = CompareFilter.initial_x.shape[0]
+
+    state_sym = sp.MatrixSymbol('state', dim_state, 1)
+    state = sp.Matrix(state_sym)
+
+    position = state[States.POSITION, :][0,:]
+    velocity = state[States.VELOCITY, :][0,:]
+
+    dt = sp.Symbol('dt')
+    state_dot = sp.Matrix(np.zeros((dim_state, 1)))
+    state_dot[States.POSITION.start, 0] = velocity
+    f_sym = state + dt * state_dot
+
+    obs_eqs = [
+      [sp.Matrix([position]), ObservationKind.POSITION, None],
+    ]
+
+    gen_code(generated_dir, name, f_sym, dt, state_sym, obs_eqs, dim_state, dim_state)
+
+  def __init__(self, generated_dir):
+    dim_state = self.initial_x.shape[0]
+    dim_state_err = self.initial_P_diag.shape[0]
+
+    # init filter
+    self.filter_py = EKF_sym_pyx(generated_dir, self.name, self.Q, self.initial_x, np.diag(self.initial_P_diag), dim_state, dim_state_err)
+    self.filter_pyx = EKF_sym2(generated_dir, self.name, self.Q, self.initial_x, np.diag(self.initial_P_diag), dim_state, dim_state_err)
+
+  def get_R(self, kind, n):
+    obs_noise = self.obs_noise[kind]
+    dim = obs_noise.shape[0]
+    R = np.zeros((n, dim, dim))
+    for i in range(n):
+      R[i, :, :] = obs_noise
+    return R
+
+
+class TestCompare:
+  def test_compare(self):
+    np.random.seed(0)
+
+    kf = CompareFilter(GENERATED_DIR)
+
+    # Simple simulation
+    dt = 0.01
+    ts = np.arange(0, 5, step=dt)
+    xs = np.empty(ts.shape)
+
+    # Simulate
+    x = 0.0
+    for i, v in enumerate(np.sin(ts * 5)):
+      xs[i] = x
+      x += v * dt
+
+    # insert late observation
+    switch = (20, 40)
+    ts[switch[0]], ts[switch[1]] = ts[switch[1]], ts[switch[0]]
+    xs[switch[0]], xs[switch[1]] = xs[switch[1]], xs[switch[0]]
+
+    for t, x in zip(ts, xs):
+      # get measurement
+      meas = np.random.normal(x, 0.1)
+      z = np.array([[meas]])
+      R = kf.get_R(ObservationKind.POSITION, 1)
+
+      # Update kf
+      kf.filter_py.predict_and_update_batch(t, ObservationKind.POSITION, z, R)
+      kf.filter_pyx.predict_and_update_batch(t, ObservationKind.POSITION, z, R)
+
+      assert kf.filter_py.get_filter_time() == pytest.approx(kf.filter_pyx.get_filter_time())
+      assert np.allclose(kf.filter_py.state(), kf.filter_pyx.state())
+      assert np.allclose(kf.filter_py.covs(), kf.filter_pyx.covs())
+
+
+if __name__ == "__main__":
+  generated_dir = sys.argv[2]
+  CompareFilter.generate_code(generated_dir)