Commit 62f9ec37 authored by Chris Jewell's avatar Chris Jewell
Browse files

Merge branch 'fix-prediction-weekday' into 'master'

Fix for weekday effect in prediction

See merge request !37
parents 45a79675 3df0dd34
...@@ -3,6 +3,7 @@ ...@@ -3,6 +3,7 @@
import numpy as np import numpy as np
import xarray import xarray
import pickle as pkl import pickle as pkl
import pandas as pd
import tensorflow as tf import tensorflow as tf
from covid import model_spec from covid import model_spec
...@@ -75,11 +76,17 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps): ...@@ -75,11 +76,17 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
origin_date = np.array(cases.coords["time"][0]) origin_date = np.array(cases.coords["time"][0])
dates = np.arange( dates = np.arange(
origin_date + np.timedelta64(initial_step, "D"), origin_date,
origin_date + np.timedelta64(initial_step + num_steps, "D"), origin_date + np.timedelta64(initial_step + num_steps, "D"),
np.timedelta64(1, "D"), np.timedelta64(1, "D"),
) )
covar_data["weekday"] = xarray.DataArray(
(pd.to_datetime(dates).weekday < 5).astype(model_spec.DTYPE),
coords=[dates],
dims=["prediction_time"],
)
estimated_init_state, predicted_events = predicted_incidence( estimated_init_state, predicted_events = predicted_incidence(
samples, initial_state, covar_data, initial_step, num_steps samples, initial_state, covar_data, initial_step, num_steps
) )
...@@ -89,7 +96,7 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps): ...@@ -89,7 +96,7 @@ def predict(data, posterior_samples, output_file, initial_step, num_steps):
coords=[ coords=[
np.arange(predicted_events.shape[0]), np.arange(predicted_events.shape[0]),
covar_data.coords["location"], covar_data.coords["location"],
dates, dates[initial_step:],
np.arange(predicted_events.shape[3]), np.arange(predicted_events.shape[3]),
], ],
dims=("iteration", "location", "time", "event"), dims=("iteration", "location", "time", "event"),
......
Markdown is supported
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