import numpy as np
from abc import ABCMeta, abstractmethod
class EKF(object):
__metaclass__ = ABCMeta
def __init__(self, n, m, pval=0.1, qval=1e-4, rval=0.1):
self.Q_k = np.eye(n) * qval
self.x = np.zeros(n)
self.P_current = None
self.P_result = np.eye(n) * pval
self.R_k = np.eye(m) * rval
self.I = np.eye(n)
def updateR(self, R):
self.R_k = R
def step(self, z):
self.x, F_k = self.stateTransitionFunction(self.x)
self.P_current = F_k.dot(self.P_result).dot(F_k.T) + self.Q_k
zz_k, H_k = self.stateToMeasurementTransitionFunction(self.x)
K_k = np.dot(self.P_current.dot(H_k.T), np.linalg.inv(H_k.dot(self.P_current).dot(H_k.T) + self.R_k))
self.x += np.dot(K_k, (np.array(z) - zz_k))
self.P_result = np.dot(self.I - np.dot(K_k, H_k), self.P_current)
return self.x
@abstractmethod
def stateTransitionFunction(self, x):
raise NotImplementedError()
@abstractmethod
def stateToMeasurementTransitionFunction(self, x):
raise NotImplementedError()
class IMUEKF(EKF):
def __init__(self, dt):
super().__init__(6, 6) # 6 states: roll, pitch, yaw, gyro_bias_x, gyro_bias_y, gyro_bias_z
self.dt = dt
def stateTransitionFunction(self, x):
roll, pitch, yaw, bias_x, bias_y, bias_z = x
F_k = np.eye(6)
F_k[0, 3] = -self.dt
F_k[1, 4] = -self.dt
F_k[2, 5] = -self.dt
return x, F_k
def stateToMeasurementTransitionFunction(self, x):
roll, pitch, yaw, bias_x, bias_y, bias_z = x
H_k = np.eye(6)
return x[:3], H_k
def main():
dt = 0.01 # 100Hz sample rate
ekf = IMUEKF(dt)
for _ in range(1000):
acc = np.random.randn(3) * 0.1 # simulate accelerometer data
gyro = np.random.randn(3) * 0.01 # simulate gyroscope data
mag = np.random.randn(3) * 0.1 # simulate magnetometer data
acc = acc / np.linalg.norm(acc)
roll_acc = np.arctan2(acc[1], acc[2])
pitch_acc = np.arctan2(-acc[0], np.sqrt(acc[1]**2 + acc[2]**2))
yaw_mag = np.arctan2(mag[1], mag[0])
z = np.array([roll_acc, pitch_acc, yaw_mag, 0, 0, 0])
state = ekf.step(z)
roll, pitch, yaw = state[:3]
print(f"Roll: {np.degrees(roll):.2f}, Pitch: {np.degrees(pitch):.2f}, Yaw: {np.degrees(yaw):.2f}")
if __name__ == "__main__":
main()
扩展滤波解决IMU数据漂移问题
最新推荐文章于 2024-09-10 13:01:50 发布