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()