From 9b055eebf03476b78fc02e00eeb50acd088580ca Mon Sep 17 00:00:00 2001 From: Scott Olesen Date: Thu, 26 Dec 2024 12:38:46 -0500 Subject: [PATCH 1/2] update pre-commit hooks --- .pre-commit-config.yaml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1f36bbf..1edfbe9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,9 +10,16 @@ repos: - id: mixed-line-ending - id: trailing-whitespace - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.8.0 + rev: v0.8.4 hooks: + # Sort imports - id: ruff + args: ["check", "--select", "I", "--fix"] + # Run the linter + - id: ruff + # Run the formatter + - id: ruff-format + ##### - repo: https://github.com/Yelp/detect-secrets rev: v1.5.0 hooks: From 3591717db53802ada822167e488f9837a75d3c07 Mon Sep 17 00:00:00 2001 From: Scott Olesen Date: Thu, 26 Dec 2024 12:39:12 -0500 Subject: [PATCH 2/2] ruff compliance --- ngm/__init__.py | 14 ++-- scripts/simulate.py | 47 +++++++----- scripts/widget.py | 169 ++++++++++++++++++++++++++++++-------------- tests/test_ngm.py | 27 +++++-- 4 files changed, 173 insertions(+), 84 deletions(-) diff --git a/ngm/__init__.py b/ngm/__init__.py index 5efa201..f421bd2 100644 --- a/ngm/__init__.py +++ b/ngm/__init__.py @@ -1,7 +1,8 @@ from collections import namedtuple -import numpy as np from typing import Any +import numpy as np + DominantEigen = namedtuple("DominantEigen", ["value", "vector"]) @@ -36,9 +37,9 @@ def run_ngm( return {"M": M_vax, "Re": eigen.value, "infection_distribution": eigen.vector} -def severity(eigenvalue: float, eigenvector: np.ndarray, p_severe: np.ndarray, G: int +def severity( + eigenvalue: float, eigenvector: np.ndarray, p_severe: np.ndarray, G: int ) -> np.ndarray: - """ Calculate cumulative severe infections up to and including the Gth generation. @@ -168,7 +169,7 @@ def distribute_vaccines( remaining_proportions = np.where( np.isin(np.arange(n_groups), target_indices, invert=True), N_i / remaining_population, - 0.0 + 0.0, ) n_vax += remaining_doses * np.array(remaining_proportions) @@ -178,6 +179,7 @@ def distribute_vaccines( return n_vax + def exp_growth_model_severity(R_e, inf_distribution, p_severe, G) -> np.ndarray: """ Get cumulative infections and severe infections in generations 0, 1, ..., G @@ -195,8 +197,8 @@ def exp_growth_model_severity(R_e, inf_distribution, p_severe, G) -> np.ndarray: [:,1] is the number of infections [:,2] is the number of severe infections """ - gens = np.arange(G+1) - infections = np.cumsum(R_e ** gens) + gens = np.arange(G + 1) + infections = np.cumsum(R_e**gens) severe = np.outer(infections, inf_distribution * p_severe).sum(axis=1) return np.stack((gens, infections, severe), 1) diff --git a/scripts/simulate.py b/scripts/simulate.py index 6d5154a..dae0bc6 100644 --- a/scripts/simulate.py +++ b/scripts/simulate.py @@ -1,9 +1,11 @@ +import griddler +import griddler.griddle import numpy as np -import ngm as ngm import polars as pl import polars.selectors as cs -import griddler -import griddler.griddle + +import ngm as ngm + def simulate_scenario(params, distributions_as_percents=False): assert sum(params["pop_props"]) == 1.0 @@ -25,16 +27,22 @@ def simulate_scenario(params, distributions_as_percents=False): params["n_vax_total"], N_i, strategy=params["vax_strategy"] ) - result = ngm.run_ngm( - M_novax=M_novax, n=N_i, n_vax=n_vax, ve=params["ve"] - ) + result = ngm.run_ngm(M_novax=M_novax, n=N_i, n_vax=n_vax, ve=params["ve"]) Re = result["Re"] ifr = np.dot(result["infection_distribution"], p_severe) - fatalities_per_prior_infection = ngm.severity(eigenvalue = Re, eigenvector = result["infection_distribution"], - p_severe = p_severe, G = 1) - fatalities_after_G_generations = ngm.severity(eigenvalue = Re, eigenvector = result["infection_distribution"], - p_severe = p_severe, G = params["G"]) + fatalities_per_prior_infection = ngm.severity( + eigenvalue=Re, + eigenvector=result["infection_distribution"], + p_severe=p_severe, + G=1, + ) + fatalities_after_G_generations = ngm.severity( + eigenvalue=Re, + eigenvector=result["infection_distribution"], + p_severe=p_severe, + G=params["G"], + ) infection_distribution_dict = { f"infections_{group}": result["infection_distribution"][i] * mult @@ -67,22 +75,27 @@ def simulate_scenario(params, distributions_as_percents=False): if __name__ == "__main__": parameter_sets = griddler.griddle.read("scripts/config.yaml") - strategy_names = {"even": "even", "0": "core first", "1": "group 1 first", "2": "group 2 first", "0_1": "core and group 1 first"} + strategy_names = { + "even": "even", + "0": "core first", + "1": "group 1 first", + "2": "group 2 first", + "0_1": "core and group 1 first", + } results_all = griddler.run_squash(simulate_scenario, parameter_sets).with_columns( pl.col("vax_strategy").replace_strict(strategy_names) ) scen = results_all.select("scenario").row(0)[0] - cols_to_select = [ - "n_vax_total", - "vax_strategy", - "Re", - "ifr"] + cols_to_select = ["n_vax_total", "vax_strategy", "Re", "ifr"] results = ( results_all.with_columns(cs.float().round(3)) - .select(cs.by_name(cols_to_select) | cs.starts_with("deaths_per_prior", "infections_")) + .select( + cs.by_name(cols_to_select) + | cs.starts_with("deaths_per_prior", "infections_") + ) .sort(["n_vax_total", "vax_strategy"]) ) diff --git a/scripts/widget.py b/scripts/widget.py index 043e690..9a74170 100644 --- a/scripts/widget.py +++ b/scripts/widget.py @@ -1,21 +1,26 @@ +import altair as alt import numpy as np import polars as pl import streamlit as st -import altair as alt + import ngm from scripts.simulate import simulate_scenario -def extract_vector(prefix: str, df: pl.DataFrame, index_name: str, sigdigs, groups=["core", "children", "adults"]): + +def extract_vector( + prefix: str, + df: pl.DataFrame, + index_name: str, + sigdigs, + groups=["core", "children", "adults"], +): assert df.shape[0] == 1 cols = [prefix + grp for grp in groups] vec = ( - df - .with_columns( + df.with_columns( total=pl.sum_horizontal(cols), ) - .select( - pl.col(col).round_sig_figs(sigdigs) for col in ["total", *cols] - ) + .select(pl.col(col).round_sig_figs(sigdigs) for col in ["total", *cols]) .with_columns( summary=pl.lit(index_name), ) @@ -24,13 +29,22 @@ def extract_vector(prefix: str, df: pl.DataFrame, index_name: str, sigdigs, grou ) return vec + def summarize_scenario( - params, - sigdigs, - groups, - display=["infections_", "deaths_per_prior_infection_", "deaths_after_G_generations_"], - display_names=["Percent of infections", "Severe infections per prior infection", "Severe infections after G generations"] - ): + params, + sigdigs, + groups, + display=[ + "infections_", + "deaths_per_prior_infection_", + "deaths_after_G_generations_", + ], + display_names=[ + "Percent of infections", + "Severe infections per prior infection", + "Severe infections after G generations", + ], +): p_vax = params["n_vax"] / (params["n_total"] * params["pop_props"]) # Run the simulation with vaccination @@ -42,20 +56,23 @@ def summarize_scenario( st.subheader("\% of each group vaccinated:", help=prop_vax_help) st.dataframe( ( - pl.DataFrame({ - grp : [prob * 100] - for grp,prob in zip(params["group_names"], p_vax) - }) - .select( + pl.DataFrame( + {grp: [prob * 100] for grp, prob in zip(params["group_names"], p_vax)} + ).select( pl.col(col).round_sig_figs(sigdigs) for col in params["group_names"] ) ) ) st.subheader("Summary of Infections:") - res = pl.concat([ - extract_vector(disp, result, disp_name, sigdigs, groups = params["group_names"]) for disp,disp_name in zip(display, display_names) - ]) + res = pl.concat( + [ + extract_vector( + disp, result, disp_name, sigdigs, groups=params["group_names"] + ) + for disp, disp_name in zip(display, display_names) + ] + ) st.dataframe(res) @@ -63,47 +80,58 @@ def summarize_scenario( st.subheader("Next Generation Matrix given vaccine scenario:") m_vax = ngm.vaccinate_M(params["M_novax"], p_vax, params["ve"]) ngm_df = ( - pl.DataFrame({ - f"from {grp}": m_vax[:,i] - for i,grp in enumerate(params["group_names"]) - }) + pl.DataFrame( + {f"from {grp}": m_vax[:, i] for i, grp in enumerate(params["group_names"])} + ) .with_columns(pl.Series("", [f"to {grp}" for grp in params["group_names"]])) .select(["", *[f"from {grp}" for grp in params["group_names"]]]) ) - st.dataframe( - ngm_df - ) + st.dataframe(ngm_df) st.write(ngm_help) re_help = "The effective reproductive number accounting for the specified administration of vaccines in this scenario." - st.subheader(f"R-effective: {result['Re'].round_sig_figs(sigdigs)[0]}", help=re_help) + st.subheader( + f"R-effective: {result['Re'].round_sig_figs(sigdigs)[0]}", help=re_help + ) - ifr_help = "The probability that a random infection will result in the severe outcome of interest, e.g. death, accounting for the specified administration of vaccines in this scenario. Here \"random\" means drawing uniformly across all infections, so the probability that one draws an infection in any class is given by the distribution specified in the summary table above." - st.subheader(f"Severe infection ratio: {result['ifr'].round_sig_figs(sigdigs)[0]}", help=ifr_help) + ifr_help = 'The probability that a random infection will result in the severe outcome of interest, e.g. death, accounting for the specified administration of vaccines in this scenario. Here "random" means drawing uniformly across all infections, so the probability that one draws an infection in any class is given by the distribution specified in the summary table above.' + st.subheader( + f"Severe infection ratio: {result['ifr'].round_sig_figs(sigdigs)[0]}", + help=ifr_help, + ) - st.subheader("Cumulative infections after G generations of infection", help="This plot shows how many infections (in total across groups) there will be, both severe and otherwise, cumulatively, up to and including G generations of infection. The first generation is the generation produced by the index case, so G = 1 includes the index infection (generation 0) and one generation of spread") + st.subheader( + "Cumulative infections after G generations of infection", + help="This plot shows how many infections (in total across groups) there will be, both severe and otherwise, cumulatively, up to and including G generations of infection. The first generation is the generation produced by the index case, so G = 1 includes the index infection (generation 0) and one generation of spread", + ) percent_infections = np.array(res.select(list(params["group_names"]))[0] / 100) growth_df = ( pl.from_numpy( - ngm.exp_growth_model_severity(result["Re"], percent_infections, params["p_severe"], params["G"],), - schema=["Generation", "All Infections", "Severe Infections"] + ngm.exp_growth_model_severity( + result["Re"], + percent_infections, + params["p_severe"], + params["G"], + ), + schema=["Generation", "All Infections", "Severe Infections"], ) .with_columns( - (pl.col("All Infections") - pl.col("Severe Infections")).alias("Non-Severe Infections") + (pl.col("All Infections") - pl.col("Severe Infections")).alias( + "Non-Severe Infections" + ) ) .drop("All Infections") .unpivot(index="Generation", variable_name="Infection Type", value_name="Count") ) # Bar plot - chart = alt.Chart(growth_df).mark_bar().encode( - x='Generation:O', - y='Count:Q', - color='Infection Type:N' - ).properties( - title='' + chart = ( + alt.Chart(growth_df) + .mark_bar() + .encode(x="Generation:O", y="Count:Q", color="Infection Type:N") + .properties(title="") ) st.altair_chart(chart, use_container_width=True) @@ -111,7 +139,9 @@ def summarize_scenario( def app(): st.title("Vaccine Allocation Widget") - st.write("Uses a Next Generation Matrix (NGM) approach to approximate the dynamics of disease spread around the disease-free equilibrium.") + st.write( + "Uses a Next Generation Matrix (NGM) approach to approximate the dynamics of disease spread around the disease-free equilibrium." + ) params_default = pl.DataFrame( { @@ -123,31 +153,62 @@ def app(): ) with st.sidebar: - st.header("Model Inputs", help="If you can't see the full matrices without scrolling, drag the sidebar to make it wider.") + st.header( + "Model Inputs", + help="If you can't see the full matrices without scrolling, drag the sidebar to make it wider.", + ) - st.subheader("Population Information", help="Edit entries in the following matrix to define the:\n - Group names\n - Numbers of people in each group\n - Number of vaccines allocated to each group\n - Probability that an infection will produce the severe outcome of interest (e.g. death) in each group") + st.subheader( + "Population Information", + help="Edit entries in the following matrix to define the:\n - Group names\n - Numbers of people in each group\n - Number of vaccines allocated to each group\n - Probability that an infection will produce the severe outcome of interest (e.g. death) in each group", + ) params = st.sidebar.data_editor(params_default) - - VE = st.slider("Vaccine Efficacy", 0.0, 1.0, value=0.74, step=0.01, help="Protection due to vaccination is assumed to be all or nothing with this efficacy.") + VE = st.slider( + "Vaccine Efficacy", + 0.0, + 1.0, + value=0.74, + step=0.01, + help="Protection due to vaccination is assumed to be all or nothing with this efficacy.", + ) m_def_np = np.array([[3.0, 0.0, 0.2], [0.10, 1.0, 0.5], [0.25, 1.0, 1.5]]) M_default = ( - pl.DataFrame({ - f"from {grp}": m_def_np[:,i] - for i,grp in enumerate(params["Group name"]) - }) + pl.DataFrame( + { + f"from {grp}": m_def_np[:, i] + for i, grp in enumerate(params["Group name"]) + } + ) .with_columns(pl.Series("", [f"to {grp}" for grp in params["Group name"]])) .select(["", *[f"from {grp}" for grp in params["Group name"]]]) ) - st.subheader("Next Generation Matrix", help="For a single new infection of category `from`, specify how many infections it will generate of category `to` by editing the corresponding entry in the matrix.") + st.subheader( + "Next Generation Matrix", + help="For a single new infection of category `from`, specify how many infections it will generate of category `to` by editing the corresponding entry in the matrix.", + ) M_df = st.data_editor(M_default, disabled=["to"], hide_index=True) M_novax = M_df.drop("").to_numpy() with st.expander("Advanced Options"): - G = st.slider("Generations", 1, 10, value=10, step=1, help="Outcomes after this many generations are summarized.") - sigdigs = st.slider("Displayed significant figures", 1, 4, value=3, step=1, help="Values are reported only to this many significant figures.") + G = st.slider( + "Generations", + 1, + 10, + value=10, + step=1, + help="Outcomes after this many generations are summarized.", + ) + sigdigs = st.slider( + "Displayed significant figures", + 1, + 4, + value=3, + step=1, + help="Values are reported only to this many significant figures.", + ) # # make and run scenarios ------------------------------------------------------------ group_names = params["Group name"] @@ -159,7 +220,7 @@ def app(): "scenario_title": "Scenario: Vaccination", "group_names": group_names, "n_total": N.sum(), - "pop_props": N/N.sum(), + "pop_props": N / N.sum(), "M_novax": M_novax, "p_severe": p_severe, "n_vax": V, diff --git a/tests/test_ngm.py b/tests/test_ngm.py index b846cee..9f026e1 100644 --- a/tests/test_ngm.py +++ b/tests/test_ngm.py @@ -1,8 +1,8 @@ import numpy as np -import ngm -from numpy.testing import assert_array_equal -from numpy.testing import assert_allclose import pytest +from numpy.testing import assert_allclose, assert_array_equal + +import ngm def test_dominant_eigen_simple(): @@ -76,14 +76,20 @@ def test_simulate(): np.array([0.44507246, 0.10853944, 0.17503951, 0.2713486]), ) + def test_severe(): p_severe = np.array([0.01, 0.0]) distribution = np.array([0.25, 0.75]) g_0 = p_severe * distribution assert (ngm.severity(10.0, distribution, p_severe, 0) == g_0).all() - assert (ngm.severity(2.0, distribution, p_severe, 3) == 15.0 * distribution * p_severe).all() - assert np.isclose(ngm.severity(0.5, distribution, p_severe, 3000), 2.0 * distribution * p_severe).all() + assert ( + ngm.severity(2.0, distribution, p_severe, 3) == 15.0 * distribution * p_severe + ).all() + assert np.isclose( + ngm.severity(0.5, distribution, p_severe, 3000), 2.0 * distribution * p_severe + ).all() + def test_ensure_positive(): assert_array_equal( @@ -126,11 +132,14 @@ def test_distribute_vaccine_even(): n_vax = ngm.distribute_vaccines(V, N_i, strategy="even") assert_allclose(n_vax, np.array([1.0 / 6, 2.0 / 6, 3.0 / 6])) + def test_distribute_vaccine_01(): N_i = np.array([10.0, 20.0, 30.0, 40.0]) V = 40.0 n_vax = ngm.distribute_vaccines(V, N_i, strategy="0_1") - assert_allclose(n_vax, np.array([10.0, 20.0, 10.0 * (30.0/70.0), 10.0 * (40.0/70.0)])) + assert_allclose( + n_vax, np.array([10.0, 20.0, 10.0 * (30.0 / 70.0), 10.0 * (40.0 / 70.0)]) + ) def test_distribute_vaccine(): @@ -159,9 +168,13 @@ def test_distribute_zero_doses(): n_vax = ngm.distribute_vaccines(V=0.0, N_i=N_i, strategy=strategy) assert_allclose(n_vax, np.array([0.0, 0.0, 0.0])) + def test_exp_growth(): r0 = 2.3 p_severe = np.array([0.02, 0.06, 0.02]) distribution = np.array([0.25, 0.25, 0.5]) G = 7 - assert ngm.severity(r0, distribution, p_severe, G).sum() == ngm.exp_growth_model_severity(r0, distribution, p_severe, G)[-1, 2] + assert ( + ngm.severity(r0, distribution, p_severe, G).sum() + == ngm.exp_growth_model_severity(r0, distribution, p_severe, G)[-1, 2] + )