def __init__(self,
fips,
N_samples,
t_list,
I_initial=1,
suppression_policy=None):
self.fips = fips
self.N_samples = N_samples
self.I_initial = I_initial
self.suppression_policy = suppression_policy
self.t_list = t_list
self._latest = combined_datasets.get_us_latest_for_fips(self.fips)
def __init__(self, fips, N_samples, t_list,
I_initial=1, suppression_policy=None):
self.fips = fips
self.agg_level = AggregationLevel.COUNTY if len(self.fips) == 5 else AggregationLevel.STATE
self.N_samples = N_samples
self.I_initial = I_initial
self.suppression_policy = suppression_policy
self.t_list = t_list
if self.agg_level is AggregationLevel.COUNTY:
self.county_metadata = load_data.load_county_metadata().set_index('fips').loc[fips].to_dict()
self.state_abbr = us.states.lookup(self.county_metadata['state']).abbr
self._latest = combined_datasets.get_us_latest_for_fips(self.fips)
else:
self.state_abbr = us.states.lookup(fips).abbr
self._latest = combined_datasets.get_us_latest_for_state(self.state_abbr)
def __init__(
self,
fips,
t0_case_count=1,
start_days_before_t0=2,
start_days_after_t0=1000,
min_days_required=5,
):
self.t0_case_count = t0_case_count
self.start_days_before_t0 = start_days_before_t0
self.start_days_after_t0 = start_days_after_t0
self.min_days_required = min_days_required
timeseries = combined_datasets.get_timeseries_for_fips(
fips, columns=[CommonFields.CASES], min_range_with_some_value=True
)
latest_values = combined_datasets.get_us_latest_for_fips(fips)
self.cases = timeseries.data
n_days = len(self.cases[CommonFields.CASES].dropna())
if n_days < min_days_required:
raise ValueError(f"Only {n_days} observations for county. Cannot fit.")
self.fips = fips
self.county = latest_values[CommonFields.COUNTY]
self.state = latest_values[CommonFields.STATE]
self.t = (self.cases.date - self.cases.date.min()).dt.days.values
self.data_start_date = self.cases.date.min()
self.y = self.cases[CommonFields.CASES].values
self.fit_predictions = None
self.t0 = None
self.t0_date = None
self.reduced_chi2 = None
self.model_params = None
def map_fips(self, fips: str) -> None:
"""
For a given fips code, for either a county or state, generate the CAN UI output format.
Parameters
----------
fips: str
FIPS code to map.
"""
log.info("Mapping output to WebUI.", state=self.state, fips=fips)
shim_log = structlog.getLogger(fips=fips)
pyseir_outputs = load_data.load_ensemble_results(fips)
try:
fit_results = load_inference_result(fips)
t0_simulation = datetime.fromisoformat(fit_results["t0_date"])
except (KeyError, ValueError):
log.error("Fit result not found for fips. Skipping...", fips=fips)
return
population = self._get_population(fips)
# We will shim all suppression policies by the same amount (since historical tracking error
# for all policies is the same).
baseline_policy = "suppression_policy__inferred" # This could be any valid policy
# We need the index in the model's temporal frame.
idx_offset = int(fit_results["t_today"] - fit_results["t0"])
# Get the latest observed values to use in calculating shims
observed_latest_dict = combined_datasets.get_us_latest_for_fips(fips)
observed_death_latest = observed_latest_dict[CommonFields.DEATHS]
observed_total_hosps_latest = observed_latest_dict[
CommonFields.CURRENT_HOSPITALIZED]
observed_icu_latest = observed_latest_dict[CommonFields.CURRENT_ICU]
# For Deaths
model_death_latest = pyseir_outputs[baseline_policy]["total_deaths"][
"ci_50"][idx_offset]
model_acute_latest = pyseir_outputs[baseline_policy]["HGen"]["ci_50"][
idx_offset]
model_icu_latest = pyseir_outputs[baseline_policy]["HICU"]["ci_50"][
idx_offset]
model_total_hosps_latest = model_acute_latest + model_icu_latest
death_shim = shim.calculate_strict_shim(
model=model_death_latest,
observed=observed_death_latest,
log=shim_log.bind(type=CommonFields.DEATHS),
)
total_hosp_shim = shim.calculate_strict_shim(
model=model_total_hosps_latest,
observed=observed_total_hosps_latest,
log=shim_log.bind(type=CommonFields.CURRENT_HOSPITALIZED),
)
# For ICU This one is a little more interesting since we often don't have ICU. In this case
# we use information from the same aggregation level (intralevel) to keep the ratios
# between general hospitalization and icu hospitalization
icu_shim = shim.calculate_intralevel_icu_shim(
model_acute=model_acute_latest,
model_icu=model_icu_latest,
observed_icu=observed_icu_latest,
observed_total_hosps=observed_total_hosps_latest,
log=shim_log.bind(type=CommonFields.CURRENT_ICU),
)
# Iterate through each suppression policy.
# Model output is interpolated to the dates desired for the API.
suppression_policies = [
key for key in pyseir_outputs.keys()
if key.startswith("suppression_policy")
]
for suppression_policy in suppression_policies:
output_for_policy = pyseir_outputs[suppression_policy]
output_model = pd.DataFrame()
t_list = output_for_policy["t_list"]
t_list_downsampled = range(0, int(max(t_list)),
self.output_interval_days)
output_model[schema.DAY_NUM] = t_list_downsampled
output_model[schema.DATE] = [
(t0_simulation + timedelta(days=t)).date().strftime("%Y-%m-%d")
for t in t_list_downsampled
]
output_model[schema.TOTAL] = population
output_model[schema.TOTAL_SUSCEPTIBLE] = np.interp(
t_list_downsampled, t_list, output_for_policy["S"]["ci_50"])
output_model[schema.EXPOSED] = np.interp(
t_list_downsampled, t_list, output_for_policy["E"]["ci_50"])
output_model[schema.INFECTED] = np.interp(
t_list_downsampled,
t_list,
np.add(output_for_policy["I"]["ci_50"],
output_for_policy["A"]["ci_50"]),
) # Infected + Asympt.
output_model[schema.INFECTED_A] = output_model[schema.INFECTED]
interpolated_model_acute_values = np.interp(
t_list_downsampled, t_list, output_for_policy["HGen"]["ci_50"])
output_model[schema.INFECTED_B] = interpolated_model_acute_values
raw_model_icu_values = output_for_policy["HICU"]["ci_50"]
interpolated_model_icu_values = np.interp(t_list_downsampled,
t_list,
raw_model_icu_values)
output_model[schema.INFECTED_C] = (
icu_shim + interpolated_model_icu_values).clip(min=0)
# General + ICU beds. don't include vent here because they are also counted in ICU
output_model[schema.ALL_HOSPITALIZED] = (
interpolated_model_acute_values +
interpolated_model_icu_values + total_hosp_shim).clip(min=0)
output_model[schema.ALL_INFECTED] = output_model[schema.INFECTED]
# Shim Deaths to Match Observed
raw_model_deaths_values = output_for_policy["total_deaths"][
"ci_50"]
interp_model_deaths_values = np.interp(t_list_downsampled, t_list,
raw_model_deaths_values)
output_model[schema.DEAD] = (interp_model_deaths_values +
death_shim).clip(min=0)
# Continue mapping
final_beds = np.mean(output_for_policy["HGen"]["capacity"])
output_model[schema.BEDS] = final_beds
output_model[schema.CUMULATIVE_INFECTED] = np.interp(
t_list_downsampled,
t_list,
np.cumsum(output_for_policy["total_new_infections"]["ci_50"]),
)
if fit_results:
output_model[schema.Rt] = np.interp(
t_list_downsampled,
t_list,
fit_results["eps2"] * fit_results["R0"] *
np.ones(len(t_list)),
)
output_model[schema.Rt_ci90] = np.interp(
t_list_downsampled,
t_list,
2 * fit_results["eps2_error"] * fit_results["R0"] *
np.ones(len(t_list)),
)
else:
output_model[schema.Rt] = 0
output_model[schema.Rt_ci90] = 0
output_model[schema.CURRENT_VENTILATED] = (
icu_shim +
np.interp(t_list_downsampled, t_list,
output_for_policy["HVent"]["ci_50"])).clip(min=0)
output_model[schema.POPULATION] = population
# Average capacity.
output_model[schema.ICU_BED_CAPACITY] = np.mean(
output_for_policy["HICU"]["capacity"])
output_model[schema.VENTILATOR_CAPACITY] = np.mean(
output_for_policy["HVent"]["capacity"])
# Truncate date range of output.
output_dates = pd.to_datetime(output_model["date"])
output_model = output_model[
(output_dates >= datetime(month=3, day=3, year=2020))
& (output_dates < datetime.today() + timedelta(days=90))]
output_model = output_model.fillna(0)
# Fill in results for the Rt indicator.
rt_results = load_Rt_result(fips)
if rt_results is not None:
rt_results.index = rt_results[
"Rt_MAP_composite"].index.strftime("%Y-%m-%d")
merged = output_model.merge(
rt_results[["Rt_MAP_composite", "Rt_ci95_composite"]],
right_index=True,
left_on="date",
how="left",
)
output_model[schema.RT_INDICATOR] = merged["Rt_MAP_composite"]
# With 90% probability the value is between rt_indicator - ci90
# to rt_indicator + ci90
output_model[schema.RT_INDICATOR_CI90] = (
merged["Rt_ci95_composite"] - merged["Rt_MAP_composite"])
else:
log.warning(
"No Rt Results found, clearing Rt in output.",
fips=fips,
suppression_policy=suppression_policy,
)
output_model[schema.RT_INDICATOR] = "NaN"
output_model[schema.RT_INDICATOR_CI90] = "NaN"
output_model[[schema.RT_INDICATOR,
schema.RT_INDICATOR_CI90]] = output_model[[
schema.RT_INDICATOR, schema.RT_INDICATOR_CI90
]].fillna("NaN")
int_columns = [
col for col in output_model.columns if col not in (
schema.DATE,
schema.Rt,
schema.Rt_ci90,
schema.RT_INDICATOR,
schema.RT_INDICATOR_CI90,
schema.FIPS,
)
]
output_model.loc[:,
int_columns] = output_model[int_columns].fillna(
0).astype(int)
output_model.loc[:, [
schema.Rt, schema.Rt_ci90, schema.RT_INDICATOR, schema.
RT_INDICATOR_CI90
]] = output_model[[
schema.Rt, schema.Rt_ci90, schema.RT_INDICATOR,
schema.RT_INDICATOR_CI90
]].fillna(0)
output_model[schema.FIPS] = fips
intervention = Intervention.from_webui_data_adaptor(
suppression_policy)
output_model[schema.INTERVENTION] = intervention.value
output_path = get_run_artifact_path(fips,
RunArtifact.WEB_UI_RESULT,
output_dir=self.output_dir)
output_path = output_path.replace("__INTERVENTION_IDX__",
str(intervention.value))
output_model.to_json(output_path, orient=OUTPUT_JSON_ORIENT)