def branch_or_leaf(dag: LocationDAG, location_id: int, sex: int, model_version_id: int,
parent_location: int, parent_sex: int,
n_sim: int, n_pool: int, upstream: List[str], tasks: List[_CascadeOperation]):
"""
Recursive function that either creates a branch (by calling itself) or a leaf fit depending
on whether or not it is at a terminal node. Determines if it's at a terminal node using
the dag.successors() method from networkx. Appends tasks onto the tasks parameter.
"""
if not dag.is_leaf(location_id=location_id):
branch = branch_fit(
model_version_id=model_version_id,
location_id=location_id, sex_id=sex,
prior_parent=parent_location, prior_sex=parent_sex,
child_locations=dag.children(location_id), child_sexes=[sex],
n_sim=n_sim, n_pool=n_pool,
upstream_commands=upstream,
ode_fit_strategy=True
)
tasks += branch
for location in dag.children(location_id):
branch_or_leaf(dag=dag, location_id=location, sex=sex, model_version_id=model_version_id,
parent_location=location_id, parent_sex=sex,
n_sim=n_sim, n_pool=n_pool, upstream=[branch[-1].command], tasks=tasks)
else:
leaf = leaf_fit(
model_version_id=model_version_id,
location_id=location_id,
sex_id=sex,
prior_parent=parent_location,
prior_sex=parent_sex,
n_sim=n_sim, n_pool=n_pool,
upstream_commands=upstream,
ode_fit_strategy=True
)
tasks += leaf
def construct_two_level_model(self, location_dag: LocationDAG, parent_location_id: int,
covariate_specs: CovariateSpecs,
weights: Optional[Dict[str, Var]] = None,
omega_df: Optional[pd.DataFrame] = None,
update_prior: Optional[Dict[str, Dict[str, np.ndarray]]] = None,
min_cv: Optional[Dict[str, Dict[str, float]]] = None,
update_mulcov_prior: Optional[Dict[Tuple[str, str, str], _Prior]] = None):
"""
Construct a Model object for a parent location and its children.
Parameters
----------
location_dag
Location DAG specifying the location hierarchy
parent_location_id
Parent location to build the model for
covariate_specs
covariate specifications, specifically will use covariate_specs.covariate_multipliers
weights
omega_df
data frame with omega values in it (other cause mortality)
update_prior
dictionary of dictionary for prior updates to rates
update_mulcov_prior
dictionary of mulcov prior updates
min_cv
dictionary (can be defaultdict) for minimum coefficient of variation
keyed by cascade level, then by rate
"""
children = location_dag.children(parent_location_id)
cascade_level = str(location_dag.depth(parent_location_id)) # min_cv lookup expects a string key
is_leaf = location_dag.is_leaf(parent_location_id)
if is_leaf:
cascade_level = MOST_DETAILED_CASCADE_LEVEL
model = Model(
nonzero_rates=self.settings.rate,
parent_location=parent_location_id,
child_location=children,
covariates=covariate_specs.covariate_list,
weights=weights
)
# First construct the rate grid, and update with prior
# information from a parent for value, dage, and dtime.
for smooth in self.settings.rate:
rate_grid = self.get_smoothing_grid(rate=smooth)
if update_prior is not None:
if smooth.rate in update_prior:
self.override_priors(rate_grid=rate_grid, update_dict=update_prior[smooth.rate])
if min_cv is not None:
self.apply_min_cv_to_prior_grid(
prior_grid=rate_grid.value, min_cv=min_cv[cascade_level][smooth.rate]
)
model.rate[smooth.rate] = rate_grid
# Second construct the covariate grids
for mulcov in covariate_specs.covariate_multipliers:
grid = smooth_grid_from_smoothing_form(
default_age_time=self.age_time_grid,
single_age_time=self.single_age_time_grid,
smooth=mulcov.grid_spec
)
if update_mulcov_prior is not None and (mulcov.group, *mulcov.key) in update_mulcov_prior:
ages = grid.ages
times = grid.times
for age, time in itertools.product(ages, times):
lb = grid.value[age, time].lower
ub = grid.value[age, time].upper
update_mulcov_prior[(mulcov.group, *mulcov.key)].lower = lb
update_mulcov_prior[(mulcov.group, *mulcov.key)].upper = ub
grid.value[age, time] = update_mulcov_prior[(mulcov.group, *mulcov.key)]
model[mulcov.group][mulcov.key] = grid
# Construct the random effect grids, based on the parent location
# specified.
if self.settings.random_effect:
random_effect_by_rate = defaultdict(list)
for smooth in self.settings.random_effect:
re_grid = smooth_grid_from_smoothing_form(
default_age_time=self.age_time_grid,
single_age_time=self.single_age_time_grid,
smooth=smooth
)
if not smooth.is_field_unset("location") and smooth.location in model.child_location:
location = smooth.location
else:
location = None
model.random_effect[(smooth.rate, location)] = re_grid
random_effect_by_rate[smooth.rate].append(location)
for rate_to_check, locations in random_effect_by_rate.items():
if locations != [None] and set(locations) != set(model.child_location):
raise RuntimeError(f"Random effect for {rate_to_check} does not have "
f"entries for all child locations, only {locations} "
f"instead of {model.child_location}.")
# Lastly, constrain omega for the parent and the random effects for the children.
if self.settings.model.constrain_omega:
LOG.info("Adding the omega constraint.")
if omega_df is None:
raise RuntimeError("Need an omega data frame in order to constrain omega.")
parent_omega = omega_df.loc[omega_df.location_id == parent_location_id].copy()
if parent_omega.empty:
raise RuntimeError(f"No omega values for location {parent_location_id}.")
omega = rectangular_data_to_var(gridded_data=parent_omega)
model.rate["omega"] = constraint_from_rectangular_data(
rate_var=omega,
default_age_time=self.age_time_grid
)
locations = set(omega_df.location_id.unique().tolist())
children_without_omega = set(children) - set(locations)
if children_without_omega:
LOG.warning(f"Children of {parent_location_id} missing omega {children_without_omega}"
f"so not including child omega constraints")
else:
for child in children:
child_omega = omega_df.loc[omega_df.location_id == child].copy()
assert not child_omega.empty
child_rate = rectangular_data_to_var(gridded_data=child_omega)
def child_effect(age, time):
return np.log(child_rate(age, time) / omega(age, time))
model.random_effect[("omega", child)] = constraint_from_rectangular_data(
rate_var=child_effect,
default_age_time=self.age_time_grid
)
return model