CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

| Download
Project: test
Views: 91872
Tweet
Share
Share
1
# -*- coding: utf-8 -*-

2
"""

3
Created on Wed Jun  4 12:33:38 2014

4


5
@author: rlabbe

6
"""

7


8
from __future__ import print_function,division

9
from filterpy.kalman import KalmanFilter

10
import numpy as np

11
import matplotlib.pyplot as plt

12
import baseball

13
from numpy.random import randn

14


15
def ball_filter6(dt,R=1., Q = 0.1):

16
    f1 = KalmanFilter(dim=6)

17
    g = 10

18


19
    f1.F = np.mat ([[1., dt, dt**2,  0,       0,  0],

20
                    [0,  1., dt,     0,       0,  0],

21
                    [0,  0,  1.,     0,       0,  0],

22
                    [0,  0,  0.,    1., dt, -0.5*dt*dt*g],

23
                    [0,  0,  0,      0, 1.,      -g*dt],

24
                    [0,  0,  0,      0, 0.,      1.]])

25


26
    f1.H = np.mat([[1,0,0,0,0,0],

27
                   [0,0,0,0,0,0],

28
                   [0,0,0,0,0,0],

29
                   [0,0,0,1,0,0],

30
                   [0,0,0,0,0,0],

31
                   [0,0,0,0,0,0]])

32


33


34
    f1.R = np.mat(np.eye(6)) * R

35


36
    f1.Q = np.zeros((6,6))

37
    f1.Q[2,2] = Q

38
    f1.Q[5,5] = Q

39
    f1.x = np.mat([0, 0 , 0, 0, 0, 0]).T

40
    f1.P = np.eye(6) * 50.

41
    f1.B = 0.

42
    f1.u = 0

43


44
    return f1

45


46


47
def ball_filter4(dt,R=1., Q = 0.1):

48
    f1 = KalmanFilter(dim=4)

49
    g = 10

50


51
    f1.F = np.mat ([[1., dt,  0, 0,],

52
                    [0,  1.,  0, 0],

53
                    [0,  0,  1., dt],

54
                    [0,  0,  0.,  1.]])

55


56
    f1.H = np.mat([[1,0,0,0],

57
                   [0,0,0,0],

58
                   [0,0,1,0],

59
                   [0,0,0,0]])

60


61


62


63
    f1.B = np.mat([[0,0,0,0],

64
                   [0,0,0,0],

65
                   [0,0,1.,0],

66
                   [0,0,0,1.]])

67


68
    f1.u = np.mat([[0],

69
                   [0],

70
                   [-0.5*g*dt**2],

71
                   [-g*dt]])

72


73
    f1.R = np.mat(np.eye(4)) * R

74


75
    f1.Q = np.zeros((4,4))

76
    f1.Q[1,1] = Q

77
    f1.Q[3,3] = Q

78
    f1.x = np.mat([0, 0 , 0, 0]).T

79
    f1.P = np.eye(4) * 50.

80
    return f1

81


82


83
def plot_ball_filter6 (f1, zs, skip_start=-1, skip_end=-1):

84
    xs, ys = [],[]

85
    pxs, pys = [],[]

86


87
    for i,z in enumerate(zs):

88
        m = np.mat([z[0], 0, 0, z[1], 0, 0]).T

89


90
        f1.predict ()

91


92
        if i == skip_start:

93
            x2 = xs[-2]

94
            x1 = xs[-1]

95
            y2 = ys[-2]

96
            y1 = ys[-1]

97


98
        if i >= skip_start and i <= skip_end:

99
            #x,y = baseball.predict (x2, y2, x1, y1, 1/30, 1/30* (i-skip_start+1))

100
            x,y = baseball.predict (xs[-2], ys[-2], xs[-1], ys[-1], 1/30, 1/30)

101


102
            m[0] = x

103
            m[3] = y

104
        #print ('skip', i, f1.x[2],f1.x[5])

105


106
        f1.update(m)

107


108


109
        '''

110
        if i >= skip_start and i <= skip_end:

111
            #f1.x[2] = -0.1

112
            #f1.x[5] = -17.

113
            pass

114
        else:

115
            f1.update (m)

116


117
        '''

118


119
        xs.append (f1.x[0,0])

120
        ys.append (f1.x[3,0])

121
        pxs.append (z[0])

122
        pys.append(z[1])

123


124
        if i > 0 and z[1] < 0:

125
            break;

126


127
    p1, = plt.plot (xs, ys, 'r--')

128
    p2, = plt.plot (pxs, pys)

129
    plt.axis('equal')

130
    plt.legend([p1,p2], ['filter', 'measurement'], 2)

131
    plt.xlim([0,xs[-1]])

132
    plt.show()

133


134


135


136
def plot_ball_filter4 (f1, zs, skip_start=-1, skip_end=-1):

137
    xs, ys = [],[]

138
    pxs, pys = [],[]

139


140
    for i,z in enumerate(zs):

141
        m = np.mat([z[0], 0, z[1], 0]).T

142


143
        f1.predict ()

144


145
        if i == skip_start:

146
            x2 = xs[-2]

147
            x1 = xs[-1]

148
            y2 = ys[-2]

149
            y1 = ys[-1]

150


151
        if i >= skip_start and i <= skip_end:

152
            #x,y = baseball.predict (x2, y2, x1, y1, 1/30, 1/30* (i-skip_start+1))

153
            x,y = baseball.predict (xs[-2], ys[-2], xs[-1], ys[-1], 1/30, 1/30)

154


155
            m[0] = x

156
            m[2] = y

157


158
        f1.update (m)

159


160


161
        '''

162
        if i >= skip_start and i <= skip_end:

163
            #f1.x[2] = -0.1

164
            #f1.x[5] = -17.

165
            pass

166
        else:

167
            f1.update (m)

168


169
        '''

170


171
        xs.append (f1.x[0,0])

172
        ys.append (f1.x[2,0])

173
        pxs.append (z[0])

174
        pys.append(z[1])

175


176
        if i > 0 and z[1] < 0:

177
            break;

178


179
    p1, = plt.plot (xs, ys, 'r--')

180
    p2, = plt.plot (pxs, pys)

181
    plt.axis('equal')

182
    plt.legend([p1,p2], ['filter', 'measurement'], 2)

183
    plt.xlim([0,xs[-1]])

184
    plt.show()

185


186


187
start_skip = 20

188
end_skip = 60

189


190
def run_6():

191
    dt = 1/30

192
    noise = 0.0

193
    f1 = ball_filter6(dt, R=.16, Q=0.1)

194
    plt.cla()

195


196
    x,y = baseball.compute_trajectory(v_0_mph = 100., theta=50., dt=dt)

197


198


199
    znoise = [(i+randn()*noise,j+randn()*noise) for (i,j) in zip(x,y)]

200


201
    plot_ball_filter6 (f1, znoise, start_skip, end_skip)

202


203


204
def run_4():

205
    dt = 1/30

206
    noise = 0.0

207
    f1 = ball_filter4(dt, R=.16, Q=0.1)

208
    plt.cla()

209


210
    x,y = baseball.compute_trajectory(v_0_mph = 100., theta=50., dt=dt)

211


212


213
    znoise = [(i+randn()*noise,j+randn()*noise) for (i,j) in zip(x,y)]

214


215
    plot_ball_filter4 (f1, znoise, start_skip, end_skip)

216


217


218
if __name__ == "__main__":

219
    run_4()

220