def trend_following_pack_compartments(self, states, orig_gt, num_forecast_steps): """Packs predictions into compartments with associated ground truth.""" (death_d_f_all_locations, death_horizon_ahead_d_f_all_locations, confirmed_f_all_locations, confirmed_horizon_ahead_d_f_all_locations, hospitalized_f_all_locations, hospitalized_increase_f_all_locations) = states # pack all results in a list of compartment dataclasses. death_d_compartment = generic_seir_model_constructor.Compartment( name=constants.DEATH, predictions=death_d_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["death"]) confirmed_compartment = generic_seir_model_constructor.Compartment( name=constants.CONFIRMED, predictions=confirmed_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["confirmed"]) hospitalized_compartment = generic_seir_model_constructor.Compartment( name=constants.HOSPITALIZED, predictions=hospitalized_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["hospitalized"]) hospitalized_increase_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HOSPITALIZED_INCREASE, predictions=hospitalized_increase_f_all_locations, num_forecast_steps=num_forecast_steps)) def create_horizon_ahead_gt(gt): """Creates incremental (1-day) ground truth values.""" horizon_ahead_gt = {} for location in gt: horizon_ahead_gt[location] = ( gt[location][num_forecast_steps - 1:] - gt[location][:-num_forecast_steps + 1]) return horizon_ahead_gt death_horizon_ahead_d_compartment = generic_seir_model_constructor.Compartment( name=constants.HORIZON_AHEAD_DEATH, predictions=death_horizon_ahead_d_f_all_locations, num_forecast_steps=1, ground_truth=create_horizon_ahead_gt(orig_gt["death"])) confirmed_horizon_ahead_d_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HORIZON_AHEAD_CONFIRMED, predictions=confirmed_horizon_ahead_d_f_all_locations, num_forecast_steps=1, ground_truth=create_horizon_ahead_gt(orig_gt["confirmed"]))) compartments = [ death_d_compartment, death_horizon_ahead_d_compartment, confirmed_compartment, confirmed_horizon_ahead_d_compartment, hospitalized_compartment, hospitalized_increase_compartment ] return [ compart for compart in compartments if compart.predictions is not None ]
def pack_compartments(self, states, ground_truth_timeseries, num_forecast_steps): """Packs predictions into compartments with associated ground truth.""" (susceptible_f_all_locations, exposed_f_all_locations, infected_d_f_all_locations, infected_ud_f_all_locations, recovered_d_f_all_locations, recovered_ud_f_all_locations, death_d_f_all_locations, death_horizon_ahead_d_f_all_locations, confirmed_f_all_locations, confirmed_horizon_ahead_d_f_all_locations, hospitalized_f_all_locations, hospitalized_increase_f_all_locations, hospitalized_cumulative_f_all_locations, icu_f_all_locations, ventilator_f_all_locations, infected_ud_increase_f_all_locations, rates) = states (_, _, _, _, orig_gt) = ground_truth_timeseries # pack all results in a list of compartment dataclasses. susceptible_compartment = generic_seir_model_constructor.Compartment( name=constants.SUSCEPTIBLE, predictions=susceptible_f_all_locations, num_forecast_steps=num_forecast_steps) exposed_compartment = generic_seir_model_constructor.Compartment( name=constants.EXPOSED, predictions=exposed_f_all_locations, num_forecast_steps=num_forecast_steps) infected_d_compartment = generic_seir_model_constructor.Compartment( name=constants.INFECTED_DOC, predictions=infected_d_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["infected"]) infected_ud_compartment = generic_seir_model_constructor.Compartment( name=constants.INFECTED_UNDOC, predictions=infected_ud_f_all_locations, num_forecast_steps=num_forecast_steps) infected_ud_increase_compartment = generic_seir_model_constructor.Compartment( name=constants.INFECTED_UNDOC_INCREASE, predictions=infected_ud_increase_f_all_locations, num_forecast_steps=num_forecast_steps) recovered_d_compartment = generic_seir_model_constructor.Compartment( name=constants.RECOVERED_DOC, predictions=recovered_d_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["recovered"]) recovered_ud_compartment = generic_seir_model_constructor.Compartment( name=constants.RECOVERED_UNDOC, predictions=recovered_ud_f_all_locations, num_forecast_steps=num_forecast_steps) death_d_compartment = generic_seir_model_constructor.Compartment( name=constants.DEATH, predictions=death_d_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["death"]) confirmed_compartment = generic_seir_model_constructor.Compartment( name=constants.CONFIRMED, predictions=confirmed_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["confirmed"]) hospitalized_compartment = generic_seir_model_constructor.Compartment( name=constants.HOSPITALIZED, predictions=hospitalized_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["hospitalized"]) hospitalized_increase_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HOSPITALIZED_INCREASE, predictions=hospitalized_increase_f_all_locations, num_forecast_steps=num_forecast_steps)) hospitalized_cumulative_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HOSPITALIZED_CUMULATIVE, predictions=hospitalized_cumulative_f_all_locations, num_forecast_steps=num_forecast_steps)) icu_compartment = generic_seir_model_constructor.Compartment( name=constants.ICU, predictions=icu_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["icu"]) ventilator_compartment = generic_seir_model_constructor.Compartment( name=constants.VENTILATOR, predictions=ventilator_f_all_locations, num_forecast_steps=num_forecast_steps, ground_truth=orig_gt["ventilator"]) def create_horizon_ahead_gt(gt): """Creates incremental (1-day) ground truth values.""" horizon_ahead_gt = {} for location in gt: horizon_ahead_gt[location] = ( gt[location][num_forecast_steps - 1:] - gt[location][:-num_forecast_steps + 1]) return horizon_ahead_gt death_horizon_ahead_d_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HORIZON_AHEAD_DEATH, predictions=death_horizon_ahead_d_f_all_locations, num_forecast_steps=1, ground_truth=create_horizon_ahead_gt(orig_gt["death"]))) confirmed_horizon_ahead_d_compartment = ( generic_seir_model_constructor.Compartment( name=constants.HORIZON_AHEAD_CONFIRMED, predictions=confirmed_horizon_ahead_d_f_all_locations, num_forecast_steps=1, ground_truth=create_horizon_ahead_gt(orig_gt["confirmed"]))) rates_compartments = [] for name, predictions in rates.items(): rates_compartments.append( generic_seir_model_constructor.Compartment( name=name, predictions=predictions, num_forecast_steps=num_forecast_steps, use_quantiles=False)) compartments = [ susceptible_compartment, exposed_compartment, infected_d_compartment, infected_ud_compartment, recovered_d_compartment, recovered_ud_compartment, death_d_compartment, death_horizon_ahead_d_compartment, confirmed_compartment, confirmed_horizon_ahead_d_compartment, hospitalized_compartment, hospitalized_increase_compartment, hospitalized_cumulative_compartment, icu_compartment, ventilator_compartment, infected_ud_increase_compartment ] compartments += rates_compartments return compartments