def write_squeezed(sqzd, location_id, year_id, sex_id, map_file): tmap = pd.read_csv(map_file) for me_id, df in sqzd.groupby(['me_id']): t_meid = tmap.query('modelable_entity_id_source == %s' % me_id) t_meid = t_meid['modelable_entity_id_target'].squeeze() try: t_meid = int(t_meid) except Exception: pass if not isinstance(t_meid, int): continue print('Writing squeezed %s to file' % t_meid) df['location_id'] = int(float(location_id)) df['year_id'] = int(float(year_id)) df['sex_id'] = int(float(sex_id)) df['measure_id'] = 5 df['age_group_id'] = df.age_group_id.astype(float).astype(int) df["modelable_entity_id"] = t_meid pusher = SuperPusher(spec={ 'file_pattern': ("{modelable_entity_id}/{location_id}/" "{measure_id}_{year_id}_{sex_id}.h5"), 'h5_tablename': 'draws' }, directory=output_dir) pusher.push(df, append=False)
def export_summary(self, component, year_type, df): if year_type == "single_year": pattern = "{measure_id}/single_year/{location_id}/{year_id}.csv" index_cols = { "cause": [ "measure_id", "year_id", "location_id", "sex_id", "age_group_id", "cause_id", "metric_id" ], "impairment": [ "measure_id", "year_id", "location_id", "sex_id", "age_group_id", "cause_id", "rei_id", "metric_id" ], "injuries": [ "measure_id", "year_id", "location_id", "sex_id", "age_group_id", "cause_id", "rei_id", "metric_id" ], "sequela": [ "measure_id", "year_id", "location_id", "sex_id", "age_group_id", "sequela_id", "metric_id" ] } if year_type == "multi_year": pattern = "{measure_id}/multi_year/{location_id}.csv" index_cols = { "cause": [ "measure_id", "year_start_id", "year_end_id", "location_id", "sex_id", "age_group_id", "cause_id", "metric_id" ], "impairment": [ "measure_id", "year_start_id", "year_end_id", "location_id", "sex_id", "age_group_id", "cause_id", "rei_id", "metric_id" ], "injuries": [ "measure_id", "year_start_id", "year_end_id", "location_id", "sex_id", "age_group_id", "cause_id", "rei_id", "metric_id" ], "sequela": [ "measure_id", "year_start_id", "year_end_id", "location_id", "sex_id", "age_group_id", "sequela_id", "metric_id" ] } df = sort_index_columns(df, index_cols[component]) df = df[index_cols[component] + ["val", "upper", "lower"]] # path to use in summaries directory = os.path.join(self.como_version.como_dir, "summaries", component) pusher = SuperPusher(directory=directory, spec={"file_pattern": pattern}) pusher.push(df, index=False)
def export_summary(self, component, year_type, df): if year_type == "single_year": pattern = "{measure_id}/single_year/{location_id}/{year_id}.csv" if year_type == "multi_year": pattern = "{measure_id}/multi_year/{location_id}.csv" # path to use in summaries directory = os.path.join(self.como_version.como_dir, "summaries", component) pusher = SuperPusher(directory=directory, spec={"file_pattern": pattern}) pusher.push(df)
class SevSplitter(object): def __init__(self, split_version_id, output_dir, decomp_step, location_id=None, year_id=None, age_group_id=None, sex_id=None, measure_id=None, location_set_id=35, gbd_round_id=gbd.GBD_ROUND_ID, n_draws=1000): # static ids self.split_version_id = split_version_id self.decomp_step = decomp_step self.location_set_id = location_set_id self.gbd_round_id = gbd_round_id # read func is derived from static values. we call it to initialize the # internal caching self._read_func = split_prop_read_func() cached_props = self._read_func(params={}, filters=self.ss_filters) # dimensions are derived unless explicit if not location_id: location_id = [ node.id for node in dbtrees.loctree(location_set_id=location_set_id, gbd_round_id=gbd_round_id).leaves() ] if not year_id: year_id = estimation_years_from_gbd_round_id(gbd_round_id) if not age_group_id: # this has the advantage of instantiating the lru cache in the main # process before multiprocessing age_group_id = get_age_group_set(12)["age_group_id"].tolist() if -1 in cached_props["age_start"].unique().tolist(): age_group_id.append(164) if not sex_id: sex_id = [gbd.sex.MALE, gbd.sex.FEMALE] if not measure_id: measure_id = [gbd.measures.PREVALENCE, gbd.measures.INCIDENCE] index_dict = { "location_id": location_id, "year_id": year_id, "age_group_id": age_group_id, "sex_id": sex_id, "measure_id": measure_id } data_dict = {"data": ['draw_{}'.format(i) for i in range(n_draws)]} self.dimensions = dimensionality.DataFrameDimensions( index_dict, data_dict) sp_formula = SevPropFormula() sp_formula.build_custom_draw_source(params={}, read_func=self._read_func) sp_formula.add_transforms() self._ss_draw_source = sp_formula.draw_source # epi draws source mvid, dstep = get_best_model_version_and_decomp_step( output_dir, int(self.parent_meid)) src = Epi.create_modelable_entity_draw_source( n_workers=1, modelable_entity_id=int(self.parent_meid), model_version_id=mvid, gbd_round_id=gbd_round_id, decomp_step=dstep) src.remove_transform(automagic_age_sex_agg) if n_draws < 1000: src.add_transform(group_and_downsample, n_draws) self._epi_draw_source = src self.pusher = SuperPusher(spec={ 'file_pattern': "{modelable_entity_id}/{location_id}.h5", 'h5_tablename': 'draws' }, directory=output_dir) @property def ss_filters(self): return {"split_version_id": self.split_version_id} @property def demo_filters(self): return { "location_id": self.dimensions.index_dim.get_level("location_id"), "age_group_id": self.dimensions.index_dim.get_level("age_group_id"), "year_id": self.dimensions.index_dim.get_level("year_id"), "sex_id": self.dimensions.index_dim.get_level("sex_id"), "measure_id": self.dimensions.index_dim.get_level("measure_id") } @property def parent_meid(self): df = self._read_func(params={}, filters=self.ss_filters) return df.parent_meid.unique()[0] @property def child_meid(self): df = self._read_func(params={}, filters=self.ss_filters) return df.child_meid.unique().tolist() def split(self): # get input draws draws = self._epi_draw_source.content(filters=self.demo_filters.copy()) # get split props filters = self.ss_filters filters.update(self.demo_filters) gprops = self._ss_draw_source.content(filters=filters) splits = merge_split(draws, gprops, group_cols=self.dimensions.index_names, value_cols=self.dimensions.data_list()) splits = splits.assign(modelable_entity_id=splits['child_meid']) splits = splits[self.dimensions.index_names + ["modelable_entity_id"] + self.dimensions.data_list()] splits = splits.fillna(0) self.pusher.push(splits, append=False) def _q_split(self, inq, outq): for location_id in iter(inq.get, sentinel): print(location_id) try: self.dimensions.index_dim.replace_level( "location_id", location_id) self.split() outq.put((False, location_id)) except Exception as e: outq.put((ExceptionWrapper(e), location_id)) def run_all_splits_mp(self, n_processes=23): inq = Queue() outq = Queue() # Create and feed sim procs split_procs = [] min_procs = min([ n_processes, self.dimensions.index_dim.cardinality("location_id") ]) for i in range(min_procs): p = Process(target=self._q_split, args=(inq, outq)) split_procs.append(p) p.start() # run the simulations for location_id in self.dimensions.index_dim.get_level("location_id"): inq.put(location_id) # make the workers die after for _ in split_procs: inq.put(sentinel) # get results results = [] for location_id in self.dimensions.index_dim.get_level("location_id"): proc_result = outq.get() results.append(proc_result) # close up the queue for p in split_procs: p.join() for exc, location_id in results: if exc: exc.re_raise()
class ExAdjust(object): def __init__(self, me_map, output_dir, location_id=[], year_id=[], age_group_id=[], sex_id=[], measure_id=[], location_set_id=35, gbd_round_id=5, n_draws=1000, copy_env_inc=False): # set static values self.me_map = me_map self.location_set_id = location_set_id self.gbd_round_id = gbd_round_id self.copy_env_inc = copy_env_inc # dimensions are derived unless explicit if not location_id: location_id = [ node.id for node in dbtrees.loctree(location_set_id=location_set_id, gbd_round_id=gbd_round_id).leaves() ] if not year_id: year_id = [1990, 1995, 2000, 2005, 2010, 2017] if not age_group_id: age_group_id = range(2, 21) + [30, 31, 32, 235] if not sex_id: sex_id = [1, 2] if not measure_id: measure_id = [5, 6] index_dict = { "location_id": location_id, "year_id": year_id, "age_group_id": age_group_id, "sex_id": sex_id, "measure_id": measure_id } data_dict = {"data": ['draw_{}'.format(i) for i in range(n_draws)]} self.dimensions = dimensionality.DataFrameDimensions( index_dict, data_dict) # draws that are imported or computed are stored here self.draws = {} # objects for reading data self._importers = {} for me_id in me_map["sub"].keys() + [me_map["env"]]: me_source = Epi.create_modelable_entity_draw_source( n_workers=1, modelable_entity_id=me_id, gbd_round_id=gbd_round_id) me_source.remove_transform(automagic_age_sex_agg) self._importers[me_id] = me_source # object for pushing results to disk self._pusher = SuperPusher(spec={ 'file_pattern': "{modelable_entity_id}/{location_id}.h5", 'h5_tablename': 'draws' }, directory=output_dir) @property def filters(self): return { "location_id": self.dimensions.index_dim.get_level("location_id"), "age_group_id": self.dimensions.index_dim.get_level("age_group_id"), "year_id": self.dimensions.index_dim.get_level("year_id"), "sex_id": self.dimensions.index_dim.get_level("sex_id"), "measure_id": self.dimensions.index_dim.get_level("measure_id") } def _import_draws(self): gbdizer = gbdize.GBDizeDataFrame(self.dimensions) # import draws for me_id in self._importers.keys(): draw_source = self._importers[me_id] draws = draw_source.content(filters=self.filters) draws = gbdizer.fill_empty_indices(draws, 0) self.draws[me_id] = draws.set_index(self.dimensions.index_names) def _calc_sigma_sub(self): """calculate the sum of the sub sequela""" # concatenate all required frames sub_dfs = [] for me_id in self.me_map["sub"].keys(): sub_dfs.append(self.draws[me_id]) sub_df = pd.concat(sub_dfs) # return the sum sub_df.reset_index(inplace=True) if self.copy_env_inc: draw_cols = self.dimensions.data_dim.get_level("data") sub_df.loc[sub_df['measure_id'] == 6, draw_cols] = 0 return sub_df.groupby(self.dimensions.index_names).sum() def _resid(self): """calculate the residual numbers""" # get the needed data sigma_sub_df = self.draws["sigma_sub"] env_df = self.draws[self.me_map["env"]] # if it is a squeeze type then we use the absolute value of the diff resid_df = (env_df - sigma_sub_df)[(sigma_sub_df <= env_df)].fillna(0) return resid_df def _excess(self, sub_me): """calculate the excess proportions""" # get the needed data sub_me_df = self.draws[sub_me] env_df = self.draws[self.me_map["env"]] sigma_sub_df = self.draws["sigma_sub"] # create a boolean dataframe for our 2 cases more = (sigma_sub_df > env_df) # now calculate the excess values excess_df = ((sigma_sub_df[more] - env_df[more]) * sub_me_df[more] / sigma_sub_df[more]).fillna(value=0) return excess_df def _squeeze(self, sub_me): """calculate the squeezed proportions""" # get the needed data sub_me_df = self.draws[sub_me] env_df = self.draws[self.me_map["env"]] sigma_sub_df = self.draws["sigma_sub"] # create a boolean dataframe for our 2 cases more = (sigma_sub_df > env_df) # get the squeezed values when squeeze_more = env_df[more] * sub_me_df[more] / sigma_sub_df[more] squeeze_less = sub_me_df[~more] squeeze_df = squeeze_more.fillna(squeeze_less) return squeeze_df def _export(self): """export all data""" # export residual me_id = self.me_map["resid"] resid_df = self.draws[me_id].reset_index() resid_df["modelable_entity_id"] = me_id self._pusher.push(resid_df, append=False) # export any subcause adjustments for sub_me in self.me_map["sub"].keys(): if "squeeze" in self.me_map["sub"][sub_me].keys(): me_id = self.me_map["sub"][sub_me]["squeeze"] squeeze_df = self.draws[me_id].reset_index() squeeze_df["modelable_entity_id"] = me_id self._pusher.push(squeeze_df, append=False) if "excess" in self.me_map["sub"][sub_me].keys(): me_id = self.me_map["sub"][sub_me]["excess"] excess_df = self.draws[me_id].reset_index() excess_df["modelable_entity_id"] = me_id self._pusher.push(excess_df, append=False) def adjust(self): """run exclusivity adjustment on all MEs""" self._import_draws() self.draws["sigma_sub"] = self._calc_sigma_sub() self.draws[self.me_map["resid"]] = self._resid() for sub_me in self.me_map["sub"].keys(): if "squeeze" in self.me_map["sub"][sub_me].keys(): self.draws[self.me_map["sub"][sub_me]["squeeze"]] = ( self._squeeze(sub_me)) if "excess" in self.me_map["sub"][sub_me].keys(): self.draws[self.me_map["sub"][sub_me]["excess"]] = ( self._excess(sub_me)) self._export() def _q_adjust(self, inq, outq): for location_id in iter(inq.get, sentinel): try: self.dimensions.index_dim.replace_level( "location_id", location_id) self.adjust() outq.put((False, location_id)) except Exception as e: outq.put((ExceptionWrapper(e), location_id)) def run_all_adjustments_mp(self, n_processes=23): inq = Queue() outq = Queue() # Create and feed sim procs adjust_procs = [] min_procs = min([ n_processes, self.dimensions.index_dim.cardinality("location_id") ]) for i in range(min_procs): p = Process(target=self._q_adjust, args=(inq, outq)) adjust_procs.append(p) p.start() # run the silulations for location_id in self.dimensions.index_dim.get_level("location_id"): inq.put(location_id) # make the workers die after for _ in adjust_procs: inq.put(sentinel) # get results results = [] for location_id in self.dimensions.index_dim.get_level("location_id"): proc_result = outq.get() results.append(proc_result) # close up the queue for p in adjust_procs: p.join() for exc, location_id in results: if exc: exc.re_raise()
class ExAdjust(object): def __init__( self, process_name, output_dir, decomp_step, location_id=None, year_id=None, age_group_id=None, sex_id=None, measure_id=None, location_set_id=35, gbd_round_id=gbd.GBD_ROUND_ID, n_draws=1000): # validate decomp_step validate_decomp_step("ExAdjust", decomp_step, gbd_round_id) this_file = os.path.realpath(__file__) this_dir = os.path.dirname(this_file) filepath = os.path.join(this_dir,"..","maps","final.json") with open(filepath, 'r') as f: emap = json.load(f) me_map = emap[process_name]["kwargs"]["me_map"] # set static values self.me_map = json_parser(json.dumps(me_map)) self.decomp_step = decomp_step self.location_set_id = location_set_id self.gbd_round_id = gbd_round_id try: self.copy_env_inc = emap[ process_name]["kwargs"].pop("copy_env_inc") except KeyError: self.copy_env_inc = False # dimensions are derived unless explicit if not location_id: location_id = [ node.id for node in dbtrees.loctree( location_set_id=location_set_id, gbd_round_id=gbd_round_id).leaves()] if not year_id: year_id = estimation_years_from_gbd_round_id(gbd_round_id) if not age_group_id: # this has the advantage of instantiating the lru cache in the main # process before multiprocessing age_group_id = get_age_group_set(12)["age_group_id"].tolist() if not sex_id: sex_id = [gbd.sex.MALE, gbd.sex.FEMALE] if not measure_id: measure_id = [gbd.measures.PREVALENCE, gbd.measures.INCIDENCE] index_dict = { "location_id": location_id, "year_id": year_id, "age_group_id": age_group_id, "sex_id": sex_id, "measure_id": measure_id } data_dict = {"data": ['draw_{}'.format(i) for i in range(n_draws)]} self.dimensions = dimensionality.DataFrameDimensions(index_dict, data_dict) # draws that are imported or computed are stored here self.draws = {} # objects for reading data self._importers = {} for me_id in list(self.me_map["sub"].keys()) + [self.me_map["env"]]: mvid, dstep = ( get_best_model_version_and_decomp_step(output_dir, me_id) ) me_source = Epi.create_modelable_entity_draw_source( n_workers=1, modelable_entity_id=me_id, model_version_id=mvid, gbd_round_id=gbd_round_id, decomp_step=dstep ) me_source.remove_transform(automagic_age_sex_agg) if n_draws < 1000: me_source.add_transform(group_and_downsample, n_draws) self._importers[me_id] = me_source # object for pushing results to disk self._pusher = SuperPusher( spec={'file_pattern': "{modelable_entity_id}/{location_id}.h5", 'h5_tablename': 'draws'}, directory=output_dir) @property def filters(self): return { "location_id": self.dimensions.index_dim.get_level("location_id"), "age_group_id": self.dimensions.index_dim.get_level("age_group_id" ), "year_id": self.dimensions.index_dim.get_level("year_id"), "sex_id": self.dimensions.index_dim.get_level("sex_id"), "measure_id": self.dimensions.index_dim.get_level("measure_id") } def _import_draws(self): gbdizer = gbdize.GBDizeDataFrame(self.dimensions) # import draws for me_id in self._importers.keys(): draw_source = self._importers[me_id] draws = draw_source.content(filters=self.filters) draws = gbdizer.fill_empty_indices(draws, 0) self.draws[me_id] = draws.set_index(self.dimensions.index_names) def _calc_sigma_sub(self): """calculate the sum of the sub sequela""" # concatenate all required frames sub_dfs = [] for me_id in self.me_map["sub"].keys(): sub_dfs.append(self.draws[me_id]) sub_df = pd.concat(sub_dfs) # return the sum sub_df.reset_index(inplace=True) if self.copy_env_inc: draw_cols = self.dimensions.data_dim.get_level("data") sub_df.loc[sub_df['measure_id'] == 6, draw_cols] = 0 return sub_df.groupby(self.dimensions.index_names).sum() def _resid(self): """calculate the residual numbers""" # get the needed data sigma_sub_df = self.draws["sigma_sub"] env_df = self.draws[self.me_map["env"]] # if it is a squeeze type then we use the absolute value of the diff resid_df = (env_df - sigma_sub_df)[(sigma_sub_df <= env_df)].fillna(0) return resid_df def _excess(self, sub_me): """calculate the excess proportions""" # get the needed data sub_me_df = self.draws[sub_me] env_df = self.draws[self.me_map["env"]] sigma_sub_df = self.draws["sigma_sub"] # create a boolean dataframe for our 2 cases more = (sigma_sub_df > env_df) # now calculate the excess values excess_df = ( (sigma_sub_df[more] - env_df[more]) * sub_me_df[more] / sigma_sub_df[more] ).fillna(value=0) return excess_df def _squeeze(self, sub_me): """calculate the squeezed proportions""" # get the needed data sub_me_df = self.draws[sub_me] env_df = self.draws[self.me_map["env"]] sigma_sub_df = self.draws["sigma_sub"] # create a boolean dataframe for our 2 cases more = (sigma_sub_df > env_df) # get the squeezed values when squeeze_more = env_df[more] * sub_me_df[more] / sigma_sub_df[more] squeeze_less = sub_me_df[~more] squeeze_df = squeeze_more.fillna(squeeze_less) return squeeze_df def _export(self): """export all data""" # export residual me_id = self.me_map["resid"] resid_df = self.draws[me_id].reset_index() resid_df["modelable_entity_id"] = me_id self._pusher.push(resid_df, append=False) # export any subcause adjustments for sub_me in self.me_map["sub"].keys(): if "squeeze" in list(self.me_map["sub"][sub_me].keys()): me_id = self.me_map["sub"][sub_me]["squeeze"] squeeze_df = self.draws[me_id].reset_index() squeeze_df["modelable_entity_id"] = me_id self._pusher.push(squeeze_df, append=False) if "excess" in list(self.me_map["sub"][sub_me].keys()): me_id = self.me_map["sub"][sub_me]["excess"] excess_df = self.draws[me_id].reset_index() excess_df["modelable_entity_id"] = me_id self._pusher.push(excess_df, append=False) def adjust(self): """run exclusivity adjustment on all MEs""" self._import_draws() self.draws["sigma_sub"] = self._calc_sigma_sub() self.draws[self.me_map["resid"]] = self._resid() for sub_me in self.me_map["sub"].keys(): if "squeeze" in list(self.me_map["sub"][sub_me].keys()): self.draws[self.me_map["sub"][sub_me]["squeeze"]] = ( self._squeeze(sub_me)) if "excess" in list(self.me_map["sub"][sub_me].keys()): self.draws[self.me_map["sub"][sub_me]["excess"]] = ( self._excess(sub_me)) self._export() def _q_adjust(self, inq, outq): for location_id in iter(inq.get, sentinel): try: self.dimensions.index_dim.replace_level("location_id", location_id) self.adjust() outq.put((False, location_id)) except Exception as e: outq.put((ExceptionWrapper(e), location_id)) def run_all_adjustments_mp(self, n_processes=23): inq = Queue() outq = Queue() # Create and feed sim process adjust_procs = [] min_procs = min( [n_processes, self.dimensions.index_dim.cardinality("location_id")] ) for i in range(min_procs): p = Process(target=self._q_adjust, args=(inq, outq)) adjust_procs.append(p) p.start() # run the simulations for location_id in self.dimensions.index_dim.get_level("location_id"): inq.put(location_id) # make the workers die after for _ in adjust_procs: inq.put(sentinel) # get results results = [] for location_id in self.dimensions.index_dim.get_level("location_id"): proc_result = outq.get() results.append(proc_result) # close up the queue for p in adjust_procs: p.join() for exc, location_id in results: if exc: exc.re_raise()
def run_squeeze(location_id, year_id, sex_id): ################################### # Prepare envelopes ################################### sequelae_map = pd.read_csv(SOURCE_TARGET_FILE) envelope_dict = create_env(location_id, year_id, sex_id) ################################### # Prepare unsqueezed prevalence ################################### # Load map of sequelae and their targets unsqueezed = get_unsqueezed(sequelae_map, location_id, year_id, sex_id) unsqueezed.loc[:, drawcols] = unsqueezed.loc[:, drawcols].clip(lower=0) ################################### # SQUEEZE ################################### # Parallelize the squeezing pool = Pool(20) ages = list(pd.unique(unsqueezed['age_group_id'])) partial_squeeze = partial(squeeze_age_group, unsqueezed=unsqueezed, env_dict=envelope_dict) squeezed = pool.map(partial_squeeze, ages, chunksize=1) pool.close() pool.join() squeezed = pd.concat(squeezed) squeezed = squeezed.groupby( ['location_id', 'year_id', 'age_group_id', 'sex_id', 'me_id']).sum() squeezed = squeezed.reset_index() ################################## # Write to files ################################## write_squeezed(squeezed, location_id, year_id, sex_id, MAP_FILE) ################################## # Allocate residuals ################################## allocate_residuals(unsqueezed, squeezed, location_id, year_id, sex_id, MAP_FILE) ########################################### # Determine the remainder of the envelopes ########################################### remains = calc_env_remainders(envelope_dict, squeezed) remain_map = { 'id_bord': 2000, 'id_mild': 1999, 'id_mod': 2001, 'id_sev': 2002, 'id_prof': 2003 } for key, meid in remain_map.iteritems(): print('Writing remainder %s to file' % meid) try: meid = int(meid) except Exception: pass df = remains[key] df['location_id'] = int(float(location_id)) df['year_id'] = int(float(year_id)) df['sex_id'] = int(float(sex_id)) df['measure_id'] = 5 df['age_group_id'] = df.age_group_id.astype(float).astype(int) df["modelable_entity_id"] = meid pusher = SuperPusher(spec={ 'file_pattern': ("{modelable_entity_id}/{location_id}/" "{measure_id}_{year_id}_{sex_id}.h5"), 'h5_tablename': 'draws' }, directory=output_dir) pusher.push(df[[ 'location_id', 'year_id', 'age_group_id', "sex_id", "modelable_entity_id", "measure_id" ] + drawcols], append=False)
def allocate_residuals(usqzd, sqzd, location_id, year_id, sex_id, map_file): tmap = pd.read_csv(map_file) resids = usqzd.merge( sqzd, on=['location_id', 'year_id', 'age_group_id', 'sex_id', 'me_id'], suffixes=('.usqzd', '.sqzd')) resids = resids[resids['resid_target_me.usqzd'].notnull()] dscols = ['draw_%s.sqzd' % d for d in range(1000)] ducols = ['draw_%s.usqzd' % d for d in range(1000)] toalloc = resids[ducols].values - resids[dscols].values toalloc = toalloc.clip(min=0) resids = resids.join( pd.DataFrame(data=toalloc, index=resids.index, columns=drawcols)) resids = resids[[ 'location_id', 'year_id', 'age_group_id', 'sex_id', 'resid_target_me.usqzd' ] + drawcols] resids.rename(columns={'resid_target_me.usqzd': 'resid_target_me'}, inplace=True) resids = resids.groupby(['resid_target_me', 'age_group_id']).sum() resids = resids.reset_index() resids = resids[['resid_target_me', 'age_group_id'] + drawcols] for me_id, resid_df in resids.groupby('resid_target_me'): t_meid = tmap.query('modelable_entity_id_source == %s' % me_id) t_meid = t_meid.modelable_entity_id_target.squeeze() try: t_meid = int(t_meid) except Exception: pass present = True try: draw_source = Epi.create_modelable_entity_draw_source( n_workers=1, modelable_entity_id=me_id, gbd_round_id=5) draw_source.remove_transform(automagic_age_sex_agg) t_df = draw_source.content( filters={ "location_id": location_id, "year_id": year_id, "sex_id": sex_id, "measure_id": 5 }) except NoBestVersionError: present = False if present: t_df = t_df.merge(resid_df, on='age_group_id', suffixes=('#base', '#resid')) newvals = (t_df.filter(like="#base").values + t_df.filter(like="#resid").values) t_df = t_df.join( pd.DataFrame(data=newvals, index=t_df.index, columns=drawcols)) print('Writing residual %s to file' % t_meid) t_df['location_id'] = int(float(location_id)) t_df['year_id'] = int(float(year_id)) t_df['sex_id'] = int(float(sex_id)) t_df['measure_id'] = 5 t_df['age_group_id'] = t_df.age_group_id.astype(float).astype(int) t_df["modelable_entity_id"] = t_meid t_df = t_df[[ 'location_id', 'year_id', 'age_group_id', "sex_id", "modelable_entity_id", "measure_id" ] + drawcols] pusher = SuperPusher(spec={ 'file_pattern': ("{modelable_entity_id}/{location_id}/" "{measure_id}_{year_id}_{sex_id}.h5"), 'h5_tablename': 'draws' }, directory=output_dir) pusher.push(t_df, append=False) else: print('ME ID %s missing' % me_id) return resids