mirror of https://github.com/commaai/rednose.git
82 lines
1.9 KiB
Python
Executable File
82 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
|
|
import numpy as np
|
|
import sympy as sp
|
|
|
|
from rednose.helpers.kalmanfilter import KalmanFilter
|
|
|
|
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
|
|
|
|
|
|
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 KinematicKalman(KalmanFilter):
|
|
name = 'kinematic'
|
|
|
|
initial_x = np.array([0.5, 0.0])
|
|
|
|
# state covariance
|
|
initial_P_diag = np.array([1.0**2, 1.0**2])
|
|
|
|
# process noise
|
|
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 = KinematicKalman.name
|
|
dim_state = KinematicKalman.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 = EKF_sym_pyx(generated_dir, self.name, self.Q, self.initial_x, np.diag(self.initial_P_diag), dim_state, dim_state_err)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
generated_dir = sys.argv[2]
|
|
KinematicKalman.generate_code(generated_dir)
|