b9f462678d
Started developing EKF chapter. Stalled - not sure what is the best initial example to develop. Air drag seems unnecessarily 'mathy', but then you need diff eq to compute Jacobians.
128 lines
3.4 KiB
Python
128 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Fri Jul 18 23:23:08 2014
|
|
|
|
@author: rlabbe
|
|
"""
|
|
from math import radians, sin, cos, sqrt, exp
|
|
import numpy.random as random
|
|
import matplotlib.pyplot as plt
|
|
import filterpy.kalman as kf
|
|
import numpy as np
|
|
|
|
def ball_kf(x, y, omega, v0, dt, r=0.5, q=0.02):
|
|
|
|
g = 9.8 # gravitational constant
|
|
f1 = kf.KalmanFilter(dim_x=5, dim_z=2)
|
|
|
|
ay = .5*dt**2
|
|
f1.F = np.array ([[1, dt, 0, 0, 0], # x = x0+dx*dt
|
|
[0, 1, 0, 0, 0], # dx = dx
|
|
[0, 0, 1, dt, ay], # y = y0 +dy*dt+1/2*g*dt^2
|
|
[0, 0, 0, 1, dt], # dy = dy0 + ddy*dt
|
|
[0, 0, 0, 0, 1]]) # ddy = -g.
|
|
|
|
f1.H = np.array([
|
|
[1, 0, 0, 0, 0],
|
|
[0, 0, 1, 0, 0]])
|
|
|
|
f1.R *= r
|
|
f1.Q *= q
|
|
|
|
omega = radians(omega)
|
|
vx = cos(omega) * v0
|
|
vy = sin(omega) * v0
|
|
|
|
f1.x = np.array([[x,vx,y,vy,-9.8]]).T
|
|
|
|
return f1
|
|
|
|
|
|
class BaseballPath(object):
|
|
def __init__(self, x0, y0, launch_angle_deg, velocity_ms, noise=(1.0,1.0)):
|
|
""" Create baseball path object in 2D (y=height above ground)
|
|
|
|
x0,y0 initial position
|
|
launch_angle_deg angle ball is travelling respective to ground plane
|
|
velocity_ms speeed of ball in meters/second
|
|
noise amount of noise to add to each reported position in (x,y)
|
|
"""
|
|
|
|
omega = radians(launch_angle_deg)
|
|
self.v_x = velocity_ms * cos(omega)
|
|
self.v_y = velocity_ms * sin(omega)
|
|
|
|
self.x = x0
|
|
self.y = y0
|
|
|
|
self.noise = noise
|
|
|
|
|
|
def drag_force (self, velocity):
|
|
""" Returns the force on a baseball due to air drag at
|
|
the specified velocity. Units are SI
|
|
"""
|
|
B_m = 0.0039 + 0.0058 / (1. + exp((velocity-35.)/5.))
|
|
return B_m * velocity
|
|
|
|
|
|
def update(self, dt, vel_wind=0.):
|
|
""" compute the ball position based on the specified time step and
|
|
wind velocity. Returns (x,y) position tuple.
|
|
"""
|
|
|
|
# Euler equations for x and y
|
|
self.x += self.v_x*dt
|
|
self.y += self.v_y*dt
|
|
|
|
# force due to air drag
|
|
v_x_wind = self.v_x - vel_wind
|
|
|
|
v = sqrt (v_x_wind**2 + self.v_y**2)
|
|
F = self.drag_force(v)
|
|
|
|
# Euler's equations for velocity
|
|
self.v_x = self.v_x - F*v_x_wind*dt
|
|
self.v_y = self.v_y - 9.81*dt - F*self.v_y*dt
|
|
|
|
return (self.x + random.randn()*self.noise[0],
|
|
self.y + random.randn()*self.noise[1])
|
|
|
|
|
|
|
|
def plot_ball():
|
|
y = 1.
|
|
x = 0.
|
|
theta = 35. # launch angle
|
|
v0 = 50.
|
|
dt = 1/10. # time step
|
|
|
|
ball = BaseballPath(x0=x, y0=y, launch_angle_deg=theta, velocity_ms=v0, noise=[.3,.3])
|
|
f1 = ball_kf(x,y,theta,v0,dt,r=1.)
|
|
f2 = ball_kf(x,y,theta,v0,dt,r=10.)
|
|
t = 0
|
|
xs = []
|
|
ys = []
|
|
xs2 = []
|
|
ys2 = []
|
|
|
|
while f1.x[2,0] > 0:
|
|
t += dt
|
|
x,y = ball.update(dt)
|
|
z = np.mat([[x,y]]).T
|
|
|
|
f1.update(z)
|
|
f2.update(z)
|
|
xs.append(f1.x[0,0])
|
|
ys.append(f1.x[2,0])
|
|
xs2.append(f2.x[0,0])
|
|
ys2.append(f2.x[2,0])
|
|
f1.predict()
|
|
f2.predict()
|
|
|
|
p1 = plt.scatter(x, y, color='green', marker='o', s=75, alpha=0.5)
|
|
|
|
p2, = plt.plot (xs, ys,lw=2)
|
|
p3, = plt.plot (xs2, ys2,lw=4, c='r')
|
|
plt.legend([p1,p2, p3], ['Measurements', 'Kalman filter(R=0.5)', 'Kalman filter(R=10)'])
|
|
plt.show() |