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
def get_individual_states(sim, order=True):
people = sim.people
# if order:
# print('Note: ordering of people for the animation is currently not supported')
# people = sorted(people, key=lambda x: x.date_exposed if x.date_exposed is not None else np.inf)
# Order these in order of precedence
# The last matching quantity will be used
states = [
{
'name': 'Healthy',
'quantity': None,
'color': '#b9d58a',
'value': 0
},
{
'name': 'Exposed',
'quantity': 'date_exposed',
'color': '#e37c30',
'value': 2
},
{
'name': 'Infectious',
'quantity': 'date_infectious',
'color': '#c35d86',
'value': 3
},
{
'name': 'Recovered',
'quantity': 'date_recovered',
'color': '#799956',
'value': 4
},
{
'name': 'Dead',
'quantity': 'date_dead',
'color': '#000000',
'value': 5
},
]
z = np.zeros((len(people), sim.npts))
for state in states:
date = state['quantity']
if date is not None:
inds = cv.defined(people[date])
for ind in inds:
z[ind, int(people[date][ind]):] = state['value']
return z, states
def get_individual_states(sim):
''' Helper function to convert people into integers '''
people = sim.people
states = [
{
'name': 'Healthy',
'quantity': None,
'color': '#b9d58a',
'value': 0
},
{
'name': 'Exposed',
'quantity': 'date_exposed',
'color': '#e37c30',
'value': 2
},
{
'name': 'Infectious',
'quantity': 'date_infectious',
'color': '#c35d86',
'value': 3
},
{
'name': 'Recovered',
'quantity': 'date_recovered',
'color': '#799956',
'value': 4
},
{
'name': 'Dead',
'quantity': 'date_dead',
'color': '#000000',
'value': 5
},
]
z = np.zeros((len(people), sim.npts))
for state in states:
date = state['quantity']
if date is not None:
inds = cv.defined(people[date])
for ind in inds:
z[ind, int(people[date][ind]):] = state['value']
return z, states
pl.fill_between([dd, dd+14], [ii-qdy]*2, [ii+qdy]*2, facecolor=diag_color, zorder=ζ['↓']-1, alpha=α)
sc.boxoff()
a[k].spines['left'].set_visible(False)
if k != 'none':
a[k].set_yticks([])
else:
a[k].set_yticks(np.arange(0,p.pop_size+1,10))
a[k].set_ylabel('Person')
pl.xlabel('Days since seed infection')
ylimmap = dict(none=1.16, test=1.16, trace=1.16)
pl.ylim([-1, p.pop_size*ylimmap[k]])
# Labels
n_inf = len(tt.infection_log)
n_diag_dir = len(np.intersect1d(cv.defined(s[k].people.date_diagnosed), cv.undefined(s[k].people.date_quarantined)))
n_diag_quar = len(np.intersect1d(cv.defined(s[k].people.date_diagnosed), cv.defined(s[k].people.date_quarantined)))
n_quar = len(np.intersect1d(cv.defined(s[k].people.date_quarantined), cv.true(s[k].people.susceptible)))
txtargs = dict(horizontalalignment='right')
ps = p.pop_size
ty1, ty2, ty25, ty3, ty4, ty5 = ps*1.02, ps*1.05, ps*1.065, ps*1.08, ps*1.11, ps*1.14
δy = 0.8
xl = pl.xlim()
tx0map = {
'none': xl[1]*0.6,
'test': xl[1]*0.6,
'trace': xl[1]*0.8,
}