diff --git a/proto/engwal.md b/proto/engwal.md index 2ce0ca2..be10c6a 100644 --- a/proto/engwal.md +++ b/proto/engwal.md @@ -11,7 +11,7 @@ The model, using the `-n` or `--nodes` parameter, runs the n-largest communities The network is based on the distances between communities (cached on disk) and optional $k$, $a$, $b$, and $c$ parameters for the function $$ -k \left( p_1^a p_2^b \over N d^c \right) +c_{ij} = k \left( P_i^a P_j^b \over N D_{ij}^c \right) $$ where $N$ is the total population (per @kfrey-idm). The $k$, $a$, $b$, and $c$ parameters may be specified on the command line. diff --git a/proto/engwal.py b/proto/engwal.py index ce0d3f0..0e539ac 100755 --- a/proto/engwal.py +++ b/proto/engwal.py @@ -5,6 +5,7 @@ from argparse import ArgumentParser from argparse import Namespace from collections import namedtuple +from datetime import datetime from functools import lru_cache from pathlib import Path from typing import List @@ -21,6 +22,8 @@ def main(args: Namespace) -> None: """Main function for the England and Wales Measles model.""" + print(f"Running England and Wales Measles model with\n{args.__dict__}") + network, population = initialize_model(args) run_model(population, network, args) generate_results(population) @@ -30,11 +33,14 @@ def main(args: Namespace) -> None: def initialize_model(args: Namespace): """Initialize the model from files and arguments.""" - places = get_places(args.nodes) + print(f"Initializing np.random with seed {args.seed}... ") + np.random.seed(args.seed) - network = load_network(args.nodes, args.g_k, args.g_a, args.g_b, args.g_c) + places = get_places(args.nodes) + network = load_network(args.nodes, args.g_k, args.g_a, args.g_b, args.g_c, args.g_max_frac) # Create a list of communities, one for each population in pops + print(f"Initializing model with {args.nodes} nodes... ") pop = Population( args.nodes, community_props=None, # ["population", "births"], @@ -59,7 +65,8 @@ def init_community(_population: Population, community: Community, index: int) -> init_pop = community.population[0] unactive = max_pop - init_pop - susceptible = np.uint32(np.round(init_pop / args.r_naught)) + # Adding a little to the susceptible population to help the initial infections take root + susceptible = np.uint32(np.round(1.0625 * init_pop / args.r_naught)) recovered = init_pop - susceptible # add groups: unactive, susceptible, exposed, infectious, recovered, and deceased print(f"{index:3} unactive: {unactive:8}, susceptible: {susceptible:8}, recovered: {recovered:8}") @@ -91,6 +98,7 @@ def run_model(population: Population, network: np.ndarray, args: Namespace) -> N """Run the model for the specified number of ticks.""" population.apply(update_report, tick=0) # Capture the initial state of the population. + print(f"Running model for {args.ticks} ticks... ") for tick in tqdm(range(args.ticks)): # 1 vital dynamics (deaths, births, and immigration) population.apply(do_vital_dynamics, tick=tick) @@ -141,7 +149,7 @@ def get_places(num_nodes: np.uint32) -> List[Tuple[str, int, int]]: return places -def load_network(num_nodes: np.uint32, k: np.float32, a: np.float32, b: np.float32, c: np.float32) -> np.ndarray: +def load_network(num_nodes: np.uint32, k: np.float32, a: np.float32, b: np.float32, c: np.float32, max_frac: np.float32) -> np.ndarray: """Load network data for England and Wales.""" # Create a list of indices for the selected places places = get_places(num_nodes) @@ -164,6 +172,10 @@ def load_network(num_nodes: np.uint32, k: np.float32, a: np.float32, b: np.float network[i, j] = k * (popi**a) * (popj**b) / (N * (distm**c)) network[j, i] = network[i, j] + outflows = network.sum(axis=0) + if (maximum := outflows.max()) > 0: + network *= max_frac / maximum + return network @@ -271,25 +283,27 @@ def update_contagion(p: Population, c: Community, i: int) -> None: def do_transmission(p: Population, c: Community, i: int, beta: np.float32, exp_mean: np.float32, exp_std: np.float32) -> None: """Do the transmission.""" - susceptibility = c.susceptible.susceptibility - etimers = c.susceptible.etimer - isusceptible = c.gmap["susceptible"] - iexposed = c.gmap["exposed"] - # TODO - iterate over groups to get N (i.e., if we change the groups, this code breaks) - N = len(c.susceptible) + len(c.exposed) + len(c.infectious) + len(c.recovered) - force = beta * p.contagion[i] * len(c.susceptible) / N - num_exposures = np.uint32(np.round(np.random.poisson(force))) - if num_exposures >= len(c.susceptible): - # raise ValueError(f"Too many exposures: {num_exposures} >= {len(c.susceptible)}") - print(f"Too many exposures: {num_exposures} >= {len(c.susceptible)}") - num_exposures = len(c.susceptible) - for _ in range(min(num_exposures, limit := len(c.susceptible))): - target = np.random.randint(limit) - if np.random.uniform() < susceptibility[target]: - susceptibility[target] = 0 - etimers[target] = max(1, int(np.round(np.random.normal(exp_mean, exp_std)))) - c.move(isusceptible, target, iexposed) - limit -= 1 + if (contagion := p.contagion[i]) > 0: + susceptibility = c.susceptible.susceptibility + etimers = c.susceptible.etimer + isusceptible = c.gmap["susceptible"] + iexposed = c.gmap["exposed"] + # TODO - iterate over groups to get N (i.e., if we change the groups, this code breaks) + N = len(c.susceptible) + len(c.exposed) + len(c.infectious) + len(c.recovered) + force = beta * contagion * len(c.susceptible) / N + num_exposures = np.uint32(np.random.poisson(force)) + if num_exposures > 0: + if num_exposures >= len(c.susceptible): + # raise ValueError(f"Too many exposures: {num_exposures} >= {len(c.susceptible)}") + print(f"Too many exposures: {num_exposures} >= {len(c.susceptible)}") + num_exposures = len(c.susceptible) + for _ in range(min(num_exposures, limit := len(c.susceptible))): + target = np.random.randint(limit) + if np.random.uniform() < susceptibility[target]: + susceptibility[target] = 0 + etimers[target] = max(1, int(np.round(np.random.normal(exp_mean, exp_std)))) + c.move(isusceptible, target, iexposed) + limit -= 1 return @@ -320,12 +334,13 @@ def parse_args() -> Namespace: TICKS = np.uint32(10 * 365) # 10 years NODES = np.uint32(32) # top 32 places by population - EXP_MEAN = np.float32(4) # 4 days - EXP_STD = np.float32(1) # 1 day - INF_MEAN = np.float32(5) # 5 days + EXP_MEAN = np.float32(7) # 7 days + EXP_STD = np.float32(2) # 2 days + INF_MEAN = np.float32(8) # 8 days INF_STD = np.float32(1) # 1 day - R_NAUGHT = np.float32(2.5) # R0 + R_NAUGHT = np.float32(10.0) # R0 SEED = np.uint32(20240227) # random seed + SEED = np.uint32(datetime.now().microsecond) # noqa: DTZ005 parser = ArgumentParser() parser.add_argument("-t", "--ticks", type=np.uint32, default=TICKS) @@ -341,11 +356,13 @@ def parse_args() -> Namespace: DEF_A = np.float32(1.0) DEF_B = np.float32(1.0) DEF_C = np.float32(2.0) + DEF_MAX_FRAC = np.float32(0.1) parser.add_argument("--g_k", type=np.float32, default=DEF_K) parser.add_argument("--g_a", type=np.float32, default=DEF_A) parser.add_argument("--g_b", type=np.float32, default=DEF_B) parser.add_argument("--g_c", type=np.float32, default=DEF_C) + parser.add_argument("--g_max_frac", type=np.float32, default=DEF_MAX_FRAC) args = parser.parse_args() args.__setattr__("beta", np.float32(args.r_naught / args.inf_mean))