def get_median_age(pop_da, gbd_round_id, year_index):
age_metadata = db.get_ages(gbd_round_id=gbd_round_id)
age_group_ids = age_metadata["age_group_id"].tolist()
days = age_metadata[[
"age_group_id", "age_group_days_start", "age_group_days_end"
]].set_index("age_group_id").to_xarray()
days["mean_age"] = (days["age_group_days_end"] -
(days["age_group_days_end"] -
days["age_group_days_start"]) / 2) / 365.25
if year_index == 2100:
pop_draw_base = pop_da.sel(scenario=0,
sex_id=3,
location_id=1,
age_group_id=age_group_ids,
year_id=year_index)
median_age = weighted_quantile_with_extra_dim(days["mean_age"],
0.5, ["age_group_id"],
pop_draw_base,
extra_dim="draw")
median_age_mean = median_age.mean("draw").values.round(2)
median_age_upper = median_age.quantile(0.975,
dim="draw").values.round(2)
median_age_lower = median_age.quantile(0.025,
dim="draw").values.round(2)
return median_age_mean, median_age_lower, median_age_upper
else:
pop_da_2017 = pop_da.sel(sex_id=3,
location_id=1,
age_group_id=age_group_ids,
year_id=year_index)
median_age = weighted_quantile(days["mean_age"], 0.5, ["age_group_id"],
pop_da_2017)
return median_age.values.round(2)
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 get_cohort_info_from_age_year(age_group_ids, years):
"""Calculates the cohort year for each age_group - year pair.
In the context of education, a cohort year is the year the cohort
turned 5 years old and not the birth year.
For a given age_group - year pair, its cohort year is calculated
in the following manner:
cohort_year = (year - age_group_lower_bound) + 5
The 5 is added at the end to force the cohort to start at age 5
and not 0.
For example:
Let the age group be 10 and the year be 1990. Then this group
will belong to the (1990 - 25) + 5 = **1970** cohort.
Args:
age_group_ids (list(int)):
List of age group ids.
years (YearRange):
The past and forecasted years.
Returns:
cohort_age_df (pandas.DataFrame):
Dataframe with (age&year) to cohort mapping.
"""
LOGGER.info("In get_cohort_info_from_age_year")
age_meta_df = db.get_ages()[[
'age_group_id', 'age_group_years_start']]
all_age_year_df = pd.MultiIndex.from_product(
[age_group_ids, years.years],
names=['age_group_id', 'year_id']
).to_frame().reset_index(drop=True)
cohort_age_df = all_age_year_df.merge(
age_meta_df, on='age_group_id').astype({'age_group_years_start': int})
# Finding the cohort from year and lower bound of age_group as described
# above.
cohort_age_df['cohort_year'] = (
cohort_age_df['year_id'] - (
cohort_age_df['age_group_years_start'] - 5))
# All cohorts don't need correction.Only those that extend into the future.
correction_cohorts_years = cohort_age_df[
cohort_age_df['year_id'] > years.past_end]['cohort_year'].unique()
cohort_age_df['need_correction'] = \
cohort_age_df['cohort_year'].isin(correction_cohorts_years)
cohort_age_df.drop('age_group_years_start', axis=1, inplace=True)
return cohort_age_df
def prep_pop_da(past_version, forecast_version, gbd_round_id, years):
forecast_pop_file = FBDPath(
f"/{gbd_round_id}/future/population/{forecast_version}/"
f"population_combined.nc")
forecast_fhs = open_xr(forecast_pop_file).data.sel(quantile='mean',
drop=True)
past_fhs_file = FBDPath(
f"/{gbd_round_id}/past/population/{past_version}/population.nc")
past_fhs = expand_dimensions(open_xr(past_fhs_file).data.sel(
year_id=years.past_years,
sex_id=forecast_fhs["sex_id"],
age_group_id=forecast_fhs["age_group_id"],
location_id=forecast_fhs["location_id"]),
scenario=forecast_fhs.scenario.values)
fhs_all_scenarios = xr.concat([past_fhs, forecast_fhs], dim="year_id")
fhs = fhs_all_scenarios.sel(scenario=[-1, 0, 1])
alt_sdg = fhs_all_scenarios.sel(scenario=[3])
alt_99 = fhs_all_scenarios.sel(scenario=[2])
ages = db.get_ages().query("age_group_id in @ALL_AGE_GROUP_IDS")
days = ages[["age_group_id", "age_group_days_start", "age_group_days_end"]]
days["mean_age"] = (days["age_group_days_end"] -
(days["age_group_days_end"] -
days["age_group_days_start"]) / 2) / 365.25
mean_age = days.set_index("age_group_id")["mean_age"].to_xarray()
data_fhs = fhs.sel(age_group_id=mean_age["age_group_id"], sex_id=SEX_IDS)
data_sdg = alt_sdg.sel(age_group_id=mean_age["age_group_id"],
sex_id=SEX_IDS)
data_99 = alt_99.sel(age_group_id=mean_age["age_group_id"], sex_id=SEX_IDS)
avg_age_fhs = (data_fhs *
mean_age).sum("age_group_id") / data_fhs.sum("age_group_id")
avg_age_sdg = (data_sdg *
mean_age).sum("age_group_id") / data_sdg.sum("age_group_id")
avg_age_99 = (data_99 *
mean_age).sum("age_group_id") / data_99.sum("age_group_id")
ds = data_fhs.rename("population").to_dataset()
ds_sdg = data_sdg.rename("population").to_dataset()
ds_99 = data_99.rename("population").to_dataset()
return avg_age_fhs, avg_age_sdg, avg_age_99, ds, ds_sdg, ds_99
def lag_scenarios(data, years):
"""Lag scenarios by age-years so that scenarios only deviate from the
reference in age-time pairs that have not finished education in the first
year of the forecast.
Args:
data (xarray.DataArray):
data that must include forecasts and can include past.
years (YearRange):
forecasting timeseries
Returns:
(xarray.DataArray):
Forecasts with adjusted/lagged scenarios.
"""
age_groups = db.get_ages()
age_group_dict = dict(
zip(age_groups["age_group_id"], age_groups["age_group_years_start"]))
forecast = data.sel(year_id=years.forecast_years)
adjusted_scenarios = []
for age_group_id in forecast["age_group_id"].values:
age_group_years_start = age_group_dict[age_group_id]
num_years_to_shift = get_num_years_to_shift(age_group_years_start)
age_slice = forecast.sel(age_group_id=age_group_id)
ref_age_slice = age_slice.sel(scenario=REFERENCE_SCENARIO, drop=True)
adjusted_scen_age_slice = []
for scenario in (-1, 1):
scen_age_slice = age_slice.sel(scenario=scenario).drop("scenario")
diff = scen_age_slice - ref_age_slice
shifted_diff = diff.shift(year_id=num_years_to_shift).fillna(0)
shifted_scen_slice = shifted_diff + ref_age_slice
shifted_scen_slice["scenario"] = scenario
adjusted_scen_age_slice.append(shifted_scen_slice)
ref_age_slice["scenario"] = 0
adjusted_scen_age_slice.append(ref_age_slice)
adjusted_scen_age_slice = xr.concat(adjusted_scen_age_slice,
dim="scenario")
adjusted_scenarios.append(adjusted_scen_age_slice)
adjusted_scenarios = xr.concat(adjusted_scenarios, dim="age_group_id")
return adjusted_scenarios.combine_first(data)
def pop_plot(avg_age_fhs,
avg_age_sdg,
avg_age_99,
ds,
ds_sdg,
ds_99,
years,
location_id=1):
location_metadata = db.get_locations_by_max_level(3)
ages = db.get_ages().query("age_group_id in @ALL_AGE_GROUP_IDS")
scenarios = [
{
"year": years.past_end,
"scenario": 0,
"name": years.past_end,
"color": "black"
},
{
"year": years.forecast_end,
"scenario": 0,
"name": "Reference",
"color": "steelblue"
},
{
"year": years.forecast_end,
"scenario": -1,
"name": "Slower Met Need and Education Pace",
"color": "firebrick"
},
{
"year": years.forecast_end,
"scenario": 1,
"name": "Faster Met Need and Education Pace",
"color": "forestgreen"
},
]
alt_scenario_sdg = [{
"year": years.forecast_end,
"scenario": 3,
"name": "SDG Met Need and Education Pace",
"color": "#ff7f00"
}]
alt_scenario_99 = [{
"year": years.forecast_end,
"scenario": 2,
"name": "Fastest Met Need and Education Pace",
"color": "#984ea3"
}]
gs = plt.GridSpec(1, 2)
ax_male = plt.subplot(gs[:, 0])
ax_female = plt.subplot(gs[:, 1])
fig = ax_female.figure
# pop plots
pop_df = ds.sel(
location_id=location_id, age_group_id=ages["age_group_id"].values
)["population"].to_dataframe().reset_index().sort_values("age_group_id")
alt_df_sdg = ds_sdg.sel(
location_id=location_id, age_group_id=ages["age_group_id"].values
)["population"].to_dataframe().reset_index().sort_values("age_group_id")
alt_df_99 = ds_99.sel(
location_id=location_id, age_group_id=ages["age_group_id"].values
)["population"].to_dataframe().reset_index().sort_values("age_group_id")
pop_df["age"] = pop_df["age_group_id"].factorize()[0]
male_pop = pop_df[pop_df.sex_id == 1]
male_max_pop = male_pop["population"].max() * 1.01
female_pop = pop_df[pop_df.sex_id == 1]
female_max_pop = female_pop["population"].max() * 1.01
# sdg
alt_df_sdg["age"] = alt_df_sdg["age_group_id"].factorize()[0]
# 99
alt_df_99["age"] = alt_df_99["age_group_id"].factorize()[0]
# comment or uncomment for fixed/non-fixed axis
max_pop = max([male_max_pop, female_max_pop])
# max_pop = pop_df["population"].max() * 1.01
if max_pop > 1e9:
max_pop = max_pop / 1e9
pop_df["population"] = pop_df["population"] / 1e9
alt_df_sdg["population"] = alt_df_sdg["population"] / 1e9
alt_df_99["population"] = alt_df_99["population"] / 1e9
label = "Population (Billions)"
elif max_pop > 1e6:
max_pop = max_pop / 1e6
pop_df["population"] = pop_df["population"] / 1e6
alt_df_sdg["population"] = alt_df_sdg["population"] / 1e6
alt_df_99["population"] = alt_df_99["population"] / 1e6
label = "Population (Millions)"
elif max_pop > 1e3:
max_pop = max_pop / 1e3
pop_df["population"] = pop_df["population"] / 1e3
alt_df_sdg["population"] = alt_df_sdg["population"] / 1e3
alt_df_99["population"] = alt_df_99["population"] / 1e3
label = "Population (Thousands)"
else:
label = "Population"
# male plot
for c in scenarios:
df = pop_df.query(
"sex_id == 1 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_male.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8,
label=c["name"])
a = avg_age_fhs.sel(location_id=location_id,
sex_id=1,
scenario=c["scenario"],
year_id=c["year"])
ax_male.plot(0,
a / 5 + 0.5,
marker="<",
color=c["color"],
markersize=20,
alpha=0.8)
# sdg
for c in alt_scenario_sdg:
df = alt_df_sdg.query(
"sex_id == 1 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_male.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8,
label=c["name"])
a = avg_age_sdg.sel(location_id=location_id,
sex_id=1,
scenario=c['scenario'],
year_id=c["year"])
ax_male.plot(0,
a / 5 + 0.5,
marker="<",
color=c["color"],
markersize=20,
alpha=0.8)
# 99
for c in alt_scenario_99:
df = alt_df_99.query(
"sex_id == 1 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_male.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8,
label=c["name"])
a = avg_age_99.sel(location_id=location_id,
sex_id=1,
scenario=c['scenario'],
year_id=c["year"])
ax_male.plot(0,
a / 5 + 0.5,
marker="<",
color=c["color"],
markersize=20,
alpha=0.8)
ax_male.set_xlim(max_pop + .1 * max_pop, 0)
ax_male.set_ylim(0, ages.shape[0])
ax_male.set_yticks(np.arange(ages.shape[0]) + 0.5)
ax_male.set_yticklabels(ages["age_group_name_short"], fontsize=14)
ax_male.set_title("Male", fontsize=18)
ax_male.set_xlabel(label)
ax_male.legend(frameon=False, loc="upper left")
sns.despine(ax=ax_male, left=True, right=False)
# female plot
for c in scenarios:
df = pop_df.query(
"sex_id == 2 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_female.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8)
a = avg_age_fhs.sel(location_id=location_id,
sex_id=2,
scenario=c["scenario"],
year_id=c["year"])
ax_female.plot(0,
a / 5 + 0.5,
marker=">",
color=c["color"],
markersize=20,
alpha=0.8)
# sdg
for c in alt_scenario_sdg:
df = alt_df_sdg.query(
"sex_id == 2 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_female.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8)
a = avg_age_sdg.sel(location_id=location_id,
sex_id=2,
scenario=c['scenario'],
year_id=c["year"])
ax_female.plot(0,
a / 5 + 0.5,
marker=">",
color=c["color"],
markersize=20,
alpha=0.8)
# 99
for c in alt_scenario_99:
df = alt_df_99.query(
"sex_id == 2 & scenario == {} & year_id == {}".format(
c["scenario"], c["year"]))
ax_female.step(x=df["population"].values.tolist() + [0],
y=df["age"].values.tolist() + [ages.shape[0]],
color=c["color"],
linewidth=2,
alpha=0.8)
a = avg_age_99.sel(location_id=location_id,
sex_id=2,
scenario=c['scenario'],
year_id=c["year"])
ax_female.plot(0,
a / 5 + 0.5,
marker=">",
color=c["color"],
markersize=20,
alpha=0.8)
ax_female.set_xlim(0, max_pop + .1 * max_pop)
ax_female.set_ylim(0, ages.shape[0])
ax_female.set_yticks(np.arange(ages.shape[0]) + 0.5)
ax_female.set_yticklabels([])
ax_female.set_title("Female", fontsize=18)
ax_female.set_xlabel(label)
sns.despine(ax=ax_female)
fig.suptitle(location_metadata.query("location_id == @location_id")
["location_name"].values[0],
fontsize=28)
fig.set_size_inches(14, 8)
plt.tight_layout()
plt.subplots_adjust(top=.85)
return fig
def main(asfr_version, past_asfr_version, location_id, gbd_round_id, years,
granularity, iterations, **kwargs):
"""
1. Read in location-specific draws of period ASFR from CCF stage
2. Add terminal age group ASFR's
3. Intercept shift asfr by holding CCF50 constant.
4. Export location-specific intercept-shifted ASFR in .nc
Args:
asfr_version (str): version name of future ccf/asfr.
past_asfr_version (str): asfr version from past.
location_id (int): location_id.
gbd_round_id (int): gbd round id.
years (YearRange): past_start:forecast_start:forecast_end
iterations (int): number of times to intercept-shift.
"""
ages_df = db.get_ages(gbd_round_id)[[
"age_group_id", "age_group_years_start", "age_group_years_end"
]]
# read the location-specific asfr .csv into dataarray
# the raw forecasted ASFR are stored in the CCF stage of the same
ccf_fbd_path = FBDPath(gbd_round_id=gbd_round_id,
past_or_future="future",
stage="ccf",
version=asfr_version)
if granularity == 1:
sub_folder = "asfr_single_year"
ccf_asfr_fbd_path = ccf_fbd_path / sub_folder
future_asfr = read_to_xr(location_id,
ccf_asfr_fbd_path,
dims=list(ASFR_NON_AGE_DIMS + ("age", )))
else:
sub_folder = "asfr"
ccf_asfr_fbd_path = ccf_fbd_path / sub_folder
future_asfr =\
read_to_xr(location_id, ccf_asfr_fbd_path,
dims=list(ASFR_NON_AGE_DIMS + ("age_group_id",)))
# we intercept-shift in 1-year ages, so convert to single years
future_asfr = _expand_5yr_age_groups_to_1yr_ages(future_asfr, ages_df)
if "sex_id" in future_asfr.dims:
raise ValueError("Found sex_id dim in future asfr")
# now etl the past asfr data
past_asfr_fbd_path = FBDPath(gbd_round_id=gbd_round_id,
past_or_future="past",
stage="asfr",
version=past_asfr_version)
past_asfr =\
open_xr(past_asfr_fbd_path /
"asfr.nc").data.sel(location_id=location_id)
if "sex_id" in past_asfr.dims:
raise ValueError("Found sex_id dim in past asfr")
# past has no scenarios, so we need to expand it for merging
past_asfr = expand_dimensions(past_asfr, scenario=future_asfr["scenario"])
# past asfr has age group ids 7-15, but future asfr in ccf only has 8-14.
# we only need age groups 8-14 for intercept shift
past_asfr_1yrs = _expand_5yr_age_groups_to_1yr_ages(
past_asfr.sel(age_group_id=range(8, 15)), ages_df)
# now ready to concat past and future together for intercept shift
asfr = xr.concat([
past_asfr_1yrs.sel(year_id=years.past_years),
future_asfr.sel(year_id=years.forecast_years)
],
dim="year_id")
del past_asfr_1yrs, future_asfr
gc.collect()
# the intercept-shift should keep ccf50 (asfr sum) constant
pre_fix_asfr_sum = asfr.sum() # sum of all asfr values before shift
asfr = ccf50_intercept_shift_lpf(asfr, gbd_round_id, years, iterations)
post_fix_asfr_sum = asfr.sum() # asfr sum post-shift should stay the same
assert np.isclose(post_fix_asfr_sum, pre_fix_asfr_sum, rtol=RTOL),\
f"The intercept shift changed total asfr sum by more than rtol={RTOL}"
# need to save years.past_end for cohort-component model
save_years = [years.past_end] + years.forecast_years.tolist()
asfr = asfr.sel(year_id=save_years) # only do forecast
# convert forecasted asfr back to 5-year age groups
asfr = _convert_ages_to_5_year_age_groups_by_mean(asfr, ages_df)
# add 10-15 (7) and 50-55 (15) age groups for forecasted asfr
asfr = extrapolate_terminal_asfr_age_groups(past_asfr,
asfr,
last_year=years.past_end)
asfr["location_id"] = location_id
asfr.name = "value"
del past_asfr
gc.collect()
LOGGER.info("Finished CCF50 intercept-shift")
asfr_fbd_path = FBDPath(gbd_round_id=gbd_round_id,
past_or_future="future",
stage="asfr",
version=asfr_version)
save_xr(asfr,
asfr_fbd_path / f"{location_id}.nc",
metric="rate",
space="identity",
version=asfr_version,
past_asfr_version=past_asfr_version,
iterations=iterations)
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()