Tracking Moving Objects

import numpy as np
import matplotlib.pyplot as plt

class KalmanFilter2D:
    def __init__(self, dt):
        """
        2D position and velocity tracker
        
        Args:
            dt: Time step between measurements
        """
        self.dt = dt
        
        # State vector: [x, y, vx, vy]
        self.x = np.zeros(4)
        
        # State transition matrix (constant velocity model)
        self.F = np.array([
            [1, 0, dt, 0],   # x = x + vx*dt
            [0, 1, 0, dt],   # y = y + vy*dt
            [0, 0, 1, 0],    # vx = vx
            [0, 0, 0, 1]     # vy = vy
        ])
        
        # Measurement matrix (we only measure position)
        self.H = np.array([
            [1, 0, 0, 0],    # Measure x
            [0, 1, 0, 0]     # Measure y
        ])
        
        # Process noise covariance
        q = 0.1  # Process noise magnitude
        self.Q = q * np.array([
            [dt**4/4, 0, dt**3/2, 0],
            [0, dt**4/4, 0, dt**3/2],
            [dt**3/2, 0, dt**2, 0],
            [0, dt**3/2, 0, dt**2]
        ])
        
        # Measurement noise covariance
        r = 5.0  # Measurement noise (meters)
        self.R = r**2 * np.eye(2)
        
        # Estimate uncertainty
        self.P = np.eye(4) * 100  # High initial uncertainty
    
    def predict(self):
        """Prediction step"""
        # Predicted state
        self.x = self.F @ self.x
        
        # Predicted covariance
        self.P = self.F @ self.P @ self.F.T + self.Q
    
    def update(self, measurement):
        """Update step"""
        # Innovation (measurement residual)
        y = measurement - self.H @ self.x
        
        # Innovation covariance
        S = self.H @ self.P @ self.H.T + self.R
        
        # Kalman gain
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # Updated state estimate
        self.x = self.x + K @ y
        
        # Updated covariance
        I = np.eye(len(self.x))
        self.P = (I - K @ self.H) @ self.P
        
        return self.x
    
    def get_position(self):
        """Get current position estimate"""
        return self.x[:2]
    
    def get_velocity(self):
        """Get current velocity estimate"""
        return self.x[2:]

# Simulation
def simulate_trajectory(num_steps, dt):
    """Generate ground truth trajectory"""
    t = np.arange(num_steps) * dt
    
    # Circular motion
    radius = 50
    omega = 0.1
    
    x_true = radius * np.cos(omega * t)
    y_true = radius * np.sin(omega * t)
    
    vx_true = -radius * omega * np.sin(omega * t)
    vy_true = radius * omega * np.cos(omega * t)
    
    return x_true, y_true, vx_true, vy_true

# Parameters
dt = 0.1  # Time step (seconds)
num_steps = 200
measurement_noise = 5.0  # GPS noise (meters)

# Generate true trajectory
x_true, y_true, vx_true, vy_true = simulate_trajectory(num_steps, dt)

# Generate noisy measurements
x_meas = x_true + np.random.normal(0, measurement_noise, num_steps)
y_meas = y_true + np.random.normal(0, measurement_noise, num_steps)

# Run Kalman filter
kf = KalmanFilter2D(dt)

# Initialize with first measurement
kf.x[:2] = [x_meas[0], y_meas[0]]

x_est = []
y_est = []
vx_est = []
vy_est = []

for i in range(num_steps):
    # Predict
    kf.predict()
    
    # Update with measurement
    measurement = np.array([x_meas[i], y_meas[i]])
    kf.update(measurement)
    
    # Store estimates
    pos = kf.get_position()
    vel = kf.get_velocity()
    x_est.append(pos[0])
    y_est.append(pos[1])
    vx_est.append(vel[0])
    vy_est.append(vel[1])

# Convert to arrays
x_est = np.array(x_est)
y_est = np.array(y_est)
vx_est = np.array(vx_est)
vy_est = np.array(vy_est)

# Plot results
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# Position plot
ax1.plot(x_true, y_true, 'g-', label='True Trajectory', linewidth=2)
ax1.plot(x_meas, y_meas, 'r.', label='Noisy Measurements', alpha=0.3)
ax1.plot(x_est, y_est, 'b-', label='Kalman Filter', linewidth=2)
ax1.set_xlabel('X Position (m)')
ax1.set_ylabel('Y Position (m)')
ax1.set_title('2D Position Tracking')
ax1.legend()
ax1.grid(True)
ax1.axis('equal')

# Velocity plot
t = np.arange(num_steps) * dt
ax2.plot(t, vx_true, 'g-', label='True Vx', linewidth=2)
ax2.plot(t, vx_est, 'b--', label='Estimated Vx')
ax2.plot(t, vy_true, 'g:', label='True Vy', linewidth=2)
ax2.plot(t, vy_est, 'r--', label='Estimated Vy')
ax2.set_xlabel('Time (s)')
ax2.set_ylabel('Velocity (m/s)')
ax2.set_title('Velocity Estimation')
ax2.legend()
ax2.grid(True)

plt.tight_layout()
plt.show()

# Calculate errors
pos_error_raw = np.sqrt((x_meas - x_true)**2 + (y_meas - y_true)**2)
pos_error_filtered = np.sqrt((x_est - x_true)**2 + (y_est - y_true)**2)

print(f"Average position error (raw): {pos_error_raw.mean():.2f} m")
print(f"Average position error (filtered): {pos_error_filtered.mean():.2f} m")
print(f"Improvement: {pos_error_raw.mean() / pos_error_filtered.mean():.1f}x")

x_new = x_old + vx * dt

measurement = H @ state  # Only observe position