扩展滤波解决IMU数据漂移问题

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()