-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathplot-forecasts-onto-truth.py
42 lines (33 loc) · 1.02 KB
/
plot-forecasts-onto-truth.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
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime
import utils
import forecasts.IHME as forecast
import truth.CSSE as truth
# Choose a location
loc = "US"
# Load truth data
truth_df = truth.load(location=loc)
# Plot forecasts onto truth data
###############################################################################
fig, ax = plt.subplots()
# Plot truth
ax.plot(truth_df['date'], truth_df['cumDeaths'],
linewidth=2, color='black', label='Actual')
# Plot forecasts
forecast_dates = forecast.list_dates()
for date in forecast_dates:
forecast_df = forecast.load(date, location=loc)
mask = np.isfinite(forecast_df['cumDeaths']) # Filter out NaN
ax.plot(forecast_df['date'][mask],
forecast_df['cumDeaths'][mask], linestyle='--')
# format the ticks
utils.xaxis_months(fig, ax)
ax.set_xlim(left=datetime(2020, 3, 1))
ax.set_ylim(bottom=0)
ax.set_title(utils.get_location_string(loc))
ax.set_xlabel('Date')
ax.set_ylabel('Deaths')
ax.legend()
plt.tight_layout()
plt.show()