def test_log_medium_issue():
logger = mock.MagicMock()
issues = []
log_issue(logger, "hi", IssueSeverity.MEDIUM, issues)
assert issues == [standard_api.Issue(description="hi", severity=5)]
logger.warning.assert_called_once_with("hi")
def test_log_low_issue():
logger = mock.MagicMock()
issues = []
log_issue(logger, "hi", IssueSeverity.LOW, issues)
assert issues == [standard_api.Issue(description="hi", severity=1)]
logger.info.assert_called_once_with("hi")
def test_log_high_issue():
logger = mock.MagicMock()
issues = []
log_issue(logger, "hi", IssueSeverity.HIGH, issues)
assert issues == [standard_api.Issue(description="hi", severity=10)]
logger.error.assert_called_once_with("hi")
def distributeContactsOverAges(
nodeState: Dict[Tuple[Age, Compartment], float],
newContacts: float,
stochastic: bool,
random_state: Optional[np.random.Generator],
issues: Optional[List[standard_api.Issue]] = None,
) -> Dict[Age, float]:
r"""
Distribute the number of new contacts across a region. There are two possible implementations:
1. Deterministic (``stochastic=False``) -- We simply allocate newContacts by proportion of people in each age groups
2. Stochastic (``stochastic=True``) -- newContacts are sampled randomly from the population from a Multinomial
distribution, with: :math:`k=\text{new contacts}`, :math:`p=\text{proportion of people in age groups}`
Note: fractional people will come out of this.
:param nodeState: The disease status of the population stratified by age.
:param newContacts: The number of new contacts to be distributed across age ranges.
:param stochastic: Whether to run the model in a stochastic or deterministic mode
:param random_state: Random number generator used for the model
:param issues: if any issues are found and a list is passed as this parameter, append the issues to it. Logging will
happen regardless.
:return: The number of new infections in each age group.
"""
ageToSus = {}
totalSus = 0.0
for age in getAges(nodeState):
sus = getSusceptibles(age, nodeState)
ageToSus[age] = sus
totalSus += sus
if totalSus < newContacts:
log_issue(
logger,
f"totalSus < incoming contacts ({totalSus} < {newContacts}) - adjusting to totalSus",
IssueSeverity.HIGH,
issues if issues is not None else [],
)
newContacts = totalSus
if totalSus > 0:
if stochastic:
newInfectionsByAge = _distributeContactsOverAgesStochastic(
ageToSus, totalSus, newContacts, random_state)
else:
newInfectionsByAge = _distributeContactsOverAgesDeterministic(
ageToSus, totalSus, newContacts)
else:
newInfectionsByAge = {age: 0 for age in ageToSus}
return newInfectionsByAge
def createNetworkOfPopulation(
compartment_transition_table: pd.DataFrame,
population_table: pd.DataFrame,
commutes_table: pd.DataFrame,
mixing_matrix_table: pd.DataFrame,
infectious_states: pd.DataFrame,
infection_prob: pd.DataFrame,
initial_infections: pd.DataFrame,
trials: pd.DataFrame,
start_end_date: pd.DataFrame,
movement_multipliers_table: pd.DataFrame = None,
stochastic_mode: pd.DataFrame = None,
) -> Tuple[NetworkOfPopulation, List[standard_api.Issue]]:
"""Create the network of the population, loading data from files.
:param compartment_transition_table: pd.Dataframe specifying the transition rates between infected compartments.
:param population_table: pd.Dataframe with the population size in each region by gender and age.
:param commutes_table: pd.Dataframe with the movements between regions.
:param mixing_matrix_table: pd.Dataframe with the age infection matrix.
:param infectious_states: States that are considered infectious
:param infection_prob: Probability that a given contact will result in an infection
:param initial_infections: Initial infections of the population at time 0
:param trials: Number of trials for the model
:param start_end_date: Starting and ending dates of the model
:param movement_multipliers_table: pd.Dataframe with the movement multipliers. This may be None, in
which case no multipliers are applied to the movements.
:param stochastic_mode: Use stochastic mode for the model
:return: A tuple with the constructed network and a list of any issues found. Although it may look convenient,
storing this list inside of NetworkOfPopulation is not a good idea, as that may give the functions that
receive it the false impression they can just append new issues there
"""
issues: List[standard_api.Issue] = []
infection_prob = loaders.readInfectionProbability(infection_prob)
infectious_states = loaders.readInfectiousStates(infectious_states)
initial_infections = loaders.readInitialInfections(initial_infections)
trials = loaders.readTrials(trials)
start_date, end_date = loaders.readStartEndDate(start_end_date)
stochastic_mode = loaders.readStochasticMode(stochastic_mode)
# diseases progression matrix
progression = loaders.readCompartmentRatesByAge(
compartment_transition_table)
# population census data
population = loaders.readPopulationAgeStructured(population_table)
# Check some requirements for this particular model to work with the progression matrix
all_states = set()
for states in progression.values():
assert SUSCEPTIBLE_STATE not in states, "progression from susceptible state is not allowed"
for state, nextStates in states.items():
for nextState in nextStates:
all_states.add(state)
all_states.add(nextState)
assert state == nextState or nextState != EXPOSED_STATE, \
"progression into exposed state is not allowed other than in self reference"
assert (set(infectious_states) - all_states) == set(), \
f"mismatched infectious states and states {infectious_states} {all_states}"
# people movement's graph
graph = loaders.genGraphFromContactFile(commutes_table)
# movement multipliers (dampening or heightening)
if movement_multipliers_table is not None:
movementMultipliers = loaders.readMovementMultipliers(
movement_multipliers_table)
else:
movementMultipliers = {}
# age-based infection matrix
mixingMatrix = loaders.MixingMatrix(mixing_matrix_table)
agesInInfectionMatrix = set(mixingMatrix)
for age in mixingMatrix:
assert agesInInfectionMatrix == set(
mixingMatrix[age]), "infection matrix columns/rows mismatch"
# Checks across datasets
assert agesInInfectionMatrix == set(
progression.keys()), "infection matrix and progression ages mismatch"
assert agesInInfectionMatrix == {age for region in population.values() for age in region}, \
"infection matrix and population ages mismatch"
disconnected_nodes = set(population.keys()) - set(graph.nodes())
if disconnected_nodes:
log_issue(
logger,
f"These nodes have no contacts in the current network: {','.join(disconnected_nodes)}",
IssueSeverity.MEDIUM,
issues,
)
state0: Dict[str, Dict[Tuple[str, str], float]] = {}
for node in list(graph.nodes()):
region = state0.setdefault(node, {})
for age, compartments in progression.items():
if node not in population:
log_issue(
logger,
f"Node {node} is not in the population table, assuming population of 0 for all ages",
IssueSeverity.MEDIUM,
issues,
)
pop = 0.0
else:
pop = population[node][age]
region[(age, SUSCEPTIBLE_STATE)] = pop
for compartment in compartments:
region[(age, compartment)] = 0
logger.info("Nodes: %s, Ages: %s, States: %s", len(state0),
agesInInfectionMatrix, all_states)
nop = NetworkOfPopulation(
progression=progression,
graph=graph,
initialState=state0,
mixingMatrix=mixingMatrix,
movementMultipliers=movementMultipliers,
infectiousStates=infectious_states,
infectionProb=infection_prob,
initialInfections=initial_infections,
trials=trials,
startDate=start_date,
endDate=end_date,
stochastic=stochastic_mode,
)
return (nop, issues)
def main(argv):
"""
Main function to run the network of populations simulation
"""
t0 = time.time()
args = build_args(argv)
setup_logger(args)
logger.info("Parameters\n%s",
"\n".join(f"\t{key}={value}"
for key, value in args._get_kwargs())) # pylint: disable=protected-access
issues: List[standard_api.Issue] = []
info = common.get_repo_info()
if not info.git_sha:
log_issue(
logger,
"Not running from a git repo, so no git_sha associated with the run",
IssueSeverity.HIGH,
issues,
)
elif info.is_dirty:
log_issue(logger, "Running out of a dirty git repo",
IssueSeverity.HIGH, issues)
with standard_api.StandardAPI.from_config(args.data_pipeline_config,
uri=info.uri,
git_sha=info.git_sha) as store:
network, new_issues = ss.createNetworkOfPopulation(
store.read_table("human/compartment-transition",
"compartment-transition"),
store.read_table("human/population", "population"),
store.read_table("human/commutes", "commutes"),
store.read_table("human/mixing-matrix", "mixing-matrix"),
store.read_table("human/infectious-compartments",
"infectious-compartments"),
store.read_table("human/infection-probability",
"infection-probability"),
store.read_table("human/initial-infections", "initial-infections"),
store.read_table("human/trials", "trials"),
store.read_table("human/start-end-date", "start-end-date"),
store.read_table("human/movement-multipliers",
"movement-multipliers")
if args.use_movement_multipliers else None,
store.read_table("human/stochastic-mode", "stochastic-mode"),
)
issues.extend(new_issues)
random_seed = loaders.readRandomSeed(
store.read_table("human/random-seed", "random-seed"))
results = runSimulation(
network,
random_seed,
issues=issues,
max_workers=None if not args.workers else args.workers)
aggregated = aggregateResults(results)
logger.info("Writing output")
store.write_table(
"output/simple_network_sim/outbreak-timeseries",
"outbreak-timeseries",
_convert_category_to_str(aggregated.output),
issues=issues,
description=aggregated.description,
)
for i, result in enumerate(results):
store.write_table(
"output/simple_network_sim/outbreak-timeseries",
f"run-{i}",
_convert_category_to_str(result.output),
issues=result.issues,
description=result.description,
)
logger.info("Took %.2fs to run the simulation.", time.time() - t0)
logger.info(
"Use `python -m simple_network_sim.network_of_populations.visualisation -h` to find out how to take "
"a peak what you just ran. You will need use the access-<hash>.yaml file that was created by this run."
)