def test_indexing():
# Definitions
farr = np.array([1.5, 0, 0, 1, 1, 0]) # Float array
barr = np.array(farr, dtype=bool) # Boolean array
darr = np.array([0, np.nan, 1, np.nan, 0,
np.nan]) # Defined/undefined array
inds = np.array([0, 10, 20, 30, 40, 50]) # Indices
inds2 = np.array([1, 2, 3, 4]) # Skip first and last index
# Test true, false, defined, and undefined
assert cv.true(farr).tolist() == [0, 3, 4]
assert cv.false(farr).tolist() == [1, 2, 5]
assert cv.defined(darr).tolist() == [0, 2, 4]
assert cv.undefined(darr).tolist() == [1, 3, 5]
# Test with indexing
assert cv.itrue(barr, inds).tolist() == [0, 30, 40]
assert cv.ifalse(barr, inds).tolist() == [10, 20, 50]
assert cv.idefined(darr, inds).tolist() == [0, 20, 40]
assert cv.iundefined(darr, inds).tolist() == [10, 30, 50]
# Test with double indexing
assert cv.itruei(barr, inds2).tolist() == [3, 4]
assert cv.ifalsei(barr, inds2).tolist() == [1, 2]
assert cv.idefinedi(darr, inds2).tolist() == [2, 4]
assert cv.iundefinedi(darr, inds2).tolist() == [1, 3]
return
import sciris as sc
import covasim as cv
states = [
'susceptible',
'naive',
'exposed',
'infectious',
'symptomatic',
'severe',
'critical',
'tested',
'diagnosed',
'recovered',
'dead',
]
sim = cv.Sim()
sim.run()
d = sc.objdict()
for state in states:
n_in = len(cv.true(sim.people[state]))
n_out = len(cv.false(sim.people[state]))
d[state] = n_in
assert n_in + n_out == sim['pop_size']
print(sim.summary)
print(d)
assert d.naive + d.exposed + d.recovered + d.dead == sim['pop_size']
def test_states():
''' Test state consistency against state_diagram.xlsx '''
filename = 'state_diagram.xlsx'
sheets = ['Without waning', 'With waning']
indexcol = 'In ↓ you can be →'
# Load state diagram
dfs = sc.odict()
for sheet in sheets:
dfs[sheet] = pd.read_excel(filename, sheet_name=sheet)
dfs[sheet] = dfs[sheet].set_index(indexcol)
# Create and run simulation
for use_waning in [False, True]:
sc.heading(f'Testing state consistency with waning = {use_waning}')
df = dfs[use_waning] # Different states are possible with or without waning
# Parameters chosen to be midway through the sim so as few states as possible are empty
pars = dict(
pop_size = 1e3,
pop_infected = 20,
n_days = 70,
use_waning = use_waning,
verbose = 0,
interventions = [
cv.test_prob(symp_prob=0.4, asymp_prob=0.01),
cv.contact_tracing(trace_probs=0.1),
cv.simple_vaccine(days=60, prob=0.1)
]
)
sim = cv.Sim(pars).run()
ppl = sim.people
# Check states
errormsg = ''
states = df.columns.values.tolist()
for s1 in states:
for s2 in states:
if s1 != s2:
relation = df.loc[s1, s2] # e.g. df.loc['susceptible', 'exposed']
print(f'Checking {s1:13s} → {s2:13s} = {relation:2n} ... ', end='')
inds = cv.true(ppl[s1])
n_inds = len(inds)
vals2 = ppl[s2][inds]
is_true = cv.true(vals2)
is_false = cv.false(vals2)
n_true = len(is_true)
n_false = len(is_false)
if relation == 1 and n_true != n_inds:
errormsg = f'Being {s1}=True implies {s2}=True, but only {n_true}/{n_inds} people are'
print(f'× {n_true}/{n_inds} error!')
elif relation == -1 and n_false != n_inds:
errormsg = f'Being {s1}=True implies {s2}=False, but only {n_false}/{n_inds} people are'
print(f'× {n_false}/{n_inds} error!')
else:
print(f'✓ {n_true}/{n_inds}')
if errormsg:
raise RuntimeError(errormsg)
return
def plot():
fig = pl.figure(num='Fig. 3: Theoretical TTQ', figsize=(24,14))
euclid = False
ay0 = 0.06
adx = 0.2
ady = 0.91
axa, axb, axc = 0.04, 0.27, 0.49
a['none'] = pl.axes([axa, ay0, adx, ady])
a['test'] = pl.axes([axb, ay0, adx, ady])
a['trace'] = pl.axes([axc, ay0, adx, ady])
xoff = 0.02
yoff = 0.05
pl.figtext(axa-xoff, ady+yoff, 'a', fontsize=40)
pl.figtext(axb-xoff, ady+yoff, 'b', fontsize=40)
pl.figtext(axc-xoff, ady+yoff, 'c', fontsize=40)
max_n = 0; ζ = {'↓':-2, '↑':10} # Configure plot zorder
# Processing trees
for k in ['none', 'test', 'trace']:
tt = t[k]
max_n = max(max_n, len(tt.infection_log))
# Get the history of each chain
hist = [None]*p.pop_size
detailed = tt.detailed.to_dict('records')
for i,entry in enumerate(detailed):
if ~np.isnan(entry['target']):
hist[i] = [i]
source = entry['source']
while ~np.isnan(source):
hist[i].insert(0, source)
source = detailed[int(source)]['source']
# Figure out order
histstrs = []
for h in hist:
if h is not None:
string = ''.join([str(n) for n in h])
histstrs.append(string)
iorder = np.arange(p.pop_size)
inf_inds = cv.false(s[k].people.susceptible)
iorder[inf_inds] = inf_inds[np.argsort(np.argsort(histstrs))] # This actually works 🤔🙈
min_inf_ind = inf_inds.min()
min_ind_ind = sc.findinds(iorder==min_inf_ind)[0]
orig_0_ind = sc.findinds(iorder==0)[0]
iorder[min_ind_ind] = 0
iorder[orig_0_ind] = min_inf_ind
# Initialization
n = p.n_days + 1
frames = [list() for i in range(n)]
quar_color = '#6495ed'
diag_color = '#80d380'
inf_color = '#b22222'
quar_color2 = '#b3ceff'
# Construct each frame of the animation
for ddict in detailed: # Loop over every person
if np.isnan(ddict['target']):
continue # Skip the 'None' node corresponding to seeded infections
frame = {}
target_ind = ddict['target']
if not np.isnan(ddict['date']): # If this person was infected
source_ind = ddict['source'] # Index of the person who infected the target
target_date = ddict['date']
if ~np.isnan(source_ind): # Seed infections and importations won't have a source
source_date = detailed[int(source_ind)]['date']
else:
source_ind = target_ind
source_date = 0
# Construct this frame
frame['x'] = [int(source_date), int(target_date)]
frame['y'] = [int(source_ind), int(target_ind)]
frame['c'] = inf_color
frames[int(target_date)].append(frame)
# Plot
pl.sca(a[k])
for frame in frames:
for entry in frame:
x0 = entry['x'][0]
x1 = entry['x'][1]
y0 = iorder[entry['y'][0]]
y1 = iorder[entry['y'][1]]
color = entry['c']
if euclid:
pl.plot([x0, x1], [y0, y1], lw=3, c=[0.7]*3)
else:
pl.plot([x0, x0, x1], [y0, y1, y1], lw=3, c=[0.7]*3)
pl.scatter([x1], [y1], s=120, zorder=ζ['↑'], c=[color])
for i in range(p.pop_size):
ii = iorder[i]
dq = s[k].people.date_quarantined[i]
dd = s[k].people.date_diagnosed[i]
de = s[k].people.date_exposed[i]
qdy = 0.2
α = 0.9
mark = ['>', 'x'][1]
if not np.isnan(dq) and not np.isnan(de) and np.isnan(dd) and verbose:
print(f'Person {i} is infected and quarantined but not diagnosed, due to presymptomatic period')
if not np.isnan(dq):
pl.plot([de, dq], [ii, ii], c=quar_color2, linestyle='--', zorder=ζ['↓']-4)
pl.fill_between([dq, dq+14], [ii-qdy]*2, [ii+qdy]*2, facecolor=quar_color2, zorder=ζ['↓']-2, alpha=α)
if np.isnan(de):
pl.scatter([dq], [ii], s=120, c='w', edgecolor=quar_color2, lw=2, alpha=1.0, zorder=ζ['↓'])
if not np.isnan(dd):
if not np.isnan(dq):
color = quar_color2
else:
color = diag_color
pl.plot([de, dd], [ii, ii], c=color, linestyle='--', zorder=ζ['↓']-4)
pl.scatter([dd], [ii], marker=mark, linewidth=3, s=120, c=[color], zorder=ζ['↓']+1)
pl.fill_between([dd, dd+14], [ii-qdy]*2, [ii+qdy]*2, facecolor=diag_color, zorder=ζ['↓']-1, alpha=α)