Sensor Fusion: Combining Intelligence

class ComplementaryFilter:
    def __init__(self, alpha=0.96):
        self.alpha = alpha
        self.angle = 0.0

    def update(self, gyro_rate, accel_angle, dt):
        """
        gyro_rate: angular velocity from gyroscope (deg/s)
        accel_angle: angle calculated from accelerometer (deg)
        dt: time step (seconds)
        """
        # Predict from gyro (integrate), correct with accel (low-pass)
        self.angle = self.alpha * (self.angle + gyro_rate * dt) + \
                     (1 - self.alpha) * accel_angle
        return self.angle

# Demo
filter = ComplementaryFilter(alpha=0.96)
angles = []
for i in range(100):
    gyro = 5.0 + np.random.normal(0, 0.5)  # 5 deg/s with noise
    accel = 45.0 + np.random.normal(0, 2.0)  # True angle ~45°
    dt = 0.01
    angle = filter.update(gyro, accel, dt)
    angles.append(angle)

print(f"Final angle: {angles[-1]:.2f}° (target: ~45°)")

class MadgwickFilter:
    """Simplified Madgwick AHRS filter."""

    def __init__(self, sample_rate=100, beta=0.1):
        self.sample_rate = sample_rate
        self.beta = beta  # Filter gain
        self.q = np.array([1.0, 0.0, 0.0, 0.0])  # Quaternion [w, x, y, z]

    def update(self, gyro, accel, mag=None, dt=None):
        """
        gyro: [gx, gy, gz] in rad/s
        accel: [ax, ay, az] normalized
        mag: [mx, my, mz] normalized (optional)
        """
        if dt is None:
            dt = 1.0 / self.sample_rate

        q = self.q.copy()

        # Normalize accelerometer
        norm = np.sqrt(accel[0]**2 + accel[1]**2 + accel[2]**2)
        if norm > 0:
            ax, ay, az = accel / norm
        else:
            ax, ay, az = 0, 0, 1

        # Gradient descent step
        f = np.array([
            2*(q[1]*q[3] - q[0]*q[2]) - ax,
            2*(q[0]*q[1] + q[2]*q[3]) - ay,
            2*(0.5 - q[1]**2 - q[2]**2) - az
        ])

        j = np.array([
            [-2*q[2],  2*q[3], -2*q[0], 2*q[1]],
            [ 2*q[1],  2*q[0],  2*q[3], 2*q[2]],
            [      0, -4*q[1], -4*q[2],      0]
        ])

        step = j.T @ f
        step = step / np.linalg.norm(step) if np.linalg.norm(step) > 0 else step

        # Quaternion rate from gyroscope
        q_dot = 0.5 * self._quaternion_multiply(q, np.array([0, *gyro])) - self.beta * step

        # Integrate
        self.q = q + q_dot * dt
        self.q = self.q / np.linalg.norm(self.q)

        return self.q

    def _quaternion_multiply(self, q1, q2):
        w = q1[0]*q2[0] - q1[1]*q2[1] - q1[2]*q2[2] - q1[3]*q2[3]
        x = q1[0]*q2[1] + q1[1]*q2[0] + q1[2]*q2[3] - q1[3]*q2[2]
        y = q1[0]*q2[2] - q1[1]*q2[3] + q1[2]*q2[0] + q1[3]*q2[1]
        z = q1[0]*q2[3] + q1[1]*q2[2] - q1[2]*q2[1] + q1[3]*q2[0]
        return np.array([w, x, y, z])

    def get_euler(self):
        """Convert quaternion to Euler angles (roll, pitch, yaw)."""
        w, x, y, z = self.q
        roll = np.arctan2(2*(w*x + y*z), 1 - 2*(x**2 + y**2))
        pitch = np.arcsin(np.clip(2*(w*y - z*x), -1, 1))
        yaw = np.arctan2(2*(w*z + x*y), 1 - 2*(y**2 + z**2))
        return np.degrees([roll, pitch, yaw])

