def prior_test(sim): old_or_sick = cv.true((sim.people.age > 60) + (sim.people.symptomatic == True)) others = cv.true((sim.people.age <= 60) * (sim.people.symptomatic == False)) inds = sim.people.uid # Everyone in the population -- equivalent to np.arange(len(sim.people)) vals = np.ones(len(sim.people)) # Create the array vals[old_or_sick] = 2 vals[others] = 0.5 output = dict(inds=inds, vals=vals) return output
def plot_schools(pop): ''' Not a formal test, but a sanity check for school distributions ''' keys = ['pk', 'es', 'ms', 'hs'] # Exclude universities for this analysis ppl_keys = ['all', 'students', 'teachers', 'staff'] xpeople = np.arange(len(ppl_keys)) # X axis for people school_types_by_ind = {} for key,vals in pop.school_types.items(): for val in vals: if key in keys: school_types_by_ind[val] = key results = {} for sc_id,sc_type in school_types_by_ind.items(): thisres = sc.objdict() sc_inds = (pop.school_id == sc_id) thisres.all = cv.true(sc_inds) thisres.students = cv.true(np.array(pop.student_flag) * sc_inds) thisres.teachers = cv.true(np.array(pop.teacher_flag) * sc_inds) thisres.staff = cv.true(np.array(pop.staff_flag) * sc_inds) results[sc_id] = thisres # Do plotting fig = pl.figure(figsize=figsize, dpi=dpi) pl.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95, hspace=0.5, wspace=0.5) n_schools = len(results) n_cols = len(ppl_keys) + 1 count = 0 for sc_id in results.keys(): count += 1 school_type = school_types_by_ind[sc_id] ax = pl.subplot(n_schools, n_cols, count) thisres = results[sc_id] thisres.people_counts = [len(thisres[k]) for k in ppl_keys] ax.bar(xpeople, thisres.people_counts) ax.set_xticks(xpeople) ax.set_xticklabels(ppl_keys) title = f'School ID {sc_id}, school type {school_type}, total size: {len(thisres.all)}' ax.set_title(title) thisres.ages = sc.objdict() for key in ppl_keys: count += 1 ax = pl.subplot(n_schools, n_cols, count) thisres.ages[key] = pop.age[thisres[key]] pl.hist(thisres.ages[key]) ax.set_title(f'Ages for {key} in school {sc_id} ({school_type})') if do_maximize: cv.maximize(fig=fig) if to_json: sc.savejson(outfile, results, indent=2) return results
def vaccinate_by_age(sim): young = cv.true(sim.people.age < 50) # cv.true() returns indices of people matching this condition, i.e. people under 50 middle = cv.true((sim.people.age >= 50) * (sim.people.age < 75)) # Multiplication means "and" here old = cv.true(sim.people.age > 75) inds = sim.people.uid # Everyone in the population -- equivalent to np.arange(len(sim.people)) vals = np.ones(len(sim.people)) # Create the array vals[young] = 0.1 # 10% probability for people <50 vals[middle] = 0.5 # 50% probability for people 50-75 vals[old] = 0.9 # 90% probbaility for people >75 output = dict(inds=inds, vals=vals) return output
def test_data_interventions(): sc.heading('Testing data interventions and other special cases') # Create sim sim = cv.Sim(pop_size=100, n_days=60, datafile=csv_file, verbose=verbose) # Intervention conversion ce = cv.InterventionDict(**{ 'which': 'clip_edges', 'pars': { 'days': [10, 30], 'changes': [0.5, 1.0] } }) print(ce) with pytest.raises(sc.KeyNotFoundError): cv.InterventionDict(**{ 'which': 'invalid', 'pars': { 'days': 10, 'changes': 0.5 } }) # Test numbers and contact tracing swab_dict = {'dist': 'uniform', 'par1': 1, 'par2': 3} tp1 = cv.test_prob(0.05, start_day=5, end_day=15, ili_prev=0.1, swab_delay=swab_dict, subtarget={ 'inds': lambda sim: cv.true(sim.people.age > 50), 'vals': 0.3 }) tn1 = cv.test_num(10, start_day=3, end_day=20, ili_prev=0.1, swab_delay=swab_dict) tn2 = cv.test_num(daily_tests='data', quar_policy=[0, 5], subtarget={ 'inds': lambda sim: cv.true(sim.people.age > 50), 'vals': 1.2 }) ct = cv.contact_tracing(presumptive=True, capacity=5) # Create and run sim['interventions'] = [ce, tp1, tn1, tn2, ct] sim.run() return
def vaccinate_by_age2(sim): # cv.true() returns indices of people matching this condition, i.e. people under 50 young_or_sick = cv.true((sim.people.age < 12) + (sim.people.susceptible == False)) old = cv.true((sim.people.age >= 65) * (sim.people.susceptible == True)) others = cv.true((sim.people.age >= 12) * (sim.people.age < 65) * (sim.people.susceptible == True)) inds = sim.people.uid # Everyone in the population -- equivalent to np.arange(len(sim.people)) vals = np.ones(len(sim.people)) # Create the array vals[young_or_sick] = 0 vals[old] = 10 vals[others] = 0.1 output = dict(inds=inds, vals=vals) return output
def apply(self, sim): if sim.t == self.day: eligible = cv.true(~np.isfinite(sim.people.date_exposed) & ~sim.people.vaccinated) self.placebo_inds = eligible[cv.choose( len(eligible), min(self.trial_size, len(eligible)))] return
def subtarget_25_30(sim): inds = cv.true((sim.people.age >= 25) & (sim.people.age < 30)) if not hasattr(sim, 'subtarget_25_30'): sim.subtarget_25_30 = cv.binomial_filter( 0.1, inds) # 90% of 25-30 years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_25_30) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget_16_17(sim): inds = cv.true((sim.people.age >= 16) & (sim.people.age < 17)) if not hasattr(sim, 'subtarget_16_17'): sim.subtarget_16_17 = cv.binomial_filter( 0.9, inds) # 10% of 16-18 years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_16_17) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
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 subtarget_18_25(sim): inds = cv.true((sim.people.age >= 18) & (sim.people.age < 25)) if not hasattr(sim, 'subtarget_18_30'): sim.subtarget_18_25 = cv.binomial_filter( 0.1, inds) # 90% of 18-25 years vaccinated inds = np.setdiff1d(inds, sim.subtarget_18_25) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget_40_45(sim): inds = cv.true((sim.people.age >= 40) & (sim.people.age < 45)) if not hasattr(sim, 'subtarget_40_45'): sim.subtarget_40_45 = cv.binomial_filter( 0.1, inds) # 90% of 40-45 years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_40_45) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget_50_60(sim): inds = cv.true((sim.people.age >= 50) & (sim.people.age < 60)) if not hasattr(sim, 'subtarget_50_60'): sim.subtarget_50_60 = cv.binomial_filter( 0.1, inds) # 90% of 50-60 years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_50_60) return {'inds': inds, 'vals': 0.005 * np.ones(len(inds))}
def subtarget_60_75(sim): inds = cv.true((sim.people.age >= 60) & (sim.people.age < 75)) if not hasattr(sim, 'subtarget_60_75'): sim.subtarget_60_75 = cv.binomial_filter( 0.05, inds) # 95% of 60+ years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_60_75) return {'inds': inds, 'vals': 0.02 * np.ones(len(inds))}
def subtarget_75_100(sim): inds = cv.true(sim.people.age >= 75) if not hasattr(sim, 'subtarget_75_100'): sim.subtarget_75_100 = cv.binomial_filter( 0.05, inds) # 95% of 70+ years olds vaccinated inds = np.setdiff1d(inds, sim.subtarget_75_100) return {'inds': inds, 'vals': 0.02 * np.ones(len(inds))}
def subtarget_30_40(sim): inds = cv.true((sim.people.age >= 30) & (sim.people.age < 40)) if not hasattr(sim, 'subtarget_30_40'): sim.subtarget_30_40 = cv.binomial_filter( 0.4, inds ) # 30% of 30-40+ years olds will not agree to get vaccinated inds = np.setdiff1d(inds, sim.subtarget_30_40) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget_45_50(sim): inds = cv.true((sim.people.age >= 45) & (sim.people.age < 50)) if not hasattr(sim, 'subtarget_45_50'): sim.subtarget_45_50 = cv.binomial_filter( 0.25, inds ) # 30% of 40-50 years olds will not agree to get vaccinated inds = np.setdiff1d(inds, sim.subtarget_45_50) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget(sim): ''' Select people who are susceptible ''' if sim.t == start_trial: eligible = cv.true(~np.isfinite(sim.people.date_exposed)) inds = eligible[cv.choose(len(eligible), min(trial_size // 2, len(eligible)))] else: inds = [] return {'vals': [1.0 for ind in inds], 'inds': inds}
def subtarget_12_17(sim): inds = cv.true((sim.people.age >= 12) & (sim.people.age < 17)) if not hasattr(sim, 'subtarget_12_17'): sim.subtarget_18_25 = cv.binomial_filter( 0.3, inds) # 30% of 18+ years olds will not agree to get vaccinated inds = np.setdiff1d(inds, sim.subtarget_12_17) return {'inds': inds, 'vals': 0.002 * np.ones(len(inds))}
def subtarget(sim, age=None, vx_scenario=None): rollout = vx_rollout[age] inds = cv.true((sim.people.age >= rollout['start_age']) * (sim.people.age < rollout['end_age'])) key = f'subtarget_{age}' if not hasattr(sim, key): prob = scendata[age][vx_scenario] vx_inds = cv.binomial_filter(prob, inds) setattr(sim, key, vx_inds) else: vx_inds = getattr(sim, key) inds = np.setdiff1d(inds, vx_inds) return {'inds': inds, 'vals': 0.65*np.ones(len(inds))}
def check_88(sim): people_who_are_88 = sim.people.age.round( ) == 88 # Find everyone who's aged 88 (to the nearest year) people_exposed = sim.people.exposed # Find everyone who's infected with COVID people_who_are_88_with_covid = cv.true( people_who_are_88 * people_exposed) # Multiplication is the same as logical "and" n = len(people_who_are_88_with_covid) # Count how many people there are if n: print( f'Oh no! {n} people aged 88 have covid on timestep {sim.t} {"🤯"*n}' ) return
def subtarget_50_60(sim): inds = cv.true((sim.people.age >= 50) & (sim.people.age < 60)) return {'inds': inds, 'vals': 0.010 * np.ones(len(inds))}
def subtarget_75_100(sim): inds = cv.true(sim.people.age >= 75) return {'inds': inds, 'vals': 0.020 * np.ones(len(inds))}
def subtarget_60_75(sim): inds = cv.true((sim.people.age >= 60) & (sim.people.age < 75)) return {'inds': inds, 'vals': 0.020 * np.ones(len(inds))}
def subtarget_40_50(sim): inds = cv.true((sim.people.age >= 40) & (sim.people.age < 50)) return {'inds': inds, 'vals': 0.005 * np.ones(len(inds))}
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 apply(self, sim): ''' Perform vaccination ''' t = sim.t if t < self.start_day: return elif self.end_day is not None and t > self.end_day: return # Check that there are still vaccines rel_t = t - self.start_day if rel_t < len(self.daily_vaccines): n_vaccines = sc.randround(self.daily_vaccines[rel_t]/sim.rescale_vec[t]) # Correct for scaling that may be applied by rounding to the nearest number of tests if not (n_vaccines and np.isfinite(n_vaccines)): # If there are no doses today, abort early return else: if sim.rescale_vec[t] != sim['pop_scale']: raise RuntimeError('bad rescale time') else: return vacc_probs = np.ones(sim.n) # Begin by assigning equal vaccine weight (converted to a probability) to everyone # Remove minors ppl = sim.people inds0 = cv.true(ppl.age <= 18) vacc_probs[inds0] *= 0 # add age priority inds1 = cv.true(ppl.age >= self.age_priority) vacc_probs[inds1] *= 10 # Handle scheduling # first dose: vacc_probs[self.vaccinations == 0] *= self.dose_priority[0] # time between first and second dose: no_dose = [sim.t < (d[0] + self.delay) if len(d) > 0 else False for d in self.vaccination_dates] vacc_probs[no_dose] *= 0 # time available for second dose: second_dose = [sim.t >= (d[0] + self.delay) if len(d) > 0 else False for d in self.vaccination_dates] vacc_probs[second_dose] *= self.dose_priority[1] # Don't give dose 2 people who have had more than 1 vacc_inds = cvu.true(self.vaccinations > 1) vacc_probs[vacc_inds] = 0.0 # Now choose who gets vaccinated and vaccinate them n_vaccines = min(n_vaccines, (vacc_probs!=0).sum()) # Don't try to vaccinate more people than have nonzero vaccination probability all_vacc_inds = cvu.choose_w(probs=vacc_probs, n=n_vaccines, unique=True) # Choose who actually tests sim.results['new_doses'][t] += len(all_vacc_inds) # Did the vaccine take? vacc_take_inds = all_vacc_inds[np.logical_or(cvu.n_binomial(self.take_prob, len(all_vacc_inds)) & (self.vaccinations[all_vacc_inds] == 0), self.vaccine_take[all_vacc_inds] & (self.vaccinations[all_vacc_inds] == 1))] self.vaccine_take[vacc_take_inds] = True # Calculate the effect per person vacc_doses = self.vaccinations[vacc_take_inds] # Calculate current doses eff_doses = np.minimum(vacc_doses, len(self.cumulative) - 1) # Convert to a valid index vacc_eff = self.cumulative[eff_doses] # Pull out corresponding effect sizes rel_trans_eff = (1.0 - vacc_eff) + vacc_eff * self.rel_trans rel_symp_eff = (1.0 - vacc_eff) + vacc_eff * self.rel_symp # Apply the vaccine to people sim.people.rel_trans[vacc_take_inds] = self.orig_rel_trans[vacc_take_inds]*rel_trans_eff sim.people.symp_prob[vacc_take_inds] = self.orig_symp_prob[vacc_take_inds]*rel_symp_eff self.vaccinations[all_vacc_inds] += 1 for v_ind in all_vacc_inds: self.vaccination_dates[v_ind].append(sim.t) return
def test_vaccine_target_eff(): sc.heading('Testing vaccine with pre-specified efficacy...') target_eff_1 = 0.7 target_eff_2 = 0.95 default_pars = cv.parameters.get_vaccine_dose_pars(default=True) test_pars = dict(doses=2, interval=21, target_eff=[target_eff_1, target_eff_2]) vacc_pars = sc.mergedicts(default_pars, test_pars) # construct analyzer to select placebo arm class placebo_arm(cv.Analyzer): def __init__(self, day, trial_size, **kwargs): super().__init__(**kwargs) self.day = day self.trial_size = trial_size return def initialize(self, sim=None): self.placebo_inds = [] self.initialized = True return def apply(self, sim): if sim.t == self.day: eligible = cv.true(~np.isfinite(sim.people.date_exposed) & ~sim.people.vaccinated) self.placebo_inds = eligible[cv.choose(len(eligible), min(self.trial_size, len(eligible)))] return pars = dict( rand_seed = 1, # Note: results may be sensitive to the random seed pop_size = 20_000, beta = 0.01, n_days = 90, verbose = -1, use_waning = True, ) # Define vaccine arm trial_size = 4_000 start_trial = 20 def subtarget(sim): ''' Select people who are susceptible ''' if sim.t == start_trial: eligible = cv.true(~np.isfinite(sim.people.date_exposed)) inds = eligible[cv.choose(len(eligible), min(trial_size // 2, len(eligible)))] else: inds = [] return {'vals': [1.0 for ind in inds], 'inds': inds} # Initialize vx = cv.vaccinate_prob(vaccine=vacc_pars, days=[start_trial], label='target_eff', prob=0.0, subtarget=subtarget) sim = cv.Sim( pars = pars, interventions = vx, analyzers = placebo_arm(day=start_trial, trial_size=trial_size // 2) ) # Run sim.run() print('Vaccine efficiency:') results = sc.objdict() vacc_inds = cv.true(sim.people.vaccinated) # Find trial arm indices, those who were vaccinated placebo_inds = sim['analyzers'][0].placebo_inds assert (len(set(vacc_inds).intersection(set(placebo_inds))) == 0) # Check that there is no overlap # Calculate vaccine efficacy against infection VE_inf = 1 - (np.isfinite(sim.people.date_exposed[vacc_inds]).sum() / np.isfinite(sim.people.date_exposed[placebo_inds]).sum()) # Calculate vaccine efficacy against symptoms VE_symp = 1 - (np.isfinite(sim.people.date_symptomatic[vacc_inds]).sum() / np.isfinite(sim.people.date_symptomatic[placebo_inds]).sum()) # Calculate vaccine efficacy against severe disease VE_sev = 1 - (np.isfinite(sim.people.date_severe[vacc_inds]).sum() / np.isfinite(sim.people.date_severe[placebo_inds]).sum()) results['inf'] = VE_inf results['symp'] = VE_symp results['sev'] = VE_sev print(f'Against: infection: {VE_inf * 100:0.2f}%, symptoms: {VE_symp * 100:0.2f}%, severity: {VE_sev * 100:0.2f}%') # Check that actual efficacy is within 6 %age points of target errormsg = f'Expected VE to be about {target_eff_2}, but it is {VE_symp}. Check different random seeds; this test is highly sensitive.' assert round(abs(VE_symp-target_eff_2),2)<=0.1, errormsg nab_init = sim['vaccine_pars']['target_eff']['nab_init'] boost = sim['vaccine_pars']['target_eff']['nab_boost'] print(f'Initial NAbs: {nab_init}') print(f'Boost: {boost}') return sim
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 subtarget_18_30(sim): inds = cv.true((sim.people.age >= 18) & (sim.people.age < 30)) return {'inds': inds, 'vals': 0.003 * np.ones(len(inds))}
def update(self, sim): ''' Called on each day to update school statistics ''' t = sim.t ppl = sim.people rescale = sim.rescale_vec[t] if self.school.ct_mgr.school_day: self.num_school_days += 1 student_uids = cv.itruei(ppl.student_flag, self.school.uids) teacher_uids = cv.itruei(ppl.teacher_flag, self.school.uids) staff_uids = cv.itruei(ppl.staff_flag, self.school.uids) infectious_ids = {} for group, ids in zip(['students', 'teachers', 'staff'], [student_uids, teacher_uids, staff_uids]): infectious_ids[group] = cv.true(ppl.infectious[ids]) self.infectious[group][t] = len(infectious_ids[group]) * rescale self.newly_exposed[group][t] = len( cv.true(ppl.date_exposed[ids] == t - 1)) * rescale self.scheduled[group][t] = len( np.intersect1d(self.school.scheduled_uids, ids)) * rescale # Scheduled self.in_person[group][t] = len( np.intersect1d(self.school.uids_passed_screening, ids)) * rescale # Post-screening # Tracing statistics to compare against previous work: # if len(self.school.uids_arriving_at_school) > 0: # school_infectious = cv.itrue(ppl.infectious[self.school.uids_arriving_at_school], self.school.uids_arriving_at_school) # else: # school_infectious = np.empty(0, dtype='int64') # Options here: (TODO - avoid code replication) # 1. Use ids of students who arrived as school (pre-screening): self.school.uids_arriving_at_school (pre-screening) # 2. Use ids of students who passed screening: self.school.uids_passed_screening # First "infectious_arrive_at_school" assumes there is a transmission risk even pre-screening (e.g. bus) n_students_at_school = len( cv.itruei(ppl.student_flag * ppl.infectious, self.school.uids_arriving_at_school)) n_teachers_at_school = len( cv.itruei(ppl.teacher_flag * ppl.infectious, self.school.uids_arriving_at_school)) n_staff_at_school = len( cv.itruei(ppl.staff_flag * ppl.infectious, self.school.uids_arriving_at_school)) for group, count in zip( ['students', 'teachers', 'staff'], [n_students_at_school, n_teachers_at_school, n_staff_at_school]): self.infectious_arrive_at_school[group][t] = count * rescale # Second "infectious_stay_at_school" effectively assumes "screen-positive" kids would be kept home from school in the first place n_students_passedscreening = len( cv.itruei(ppl.student_flag * ppl.infectious, self.school.uids_passed_screening)) n_teachers_passedscreening = len( cv.itruei(ppl.teacher_flag * ppl.infectious, self.school.uids_passed_screening)) n_staff_passedscreening = len( cv.itruei(ppl.staff_flag * ppl.infectious, self.school.uids_passed_screening)) for group, count in zip(['students', 'teachers', 'staff'], [ n_students_passedscreening, n_teachers_passedscreening, n_staff_passedscreening ]): self.infectious_stay_at_school[group][t] = count * rescale