def gravity_model_contact_events(agents: List[Agent],
positions: np.array,
env: simpy.Environment,
rng: np.random.Generator):
tree = KDTree(data=positions)
close_pairs = list(tree.query_pairs(r=contact_distance_upper_bound))
inverse_distances = np.array(
[np.linalg.norm(positions[idx1] - positions[idx2]) ** -contact_rate_gravity_exponent
for idx1, idx2 in close_pairs])
inverse_distances /= inverse_distances.sum()
while True:
choices = rng.choice(a=close_pairs, p=inverse_distances, size=len(agents)).tolist()
for choice in choices:
yield env.timeout(delay=rng.exponential(scale=1 / len(agents) / contact_rate_per_individual))
contact_agents = [agents[idx]
for idx in choice
if not agents[idx].state.symptomatic()
or rng.uniform() > p_symptomatic_individual_isolates]
if len(contact_agents) < 2:
# Symptomatic self-isolation means this is no longer a contact event and doesn't need
# recording. Skip to the next event.
continue
infected = get_infected(contact_agents, rng=rng)
for i in infected:
env.process(generator=infection_events(env=env, infected=i, rng=rng))
def sample_forwards(
fin_t: float,
init_x: Optional[int],
gen: np.ndarray,
ome: np.random.Generator = np.random.default_rng()
) -> Skeleton:
"""
:param fin_t:
:param init_x:
:param gen:
:param ome:
:return:
>>> # generate fixture
>>> dim = 3
>>> fin_t = 1e5
>>> gen = sample_trial_generator(dim)
>>> stat = get_stat(gen)
>>> # test stationary distribution
>>> skel = sample_forwards(fin_t, None, gen)
>>> probs = [np.sum(np.ediff1d(skel.t, to_end=(fin_t - skel.t[-1],))[skel.xt == i]) / fin_t for i in range(dim)]
>>> np.allclose(probs, stat, 1e-2)
True
"""
dim = gen.shape[0]
if init_x is None:
init_x = ome.choice(dim, 1, p=get_stat(gen)).item()
t, x = [0], [init_x]
while t[-1] < fin_t:
hold = ome.exponential(1 / np.delete(gen[x[-1]], x[-1]))
move = np.argmin(hold)
t.append(t[-1] + hold[move])
x.append(move + int(move >= x[-1]))
return Skeleton(np.array(t[:-1]), np.array(x[:-1]), fin_t)
def sampling_fn(rng: np.random.Generator):
return rng.exponential(scale=scale, size=1)