diff --git a/ogusa/demographics.py b/ogusa/demographics.py index 57dd9878..7775152f 100644 --- a/ogusa/demographics.py +++ b/ogusa/demographics.py @@ -86,8 +86,9 @@ def get_un_data( # keep just what is needed from data df = df[df.variant == "Median"] df = df[df.sex == "Both sexes"][["timeLabel", "ageLabel", "value"]] - df.rename({"timeLabel": "year", "ageLabel": "age"}, - axis=1, inplace=True) + df.rename( + {"timeLabel": "year", "ageLabel": "age"}, axis=1, inplace=True + ) df.loc[df.age == "100+", "age"] = 100 df.age = df.age.astype(int) df.year = df.year.astype(int) @@ -101,8 +102,15 @@ def get_un_data( return df -def get_fert(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_CODE, start_year=START_YEAR, - end_year=END_YEAR, graph=False): +def get_fert( + totpers=100, + min_age=0, + max_age=100, + country_id=UN_COUNTRY_CODE, + start_year=START_YEAR, + end_year=END_YEAR, + graph=False, +): """ This function generates a vector of fertility rates by model period age that corresponds to the fertility rate data by age in years. @@ -120,7 +128,9 @@ def get_fert(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_CODE, st """ # Read UN data - df = get_un_data("68", start_year=start_year, end_year=end_year,country_id=country_id) + df = get_un_data( + "68", start_year=start_year, end_year=end_year, country_id=country_id + ) # put in vector fert_rates = df.value.values # fill in with zeros for ages < 15 and > 49 @@ -162,8 +172,15 @@ def get_fert(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_CODE, st return fert_rates -def get_mort(totpers=100, min_age=0, max_age=100, start_year=START_YEAR, - end_year=END_YEAR, country_id=UN_COUNTRY_CODE, graph=True): +def get_mort( + totpers=100, + min_age=0, + max_age=100, + start_year=START_YEAR, + end_year=END_YEAR, + country_id=UN_COUNTRY_CODE, + graph=True, +): """ This function generates a vector of mortality rates by model period age. @@ -182,8 +199,9 @@ def get_mort(totpers=100, min_age=0, max_age=100, start_year=START_YEAR, """ # Read UN data - df = get_un_data("80", start_year=start_year, - end_year=end_year,country_id=country_id) + df = get_un_data( + "80", start_year=start_year, end_year=end_year, country_id=country_id + ) # put in vector mort_rates_data = df.value.values # In UN data, mortality rates for 0 year olds are the infant @@ -266,8 +284,14 @@ def pop_rebin(curr_pop_dist, totpers_new): return curr_pop_new -def get_imm_rates(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_CODE, start_year=START_YEAR, - end_year=END_YEAR): +def get_imm_rates( + totpers=100, + min_age=0, + max_age=100, + country_id=UN_COUNTRY_CODE, + start_year=START_YEAR, + end_year=END_YEAR, +): """ Calculate immigration rates by age as a residual given population levels in different periods, then output average calculated @@ -288,9 +312,18 @@ def get_imm_rates(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_COD """ # Read UN data num_years = 4 # note that code below only uses four years, in future, this should be more flexbile - imm_start_year = start_year - num_years # year to start finding residual imm rate - imm_end_year = start_year + num_years - 1 # last year to find residual imm rate - df = get_un_data("47", country_id=country_id, start_year=imm_start_year, end_year=imm_end_year) + imm_start_year = ( + start_year - num_years + ) # year to start finding residual imm rate + imm_end_year = ( + start_year + num_years - 1 + ) # last year to find residual imm rate + df = get_un_data( + "47", + country_id=country_id, + start_year=imm_start_year, + end_year=imm_end_year, + ) # separate pop dist by year and put into dictionary of arrays pop_dict = {} @@ -309,8 +342,12 @@ def get_imm_rates(totpers=100, min_age=0, max_age=100, country_id=UN_COUNTRY_COD pop_list.append(pop_dict[t][0]) pop11vec = np.array(pop_list[:-1]) pop21vec = np.array(pop_list[1:]) - fert_rates = get_fert(totpers, min_age, max_age, country_id, start_year, end_year, False) - mort_rates, infmort_rate = get_mort(totpers, min_age, max_age, country_id, start_year, end_year, False) + fert_rates = get_fert( + totpers, min_age, max_age, country_id, start_year, end_year, False + ) + mort_rates, infmort_rate = get_mort( + totpers, min_age, max_age, country_id, start_year, end_year, False + ) newbornvec = np.dot( fert_rates, np.vstack((pop_dict[0], pop_dict[1], pop_dict[2])).T ) @@ -447,10 +484,16 @@ def get_pop_objs( # Get fertility, mortality, and immigration rates # will be used to generate population distribution in future years - fert_rates = get_fert(E + S, min_age, max_age, country_id, data_year, data_year) - mort_rates, infmort_rate = get_mort(E + S, min_age, max_age, country_id, data_year, data_year) + fert_rates = get_fert( + E + S, min_age, max_age, country_id, data_year, data_year + ) + mort_rates, infmort_rate = get_mort( + E + S, min_age, max_age, country_id, data_year, data_year + ) mort_rates_S = mort_rates[-S:] - imm_rates_orig = get_imm_rates(E + S, min_age, max_age, country_id, data_year, data_year) + imm_rates_orig = get_imm_rates( + E + S, min_age, max_age, country_id, data_year, data_year + ) OMEGA_orig = np.zeros((E + S, E + S)) OMEGA_orig[0, :] = (1 - infmort_rate) * fert_rates + np.hstack( (imm_rates_orig[0], np.zeros(E + S - 1)) @@ -470,7 +513,9 @@ def get_pop_objs( # Generate time path of the nonstationary population distribution omega_path_lev = np.zeros((E + S, T + S)) - pop_data = get_un_data("47", country_id=country_id, start_year=data_year, end_year=data_year) + pop_data = get_un_data( + "47", country_id=country_id, start_year=data_year, end_year=data_year + ) # TODO: allow one to read in multiple years of UN forecast then # extrapolate from the end of that pop_data_sample = pop_data[ diff --git a/ogusa/tests/test_demographics.py b/ogusa/tests/test_demographics.py index 25c4b27a..cf2f6ad8 100644 --- a/ogusa/tests/test_demographics.py +++ b/ogusa/tests/test_demographics.py @@ -12,7 +12,9 @@ def test_get_pop_objs(): T = int(round(4.0 * S)) start_year = 2019 - pop_dict = demographics.get_pop_objs(E, S, T, 1, 100, start_year - 1, start_year, False) + pop_dict = demographics.get_pop_objs( + E, S, T, 1, 100, start_year - 1, start_year, False + ) assert np.allclose(pop_dict["omega_SS"], pop_dict["omega"][-1, :]) @@ -26,7 +28,9 @@ def test_pop_smooth(): T = int(round(4.0 * S)) start_year = 2019 - pop_dict = demographics.get_pop_objs(E, S, T, 1, 100, start_year - 1, start_year, False) + pop_dict = demographics.get_pop_objs( + E, S, T, 1, 100, start_year - 1, start_year, False + ) assert np.any( np.absolute(pop_dict["omega"][:-1, :] - pop_dict["omega"][1:, :]) @@ -46,7 +50,9 @@ def test_imm_smooth(): T = int(round(4.0 * S)) start_year = 2019 - pop_dict = demographics.get_pop_objs(E, S, T, 1, 100, start_year - 1, start_year, False) + pop_dict = demographics.get_pop_objs( + E, S, T, 1, 100, start_year - 1, start_year, False + ) assert np.any( np.absolute( diff --git a/ogusa/utils.py b/ogusa/utils.py index aa5317bb..8520e37f 100644 --- a/ogusa/utils.py +++ b/ogusa/utils.py @@ -248,4 +248,4 @@ def get_legacy_session(): ctx.options |= 0x4 # OP_LEGACY_SERVER_CONNECT #in Python 3.12 you will be able to switch from 0x4 to ssl.OP_LEGACY_SERVER_CONNECT. session = requests.session() session.mount("https://", CustomHttpAdapter(ctx)) - return session \ No newline at end of file + return session