def _distributeContactsOverAgesStochastic(
ageToSusceptibles: Dict[Age, float], totalSusceptibles: float,
newContacts: float,
random_state: np.random.Generator) -> Dict[Age, float]:
"""
Distribute the number of new contacts across a region using the stochastic algorithm.
:param ageToSusceptibles: The number of susceptibles in each age group
:param totalSusceptibles: Total number of susceptible people in all age groups
:param newContacts: The number of new contacts to be distributed across age ranges.
:param random_state: Random number generator used for the model
:return: The number of new infections in each age group.
"""
assert isinstance(newContacts, int) or newContacts.is_integer()
ageProbabilities = np.array(list(
ageToSusceptibles.values())) / totalSusceptibles
allocationsByAge = random_state.multinomial(newContacts, ageProbabilities)
return dict(zip(ageToSusceptibles.keys(), allocationsByAge))
def _internalStateDiseaseUpdate(age: Age, state: Compartment, people: float,
outTransitions: Dict[Compartment, float],
newStates: Dict[Tuple[Age, Compartment],
float], stochastic: bool,
random_state: np.random.Generator):
"""
Returns the status of exposed individuals, moving them into the next disease state with a
probability defined in the given progression matrix.
:param age: Current age
:param age: Current state
:param people: Number of people in current age, state
:param outTransitions: Transition probabilities out of the current age, state,
to next potential ones
:param newStates: Distribution of people in next states
:param stochastic: Use the model in stochastic mode?
:param random_state: Random number generator used for the model
:return: the numbers in each exposed state stratified by age
"""
if stochastic:
if state == SUSCEPTIBLE_STATE:
return
assert isinstance(people, int) or people.is_integer()
outRepartitions = random_state.multinomial(
people, list(outTransitions.values()))
outRepartitions = dict(zip(outTransitions.keys(), outRepartitions))
for nextState, nextStateCases in outRepartitions.items():
newStates.setdefault((age, nextState), 0.0)
newStates[(age, nextState)] += nextStateCases
else:
for nextState, transition in outTransitions.items():
newStates.setdefault((age, nextState), 0.0)
newStates[(age, nextState)] += transition * people