-
Notifications
You must be signed in to change notification settings - Fork 37
/
example.py
75 lines (58 loc) · 2.42 KB
/
example.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
"""
Smoothing and prediction examples with figures. Needs matplotlib
to display the graphics.
Example output:
https://github.com/oseiskar/simdkalman/blob/master/doc/example.png
"""
import simdkalman
import numpy as np
import numpy.random as random
kf = simdkalman.KalmanFilter(
state_transition = np.array([[1,1],[0,1]]),
process_noise = np.diag([0.1, 0.01]),
observation_model = np.array([[1,0]]),
observation_noise = 1.0)
# simulate 100 random walk time series
rand = lambda: random.normal(size=(100, 200))
data = np.cumsum(np.cumsum(rand()*0.02, axis=1) + rand(), axis=1) + rand()*3
# introduce 10% of NaNs denoting missing values
data[random.uniform(size=data.shape) < 0.1] = np.nan
# fit noise parameters to data with the EM algorithm (optional)
kf = kf.em(data, n_iter=10)
# smooth and explain existing data
smoothed = kf.smooth(data)
# predict new data
pred = kf.predict(data, 15)
# could be also written as
# r = kf.compute(data, 15); smoothed = r.smoothed; pred = r.predicted
import matplotlib.pyplot as plt
# show the first 3 smoothed time series
for i in range(3):
_, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
plt.title("time series %d" % (i+1))
x = np.arange(0, data.shape[1])
ax1.plot(x, data[i,:], 'b.', label="data")
smoothed_obs = smoothed.observations.mean[i,:]
obs_stdev = np.sqrt(smoothed.observations.cov[i,:])
ax1.plot(x, smoothed_obs, 'r-', label="smoothed")
ax1.plot(x, smoothed_obs - obs_stdev, 'k--', label="67% confidence")
ax1.plot(x, smoothed_obs + obs_stdev, 'k--')
x_pred = np.arange(data.shape[1], data.shape[1]+pred.observations.mean.shape[1])
y_pred = pred.observations.mean[i,:]
pred_stdev = np.sqrt(pred.observations.cov[i,:])
ax1.plot(x_pred, y_pred, 'b-', label="predicted")
ax1.plot(x_pred, y_pred + pred_stdev, 'k--')
ax1.plot(x_pred, y_pred - pred_stdev, 'k--')
ax1.legend()
trend = smoothed.states.mean[i,:,1]
trend_stdev = np.sqrt(smoothed.states.cov[i,:,1,1])
ax2.plot(x, trend, 'g-', label="trend")
ax2.plot(x, trend - trend_stdev, 'k--', label="67% confidence")
ax2.plot(x, trend + trend_stdev, 'k--')
trend_pred = pred.states.mean[i,:,1]
trend_pred_stdev = np.sqrt(pred.states.cov[i,:,1,1])
ax2.plot(x_pred, trend_pred, 'b-', label='predicted')
ax2.plot(x_pred, trend_pred + trend_pred_stdev, 'k--')
ax2.plot(x_pred, trend_pred - trend_pred_stdev, 'k--')
ax2.legend()
plt.show()