mirror of https://github.com/commaai/rednose.git
126 lines
3.3 KiB
Python
Executable File
126 lines
3.3 KiB
Python
Executable File
#!/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)
|