def initialize(self, sim):
# Create schools, stealing 's' edges into the School class instances upon *initialize*
self.school_types = sim.people.school_types # Dict with keys of school types (e.g. 'es') and values of list of school ids (e.g. [1,5])
sdf = sim.people.contacts['s'].to_df()
sim.school_stats = {}
for school_type, scids in self.school_types.items():
for school_id in scids:
uids = sim.people.schools[
school_id] # Dict with keys of school_id and values of uids in that school
if self.scenario[school_type] is not None:
stats = {
'type': school_type,
'scenario': self.scenario[school_type],
}
sim.school_stats[int2key(school_id)] = stats
# Extract 's'-layer associated with this school
rows = (sdf['p1'].isin(uids)) | (sdf['p2'].isin(uids))
s_subset = sdf.loc[rows]
sdf = sdf.loc[~rows] # Remove rows from the 's' contacts
this_school_layer = cv.Layer().from_df(s_subset)
sch = School(sim, school_id, school_type, uids,
this_school_layer,
**self.scenario[school_type])
self.schools.append(sch)
# Configure the new layer
sim['beta_layer'][
sch.sid] = self.scenario[school_type]['beta_s']
sim['iso_factor'][sch.sid] = sim['iso_factor']['s']
sim['quar_factor'][sch.sid] = sim['quar_factor']['s']
# Delete remaining entries in sim.people.contacts['s'], these were associated with schools that will not open, e.g. pk and uv
sim.people.contacts['s'] = cv.Layer()
self.initialized = True
def begin_day(self, date):
''' Called at the beginning of each day to configure the school layer '''
dayname = sc.readdate(date).strftime('%A')
group = self.schedule[dayname]
self.school_day = group == 'all'
# Could modify layer based on group
if group == 'all':
# Start with the original layer, will remove uids at home later
self.layer = sc.dcp(self.base_layer) # needed?
uids = self.uids
else:
self.layer = cv.Layer() # Empty
uids = np.empty(0, dtype='int64')
return uids # Everyone is scheduled for school today, unless it's a weekend
def begin_day(self, date):
''' Called at the beginning of each day to configure the school layer '''
dayname = sc.readdate(date).strftime('%A')
group = self.schedule[dayname]
self.school_day = group in ['A', 'B']
# Could modify layer based on group
if group == 'A':
self.layer = sc.dcp(self.A_base_layer) # needed?
uids = self.A_group
elif group == 'B':
self.layer = sc.dcp(self.B_base_layer) # needed?
uids = self.B_group
else:
uids = np.empty(0, dtype='int64')
self.layer = cv.Layer() # Empty
return uids # Hybrid scheduling
def test_basepeople():
sc.heading('Testing base.py people and contacts...')
# Create a small sim for later use
sim = cv.Sim(pop_size=100, verbose=verbose)
sim.initialize()
# BasePeople methods
ppl = sim.people
ppl.get(['susceptible', 'infectious'])
ppl.keys()
ppl.person_keys()
ppl.state_keys()
ppl.date_keys()
ppl.dur_keys()
ppl.indices()
ppl._resize_arrays(
new_size=200
) # This only resizes the arrays, not actually create new people
ppl._resize_arrays(new_size=100) # Change back
ppl.to_df()
ppl.to_arr()
ppl.person(50)
people = ppl.to_people()
ppl.from_people(people)
ppl.make_edgelist([{'new_key': [0, 1, 2]}])
ppl.brief()
# Contacts methods
contacts = ppl.contacts
df = contacts['a'].to_df()
ppl.remove_duplicates(df)
with pytest.raises(sc.KeyNotFoundError):
contacts['invalid_key']
contacts.values()
len(contacts)
print(contacts)
print(contacts['a'])
# Layer methods
hospitals_layer = cv.Layer()
contacts.add_layer(hospitals=hospitals_layer)
contacts.pop_layer('hospitals')
df = hospitals_layer.to_df()
hospitals_layer.from_df(df)
# Generate an average of 10 contacts for 1000 people
n = 10_000
n_people = 1000
p1 = np.random.randint(n_people, size=n)
p2 = np.random.randint(n_people, size=n)
beta = np.ones(n)
layer = cv.Layer(p1=p1, p2=p2, beta=beta)
# Convert one layer to another with extra columns
index = np.arange(n)
self_conn = p1 == p2
layer2 = cv.Layer(**layer, index=index, self_conn=self_conn)
assert len(layer2) == n
assert len(layer2.keys()) == 5
# Test dynamic layers, plotting, and stories
pars = dict(pop_size=100,
n_days=10,
verbose=verbose,
pop_type='hybrid',
beta=0.02)
s1 = cv.Sim(pars, dynam_layer={'c': 1})
s1.run()
s1.people.plot()
for person in [0, 50]:
s1.people.story(person)
# Run without dynamic layers and assert that the results are different
s2 = cv.Sim(pars, dynam_layer={'c': 0})
s2.run()
assert cv.diff_sims(s1, s2, output=True)
# Create a bare People object
ppl = cv.People(100)
with pytest.raises(sc.KeyNotFoundError): # Need additional parameters
ppl.initialize()
return
# Create the first sim
orig_sim = cv.Sim(pop_type='hybrid',
n_days=120,
label='Default hybrid population')
orig_sim.initialize() # Initialize the population
# Create the second sim
sim = orig_sim.copy()
# Define the new layer, 'transport'
n_people = len(sim.people)
n_contacts_per_person = 0.5
n_contacts = int(n_contacts_per_person * n_people)
contacts_p1 = cv.choose(max_n=n_people, n=n_contacts)
contacts_p2 = cv.choose(max_n=n_people, n=n_contacts)
beta = np.ones(n_contacts)
layer = cv.Layer(p1=contacts_p1, p2=contacts_p2,
beta=beta) # Create the new layer
# Add this layer in and re-initialize the sim
sim.people.contacts.add_layer(transport=layer)
sim.reset_layer_pars(
) # Automatically add layer 'q' to the parameters using default values
sim.initialize() # Reinitialize
sim.label = f'Transport layer with {n_contacts_per_person} contacts/person'
# Run and compare
msim = cv.MultiSim([orig_sim, sim])
msim.run()
msim.plot()
def broken_bubbles(sim, scen, test):
frac_edges_to_rewire = 0.5
np.random.seed(1)
# Modify scen with test
for stype, spec in scen.items():
if spec is not None:
spec[
'testing'] = test # dcp probably not needed because deep copied in new_schools
school_contacts = []
sdf = sim.people.contacts['s'].to_df()
student_flag = np.array(sim.people.student_flag, dtype=bool)
sdf['p1_student'] = student_flag[sdf['p1']]
sdf['p2_student'] = student_flag[sdf['p1']]
school_types = sim.people.school_types
for school_type, scids in school_types.items():
for school_id in scids:
uids = sim.people.schools[
school_id] # Dict with keys of school_id and values of uids in that school
edges_this_school = sdf.loc[((sdf['p1'].isin(uids)) |
(sdf['p2'].isin(uids)))]
if scen[school_type] is None:
school_contacts.append(edges_this_school)
else:
student_to_student_edge_bool = (
edges_this_school['p1_student']
& edges_this_school['p2_student'])
student_to_student_edges = edges_this_school.loc[
student_to_student_edge_bool]
inds_to_rewire = np.random.choice(
student_to_student_edges.index,
size=int(frac_edges_to_rewire *
student_to_student_edges.shape[0]),
replace=False)
inds_to_keep = np.setdiff1d(student_to_student_edges.index,
inds_to_rewire)
edges_to_rewire = student_to_student_edges.loc[inds_to_rewire]
stublist = np.concatenate(
(edges_to_rewire['p1'], edges_to_rewire['p2']))
p1_inds = np.random.choice(len(stublist),
size=len(stublist) // 2,
replace=False)
p2_inds = np.setdiff1d(range(len(stublist)), p1_inds)
p1 = stublist[p1_inds]
p2 = stublist[p2_inds]
new_edges = pd.DataFrame({'p1': p1, 'p2': p2})
new_edges['beta'] = cv.defaults.default_float(1.0) #1.0
# Remove self loops
new_edges = new_edges.loc[new_edges['p1'] != new_edges['p2']]
rewired_student_to_student_edges = pd.concat([
student_to_student_edges.loc[
inds_to_keep, ['p1', 'p2', 'beta']], # Keep these
new_edges
])
print(
f'During rewiring, the number of student-student edges went from {student_to_student_edges.shape[0]} to {rewired_student_to_student_edges.shape[0]}'
)
other_edges = edges_this_school.loc[
(~edges_this_school['p1_student']) |
(~edges_this_school['p2_student'])]
rewired_edges_this_school = pd.concat(
[rewired_student_to_student_edges, other_edges])
school_contacts.append(rewired_edges_this_school)
all_school_contacts = pd.concat(school_contacts)
sim.people.contacts['s'] = cv.Layer().from_df(all_school_contacts)
sm = cvsch.schools_manager(scen)
sim['interventions'] += [sm]
def begin_day(self, date):
''' No students, so return an empty layer '''
self.layer = cv.Layer() # Empty
uids = np.empty(0, dtype='int64')
return uids
def __init__(self, sim, uids, layer):
super().__init__(uids, layer)
self.base_layer = cv.Layer(
) # Empty base layer (ignore the passed-in layer)
return
def __init__(self,
sim,
school_id,
school_type,
uids,
layer,
start_day,
screen_prob,
screen2pcr,
test_prob,
trace_prob,
quar_prob,
schedule,
beta_s,
ili_prob,
testing,
verbose=False,
**kwargs):
'''
Initialize the School
sim (covasim Sim) : Pointer to the simulation object
school_id (int) : ID of this school
school_type (str) : Type of this school in pk, es, ms, hs, uv
uids (array) : Array of ids of individuals associated with this school
layer (Layer) : The fragment of the original 's' network associated with this school
start_day (str) : Opening day for school
screen_prob (float) : Coverage of screening
test_prob (float) : Probability of PCR testing on screen +
screen2pcr (int) : Days between positive screening receiving PCR results, for those testing
trace_prob (float) : Probability of tracing from PCR+
quar_prob (float) : Probability school contacts quarantine on trace
schedule (str) : Full, Hybrid, or Remote
beta_s (float) : beta for this school
ili_prob (float) : Daily probability of ILI
testing (struct) : List of dictionaries of parameters for SchoolTesting
'''
self.sid = int2key(school_id) # Convert to an string
self.stype = school_type
self.uids = np.array(uids)
self.start_day = sim.day(start_day)
self.screen_prob = screen_prob
self.screen2pcr = screen2pcr
self.test_prob = test_prob
self.trace_prob = trace_prob
self.quar_prob = quar_prob
self.schedule = schedule
self.beta_s = beta_s # Not currently used here, but rather in the school_intervention
self.ili_prob = ili_prob
self.verbose = verbose
# TODO: these flags should have been arrays in population.py, but weren't. Convert here for performance.
sim.people.student_flag = np.array(sim.people.student_flag, dtype=bool)
sim.people.teacher_flag = np.array(sim.people.teacher_flag, dtype=bool)
sim.people.staff_flag = np.array(sim.people.staff_flag, dtype=bool)
self.is_open = False # Schools start closed
self.uids_at_home = {} # Dict from uid to release date
if self.schedule.lower() == 'hybrid':
self.ct_mgr = HybridContactManager(sim, self.uids, layer)
elif self.schedule.lower() == 'full':
self.ct_mgr = FullTimeContactManager(sim, self.uids, layer)
elif self.schedule.lower() == 'remote':
self.ct_mgr = RemoteContactManager(sim, self.uids, layer)
else:
print(
f'Warning: Unrecognized schedule ({self.schedule}) passed to School class.'
)
self.stats = SchoolStats(self, sim)
self.testing = SchoolTesting(self, testing, sim)
self.empty_layer = cv.Layer() # Cache an empty layer
return
def test_base():
sc.heading('Testing base.py...')
json_path = 'base_tests.json'
sim_path = 'base_tests.sim'
# Create a small sim for later use
sim = cv.Sim(pop_size=100, verbose=verbose)
sim.run()
# Check setting invalid key
with pytest.raises(sc.KeyNotFoundError):
po = cv.ParsObj(pars={'a':2, 'b':3})
po.update_pars({'c':4})
# Printing result
r = cv.Result()
print(r)
print(r.npts)
# Day and date conversion
daystr = '2020-04-04'
sim.day(daystr)
sim.day(sc.readdate(daystr))
with pytest.raises(ValueError):
sim.day('not a date')
sim.date(34)
sim.date([34, 54])
sim.date(34, 54, as_date=True)
# BaseSim methods
sim.copy()
sim.export_results(filename=json_path)
sim.export_pars(filename=json_path)
sim.shrink(in_place=False)
for keep_people in [True, False]:
sim.save(filename=sim_path, keep_people=keep_people)
cv.Sim.load(sim_path)
# BasePeople methods
ppl = sim.people
ppl.get(['susceptible', 'infectious'])
ppl.keys()
ppl.person_keys()
ppl.state_keys()
ppl.date_keys()
ppl.dur_keys()
ppl.indices()
ppl._resize_arrays(pop_size=200) # This only resizes the arrays, not actually create new people
ppl._resize_arrays(pop_size=100) # Change back
ppl.to_df()
ppl.to_arr()
ppl.person(50)
people = ppl.to_people()
ppl.from_people(people)
ppl.make_edgelist([{'new_key':[0,1,2]}])
ppl.brief()
# Contacts methods
contacts = ppl.contacts
df = contacts['a'].to_df()
ppl.remove_duplicates(df)
with pytest.raises(sc.KeyNotFoundError):
contacts['invalid_key']
contacts.values()
len(contacts)
print(contacts)
print(contacts['a'])
# Layer methods
hospitals_layer = cv.Layer()
contacts.add_layer(hospitals=hospitals_layer)
contacts.pop_layer('hospitals')
df = hospitals_layer.to_df()
hospitals_layer.from_df(df)
# Tidy up
remove_files(json_path, sim_path)
return