overall_rt.py 1020 Bytes
Newer Older
1
2
3
"""Calculates overall Rt given a posterior next generation matix"""

import numpy as np
4
import xarray
5
6
7
8
9
10
11
12
import pandas as pd

from covid.summary import (
    rayleigh_quotient,
    power_iteration,
)


Chris Jewell's avatar
Chris Jewell committed
13
14
15
16
17
18
def overall_rt(inference_data, output_file):

    r_t = xarray.open_dataset(inference_data, group="posterior_predictive")[
        "R_t"
    ]

19
    q = np.arange(0.05, 1.0, 0.05)
Chris Jewell's avatar
Chris Jewell committed
20
21
22
23
24
    quantiles = r_t.isel(time=-1).quantile(q=q)
    quantiles.to_dataframe().T.to_excel(output_file)
    # pd.DataFrame({"Rt": np.quantile(r_t, q, axis=-1)}, index=q).T.to_excel(
    #     output_file
    # )
25
26
27
28
29
30
31
32
33


if __name__ == "__main__":

    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument(
        "input_file",
34
35
        type=str,
        help="The input .pkl file containing the next generation matrix",
36
37
    )
    parser.add_argument(
38
        "output_file", type=str, help="The name of the output .xlsx file"
39
40
41
42
    )

    args = parser.parse_args()
    overall_rt(args.input_file, args.output_file)