Commit 23149586 authored by Chris Jewell's avatar Chris Jewell
Browse files

Added full quantiles for longformat summary output.

parent e82b6c22
......@@ -31,7 +31,9 @@ def _expand_quantiles(q_dict):
return quantiles
for k, v in q_dict.items():
q_key = f"Quantile {k}"
q_key = (
f"Quantile {float(k)}" # Coerce back to float to strip trailing 0s
)
if q_key not in quantiles.keys():
raise KeyError(f"quantile '{k}' not compatible with template form")
quantiles[q_key] = v
......
......@@ -14,12 +14,9 @@ from covid.formats import make_dstl_template
def xarray2summarydf(arr):
mean = arr.mean(dim="iteration").to_dataset(name="value")
quantiles = arr.quantile(q=[0.05, 0.5, 0.95], dim="iteration").to_dataset(
dim="quantile"
)
ds = mean.merge(quantiles).rename_vars(
{0.05: "0.05", 0.5: "0.5", 0.95: "0.95"}
)
q = np.arange(start=0.05, stop=1.0, step=0.05)
quantiles = arr.quantile(q=q, dim="iteration").to_dataset(dim="quantile")
ds = mean.merge(quantiles).rename_vars({qi: f"{qi:.2f}" for qi in q})
return ds.to_dataframe().reset_index()
......@@ -93,6 +90,8 @@ def summary_longformat(input_files, output_file):
rt_summary["time"] = data["date_range"][1]
df = pd.concat([df, rt_summary], axis="index")
quantiles = df.columns[df.columns.str.startswith("0.")]
return make_dstl_template(
group="Lancaster",
model="SpatialStochasticSEIR",
......@@ -103,9 +102,5 @@ def summary_longformat(input_files, output_file):
geography=df["location"],
value_date=df["time"],
value_type=df["value_name"],
quantiles={
"0.05": df["0.05"],
"0.5": df["0.5"],
"0.95": df["0.95"],
},
quantiles={q: df[q] for q in quantiles},
).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