def plot_met_demand(forecasts):
''' Plot forecasted met_demand.
'''
loc_map = get_modeled_locations().\
set_index('location_id')['location_name'].to_dict()
age_map = get_ages().set_index('age_group_id')['age_group_name'].to_dict()
locations = [l for l in forecasts.location_id.values if l in loc_map.keys()]
outfile = os.path.join(settings.MET_DEMAND_FORECAST_DIR, 'graphs/met_demand.pdf')
if not os.path.exists(os.path.dirname(outfile)):
os.makedirs(os.path.dirname(outfile))
with PdfPages(outfile) as pp:
for location_id in locations:
location = loc_map[location_id]
nrow = 3
ncol = 3
fig = plt.figure(figsize=(12, 10))
grid = gridspec.GridSpec(nrow, ncol)
for ix, age_group_id in enumerate(np.arange(8, 15)):
ax = fig.add_subplot(grid[ix])
tmp = forecasts.loc[{'location_id': location_id, 'age_group_id': age_group_id}]
tmp_past = tmp.loc[{'year_id': np.arange(1990, 2017)}]
tmp_future = tmp.loc[{'year_id': np.arange(2017, 2041)}]
ax.plot(tmp_past.year_id, tmp_past.values[0], color='b')
ax.plot(tmp_future.year_id, tmp_future.values[0], color='g')
ax.text(0.7, 0.95, age_map[age_group_id],
verticalalignment='top',
horizontalalignment='left', transform=ax.transAxes,
color='black', fontsize=12)
fig.suptitle(location, fontsize=15)
pp.savefig(bbox_inches='tight')
def demographic_coords(gbd_round_id, years):
"""
Creates and caches an OrderedDict of demographic indices
Args:
gbd_round_id (int): gbd round id.
years (YearRange): [past_start, forecast_start, forecast_end] years.
Returns:
OrderedDict: ordered dict of all non-draw dimensions and their
coordinates.
"""
if not (hasattr(demographic_coords, "coords") and
demographic_coords.gbd_round_id == gbd_round_id and
demographic_coords.years == years):
location_ids =\
db.get_modeled_locations(gbd_round_id=gbd_round_id).\
location_id.values.tolist()
age_group_ids =\
db.get_ages(gbd_round_id=gbd_round_id)[AGE_DIM].unique().tolist()
_coords =\
collections.OrderedDict([(LOCATION_DIM, location_ids),
(AGE_DIM, age_group_ids),
(SEX_DIM, list(SEXES)),
(YEAR_DIM, years.years),
(SCENARIO_DIM, list(SCENARIOS))])
demographic_coords.coords = _coords
demographic_coords.gbd_round_id = gbd_round_id
demographic_coords.years = years
return demographic_coords.coords
def plot_paf(df, acause, risk, date, paf_cols, demography_cols):
''' Plot cause specific scalars.'''
age_map = get_ages().\
set_index('age_group_id')['age_group_name'].to_dict()
location_map = get_modeled_locations().\
set_index('location_id')['location_name'].to_dict()
df['mean'] = np.mean(df.loc[:, paf_cols].values, axis=1)
df['lower'] = np.percentile(df.loc[:, paf_cols].values, 2.5, axis=1)
df['upper'] = np.percentile(df.loc[:, paf_cols].values, 97.5, axis=1)
df = df[demography_cols + ['mean', 'lower', 'upper']]
df = df.loc[df.scenario == 0]
outfile = ('/ihme/forecasting/data/paf/{date}/plots/'
'{acause}_{risk}_2.pdf'.format(date=date,
risk=risk,
acause=acause))
if not os.path.exists(os.path.dirname(outfile)):
os.makedirs(os.path.dirname(outfile))
location_ids = get_gbd_demographics().location_id.unique()
with PdfPages(outfile) as pp:
for location_id in location_ids: # [6, 102]
test = df.loc[df.location_id == location_id]
nrow = 4
ncol = 3
fig = plt.figure(figsize=(12, 10))
grid = gridspec.GridSpec(nrow, ncol)
for ix, age_group_id in enumerate(xrange(2, 14)):
ax = fig.add_subplot(grid[ix])
tmp = test.loc[test.age_group_id == age_group_id]
tmp = tmp.drop_duplicates(demography_cols)
ax.plot(tmp.year_id.unique(),
tmp.loc[tmp.sex_id == 2]['mean'].values,
'b',
label='Female')
ax.fill_between(tmp.year_id.unique(),
tmp.loc[tmp.sex_id == 2]['lower'].values,
tmp.loc[tmp.sex_id == 2]['upper'].values)
ax.text(0.7,
0.95,
age_map[age_group_id],
verticalalignment='top',
horizontalalignment='left',
transform=ax.transAxes,
color='black',
fontsize=12)
location = location_map[location_id]
suptitle = '{location}, {acause}'.format(location=location,
acause=acause)
fig.suptitle(suptitle, fontsize=15)
pp.savefig()
def _draw_epsilons(epsilon_past, draws, smoothing, years, acause, decay,
gbd_round_id):
"""Draws forecasts for epsilons. For all-cause, this is done by running an
attenuated drift model first to remove some of the time trend from the
epsilons. Then for all causes (including all-cause) except NTD's, a Pooled
random walk model is used to forecast the remaining residuals and generate
expanding uncertainty.
:param xarray.DataArray epsilon_past: past epsilons (error of predictions
based on data)
:param int draws: number of draws to grab.
:param list[str] smoothing: what dimensions to smooth over in ARIMA
modeling.
:param fbd_core.argparse.YearRange years: a container for the three years
which define our forecast.
:param str acause: the cause to forecast epsilons for
"""
logger.info(("Sampling epsilon_hat from ARIMA with cross sections by {}"
"{}-{}").format(smoothing, years.forecast_start,
years.forecast_end))
# for all-cause, remove drift first, then run random walk on the remainder
if acause == "_all":
drift_component = remove_drift.get_decayed_drift_preds(
epsilon_past, years, decay)
remainder = epsilon_past - drift_component.sel(
year_id=years.past_years)
ds = xr.Dataset(dict(y=remainder.copy()))
# if not all-cause, directly model epsilons with random walk
else:
ds = xr.Dataset(dict(y=epsilon_past.copy()))
regdf = db.get_modeled_locations(gbd_round_id)[[
"location_id", "region_id", "super_region_id"
]]
regdf.set_index("location_id", inplace=True)
ds.update(xr.Dataset(regdf))
if acause == "_all":
rw_model = rw.RandomWalk(ds, years.past_start, years.forecast_start,
years.forecast_end, draws)
rw_model.fit()
rw_preds = rw_model.predict()
epsilon_hat = drift_component + rw_preds
else:
# if not all-cause, use a pooled random walk whose location pooling dimension is based on cause levelß
pooled_rw = PooledRandomWalk(ds,
draws=draws,
holdout=years.forecast_start,
forecast=years.forecast_end,
dims=smoothing)
pooled_rw.fit()
rw_preds = pooled_rw.predict()
epsilon_hat = rw_preds
return epsilon_hat
def plot_scalars(scalar_ds,
acause,
sex_id,
location_ids,
scenario,
date,
outdir,
start_age_group_id=10,
end_age_group_id=22):
''' Plot scalars with uncertainties.
Parameters
----------
scalar_ds: scalar in xarray.Dataset format.
acause: str
sex_id: int, 1 or 2
location_ids: list
scenario: -1 or 0 or 1
date: date to version-control plots
outdir: output directory
start_age_group_id: int, default 10
end_age_group_id: int, default 22
'''
outfile = os.path.join(outdir.format(d=date),
'{}_{}.pdf'.format(acause, scenario))
if not os.path.exists(os.path.dirname(outfile)):
os.makedirs(os.path.dirname(outfile))
ds = scalar_ds.loc[{'scenario': scenario, 'sex_id': sex_id}]
# Calculate statistics.
ds['mean'] = ds['value'].mean(dim='draw')
ds['lower'] = ds['value'].quantile(0.025, dim='draw')
ds['upper'] = ds['value'].quantile(0.975, dim='draw')
age_map = get_ages().set_index('age_group_id')['age_group_name'].to_dict()
with PdfPages(outfile) as pp:
for location_id in location_ids:
loc_ds = ds.loc[{"location_id": location_id}]
nrow = 4
ncol = 3
fig = plt.figure(figsize=(12, 10))
grid = gridspec.GridSpec(nrow, ncol)
for ix, age_group_id in enumerate(
xrange(start_age_group_id, end_age_group_id)):
ax = fig.add_subplot(grid[ix])
age_ds = loc_ds.loc[{"age_group_id": age_group_id}]
ax.plot(age_ds['year_id'], age_ds['mean'])
ax.fill_between(age_ds['year_id'], age_ds['lower'],
age_ds['upper'])
ax.text(0.7,
0.95,
age_map[age_group_id],
verticalalignment='top',
horizontalalignment='left',
transform=ax.transAxes,
color='black',
fontsize=12)
location_map = get_modeled_locations().\
set_index('location_id')['location_name'].to_dict()
location = location_map[location_id]
suptitle = "{location}, {acause}; Scenario: {scenario}".format(
location=location, acause=acause, scenario=scenario)
fig.suptitle(suptitle, fontsize=15)
pp.savefig(bbox_inches='tight')
logger.info('Scalars plotting finished: {}'.format(acause))
def plot_risk_attributable(ds,
acause,
sex_id,
location_ids,
risk,
outdir,
start_age_group_id=10,
end_age_group_id=22):
''' Plot risk attributable mortality across scenarios.
# NOTE: this is called within plot_risk_attr_mort()
Parameters
----------
ds: risk attributable mortality in xarray format.
acause: str
risk: str
sex_id: int, 1 or 2
location_ids: list
outdir: output directory
start_age_group_id: int, default 10
end_age_group_id: int, default 22
'''
# TODO Kendrick said: This is kind of weird (I think), because if I call
# a different plotting function plot_thing(), and then this plotting
# function, and then call plot_thing() again,
# it could potentially change plot_thing() because of the sns stuff.
sns.despine()
sns.set_style('ticks')
age_map = get_ages().\
set_index('age_group_id')['age_group_name'].to_dict()
location_map = get_modeled_locations().\
set_index('location_id')['location_name'].to_dict()
color_map = ['r', 'b', 'g']
sexn = 'male' if sex_id == 1 else 'female'
outfile = os.path.join(outdir, '{}_{}_{}.pdf'.format(acause, risk, sexn))
if not os.path.exists(os.path.dirname(outfile)):
os.makedirs(os.path.dirname(outfile))
# Calculate statistics.
ds['mean'] = ds['value'].mean(dim='draw')
ds = ds.loc[dict(sex_id=sex_id)]
with PdfPages(outfile) as pp:
for location_id in location_ids:
loc_ds = ds.loc[{"location_id": location_id}]
nrow = 4
ncol = 3
fig = plt.figure(figsize=(15, 12))
grid = gridspec.GridSpec(nrow, ncol)
for ix, age_group_id in enumerate(
xrange(start_age_group_id, end_age_group_id)):
ax = fig.add_subplot(grid[ix])
age_ds = loc_ds.loc[{"age_group_id": age_group_id}]
for scenario in [-1, 1, 0]:
scenario_ds = age_ds.loc[dict(scenario=scenario)]
ax.plot(scenario_ds['year_id'],
scenario_ds['mean'],
color=color_map[scenario + 1])
ax.text(0.7,
0.95,
age_map[age_group_id],
verticalalignment='top',
horizontalalignment='left',
transform=ax.transAxes,
color='black',
fontsize=12)
ax.axvline(x=2015, color='k')
location = location_map[location_id]
suptitle = "{location}, {acause}, {risk}".format(location=location,
acause=acause,
scenario=scenario,
risk=risk)
fig.suptitle(suptitle, fontsize=15)
pp.savefig()