import matplotlib.pyplot as plt
# ============================================================
# IMU Orientation Estimator: Complementary + Kalman Filter
# ============================================================
# Simulates accelerometer and gyroscope data with realistic
# noise, then fuses them using two different algorithms.
# --- Simulation parameters ---
dt = 0.01 # 100 Hz sample rate
t_end = 60.0 # 60 seconds of data
t = np.arange(0, t_end, dt)
# --- True angle trajectory ---
# Slow sinusoidal motion (like a robot arm or gimbal)
true_angle = 30.0 * np.sin(0.2 * t) + 10.0 * np.sin(0.05 * t)
true_rate = np.gradient(true_angle, dt) # true angular rate (deg/s)
# --- Simulated accelerometer ---
# Good for static angle, but noisy and affected by vibrations
accel_noise_std = 3.0 # degrees
vibration_noise = 5.0 * np.sin(50 * t) * np.exp(-0.05 * t) # decaying vibration
accel_angle = true_angle + accel_noise_std * np.random.randn(N) + vibration_noise
# Add occasional impulse disturbances (bumps)
bump_indices = np.random.choice(N, size=20, replace=False)
accel_angle[bump_indices] += np.random.uniform(-15, 15, size=20)
# --- Simulated gyroscope ---
# Smooth but drifts over time
gyro_noise_std = 0.3 # deg/s
gyro_bias_rate = 0.02 # deg/s drift accumulation rate
gyro_bias = np.cumsum(gyro_bias_rate * dt * np.random.randn(N))
gyro_rate = true_rate + gyro_noise_std * np.random.randn(N) + gyro_bias
# --- Raw gyro integration (to show drift) ---
gyro_angle[0] = accel_angle[0]
gyro_angle[i] = gyro_angle[i - 1] + gyro_rate[i] * dt
# ============================================================
# ============================================================
def complementary_filter(accel_angle, gyro_rate, dt, alpha=0.98):
Fuse accelerometer angle and gyroscope rate.
alpha: weight for gyroscope (0.95 to 0.99 typical).
angle[0] = accel_angle[0]
# Gyro prediction + accel correction
angle[i] = alpha * (angle[i - 1] + gyro_rate[i] * dt) \
+ (1 - alpha) * accel_angle[i]
# ============================================================
# Kalman Filter for IMU Fusion
# ============================================================
def kalman_filter_imu(accel_angle, gyro_rate, dt,
Q_angle=0.001, Q_bias=0.003, R_measure=3.0):
2-state Kalman filter: estimates angle and gyro bias.
State: x = [angle, bias]^T
Prediction: angle_new = angle + (gyro - bias) * dt
Measurement: z = accel_angle
Q_angle: process noise for angle state
Q_bias: process noise for bias state
R_measure: measurement noise (accelerometer variance)
angle_est[0] = accel_angle[0]
# Error covariance matrix P (2x2)
P = np.array([[1.0, 0.0],
# Process noise covariance
Q = np.array([[Q_angle, 0.0],
# Measurement matrix: we only measure angle
H = np.array([[1.0, 0.0]])
R = np.array([[R_measure]])
rate = gyro_rate[i] - bias_est[i - 1]
angle_pred = angle_est[i - 1] + rate * dt
bias_pred = bias_est[i - 1]
A = np.array([[1.0, -dt],
# Innovation (measurement residual)
y = z - angle_pred # H @ x_pred = angle_pred
K = np.array([[P[0, 0] / S],
angle_est[i] = angle_pred + K[0, 0] * y
bias_est[i] = bias_pred + K[1, 0] * y
I_KH = np.array([[1.0 - K[0, 0], 0.0],
return angle_est, bias_est
# ============================================================
# ============================================================
comp_angle = complementary_filter(accel_angle, gyro_rate, dt, alpha=0.98)
kalman_angle, kalman_bias = kalman_filter_imu(
accel_angle, gyro_rate, dt,
Q_angle=0.001, Q_bias=0.003, R_measure=4.0
# ============================================================
# ============================================================
accel_rmse = np.sqrt(np.mean((accel_angle - true_angle) ** 2))
gyro_rmse = np.sqrt(np.mean((gyro_angle - true_angle) ** 2))
comp_rmse = np.sqrt(np.mean((comp_angle - true_angle) ** 2))
kalman_rmse = np.sqrt(np.mean((kalman_angle - true_angle) ** 2))
print(" IMU Orientation Estimation: Error Comparison")
print(f" Raw Accelerometer RMSE: {accel_rmse:8.3f} deg")
print(f" Raw Gyro (integrated) : {gyro_rmse:8.3f} deg")
print(f" Complementary Filter : {comp_rmse:8.3f} deg")
print(f" Kalman Filter : {kalman_rmse:8.3f} deg")
print(f"\n Complementary filter is {accel_rmse / comp_rmse:.1f}x better than raw accel")
print(f" Kalman filter is {accel_rmse / kalman_rmse:.1f}x better than raw accel")
# ============================================================
# ============================================================
fig, axes = plt.subplots(4, 1, figsize=(12, 14), sharex=True)
# Plot 1: Raw sensor data vs truth
ax.plot(t, true_angle, 'k-', linewidth=2, label='True angle')
ax.plot(t, accel_angle, 'r.', markersize=1, alpha=0.3, label=f'Accel (RMSE={accel_rmse:.1f})')
ax.plot(t, gyro_angle, 'b-', alpha=0.5, label=f'Gyro integrated (RMSE={gyro_rmse:.1f})')
ax.set_ylabel('Angle (deg)')
ax.set_title('Raw Sensor Data: Neither Is Good Enough')
ax.legend(loc='upper right', fontsize=9)
# Plot 2: Complementary filter
ax.plot(t, true_angle, 'k-', linewidth=2, label='True angle')
ax.plot(t, comp_angle, 'g-', linewidth=1.5, label=f'Complementary (RMSE={comp_rmse:.2f})')
ax.set_ylabel('Angle (deg)')
ax.set_title('Complementary Filter: Simple but Effective')
ax.legend(loc='upper right', fontsize=9)
ax.plot(t, true_angle, 'k-', linewidth=2, label='True angle')
ax.plot(t, kalman_angle, 'm-', linewidth=1.5, label=f'Kalman (RMSE={kalman_rmse:.2f})')
ax.set_ylabel('Angle (deg)')
ax.set_title('Kalman Filter: Optimal Estimation')
ax.legend(loc='upper right', fontsize=9)
# Plot 4: Estimated gyro bias
ax.plot(t, gyro_bias, 'b-', linewidth=1.5, label='True gyro bias')
ax.plot(t, kalman_bias, 'r--', linewidth=1.5, label='Kalman bias estimate')
ax.set_xlabel('Time (s)')
ax.set_ylabel('Bias (deg/s)')
ax.set_title('Kalman Filter Tracks Gyroscope Bias')
ax.legend(loc='upper right', fontsize=9)
plt.savefig('imu_sensor_fusion.png', dpi=150, bbox_inches='tight')
print("\nPlot saved: imu_sensor_fusion.png")
Comments