Commit 32ffa1f1 authored by Chris Jewell's avatar Chris Jewell
Browse files

Changes to summary_longformat for DSTL

parent d40d516b
...@@ -10,6 +10,8 @@ from covid.util import impute_previous_cases ...@@ -10,6 +10,8 @@ from covid.util import impute_previous_cases
import covid.data as data import covid.data as data
tfd = tfp.distributions tfd = tfp.distributions
VERSION = 0.3
DTYPE = np.float64 DTYPE = np.float64
STOICHIOMETRY = np.array([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]]) STOICHIOMETRY = np.array([[-1, 1, 0, 0], [0, -1, 1, 0], [0, 0, -1, 1]])
......
...@@ -61,7 +61,7 @@ def predicted_incidence(posterior_samples, covar_data, init_step, num_steps): ...@@ -61,7 +61,7 @@ def predicted_incidence(posterior_samples, covar_data, init_step, num_steps):
), ),
fn_output_signature=(tf.float64), fn_output_signature=(tf.float64),
) )
return events return init_state, events
def read_pkl(filename): def read_pkl(filename):
...@@ -77,22 +77,27 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps): ...@@ -77,22 +77,27 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
if initial_step < 0: if initial_step < 0:
initial_step = samples["seir"].shape[-2] + initial_step initial_step = samples["seir"].shape[-2] + initial_step
dates = np.arange(covar_data["date_range"][0] + np.timedelta64(initial_step, "D"),
covar_data["date_range"][0] + np.timedelta64(initial_step + num_steps, "D"),
np.timedelta64(1, "D"))
del covar_data["date_range"] del covar_data["date_range"]
prediction = predicted_incidence( estimated_init_state, predicted_events = predicted_incidence(
samples, covar_data, initial_step, num_steps samples, covar_data, initial_step, num_steps
) )
prediction = xarray.DataArray( prediction = xarray.DataArray(
prediction, predicted_events,
coords=[ coords=[
np.arange(prediction.shape[0]), np.arange(predicted_events.shape[0]),
covar_data["locations"]["lad19cd"], covar_data["locations"]["lad19cd"],
np.arange(prediction.shape[2]), dates,
np.arange(prediction.shape[3]), np.arange(predicted_events.shape[3]),
], ],
dims=("iteration", "location", "time", "event"), dims=("iteration", "location", "time", "event"),
) )
prediction.attrs["initial_state"] = estimated_init_state
with open(output_file, "wb") as f: with open(output_file, "wb") as f:
pkl.dump(prediction, f) pkl.dump(prediction, f)
......
...@@ -80,7 +80,7 @@ def summary_longformat(input_files, output_file): ...@@ -80,7 +80,7 @@ def summary_longformat(input_files, output_file):
# Medium term prevalence # Medium term prevalence
prev_df = prevalence(medium_term, data["N"]) prev_df = prevalence(medium_term, data["N"])
prev_df["value_name"] = "Prevalence" prev_df["value_name"] = "prevalence"
df = pd.concat([df, prev_df], axis="index") df = pd.concat([df, prev_df], axis="index")
# Rt # Rt
...@@ -96,6 +96,7 @@ def summary_longformat(input_files, output_file): ...@@ -96,6 +96,7 @@ def summary_longformat(input_files, output_file):
return make_dstl_template( return make_dstl_template(
group="Lancaster", group="Lancaster",
model="SpatialStochasticSEIR", model="SpatialStochasticSEIR",
scenario="Nowcast",
creation_date=date.today(), creation_date=date.today(),
version=model_spec.VERSION, version=model_spec.VERSION,
age_band="All", age_band="All",
...@@ -107,4 +108,4 @@ def summary_longformat(input_files, output_file): ...@@ -107,4 +108,4 @@ def summary_longformat(input_files, output_file):
"0.5": df["0.5"], "0.5": df["0.5"],
"0.95": df["0.95"], "0.95": df["0.95"],
}, },
).to_excel(output_file) ).to_excel(output_file, index=False)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment