def _survivorship(lx_year, ruler, ruler_extended, **kwargs):
"""Iterate over survivorship by week.
The survivorship is used to age the population week-by-week, starting
from the last past year.
Because pop is reported at the mid-year point, we first age the last
past year by 1/2 year (using last past year's lx). Then, for all forecast
years, we age all weeks of each year by that year's lx. For example,
if 2017 is the last past year, we use lx from 2017 to age the 2017 pop for
half a year (26 weeks) to reach 01-01-2018. Then we use lx from 2018 to
age the pop for an entire year (52 weeks) to reach 01-01-2019, and so on.
Args:
lx_year (np.array): lx in shape (sex_id, age_group_id).
ruler (Ruler): ruler object that maps fbd age groups to weeks.
ruler_extended (Ruler): identical to ruler, but with 5-year age groups
within the fhs terminal age group.
Returns:
(np.array): for every week, returns np.array in the shape of
(sex_id, 1 + len(ruler.x_ww)), where 1 corresponds to nL0 and
len(ruler.x_ww) corresponds to the survivorship elements of the
Leslie matrix.
"""
nLx = lx_to_nLx_weeks(lx_year, ruler, ruler_extended)
traceplot("nLx", nLx[1])
survivorship = nLx_to_survivorship(nLx)
traceplot("surv", survivorship[1])
return survivorship
def _migration(migration_year, ruler, **kwargs):
"""Iterate over migration by week, analogous to method in _survivorship().
Args:
migration_year (np.array): in shape of (year_id, sex_id, age_group_id).
ruler (Ruler): maps migration data age groups to weeks.
Returns:
(np.array): weekly iterated migration.
"""
migration_weekly = _migration_to_weeks(migration_year, ruler)
traceplot("migrationw", migration_weekly[0])
return migration_weekly
def population_to_weeks_by_projection(population, ruler):
"""
Converts population on age group Ids to weekly populations
using a simple projection matrix, so you get a stairwise step
function of population.
"""
assert population.shape[0] == 2
averaging = uneven_average(ruler.nx_gw)
LOGGER.debug("averaging shape {}".format(averaging.shape))
pop_weeks = np.einsum("wa,sa->sw", averaging, population)
traceplot("pop_weeks", pop_weeks[0])
assert pop_weeks.shape[0] == 2
return pop_weeks
def _asfr(asfr_year, ruler, **kwargs):
"""Iterate over ASFR by week, analogous to method in _survivorship().
Args:
asfr_year (np.array): array in shape of (year_id, age_group_id).
ruler (Ruler): maps asfr age groups to weeks.
Returns:
(np.array): weekly iterated asfr.
(np.array): fertile age group ids.
"""
projected = asfr_to_weeks(asfr_year)
traceplot("asfr_projected", projected)
return projected, ruler.fertile
def lx_to_nLx_weeks(lx, ruler, ruler_extended):
"""
Integrates a fitted lx curve to compute nLx, which would be used to compute
survivorship downstream.
Args:
lx (np.array): lx with extrapolated values beyond fhs terminal age
group start.
ruler (Ruler): ruler for the fhs age groups. This ruler determines
#the age group starts that need to be returned2
ruler_extended (Ruler): ruler that matches lx. This ruler, along with
lx, provide the points for interpolation.
Returns:
(np.array): nLx, in shape of (sex, age), where age corresponds to
fhs age group ids.
"""
assert has_dims(lx, "sex age")
LOGGER.debug("lx dims {}".format(lx.shape))
# We need 2 extra data points beyond the terminal age group start to
# compute its survivorship: terminal_age_start + 1 week, and 110 years.
x_all_weeks = np.concatenate(
[ruler.x_ww, [ruler.x_ww[-1] + 1,
ruler_extended.x_ww[-1]]]) # all in weeks
integrated_lx = list()
for sex in [0, 1]:
# first interpolate lx over the "extended" fhs age groups (in weeks)
fit = PchipInterpolator(ruler_extended.x_gw, lx[sex], extrapolate=True)
# now integrate lx and evaluate over all weeks
integrate_once = 1
integrated = fit.antiderivative(integrate_once)(x_all_weeks)
LOGGER.debug("integrated lx end {}".format(integrated[-5:]))
traceplot("integrated_lx", x_all_weeks, integrated)
integrated_lx.append(integrated)
lx_int_sexed = np.vstack(integrated_lx)
# Will be one less than number of weeks
return np.diff(lx_int_sexed, 1, axis=1)
def project_draw(asfr, lx, migration, pop_start_year, srb, ruler,
ruler_extended, migration_ruler, forecast_years):
"""
Take datasets and produce future populations.
The datasets have to start with the year for the observed population.
This data is all Numpy arrays, all for the same
location, scenario, and draw. That means it will be organized by
(year_id, sex_id, age_group_id) with exceptions as noted below.
Args:
asfr (np.ndarray): Age-specific fertility rate, organized as
(year_id, age_group_id).
lx (np.ndarray): :math:`l_x`, organized as
(year_id, sex_id, age_group_id).
migration (np.ndarray): Migration, organized as
(year_id, sex_id, age_group_id).
pop_start_year (np.ndarray): Starting population, organized as
(sex_id, age_group_id).
srb (np.ndarray): Sex ratio at birth with just one value
for the last past year
ruler (Ruler): Time scales for conversion to weeks.
ruler_extended (Ruler): identical to ruler, but with 5-year age groups
within the fhs terminal age group.
migration_ruler (Ruler): Time scales for conversion from years to
weeks.
forecast_years (list): The list of all years in the forecast
Returns:
np.ndarray: Population at starting and all future years.
"""
# All ndarrays at this point.
assert has_dims(asfr, "year age")
assert has_dims(lx, "year, sex, age")
assert has_dims(migration, "year sex age")
assert has_dims(pop_start_year, "sex, age")
# Project from middle of first year to middle of last year.
year_cnt = len(forecast_years)
assert year_cnt <= lx.shape[0] - 1, "Not enough years in lx"
assert year_cnt <= asfr.shape[0] - 1, "Not enough years in asfr"
pop_save = np.zeros((year_cnt, ) + pop_start_year.shape, dtype=np.double)
pop_buffer = population_to_weeks(pop_start_year, ruler)
survivorship_iter = _pop_aging_iterator(_survivorship,
lx,
ruler,
ruler_extended=ruler_extended)
asfr_iter = _pop_aging_iterator(_asfr, asfr, ruler)
migration_iter = _pop_aging_iterator(_migration, migration,
migration_ruler)
save_year = 0
for step_idx in range(ruler.ratio * year_cnt):
asfr, age_limits = next(asfr_iter)
survivorship = next(survivorship_iter)
migration = next(migration_iter)
traceplot("pop_buffer_in", pop_buffer[1])
leslie = LeslieMatrix(asfr, age_limits[0], survivorship, float(srb))
pop_t_plus_one = leslie.increment(pop_buffer)
pop_buffer = pop_t_plus_one + migration
if np.any(pop_buffer < 0):
leslie_error_cnt = np.sum(pop_t_plus_one < 0)
survivorship_neg = np.sum(survivorship < 0)
asfr_neg = np.sum(asfr < 0)
if leslie_error_cnt > 0:
LOGGER.error("Leslie matrix negative values {} {} {}".format(
leslie_error_cnt, survivorship_neg, asfr_neg))
pop_buffer[pop_buffer < 0] = 0
LOGGER.debug("pop_mig young F {}".format(pop_buffer[1, :10]))
traceplot("pop_buffer_out", pop_buffer[1])
LOGGER.debug("pop {:3g}".format(pop_buffer.sum()))
year_idx = (step_idx + 1) // ruler.ratio
if save_year != year_idx:
pop_save[year_idx - 1] = \
population_to_groups(pop_buffer, ruler.nx_gw)
LOGGER.debug("pop save step {} {:3g}".format(
step_idx, pop_save[year_idx - 1].sum()))
save_year = year_idx
if POP_MINIMUM is not None:
pop_save[pop_save < POP_MINIMUM] = POP_MINIMUM
return pop_save