# -*- coding: utf-8 -*- """ Created on Sat Jul 5 16:07:29 2014 @author: rlabbe """ import numpy as np from KalmanFilter import KalmanFilter from math import radians, sin, cos, sqrt, exp import numpy.random as random import matplotlib.markers as markers import matplotlib.pyplot as plt from RungeKutta import * class BallPath(object): def __init__(self, x0, y0, omega_deg, velocity, g=9.8, noise=[1.0,1.0]): omega = radians(omega_deg) self.vx0 = velocity * cos(omega) self.vy0 = velocity * sin(omega) self.x0 = x0 self.y0 = y0 self.g = g self.noise = noise def pos_at_t(self, t): """ returns (x,y) tuple of ball position at time t""" x = self.vx0*t + self.x0 y = -0.5*self.g*t**2 + self.vy0*t + self.y0 return (x +random.randn()*self.noise[0], y +random.randn()*self.noise[1]) def ball_kf(x, y, omega, v0, dt, r=0.5, q=0.02): g = 9.8 # gravitational constant f1 = KalmanFilter(dim_x=5, dim_z=2) ay = .5*dt**2 f1.F = np.mat ([[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.mat([ [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.mat([x,vx,y,vy,-9.8]).T return f1 def ball_kf_noay(x, y, omega, v0, dt, r=0.5, q=0.02): g = 9.8 # gravitational constant f1 = KalmanFilter(dim_x=5, dim_z=2) ay = .5*dt**2 f1.F = np.mat ([[1, dt, 0, 0, 0], # x = x0+dx*dt [0, 1, 0, 0, 0], # dx = dx [0, 0, 1, dt, 0], # y = y0 +dy*dt [0, 0, 0, 1, dt], # dy = dy0 + ddy*dt [0, 0, 0, 0, 1]]) # ddy = -g. f1.H = np.mat([ [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.mat([x,vx,y,vy,-9.8]).T return f1 def test_kf(): dt = 0.1 t = 0 f1 = ball_kf (0,1, 35, 50, 0.1) f2 = ball_kf_noay (0,1, 35, 50, 0.1) path = BallPath( 0, 1, 35, 50, noise=(0,0)) path_rk = BallRungeKutta(0, 1, 50, 35) xs = [] ys = [] while f1.x[2,0]>= 0: t += dt f1.predict() f2.predict() #x,y = path.pos_at_t(t) x,y = path_rk.step(dt) xs.append(x) ys.append(y) plt.scatter (f1.x[0,0], f1.x[2,0], color='blue',alpha=0.6) plt.scatter (f2.x[0,0], f2.x[2,0], color='red', alpha=0.6) plt.plot(xs, ys, c='g') 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 test_baseball_path(): ball = BaseballPath (0, 1, 35, 50) while ball.y > 0: ball.update (0.1, 0.) plt.scatter (ball.x, ball.y) def test_ball_path(): y = 15 x = 0 omega = 0. noise = [1,1] v0 = 100. ball = BallPath (x0=x, y0=y, omega_deg=omega, velocity=v0, noise=noise) dt = 1 f1 = KalmanFilter(dim_x=6, dim_z=2) dt = 1/30. # time step ay = -.5*dt**2 f1.F = np.mat ([[1, dt, 0, 0, 0, 0], # x=x0+dx*dt [0, 1, dt, 0, 0, 0], # dx = dx [0, 0, 0, 0, 0, 0], # ddx = 0 [0, 0, 0, 1, dt, ay], # y = y0 +dy*dt+1/2*g*dt^2 [0, 0, 0, 0, 1, dt], # dy = dy0 + ddy*dt [0, 0, 0, 0, 0, 1]]) # ddy = -g f1.H = np.mat([ [1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0]]) f1.R = np.eye(2) * 5 f1.Q = np.eye(6) * 0. omega = radians(omega) vx = cos(omega) * v0 vy = sin(omega) * v0 f1.x = np.mat([x,vx,0,y,vy,-9.8]).T f1.P = np.eye(6) * 500. z = np.mat([[0,0]]).T count = 0 markers.MarkerStyle(fillstyle='none') np.set_printoptions(precision=4) while f1.x[3,0] > 0: count += 1 #f1.update (z) f1.predict() plt.scatter(f1.x[0,0],f1.x[3,0], color='green') def drag_force (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_drag(f, dt): vx = f.x[1,0] vy = f.x[3,0] v = sqrt(vx**2 + vy**2) F = -drag_force(v) print F f.u[0,0] = -drag_force(vx) f.u[1,0] = -drag_force(vy) #f.x[2,0]=F*vx #f.x[5,0]=F*vy def test_kf_drag(): y = 1 x = 0 omega = 35. noise = [0,0] v0 = 50. ball = BaseballPath (x0=x, y0=y, launch_angle_deg=omega, velocity_ms=v0, noise=noise) #ball = BallPath (x0=x, y0=y, omega_deg=omega, velocity=v0, noise=noise) dt = 1 f1 = KalmanFilter(dim_x=6, dim_z=2) dt = 1/30. # time step ay = -.5*dt**2 ax = .5*dt**2 f1.F = np.mat ([[1, dt, ax, 0, 0, 0], # x=x0+dx*dt [0, 1, dt, 0, 0, 0], # dx = dx [0, 0, 1, 0, 0, 0], # ddx = 0 [0, 0, 0, 1, dt, ay], # y = y0 +dy*dt+1/2*g*dt^2 [0, 0, 0, 0, 1, dt], # dy = dy0 + ddy*dt [0, 0, 0, 0, 0, 1]]) # ddy = -g # We will inject air drag using Bu f1.B = np.mat([[0., 0., 1., 0., 0., 0.], [0., 0., 0., 0., 0., 1.]]).T f1.u = np.mat([[0., 0.]]).T f1.H = np.mat([ [1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0]]) f1.R = np.eye(2) * 5 f1.Q = np.eye(6) * 0. omega = radians(omega) vx = cos(omega) * v0 vy = sin(omega) * v0 f1.x = np.mat([x,vx,0,y,vy,-9.8]).T f1.P = np.eye(6) * 500. z = np.mat([[0,0]]).T markers.MarkerStyle(fillstyle='none') np.set_printoptions(precision=4) t=0 while f1.x[3,0] > 0: t+=dt #f1.update (z) x,y = ball.update(dt) #x,y = ball.pos_at_t(t) update_drag(f1, dt) f1.predict() print f1.x.T plt.scatter(f1.x[0,0],f1.x[3,0], color='red', alpha=0.5) plt.scatter (x,y) return f1 if __name__ == '__main__': #test_baseball_path() #test_ball_path() #test_kf() f1 = test_kf_drag()