overall_rt.py 977 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
"""Calculates overall Rt given a posterior next generation matix"""

import numpy as np
import pickle as pkl
import pandas as pd

from covid.summary import (
    rayleigh_quotient,
    power_iteration,
)


def overall_rt(next_generation_matrix, output_file):

    with open(next_generation_matrix, "rb") as f:
        ngms = pkl.load(f)

    b, _ = power_iteration(ngms)
    rt = rayleigh_quotient(ngms, b)
    q = np.arange(0.05, 1.0, 0.05)
    rt_quantiles = pd.DataFrame(
        {"Rt": np.quantile(rt, q, axis=-1)}, index=q
    ).T.to_excel(output_file)


if __name__ == "__main__":

    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument(
        "input_file",
        description="The input .pkl file containing the next generation matrix",
    )
    parser.add_argument(
        "output_file", description="The name of the output .xlsx file"
    )

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