# Demo
madgwick = MadgwickFilter(sample_rate=100, beta=0.1)
for i in range(100):
    gyro = np.array([0.1, 0.05, 0.02])  # rad/s
    accel = np.array([0.0, 0.0, 9.81])   # gravity
    q = madgwick.update(gyro, accel)

roll, pitch, yaw = madgwick.get_euler()
print(f"Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°")

class EKF_AHRS:
    """6-state Extended Kalman Filter for AHRS."""

    def __init__(self, dt=0.01):
        self.dt = dt

        # State: [roll, pitch, yaw, gyro_bias_x, gyro_bias_y, gyro_bias_z]
        self.x = np.zeros(6)

        # State covariance
        self.P = np.eye(6) * 0.1

        # Process noise
        self.Q = np.eye(6) * 0.001
        self.Q[3:, 3:] *= 0.01  # Gyro bias process noise

        # Measurement noise
        self.R_accel = np.eye(2) * 0.5
        self.R_mag = np.eye(1) * 2.0

    def predict(self, gyro):
        """Predict step using gyroscope."""
        roll, pitch, yaw, bx, by, bz = self.x

        # Remove estimated bias from gyro
        gx, gy, gz = gyro - np.array([bx, by, bz])

        # State transition (simplified)
        self.x[0] += gx * self.dt  # roll
        self.x[1] += gy * self.dt  # pitch
        self.x[2] += gz * self.dt  # yaw
        # Biases remain constant (random walk)

        # Jacobian
        F = np.eye(6)
        F[0, 3] = -self.dt
        F[1, 4] = -self.dt
        F[2, 5] = -self.dt

        # Update covariance
        self.P = F @ self.P @ F.T + self.Q

    def update_accel(self, accel):
        """Update with accelerometer (roll and pitch)."""
        ax, ay, az = accel / np.linalg.norm(accel)

        # Expected measurement from current state
        roll, pitch = self.x[0], self.x[1]
        h = np.array([
            np.arctan2(ay, az) - roll,
            np.arctan2(-ax, np.sqrt(ay**2 + az**2)) - pitch
        ])

        # Jacobian
        H = np.zeros((2, 6))
        H[0, 0] = -1
        H[1, 1] = -1

        # Kalman gain
        S = H @ self.P @ H.T + self.R_accel
        K = self.P @ H.T @ np.linalg.inv(S)

        # Update
        self.x += K @ h
        self.P = (np.eye(6) - K @ H) @ self.P

    def update_mag(self, mag, declination=0.0):
        """Update with magnetometer (yaw only)."""
        mx, my = mag[0], mag[1]
        expected_yaw = np.arctan2(my, mx) - declination
        error = expected_yaw - self.x[2]
        error = np.arctan2(np.sin(error), np.cos(error))  # Wrap to [-π, π]

        H = np.zeros((1, 6))
        H[0, 2] = 1

        S = H @ self.P @ H.T + self.R_mag
        K = self.P @ H.T @ np.linalg.inv(S)

        self.x += (K @ np.array([error])).flatten()
        self.P = (np.eye(6) - K @ H) @ self.P

    def get_orientation(self):
        """Get current orientation in degrees."""
        return np.degrees(self.x[:3])

# Demo
ekf = EKF_AHRS(dt=0.01)
np.random.seed(42)

for i in range(100):
    gyro = np.array([0.1, 0.05, 0.02]) + np.random.normal(0, 0.01, 3)
    accel = np.array([0.0, 0.0, 9.81]) + np.random.normal(0, 0.1, 3)
    mag = np.array([20.0, 5.0, 0.0]) + np.random.normal(0, 0.5, 3)

    ekf.predict(gyro)
    ekf.update_accel(accel)
    if i % 5 == 0:  # Magnetometer slower
        ekf.update_mag(mag)

roll, pitch, yaw = ekf.get_orientation()
print(f"Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°")