diff --git a/epimargin/estimators.py b/epimargin/estimators.py index 6e4fe75..b6d4b5b 100644 --- a/epimargin/estimators.py +++ b/epimargin/estimators.py @@ -22,7 +22,7 @@ def rollingOLS(totals: pd.DataFrame, window: int = 3, infectious_period: float = model = RollingOLS.from_formula(formula = "logdelta ~ time", window = window, data = totals) rolling = model.fit(method = "lstsq") - growthrates = rolling.params.join(rolling.bse, rsuffix="_stderr") + growthrates = pd.DataFrame(rolling.params).join(rolling.bse, rsuffix="_stderr") growthrates["rsq"] = rolling.rsquared growthrates.rename(lambda s: s.replace("time", "gradient").replace("const", "intercept"), axis = 1, inplace = True) diff --git a/epimargin/etl/__init__.py b/epimargin/etl/__init__.py index 58021f8..7303a26 100644 --- a/epimargin/etl/__init__.py +++ b/epimargin/etl/__init__.py @@ -1,2 +1,2 @@ from .commons import download_data -__all__ = [download_data] \ No newline at end of file +__all__ = ["download_data"] \ No newline at end of file diff --git a/epimargin/models.py b/epimargin/models.py index 469fe31..a9e3005 100644 --- a/epimargin/models.py +++ b/epimargin/models.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Dict, Iterator, Optional, Sequence, Tuple, Union +from typing import Dict, Iterator, Optional, Sequence, Tuple, Union, List import geopandas as gpd import numpy as np @@ -124,7 +124,7 @@ def forward_epi_step(self, dB: int = 0): # parallel poisson draws for infection def parallel_forward_epi_step(self, dB: int = 0, num_sims = 10000): # get previous state - S, I, R, D, N = (vector[-1].copy() for vector in (self.S, self.I, self.R, self.D, self.N)) + S, I, R, D, N = (vector[-1] for vector in (self.S, self.I, self.R, self.D, self.N)) # update state Rt = self.Rt0 * S/N @@ -148,9 +148,9 @@ def parallel_forward_epi_step(self, dB: int = 0, num_sims = 10000): I -= (num_dead + num_recov) - S = S.clip(0) - I = I.clip(0) - D = D.clip(0) + S = max(0, S) + I = max(0, I) + D = max(0, D) N = S + I + R beta = (num_cases * N)/(b * S * I) @@ -172,15 +172,15 @@ def parallel_forward_epi_step(self, dB: int = 0, num_sims = 10000): # parallel binomial draws for infection def parallel_forward_binom_step(self, dB: int = 0, num_sims = 10000): # get previous state - S, I, R, D, N = (vector[-1].copy() for vector in (self.S, self.I, self.R, self.D, self.N)) + S, I, R, D, N = (vector[-1] for vector in (self.S, self.I, self.R, self.D, self.N)) # update state Rt = self.Rt0 * S/N p = self.gamma * Rt * I/N - num_cases = binom.rvs(n = S.astype(int), p = p, size = num_sims) - self.upper_CI.append(binom.ppf(self.CI, n = S.astype(int), p = p)) - self.lower_CI.append(binom.ppf(1 - self.CI, n = S.astype(int), p = p)) + num_cases = binom.rvs(n = S, p = p, size = num_sims) + self.upper_CI.append(binom.ppf(self.CI, n = S, p = p)) + self.lower_CI.append(binom.ppf(1 - self.CI, n = S, p = p)) I += num_cases S -= num_cases @@ -195,12 +195,11 @@ def parallel_forward_binom_step(self, dB: int = 0, num_sims = 10000): I -= (num_dead + num_recov) - S = S.clip(0) - I = I.clip(0) - D = D.clip(0) + S = max(0, S) + I = max(0, I) + D = max(0, D) N = S + I + R - # beta = (num_cases * N)/(b * S * I) # update state vectors self.Rt.append(Rt) @@ -222,17 +221,20 @@ def run(self, days: int): def __repr__(self) -> str: return f"[{self.name}]" -class Age_SIRVD(SIR): - """ age-structured compartmental model with a vaccinated class for each age bin """ +class Age_SIRVD(): + """ age-structured compartmental model with a vaccinated class for each age bin + note that the underlying parallelizing mechanism is different from that of SIR and NetworkedSIR + + """ def __init__(self, name: str, # name of unit population: int, # unit population dT0: Optional[int] = None, # last change in cases, None -> Poisson random intro Rt0: float = 1.9, # initial reproductive rate, - S0: int = 0, # initial susceptible - I0: int = 0, # initial infected - R0: int = 0, # initial recovered - D0: int = 0, # initial dead + S0: np.array = np.array(0), # initial susceptibles + I0: np.array = np.array(0), # initial infected + R0: np.array = np.array(0), # initial recovered + D0: np.array = np.array(0), # initial dead infectious_period: int = 5, # how long disease is communicable in days introduction_rate: float = 5.0, # parameter for new community transmissions (lambda) mortality: float = 0.02, # I -> D transition probability @@ -245,7 +247,36 @@ def __init__(self, ve: float = 0.7, # vaccine effectiveness random_seed: int = 0 # random seed, ): - super().__init__(name, population, dT0=dT0, Rt0=Rt0, I0=I0, R0=R0, D0=D0, infectious_period=infectious_period, introduction_rate=introduction_rate, mortality=mortality, mobility=mobility, upper_CI=upper_CI, lower_CI=lower_CI, CI=CI, random_seed=random_seed) + self.name = name + self.pop0 = population + self.gamma = 1.0/infectious_period + self.ll = introduction_rate + self.m = mortality + self.mu = mobility + self.Rt0 = Rt0 + self.CI = CI + + # state and delta vectors + if dT0 is None: + dT0 = np.random.poisson(self.ll) # initial number of new cases + self.dT = [dT0] # case change rate, initialized with the first introduction, if any + self.Rt = [Rt0] + self.b = [np.exp(self.gamma * (Rt0 - 1.0))] + self.S = [S0 if S0 is not None else population - R0 - D0 - I0] + self.I = [I0] + self.R = [R0] + self.D = [D0] + self.dR = [0] + self.dD = [0] + self.N = [population - D0] # total population = S + I + R + self.beta = [Rt0 * self.gamma] # initial contact rate + self.total_cases = [I0] # total cases + self.upper_CI = [upper_CI] + self.lower_CI = [lower_CI] + + np.random.seed(random_seed) + + self.N = [S0 + I0 + R0] shape = (sims, bins) = S0.shape @@ -276,7 +307,7 @@ def __init__(self, self.dT_total = [np.zeros(sims)] self.dD_total = [np.zeros(sims)] - self.dV = [] + self.dV: List[np.array] = [] self.rng = np.random.default_rng(random_seed) @@ -440,21 +471,12 @@ def set_parameters(self, **kwargs): unit.__setattr__(attr, val) return self - def aggregate(self, curves: Union[Sequence[str], str] = ["Rt", "b", "S", "I", "R", "D", "P", "beta"]) -> Union[Dict[str, Sequence[float]], Sequence[float]]: - if isinstance(curves, str): - single_curve = curves - curves = [curves] - else: - single_curve = False - aggs = { + def aggregate(self, curves: Union[Sequence[str], str] = ["Rt", "b", "S", "I", "R", "D", "P", "beta"]) -> Dict[str, Sequence[float]]: + return { curve: list(map(sum, zip(*(unit.__getattribute__(curve) for unit in self.units)))) for curve in curves } - if single_curve: - return aggs[single_curve] - return aggs - class SEIR(): """ stochastic SEIR model without external introductions """ def __init__(self, diff --git a/epimargin/plots.py b/epimargin/plots.py index 2c69d08..774a44c 100644 --- a/epimargin/plots.py +++ b/epimargin/plots.py @@ -1,7 +1,8 @@ +from build.lib.epimargin.models import NetworkedSIR import datetime from collections import namedtuple from pathlib import Path -from typing import Optional, Sequence, Tuple +from typing import Optional, Sequence, Tuple, List, Dict import matplotlib as mpl import matplotlib.dates as mdates @@ -11,7 +12,6 @@ import seaborn as sns import tikzplotlib from matplotlib.patheffects import Normal, Stroke -from matplotlib.pyplot import * from .models import SIR @@ -56,15 +56,6 @@ def rebuild_font_cache(): import matplotlib.font_manager matplotlib.font_manager._rebuild() -def despine(**kwargs): - pass - -def grid(flag): - if flag: - pass - else: - pass - # container class for different theme Aesthetics = namedtuple( "Aesthetics", @@ -301,7 +292,7 @@ def show(self, **kwargs): plt.show(**kwargs) return self -def plot_SIRD(model: SIR, layout = (1, 4)) -> PlotDevice: +def plot_SIRD(model: NetworkedSIR, layout = (1, 4)) -> PlotDevice: """ plot all 4 available curves (S, I, R, D) for a given SIR model """ fig, axes = plt.subplots(layout[0], layout[1], sharex = True, sharey = True) t = list(range(len(model.Rt))) @@ -315,7 +306,7 @@ def plot_SIRD(model: SIR, layout = (1, 4)) -> PlotDevice: fig.legend([s, i, r, d], labels = ["S", "I", "R", "D"], loc = "center right", borderaxespad = 0.1) return PlotDevice(fig) -def plot_curve(models: Sequence[SIR], labels: Sequence[str], curve: str = "I"): +def plot_curve(models: Sequence[NetworkedSIR], labels: Sequence[str], curve: str = "I"): """ plot specific epidemic curve """ fig = plt.figure() for (model, label) in zip(models, labels): @@ -368,7 +359,7 @@ def predictions(date_range, model, color, bounds = [2.5, 97.5], curve = "dT"): return [(range_marker, median_marker), model.name] def simulations( - simulation_results: Sequence[Tuple[SIR]], + simulation_results: Sequence[Tuple[NetworkedSIR]], labels: Sequence[str], historical: Optional[pd.Series] = None, historical_label: str = "Empirical Case Data", @@ -384,18 +375,18 @@ def simulations( num_sims = len(simulation_results) total_time = len(policy_outcomes[0][0]) - ranges = [{"max": [], "min": [], "mdn": [], "avg": []} for _ in range(len(policy_outcomes))] + ranges: List[Dict[str, List]] = [{"max": [], "min": [], "mdn": [], "avg": []} for _ in range(len(policy_outcomes))] for (i, policy) in enumerate(policy_outcomes): - for t in range(total_time): - curve_sorted = sorted([curve[t] for curve in policy]) + for _ in range(total_time): + curve_sorted = sorted([curve[_] for curve in policy]) ranges[i]["min"].append(curve_sorted[0]) ranges[i]["max"].append(curve_sorted[-1]) ranges[i]["mdn"].append(curve_sorted[num_sims//2]) ranges[i]["avg"].append(np.mean(curve_sorted)) legends = [] - legend_labels = [] + legend_labels = [] if historical is not None: p, = plt.plot(historical.index, historical, 'k-', alpha = 0.8, zorder = 10) t = [historical.index.max() + datetime.timedelta(days = n) for n in range(total_time)] @@ -478,7 +469,7 @@ def daily_cases(dates, T_pred, T_CI_upper, T_CI_lower, new_cases_ts, anomaly_dat plt.ylim(bottom = 0, top = top) legends += [anomalies_marker] labels += ["anomalies"] - xlim(left = dates[0], right = end) + plt.xlim(left = dates[0], right = end) plt.legend(legends, labels, prop = {'size': 14}, framealpha = theme.framealpha, handlelength = theme.handlelength, loc = "best") plt.gca().xaxis.set_major_formatter(DATE_FMT) plt.gca().xaxis.set_minor_formatter(DATE_FMT) @@ -544,5 +535,5 @@ def double_choropleth_v(*args, **kwargs): kwargs["arrangement"] = (2, 1) return double_choropleth(*args, **kwargs) -double_choropleth.horizontal = double_choropleth -double_choropleth.vertical = double_choropleth_v +double_choropleth.horizontal = double_choropleth # type: ignore +double_choropleth.vertical = double_choropleth_v # type: ignore diff --git a/epimargin/policy.py b/epimargin/policy.py index 4d7acbe..a0e03cb 100644 --- a/epimargin/policy.py +++ b/epimargin/policy.py @@ -1,20 +1,20 @@ from abc import abstractmethod from itertools import product -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Set, Tuple import numpy as np -from sklearn.metrics import auc - from scipy.stats import multinomial as Multinomial +from sklearn.metrics import auc -from .models import SIR +from .models import SIR, NetworkedSIR, Age_SIRVD from .utils import weeks + def AUC(curve): return auc(x = range(len(curve)), y = curve) # Adaptive Control policies -def simulate_lockdown(model: SIR, lockdown_period: int, total_time: int, Rt0_mandatory: Dict[str, float], Rt0_voluntary: Dict[str, float], lockdown: np.matrix, migrations: np.matrix) -> SIR: +def simulate_lockdown(model: NetworkedSIR, lockdown_period: int, total_time: int, Rt0_mandatory: Dict[str, float], Rt0_voluntary: Dict[str, float], lockdown: np.matrix, migrations: np.matrix) -> NetworkedSIR: """ simulate a lockdown where stringency is assumed to be reflected by Rt0_mandatory, with a return to Rt0_voluntary (and possible migration) once lockdown is over """ return model.set_parameters(Rt0 = Rt0_mandatory)\ .run(lockdown_period, migrations = lockdown)\ @@ -22,7 +22,7 @@ def simulate_lockdown(model: SIR, lockdown_period: int, total_time: int, Rt0_man .run(total_time - lockdown_period, migrations = migrations) def simulate_adaptive_control( - model: SIR, + model: NetworkedSIR, initial_run: int, total_time: int, lockdown: np.matrix, @@ -31,7 +31,7 @@ def simulate_adaptive_control( beta_v: Dict[str, float], beta_m: Dict[str, float], evaluation_period: int = 2*weeks, - adjacency: Optional[np.matrix] = None) -> SIR: + adjacency: Optional[np.matrix] = None) -> NetworkedSIR: """ simulate the Malani et al. adaptive lockdown policy where districts are assigned to lockdown stringency buckets based on Rt """ n = len(model) model.set_parameters(Rt0 = R_m)\ @@ -40,7 +40,10 @@ def simulate_adaptive_control( gantt = [] last_category = dict() while days_run < total_time: - Gs, Ys, Os, Rs = set(), set(), set(), set() + Gs: Set[int] = set() + Ys: Set[int] = set() + Os: Set[int] = set() + Rs: Set[int] = set() categories = dict(enumerate([Gs, Ys, Os, Rs])) category_transitions = {} for (i, unit) in enumerate(model): @@ -88,10 +91,10 @@ def simulate_adaptive_control( model.run(evaluation_period, phased_migration) days_run += evaluation_period - model.gantt = gantt + model.gantt = gantt # type: ignore return model -def simulate_adaptive_control_MHA(model: SIR, initial_run: int, total_time: int, lockdown: np.matrix, migrations: np.matrix, R_m: Dict[str, float], beta_v: Dict[str, float], beta_m: Dict[str, float], evaluation_period = 2*weeks): +def simulate_adaptive_control_MHA(model: NetworkedSIR, initial_run: int, total_time: int, lockdown: np.matrix, migrations: np.matrix, R_m: Dict[str, float], beta_v: Dict[str, float], beta_m: Dict[str, float], evaluation_period = 2*weeks): """ simulates the version of adaptive control suggested by the Indian Ministry of Home Affairs, where the trigger is based on infection count doubling time """ n = len(model) model.set_parameters(Rt0 = R_m)\ @@ -100,7 +103,10 @@ def simulate_adaptive_control_MHA(model: SIR, initial_run: int, total_time: int, gantt = [] last_category = dict() while days_run < total_time: - Gs, Ys, Os, Rs = set(), set(), set(), set() + Gs: Set[int] = set() + Ys: Set[int] = set() + Os: Set[int] = set() + Rs: Set[int] = set() categories = dict(enumerate([Gs, Ys, Os, Rs])) category_transitions = {} for (i, unit) in enumerate(model): @@ -151,26 +157,26 @@ def simulate_adaptive_control_MHA(model: SIR, initial_run: int, total_time: int, model.run(evaluation_period, phased_migration) days_run += evaluation_period - model.gantt = gantt + model.gantt = gantt # type: ignore return model def simulate_PID_controller( - model: SIR, + model: NetworkedSIR, initial_run: int, total_time: int, Rtarget: float = 0.9, kP: float = 0.05, kI: float = 0.5, kD: float = 0, - Dt: float = 1.0) -> SIR: + Dt: float = 1.0) -> NetworkedSIR: """ implements a hypothetical proportional-integral-derivative control policy where the error term is Rt magnitude above 1 """ # initial run without PID model.run(initial_run) # set up PID running variables - integral = 0 - derivative = 0 - u = 0 + integral = 0.0 + derivative = 0.0 + u = 0.0 prev_error = model[0].Rt[-1] z = np.zeros((len(model), len(model))) @@ -191,14 +197,16 @@ def simulate_PID_controller( # Vaccination policies class VaccinationPolicy(): """ parent class to hold vaccination policy state """ - def __init__(self, bin_populations: np.array) -> None: + def __init__(self, bin_populations: np.array, effectiveness: float, daily_doses: int) -> None: self.bin_populations = bin_populations + self.daily_doses = daily_doses + self.effectiveness = effectiveness def name(self) -> str: return self.__class__.__name__.lower() @abstractmethod - def distribute_doses(self, model: SIR, *kwargs) -> Tuple[np.array]: + def distribute_doses(self, model: Age_SIRVD, num_sims: int) -> Tuple[np.array, ...]: pass def exhausted(self, model) -> bool: @@ -213,17 +221,15 @@ def get_mortality(self, base_IFRs) -> float: class RandomVaccineAssignment(VaccinationPolicy): """ assigns vaccines to members of the population randomly """ def __init__(self, daily_doses: int, effectiveness: float, bin_populations: np.array, age_ratios: np.array): - self.daily_doses = daily_doses + super().__init__(bin_populations, effectiveness, daily_doses) self.age_ratios = age_ratios - self.effectiveness = effectiveness - self.bin_populations = bin_populations - def distribute_doses(self, model: SIR, num_sims: int = 10000) -> Tuple[np.array]: + def distribute_doses(self, model: Age_SIRVD, num_sims: int = 10000) -> Tuple[np.array, ...]: if self.exhausted(model): return (np.zeros(self.age_ratios.shape), np.zeros(self.age_ratios.shape), np.zeros(self.age_ratios.shape)) dV = (model.S[-1]/model.N[-1]) * self.daily_doses * self.effectiveness model.S[-1] -= dV - model.parallel_forward_epi_step(num_sims = num_sims) + model.parallel_forward_epi_step(0, num_sims = num_sims) distributed_doses = Multinomial.rvs(self.daily_doses, self.age_ratios) effective_doses = self.effectiveness * distributed_doses immunizing_doses = (model.S[-1].mean()/model.N[-1].mean()) * effective_doses @@ -236,26 +242,23 @@ def name(self) -> str: class PrioritizedAssignment(VaccinationPolicy): """ assigns vaccines to members of the population based on a prioritized ordering of subpopulations """ def __init__(self, daily_doses: int, effectiveness: float, bin_populations: np.array, prioritization: List[int], label: str): - self.daily_doses = daily_doses - self.bin_populations = bin_populations - self.prioritization = prioritization - self.effectiveness = effectiveness + super().__init__(bin_populations, effectiveness, daily_doses) + self.prioritization = prioritization self.label = label def name(self) -> str: return f"{self.label}prioritized" - def distribute_doses(self, model: SIR, num_sims: int = 10_000) -> Tuple[np.array]: + def distribute_doses(self, model: Age_SIRVD, num_sims: int = 10_000) -> Tuple[np.array, ...]: if self.exhausted(model): return (None, None, None) - # return (np.zeros(self.age_ratios.shape), np.zeros(self.age_ratios.shape), np.zeros(self.age_ratios.shape)) dV = (model.S[-1]/model.N[-1]) * self.daily_doses * self.effectiveness model.S[-1] -= dV - model.parallel_forward_epi_step(num_sims = num_sims) + model.parallel_forward_epi_step(0, num_sims = num_sims) dVx = np.zeros(self.bin_populations.shape) bin_idx, age_bin = next(((i, age_bin) for (i, age_bin) in enumerate(self.prioritization) if self.bin_populations[age_bin] > 0), (None, None)) - if age_bin is not None: + if age_bin is not None and bin_idx is not None: if self.bin_populations[age_bin] > self.daily_doses: self.bin_populations[age_bin] -= self.daily_doses dVx[age_bin] = self.daily_doses diff --git a/epimargin/utils.py b/epimargin/utils.py index 7284d0b..ef246b3 100644 --- a/epimargin/utils.py +++ b/epimargin/utils.py @@ -32,7 +32,7 @@ def mkdir(p: Path, exist_ok: bool = True) -> Path: p.mkdir(exist_ok = exist_ok) return p -def setup(**kwargs) -> Tuple[Path]: +def setup(**kwargs) -> Tuple[Path, ...]: root = cwd() if len(sys.argv) > 2: parser = argparse.ArgumentParser() diff --git a/requirements.txt b/requirements.txt index bbc53b6..210bd5c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,8 @@ astroid==2.4.1 attrs==19.3.0 backcall==0.1.0 beautifulsoup4==4.9.3 +bleach==3.3.0 +bump2version==1.0.1 bumpversion==0.6.0 certifi==2020.6.20 cftime==1.1.3 @@ -11,23 +13,27 @@ chardet==3.0.4 click==7.1.2 click-plugins==1.1.1 cligj==0.5.0 +cloudpickle==1.6.0 colorama==0.4.3 commonmark==0.9.1 cycler==0.10.0 Cython==0.29.21 +dask==2021.4.1 decorator==4.4.2 descartes==1.1.0 +docutils==0.17.1 fastprogress==0.2.3 Fiona==1.8.17 flake8==3.8.3 Flask==1.1.2 flat-table==1.1.1 +fsspec==2021.4.0 GDAL==3.1.3 geopandas==0.8.1 geos==0.2.2 h5py==2.10.0 idna==2.9 -importlib-metadata==1.6.1 +importlib-metadata==4.5.0 ipython==7.13.0 ipython-genutils==0.2.0 isort==4.3.21 @@ -35,24 +41,31 @@ itsdangerous==1.1.0 jedi==0.17.0 Jinja2==2.11.3 joblib==0.14.1 +keyring==23.0.1 kiwisolver==1.2.0 lazy-object-proxy==1.4.3 linearmodels==4.19 +locket==0.2.1 lxml==4.6.3 +mapclassify==2.4.2 MarkupSafe==1.1.1 matplotlib==3.2.1 mccabe==0.6.1 munch==2.5.0 +mypy==0.910 mypy-extensions==0.4.3 netCDF4==1.5.3 +networkx==2.5.1 numpy==1.18.2 packaging==20.4 pandas==1.0.3 parso==0.7.0 +partd==1.2.0 patsy==0.5.1 pexpect==4.8.0 pickleshare==0.7.5 Pillow==8.2.0 +pkginfo==1.7.0 pprintpp==0.4.0 prompt-toolkit==3.0.5 property-cached==1.6.4 @@ -68,7 +81,11 @@ pyproj==2.6.0 pyreadr==0.4.0 python-dateutil==2.8.1 pytz==2019.3 +PyYAML==5.4.1 +readme-renderer==29.0 requests==2.23.0 +requests-toolbelt==0.9.1 +rfc3986==1.5.0 rich==2.3.0 rope==0.17.0 scikit-learn==0.22.2.post1 @@ -82,13 +99,16 @@ statsmodels==0.11.1 Theano==1.0.5 tikzplotlib==0.9.4 toml==0.10.1 +toolz==0.11.1 tqdm==4.45.0 traitlets==4.3.3 +twine==3.4.1 typed-ast==1.4.1 typing-extensions==3.7.4.2 urllib3==1.26.5 urlpath==1.1.7 wcwidth==0.1.9 +webencodings==0.5.1 Werkzeug==1.0.1 wrapt==1.12.1 xarray==0.15.1