more fixes for code directory

code/ekf_internal.py

@@ -0,0 +1,164 @@
+# -*- 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()
+    
+    
+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('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()
+    
+show_radar_chart()

code/nonlinear_plots.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sun May 18 11:09:23 2014
+
+@author: rlabbe
+"""
+
+from __future__ import division
+import numpy as np
+import matplotlib.pyplot as plt
+from numpy.random import normal
+
+
+def plot_transfer_func(data, f, lims,num_bins=1000):
+    ys = f(data)
+
+    #plot output
+    plt.subplot(2,2,1)
+    plt.hist(ys, num_bins, orientation='horizontal',histtype='step')
+    plt.ylim(lims)
+    plt.gca().xaxis.set_ticklabels([])
+    plt.title('output')
+
+    # plot transfer function
+    plt.subplot(2,2,2)
+    x = np.arange(lims[0], lims[1],0.1)
+    y = f(x)
+    plt.plot (x,y)
+    isct = f(0)
+    plt.plot([0,0,lims[0]],[lims[0],isct,isct],c='r')
+    plt.xlim(lims)
+    plt.ylim(lims)
+    plt.title('transfer function')
+
+    # plot input
+    plt.subplot(2,2,4)
+    plt.hist(data, num_bins, histtype='step')
+    plt.xlim(lims)
+    plt.gca().yaxis.set_ticklabels([])
+    plt.title('input')
+
+    plt.show()