def test_samples(doplot=False):
sc.heading('Samples distribution')
n = 10000
nbins = 40
# Warning, must match utils.py!
choices = [
'uniform',
'normal',
'lognormal',
'normal_pos',
'normal_int',
'lognormal_int',
'neg_binomial'
]
if doplot:
pl.figure(figsize=(20,14))
# Run the samples
nchoices = len(choices)
nsqr = np.ceil(np.sqrt(nchoices))
results = {}
for c,choice in enumerate(choices):
if choice == 'neg_binomial':
par1 = 10
par2 = 0.5
elif choice in ['lognormal', 'lognormal_int']:
par1 = 1
par2 = 0.5
else:
par1 = 0
par2 = 5
results[choice] = cova.sample(dist=choice, par1=par1, par2=par2, size=n)
if doplot:
pl.subplot(nsqr, nsqr, c+1)
pl.hist(x=results[choice], bins=nbins)
pl.title(f'dist={choice}, par1={par1}, par2={par2}')
with pytest.raises(NotImplementedError):
cova.sample(dist='not_found')
return results
def run(self,
seed_infections=1,
verbose=None,
calc_likelihood=False,
do_plot=False,
**kwargs):
''' Run the simulation '''
T = sc.tic()
# Reset settings and results
if verbose is None:
verbose = self['verbose']
self.init_results()
self.init_people(
seed_infections=seed_infections) # Actually create the people
daily_tests = self.data[
'new_tests'] # Number of tests each day, from the data
evacuated = self.data['evacuated'] # Number of people evacuated
# Main simulation loop
for t in range(self.npts):
# Print progress
if verbose >= 1:
string = f' Running day {t:0.0f} of {self["n_days"]}...'
if verbose >= 2:
sc.heading(string)
else:
print(string)
test_probs = {
} # Store the probability of each person getting tested
# Update each person
for person in self.people.values():
# Count susceptibles
if person.susceptible:
self.results['n_susceptible'][t] += 1
continue # Don't bother with the rest of the loop
# Handle testing probability
if person.infectious:
test_probs[person.uid] = self[
'symptomatic'] # They're infectious: high probability of testing
else:
test_probs[person.uid] = 1.0
# If exposed, check if the person becomes infectious
if person.exposed:
self.results['n_exposed'][t] += 1
if not person.infectious and t >= person.date_infectious: # It's the day they become infectious
person.infectious = True
if verbose >= 2:
print(
f' Person {person.uid} became infectious!'
)
# If infectious, check if anyone gets infected
if person.infectious:
# First, check for recovery
if person.date_recovered and t >= person.date_recovered: # It's the day they become infectious
person.exposed = False
person.infectious = False
person.recovered = True
self.results['recoveries'][t] += 1
else:
self.results['n_infectious'][
t] += 1 # Count this person as infectious
n_contacts = cv.pt(
person.contacts
) # Draw the number of Poisson contacts for this person
contact_inds = cv.choose_people(
max_ind=len(self.people),
n=n_contacts) # Choose people at random
for contact_ind in contact_inds:
exposure = cv.bt(self['r_contact']
) # Check for exposure per person
if exposure:
target_person = self.people[contact_ind]
if target_person.susceptible: # Skip people who are not susceptible
self.results['infections'][t] += 1
target_person.susceptible = False
target_person.exposed = True
target_person.date_exposed = t
incub_pars = dict(dist='normal_int',
par1=self['incub'],
par2=self['incub_std'])
dur_pars = dict(dist='normal_int',
par1=self['dur'],
par2=self['dur_std'])
incub_dist = cv.sample(**incub_pars)
dur_dist = cv.sample(**dur_pars)
target_person.date_infectious = t + incub_dist
target_person.date_recovered = target_person.date_infectious + dur_dist
if verbose >= 2:
print(
f' Person {person.uid} infected person {target_person.uid}!'
)
# Count people who recovered
if person.recovered:
self.results['n_recovered'][t] += 1
# Implement testing -- this is outside of the loop over people, but inside the loop over time
if t < len(
daily_tests
): # Don't know how long the data is, ensure we don't go past the end
n_tests = daily_tests.iloc[t] # Number of tests for this day
if n_tests and not pl.isnan(
n_tests): # There are tests this day
self.results['tests'][
t] = n_tests # Store the number of tests
test_probs = pl.array(list(test_probs.values()))
test_probs /= test_probs.sum()
test_inds = cv.choose_people_weighted(probs=test_probs,
n=n_tests)
uids_to_pop = []
for test_ind in test_inds:
tested_person = self.people[test_ind]
if tested_person.infectious and cv.bt(
self['sensitivity']
): # Person was tested and is true-positive
self.results['diagnoses'][t] += 1
tested_person.diagnosed = True
if self['evac_positives']:
uids_to_pop.append(tested_person.uid)
if verbose >= 2:
print(
f' Person {person.uid} was diagnosed!'
)
for uid in uids_to_pop: # Remove people from the ship once they're diagnosed
self.off_ship[uid] = self.people.pop(uid)
# Implement quarantine
if t == self['quarantine']:
if verbose >= 1:
print(f'Implementing quarantine on day {t}...')
for person in self.people.values():
if 'quarantine_eff' in self.pars.keys():
quarantine_eff = self['quarantine_eff'] # Both
else:
if person.crew:
quarantine_eff = self['quarantine_eff_c'] # Crew
else:
quarantine_eff = self['quarantine_eff_g'] # Guests
person.contacts *= quarantine_eff
# Implement testing change
if t == self['testing_change']:
if verbose >= 1:
print(f'Implementing testing change on day {t}...')
self['symptomatic'] *= self[
'testing_symptoms'] # Reduce the proportion of symptomatic testing
# Implement evacuations
if t < len(evacuated):
n_evacuated = evacuated.iloc[
t] # Number of evacuees for this day
if n_evacuated and not pl.isnan(
n_evacuated
): # There are evacuees this day # TODO -- refactor with n_tests
if verbose >= 1:
print(f'Implementing evacuation on day {t}')
evac_inds = cv.choose_people(max_ind=len(self.people),
n=n_evacuated)
uids_to_pop = []
for evac_ind in evac_inds:
evac_person = self.people[evac_ind]
if evac_person.infectious and cv.bt(
self['sensitivity']):
self.results['evac_diagnoses'][t] += 1
uids_to_pop.append(evac_person.uid)
for uid in uids_to_pop: # Remove people from the ship once they're diagnosed
self.off_ship[uid] = self.people.pop(uid)
# Compute cumulative results
self.results['cum_exposed'] = pl.cumsum(self.results['infections'])
self.results['cum_tested'] = pl.cumsum(self.results['tests'])
self.results['cum_diagnosed'] = pl.cumsum(self.results['diagnoses'])
# Compute likelihood
if calc_likelihood:
self.likelihood()
# Tidy up
self.results['ready'] = True
elapsed = sc.toc(T, output=True)
if verbose >= 1:
print(f'\nRun finished after {elapsed:0.1f} s.\n')
summary = self.summary_stats()
print(f"""Summary:
{summary['n_susceptible']:5.0f} susceptible
{summary['n_exposed']:5.0f} exposed
{summary['n_infectious']:5.0f} infectious
""")
if do_plot:
self.plot(**kwargs)
return self.results
def test_samples(do_plot=False, verbose=True):
sc.heading('Samples distribution')
n = 200_000
nbins = 100
# Warning, must match utils.py!
choices = [
'uniform', 'normal', 'lognormal', 'normal_pos', 'normal_int',
'lognormal_int', 'poisson', 'neg_binomial'
]
if do_plot:
pl.figure(figsize=(20, 14))
# Run the samples
nchoices = len(choices)
nsqr, _ = sc.get_rows_cols(nchoices)
results = sc.objdict()
mean = 11
std = 7
low = 3
high = 9
normal_dists = [
'normal', 'normal_pos', 'normal_int', 'lognormal', 'lognormal_int'
]
for c, choice in enumerate(choices):
kw = {}
if choice in normal_dists:
par1 = mean
par2 = std
elif choice == 'neg_binomial':
par1 = mean
par2 = 1.2
kw['step'] = 0.1
elif choice == 'poisson':
par1 = mean
par2 = 0
elif choice == 'uniform':
par1 = low
par2 = high
else:
errormsg = f'Choice "{choice}" not implemented'
raise NotImplementedError(errormsg)
# Compute
results[choice] = cv.sample(dist=choice,
par1=par1,
par2=par2,
size=n,
**kw)
# Optionally plot
if do_plot:
pl.subplot(nsqr, nsqr, c + 1)
plotbins = np.unique(results[choice]) if (
choice == 'poisson' or '_int' in choice) else nbins
pl.hist(x=results[choice], bins=plotbins, width=0.8)
pl.title(f'dist={choice}, par1={par1}, par2={par2}')
with pytest.raises(NotImplementedError):
cv.sample(dist='not_found')
# Do statistical tests
tol = 1 / np.sqrt(
n / 50 / len(choices)
) # Define acceptable tolerance -- broad to avoid false positives
def isclose(choice, tol=tol, **kwargs):
key = list(kwargs.keys())[0]
ref = list(kwargs.values())[0]
npfunc = getattr(np, key)
value = npfunc(results[choice])
msg = f'Test for {choice:14s}: expecting {key:4s} = {ref:8.4f} ± {tol*ref:8.4f} and got {value:8.4f}'
if verbose:
print(msg)
assert np.isclose(value, ref, rtol=tol), msg
return True
# Normal
for choice in normal_dists:
isclose(choice, mean=mean)
if all([k not in choice
for k in ['_pos', '_int']]): # These change the variance
isclose(choice, std=std)
# Negative binomial
isclose('neg_binomial', mean=mean)
# Poisson
isclose('poisson', mean=mean)
isclose('poisson', var=mean)
# Uniform
isclose('uniform', mean=(low + high) / 2)
return results
def run(self,
seed_infections=1,
verbose=None,
calc_likelihood=False,
do_plot=False,
**kwargs):
''' Run the simulation '''
T = sc.tic()
# Reset settings and results
if verbose is None:
verbose = self['verbose']
self.init_results()
self.init_people(
seed_infections=seed_infections) # Actually create the people
# Main simulation loop
for t in range(self.npts):
# Print progress
if verbose >= 1:
string = f' Running day {t:0.0f} of {self["n_days"]}...'
if verbose >= 2:
sc.heading(string)
else:
print(string)
test_probs = {
} # Store the probability of each person getting tested
# Update each person
for person in self.people.values():
# Count susceptibles
if person.susceptible:
self.results['n_susceptible'][t] += 1
continue # Don't bother with the rest of the loop
# Handle testing probability
if person.infectious:
test_probs[person.uid] = self[
'symptomatic'] # They're infectious: high probability of testing
else:
test_probs[person.uid] = 1.0
# If exposed, check if the person becomes infectious
if person.exposed:
self.results['n_exposed'][t] += 1
if not person.infectious and t >= person.date_infectious: # It's the day they become infectious
person.infectious = True
if verbose >= 2:
print(
f' Person {person.uid} became infectious!'
)
if t == person.date_infectious:
# print(type(person.uid))
if int(person.uid) != 0:
# give 50% chance of quaratine?
q_per = .5
quar = np.random.choice([True, False],
p=[q_per, 1 - q_per])
#person.quarantine = np.copy(quar)
person.quarantine = False
else:
person.quarantine = False
# If infectious, check if anyone gets infected
if person.infectious:
# print(person)
# First, check for recovery
if person.date_recovered and t >= person.date_recovered: # It's the day they become infectious
person.exposed = False
person.infectious = False
person.recovered = True
self.results['recoveries'][t] += 1
else:
self.results['n_infectious'][
t] += 1 # Count this person as infectious
cnum = 11
topull = np.random.poisson(cnum)
s1_cons = person.contacts['s1']
s2_cons = person.contacts['s2']
s3_cons = person.contacts['s3']
# Draw the numbergit of Poisson contacts for this person
# this should instead initiate the random contacts?
rng = np.random.default_rng()
shuffled = np.array(range(len(self.people)))
rng.shuffle(shuffled)
c_contacts = np.array(
shuffled)[:topull] # Choose people at random
contypes = np.concatenate(
(["s1" for _ in range(len(s2_cons))
], ["s2" for _ in range(len(s1_cons))
], ["s3" for _ in range(len(s3_cons))],
["c" for _ in range(topull)]))
#contact_inds = np.concatenate((s1_cons, s2_cons, s3_cons, c_contacts))
# uncomment above to add back community contacts
contact_inds = np.concatenate(
(s1_cons, s2_cons, s3_cons))
z = -1
for contact_ind in contact_inds:
if person.quarantine == False:
exposure = bt(
self['r_contact']
) # Check for exposure per person
z += 1
if exposure:
target_person = self.people[contact_ind]
if target_person.susceptible: # Skip people who are not susceptible
self.results['infections'][t] += 1
target_person.susceptible = False
target_person.exposed = True
target_person.exposed_type = contypes[
z]
target_person.date_exposed = t
incub_pars = dict(
dist='normal_int',
par1=self['incub'],
par2=self['incub_std'])
dur_pars = dict(dist='normal_int',
par1=self['dur'],
par2=self['dur_std'])
incub_dist = cv.sample(**incub_pars)
dur_dist = cv.sample(**dur_pars)
target_person.date_infectious = t + incub_dist
target_person.date_recovered = target_person.date_infectious + dur_dist
if verbose >= 2:
print(
f' Person {person.uid} infected person {target_person.uid}!'
)
# Count people who recovered
if person.recovered:
self.results['n_recovered'][t] += 1