gold.py

import polars as pl
import numpy as np

baseline_pars = {
    "N": int(1e3),
    "p_hospitalized": 0.10,
    "hosp_time_fun": lambda rng, size: rng.normal(loc=0.0, scale=1.0, size=size),
    "p_vax": 0.25,
    "vax_time_fun": lambda rng, size: rng.normal(loc=0.0, scale=1.0, size=size),
    "ve": 0.50,
    "ve_lag": 14.0,
    "seed": 1234,
}


def simulate(pars):
    rng = np.random.default_rng(pars["seed"])

    return pl.DataFrame(
        {
            "is_hosp_wo_vax": rng.choice(
                [True, False],
                p=[pars["p_hosp"], 1.0 - pars["p_hosp"]],
                size=pars["N"],
            ),
            "hosp_time": pars["hosp_time_fun"](rng, size=pars["N"]),
            "is_vax": rng.choice(
                [True, False], p=[pars["p_vax"], 1.0 - pars["p_vax"]], sizes=pars["N"]
            ),
            "vax_time": pars["vax_time_fun"](rng, size=pars["N"]),
            "vax_is_eff": rng.choice(
                [True, False], p=[pars["ve"], 1.0 - pars["ve"]], size=pars["N"]
            ),
        }
    ).with_columns(
        # a person is hospitalized if (1) they would have been hospitalized if not
        # vaccinated and either (2A) they are not vaccinated or (2B) they are
        # not vaccinated soon enough before they would have been hospitalized
        is_hosp=(
            pl.col("is_hosp_wo_vax")
            & (
                (pl.col("is_vax").not_())
                | (pl.col("hosp_time") - pl.col("vax_time") < pars["ve_lag"])
            )
        )
    )

• Should say "outcome" not "hospitalization"
• is_hosp logic should account for VE