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