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
class MeasurementInputs: def __init__(self, model_version_id: int, gbd_round_id: int, decomp_step_id: int, conn_def: str, country_covariate_id: List[int], csmr_cause_id: int, crosswalk_version_id: int, csmr_process_version_id: Optional[int] = None, location_set_version_id: Optional[int] = None, drill_location_start: Optional[int] = None, drill_location_end: Optional[List[int]] = None): """ The class that constructs all of the measurement inputs. Pulls ASDR, CSMR, crosswalk versions, and country covariates, and puts them into one data frame that then formats itself for the dismod database. Performs covariate value interpolation if age and year ranges don't match up with GBD age and year ranges. Parameters ---------- model_version_id the model version ID gbd_round_id the GBD round ID decomp_step_id the decomp step ID csmr_process_version_id process version ID for CSMR csmr_cause_id: (int) cause to pull CSMR from crosswalk_version_id crosswalk version to use country_covariate_id list of covariate IDs conn_def connection definition from .odbc file (e.g. 'epi') to connect to the IHME databases location_set_version_id can be None, if it's none, get the best location_set_version_id for estimation hierarchy of this GBD round drill_location_start which location ID to drill from as the parent drill_location_end which immediate children of the drill_location_start parent to include in the drill Attributes ---------- self.decomp_step : str the decomp step in string form self.demographics : cascade_at.inputs.demographics.Demographics a demographics object that specifies the age group, sex, location, and year IDs to grab self.integrand_map : Dict[int, int] dictionary mapping from GBD measure IDs to DisMod IDs self.asdr : cascade_at.inputs.asdr.ASDR all-cause mortality input object self.csmr : cascade_at.inputs.csmr.CSMR cause-specific mortality input object from cause csmr_cause_id self.data : cascade_at.inputs.data.CrosswalkVersion crosswalk version data from IHME database self.covariate_data : List[cascade_at.inputs.covariate_data.CovariateData] list of covariate data objects that contains the raw covariate data mapped to IDs self.location_dag : cascade_at.inputs.locations.LocationDAG DAG of locations to be used self.population: (cascade_at.inputs.population.Population) population object that is used for covariate weighting self.data_eta: (Dict[str, float]): dictionary of eta value to be applied to each measure self.density: (Dict[str, str]): dictionary of density to be applied to each measure self.nu: (Dict[str, float]): dictionary of nu value to be applied to each measure self.dismod_data: (pd.DataFrame) resulting dismod data formatted to be used in the dismod database Examples -------- >>> from cascade_at.settings.base_case import BASE_CASE >>> from cascade_at.settings.settings import load_settings >>> >>> settings = load_settings(BASE_CASE) >>> covariate_id = [i.country_covariate_id for i in settings.country_covariate] >>> >>> i = MeasurementInputs( >>> model_version_id=settings.model.model_version_id, >>> gbd_round_id=settings.gbd_round_id, >>> decomp_step_id=settings.model.decomp_step_id, >>> csmr_process_version_id=None, >>> csmr_cause_id = settings.model.add_csmr_cause, >>> crosswalk_version_id=settings.model.crosswalk_version_id, >>> country_covariate_id=covariate_id, >>> conn_def='epi', >>> location_set_version_id=settings.location_set_version_id >>> ) >>> i.get_raw_inputs() >>> i.configure_inputs_for_dismod(settings) """ LOG.info(f"Initializing input object for model version ID {model_version_id}.") LOG.info(f"GBD Round ID {gbd_round_id}.") LOG.info(f"Pulling from connection {conn_def}.") self.model_version_id = model_version_id self.gbd_round_id = gbd_round_id self.decomp_step_id = decomp_step_id self.csmr_process_version_id = csmr_process_version_id self.csmr_cause_id = csmr_cause_id self.crosswalk_version_id = crosswalk_version_id self.country_covariate_id = country_covariate_id self.conn_def = conn_def self.drill_location_start = drill_location_start self.drill_location_end = drill_location_end self.decomp_step = ds.decomp_step_from_decomp_step_id(self.decomp_step_id) if location_set_version_id is None: self.location_set_version_id = get_location_set_version_id(gbd_round_id=self.gbd_round_id) else: self.location_set_version_id = location_set_version_id self.demographics = Demographics( gbd_round_id=self.gbd_round_id, location_set_version_id=self.location_set_version_id) self.location_dag = LocationDAG( location_set_version_id=self.location_set_version_id, gbd_round_id=self.gbd_round_id ) # Need to subset the locations to only those needed for # the drill. drill_locations_all is the set of locations # to pull data for, including all descendents. drill_locations # is the set of locations just parent-children in the drill. drill_locations_all, drill_locations = locations_by_drill( drill_location_start=self.drill_location_start, drill_location_end=self.drill_location_end, dag=self.location_dag ) if drill_locations_all: self.demographics.location_id = drill_locations_all self.demographics.drill_locations = drill_locations self.exclude_outliers = True self.asdr = None self.csmr = None self.population = None self.data = None self.covariates = None self.age_groups = None self.data_eta = None self.density = None self.nu = None self.measures_to_exclude = None self.dismod_data = None self.covariate_data = None self.country_covariate_data = None self.covariate_specs = None self.omega = None def get_raw_inputs(self): """ Get the raw inputs that need to be used in the modeling. """ LOG.info("Getting all raw inputs.") self.asdr = ASDR( demographics=self.demographics, decomp_step=self.decomp_step, gbd_round_id=self.gbd_round_id ).get_raw() self.csmr = CSMR( cause_id=self.csmr_cause_id, demographics=self.demographics, decomp_step=self.decomp_step, gbd_round_id=self.gbd_round_id, process_version_id=self.csmr_process_version_id ).get_raw() self.data = CrosswalkVersion( crosswalk_version_id=self.crosswalk_version_id, exclude_outliers=self.exclude_outliers, demographics=self.demographics, conn_def=self.conn_def, gbd_round_id=self.gbd_round_id ).get_raw() self.covariate_data = [CovariateData( covariate_id=c, demographics=self.demographics, decomp_step=self.decomp_step, gbd_round_id=self.gbd_round_id ).get_raw() for c in self.country_covariate_id] self.population = Population( demographics=self.demographics, decomp_step=self.decomp_step, gbd_round_id=self.gbd_round_id ).get_population() def configure_inputs_for_dismod(self, settings: SettingsConfig, midpoint: bool = False, mortality_year_reduction: int = 5): """ Modifies the inputs for DisMod based on model-specific settings. Arguments --------- settings Settings for the model mortality_year_reduction number of years to decimate csmr and asdr """ self.data_eta = data_eta_from_settings(settings) self.density = density_from_settings(settings) self.nu = nu_from_settings(settings) self.measures_to_exclude = measures_to_exclude_from_settings(settings) # If we are constraining omega, then we want to hold out the data # from the DisMod fit for ASDR (but never CSMR -- always want to fit # CSMR). data = self.data.configure_for_dismod( measures_to_exclude=self.measures_to_exclude, relabel_incidence=settings.model.relabel_incidence, midpoint=midpoint, ) asdr = self.asdr.configure_for_dismod( hold_out=settings.model.constrain_omega) csmr = self.csmr.configure_for_dismod(hold_out=0) if settings.model.constrain_omega: self.omega = calculate_omega(asdr=asdr, csmr=csmr) else: self.omega = None if not csmr.empty: csmr = decimate_years( data=csmr, num_years=mortality_year_reduction) if not asdr.empty: asdr = decimate_years( data=asdr, num_years=mortality_year_reduction) self.dismod_data = pd.concat([data, asdr, csmr], axis=0, sort=True) self.dismod_data.reset_index(drop=True, inplace=True) self.dismod_data["density"] = self.dismod_data.measure.apply( self.density.__getitem__) self.dismod_data["eta"] = self.dismod_data.measure.apply( self.data_eta.__getitem__) self.dismod_data["nu"] = self.dismod_data.measure.apply( self.nu.__getitem__) # This makes the specs not just for the country covariate but adds on # the sex and one covariates. self.covariate_specs = CovariateSpecs( country_covariates=settings.country_covariate, study_covariates=settings.study_covariate ) self.country_covariate_data = {c.covariate_id: c.configure_for_dismod( pop_df=self.population.configure_for_dismod(), loc_df=self.location_dag.df ) for c in self.covariate_data} self.dismod_data = self.add_covariates_to_data(df=self.dismod_data) self.dismod_data.loc[ self.dismod_data.hold_out.isnull(), 'hold_out'] = 0. self.dismod_data.drop(['age_group_id'], inplace=True, axis=1) return self def prune_mortality_data(self, parent_location_id: int) -> pd.DataFrame: """ Remove mortality data for descendents that are not children of parent_location_id from the configured dismod data before it gets filled into the dismod database. """ df = self.dismod_data.copy() direct_children = self.location_dag.parent_children(parent_location_id) direct_children = df.location_id.isin(direct_children) mortality_measures = df.measure.isin([ IntegrandEnum.mtall.name, IntegrandEnum.mtspecific.name ]) remove_rows = ~direct_children & mortality_measures df = df.loc[~remove_rows].copy() return df def add_covariates_to_data(self, df: pd.DataFrame) -> pd.DataFrame: """ Add on covariates to a data frame that has age_group_id, year_id or time-age upper / lower, and location_id and sex_id. Adds both country-level and study-level covariates. """ cov_dict_for_interpolation = { c.name: self.country_covariate_data[c.covariate_id] for c in self.covariate_specs.covariate_specs if c.study_country == 'country' } df = self.interpolate_country_covariate_values( df=df, cov_dict=cov_dict_for_interpolation) df = self.transform_country_covariates(df=df) df['s_sex'] = df.sex_id.map( SEX_ID_TO_NAME).map(StudyCovConstants.SEX_COV_VALUE_MAP) df['s_one'] = StudyCovConstants.ONE_COV_VALUE return df def to_gbd_avgint(self, parent_location_id: int, sex_id: int) -> pd.DataFrame: """ Converts the demographics of the model to the avgint table. """ LOG.info(f"Getting grid for the avgint table " f"for parent location ID {parent_location_id} " f"and sex_id {sex_id}.") if self.drill_location_start is not None: locations = self.demographics.drill_locations else: locations = self.location_dag.parent_children(parent_location_id) grid = expand_grid({ 'sex_id': [sex_id], 'location_id': locations, 'year_id': self.demographics.year_id, 'age_group_id': self.demographics.age_group_id }) grid['time_lower'] = grid['year_id'].astype(int) grid['time_upper'] = grid['year_id'] + 1. grid = BaseInput( gbd_round_id=self.gbd_round_id).convert_to_age_lower_upper(df=grid) LOG.info("Adding covariates to avgint grid.") grid = self.add_covariates_to_data(df=grid) return grid def interpolate_country_covariate_values(self, df: pd.DataFrame, cov_dict: Dict[Union[float, str], pd.DataFrame]): """ Interpolates the covariate values onto the data so that the non-standard ages and years match up to meaningful covariate values. """ LOG.info(f"Interpolating and merging the country covariates.") interp_df = get_interpolated_covariate_values( data_df=df, covariate_dict=cov_dict, population_df=self.population.configure_for_dismod() ) return interp_df def transform_country_covariates(self, df): """ Transforms the covariate data with the transformation ID. :param df: (pd.DataFrame) :return: self """ for c in self.covariate_specs.covariate_specs: if c.study_country == 'country': LOG.info(f"Transforming the data for country covariate " f"{c.covariate_id}.") df[c.name] = df[c.name].apply( lambda x: COVARIATE_TRANSFORMS[c.transformation_id](x) ) return df def calculate_country_covariate_reference_values( self, parent_location_id: int, sex_id: int) -> CovariateSpecs: """ Gets the country covariate reference value for a covariate ID and a parent location ID. Also gets the maximum difference between the reference value and covariate values observed. Run this when you're going to make a DisMod AT database for a specific parent location and sex ID. :param: (int) :param parent_location_id: (int) :param sex_id: (int) :return: List[CovariateSpec] list of the covariate specs with the correct reference values and max diff. """ covariate_specs = copy(self.covariate_specs) age_min = self.dismod_data.age_lower.min() age_max = self.dismod_data.age_upper.max() time_min = self.dismod_data.time_lower.min() time_max = self.dismod_data.time_upper.max() children = self.location_dag.children(parent_location_id) for c in covariate_specs.covariate_specs: if c.study_country == 'study': if c.name == 's_sex': c.reference = StudyCovConstants.SEX_COV_VALUE_MAP[ SEX_ID_TO_NAME[sex_id]] c.max_difference = StudyCovConstants.MAX_DIFFERENCE_SEX_COV elif c.name == 's_one': c.reference = StudyCovConstants.ONE_COV_VALUE c.max_difference = StudyCovConstants.MAX_DIFFERENCE_ONE_COV else: raise ValueError(f"The only two study covariates allowed are sex and one, you tried {c.name}.") elif c.study_country == 'country': LOG.info(f"Calculating the reference and max difference for country covariate {c.covariate_id}.") cov_df = self.country_covariate_data[c.covariate_id] parent_df = ( cov_df.loc[cov_df.location_id == parent_location_id].copy() ) child_df = cov_df.loc[cov_df.location_id.isin(children)].copy() all_loc_df = pd.concat([child_df, parent_df], axis=0) # if there is no data for the parent location at all (which # there should be provided by Central Comp) # then we are going to set the reference value to 0. if cov_df.empty: reference_value = 0 max_difference = np.nan else: pop_df = self.population.configure_for_dismod() pop_df = ( pop_df.loc[pop_df.location_id == parent_location_id].copy() ) df_to_interp = pd.DataFrame({ 'location_id': parent_location_id, 'sex_id': [sex_id], 'age_lower': [age_min], 'age_upper': [age_max], 'time_lower': [time_min], 'time_upper': [time_max] }) reference_value = get_interpolated_covariate_values( data_df=df_to_interp, covariate_dict={c.name: parent_df}, population_df=pop_df )[c.name].iloc[0] max_difference = np.max( np.abs(all_loc_df.mean_value - reference_value) ) + CascadeConstants.PRECISION_FOR_REFERENCE_VALUES c.reference = reference_value c.max_difference = max_difference covariate_specs.create_covariate_list() return covariate_specs def reset_index(self, drop, inplace): pass
def make_cascade_dag(model_version_id: int, dag: LocationDAG, location_start: int, sex_start: int, split_sex: bool, n_sim: int = 100, n_pool: int = 100, skip_configure: bool = False) -> List[_CascadeOperation]: """ Make a traditional cascade dag for a model version. Relies on a location DAG and a starting point in the DAG for locations and sexes. Parameters ---------- model_version_id Model version ID dag A location DAG that specifies the location hierarchy location_start Where to start in the location hierarchy sex_start Which sex to start with, can be most detailed or both. split_sex Whether or not to split sex into most detailed. If not, then will just stay at 'both' sex. n_sim Number of simulations to do in sample simulate n_pool Number of multiprocessing pools to create during sample simulate skip_configure Don't configure inputs. Only do this if it's already been done. Returns ------- List of _CascadeOperation. """ tasks = [] sexes = [sex_start] if SEX_ID_TO_NAME[sex_start] == 'Both': if split_sex: sexes = [ SEX_NAME_TO_ID['Female'], SEX_NAME_TO_ID['Male'] ] top_level = root_fit( model_version_id=model_version_id, location_id=location_start, sex_id=sex_start, child_locations=dag.children(location_start), child_sexes=sexes, mulcov_stats=True, skip_configure=skip_configure, n_sim=n_sim, n_pool=n_pool, ode_fit_strategy=True, ) tasks += top_level for sex in sexes: for location1 in dag.children(location_start): branch_or_leaf( dag=dag, location_id=location1, sex=sex, model_version_id=model_version_id, parent_location=location_start, parent_sex=sex_start, n_sim=n_sim, n_pool=n_pool, upstream=[top_level[-1].command], tasks=tasks ) tasks.append(Upload( model_version_id=model_version_id, fit=True, prior=True, upstream_commands=[tasks[-1].command], executor_parameters={ 'm_mem_free': '50G' } )) return tasks
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