232 lines
6.5 KiB
Python
232 lines
6.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
"""Copyright 2015 Roger R Labbe Jr.
|
|
|
|
|
|
Code supporting the book
|
|
|
|
Kalman and Bayesian Filters in Python
|
|
https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python
|
|
|
|
|
|
This is licensed under an MIT license. See the LICENSE.txt file
|
|
for more information.
|
|
"""
|
|
|
|
from __future__ import (absolute_import, division, print_function,
|
|
unicode_literals)
|
|
|
|
import book_plots as bp
|
|
import filterpy.kalman as kf
|
|
from math import radians, sin, cos, sqrt, exp
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import numpy.random as random
|
|
|
|
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
|
|
|
|
|
|
def plot_radar(xs, track, time):
|
|
|
|
plt.figure()
|
|
bp.plot_track(time, track[:, 0])
|
|
bp.plot_filter(time, xs[:, 0])
|
|
plt.legend(loc=4)
|
|
plt.xlabel('time (sec)')
|
|
plt.ylabel('position (m)')
|
|
|
|
plt.figure()
|
|
bp.plot_track(time, track[:, 1])
|
|
bp.plot_filter(time, xs[:, 1])
|
|
plt.legend(loc=4)
|
|
plt.xlabel('time (sec)')
|
|
plt.ylabel('velocity (m/s)')
|
|
|
|
plt.figure()
|
|
bp.plot_track(time, track[:, 2])
|
|
bp.plot_filter(time, xs[:, 2])
|
|
plt.ylabel('altitude (m)')
|
|
plt.legend(loc=4)
|
|
plt.xlabel('time (sec)')
|
|
plt.ylim((900, 1600))
|
|
plt.show()
|
|
|
|
def plot_bicycle():
|
|
plt.clf()
|
|
plt.axes()
|
|
ax = plt.gca()
|
|
#ax.add_patch(plt.Rectangle((10,0), 10, 20, fill=False, ec='k')) #car
|
|
ax.add_patch(plt.Rectangle((21,1), .75, 2, fill=False, ec='k')) #wheel
|
|
ax.add_patch(plt.Rectangle((21.33,10), .75, 2, fill=False, ec='k', angle=20)) #wheel
|
|
ax.add_patch(plt.Rectangle((21.,4.), .75, 2, fill=True, ec='k', angle=5, ls='dashdot', alpha=0.3)) #wheel
|
|
|
|
plt.arrow(0, 2, 20.5, 0, fc='k', ec='k', head_width=0.5, head_length=0.5)
|
|
plt.arrow(0, 2, 20.4, 3, fc='k', ec='k', head_width=0.5, head_length=0.5)
|
|
plt.arrow(21.375, 2., 0, 8.5, fc='k', ec='k', head_width=0.5, head_length=0.5)
|
|
plt.arrow(23, 2, 0, 2.5, fc='k', ec='k', head_width=0.5, head_length=0.5)
|
|
|
|
#ax.add_patch(plt.Rectangle((10,0), 10, 20, fill=False, ec='k'))
|
|
plt.text(11, 1.0, "R", fontsize=18)
|
|
plt.text(8, 2.2, r"$\beta$", fontsize=18)
|
|
plt.text(20.4, 13.5, r"$\alpha$", fontsize=18)
|
|
plt.text(21.6, 7, "w (wheelbase)", fontsize=18)
|
|
plt.text(0, 1, "C", fontsize=18)
|
|
plt.text(24, 3, "d", fontsize=18)
|
|
plt.plot([21.375, 21.25], [11, 14], color='k', lw=1)
|
|
plt.plot([21.375, 20.2], [11, 14], color='k', lw=1)
|
|
plt.axis('scaled')
|
|
plt.xlim(0,25)
|
|
plt.axis('off')
|
|
plt.show()
|
|
|
|
#plot_bicycle()
|
|
|
|
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()
|
|
|
|
|
|
def show_radar_chart():
|
|
plt.xlim([0.9,2.5])
|
|
plt.ylim([0.5,2.5])
|
|
|
|
plt.scatter ([1,2],[1,2])
|
|
#plt.scatter ([2],[1],marker='o')
|
|
ax = plt.axes()
|
|
|
|
ax.annotate('', xy=(2,2), xytext=(1,1),
|
|
arrowprops=dict(arrowstyle='->', ec='r',shrinkA=3, shrinkB=4))
|
|
|
|
ax.annotate('', xy=(2,1), xytext=(1,1),
|
|
arrowprops=dict(arrowstyle='->', ec='b',shrinkA=0, shrinkB=0))
|
|
|
|
ax.annotate('', xy=(2,2), xytext=(2,1),
|
|
arrowprops=dict(arrowstyle='->', ec='b',shrinkA=0, shrinkB=4))
|
|
|
|
|
|
|
|
ax.annotate('$\Theta$', xy=(1.2, 1.1), color='b')
|
|
ax.annotate('Aircraft', xy=(2.04,2.), color='b')
|
|
ax.annotate('altitude', xy=(2.04,1.5), color='k')
|
|
ax.annotate('X', xy=(1.5, .9))
|
|
ax.annotate('Radar', xy=(.95, 0.9))
|
|
ax.annotate('Slant\n (r)', xy=(1.5,1.62), color='r')
|
|
|
|
plt.title("Radar Tracking")
|
|
ax.xaxis.set_ticklabels([])
|
|
ax.yaxis.set_ticklabels([])
|
|
ax.xaxis.set_ticks([])
|
|
ax.yaxis.set_ticks([])
|
|
|
|
|
|
plt.show() |