code: import numpy as np
  import networkx as nx
  from math import gamma
  from scipy.optimize import root_scalar
  
  from mesa import Model
  from mesa.time import RandomActivation, SimultaneousActivation
  from mesa.datacollection import DataCollector
  
  from scseirx.testing_strategy import Testing
  
  ## data collection functions ##
  def get_N_diagnostic_tests(model):
      return model.number_of_diagnostic_tests
  
  
  def get_N_preventive_screening_tests(model):
      return model.number_of_preventive_screening_tests
  
  
  def get_infection_state(agent):
      if agent.exposed == True: return 'exposed'
      elif agent.infectious == True: return 'infectious'
      elif agent.recovered == True: return 'recovered'
      else: return 'susceptible'
  
  def get_quarantine_state(agent):
      if agent.quarantined == True: return True
      else: return False
  
  
  def get_undetected_infections(model):
      return model.undetected_infections
  
  
  def get_predetected_infections(model):
      return model.predetected_infections
  
  
  def get_pending_test_infections(model):
      return model.pending_test_infections
  
  
  def get_diagnostic_test_detected_infections_student(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['student']
  def get_diagnostic_test_detected_infections_teacher(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['teacher']
  def get_diagnostic_test_detected_infections_family_member(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['family_member']
  def get_diagnostic_test_detected_infections_resident(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['resident']
  def get_diagnostic_test_detected_infections_employee(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['employee']
  def get_diagnostic_test_detected_infections_unistudent(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['unistudent']
  def get_diagnostic_test_detected_infections_lecturer(model):
      return model.positive_tests[model.Testing.diagnostic_test_type]['lecturer']
  
  diagnostic_test_detected_infections_funcs = {
      'student':get_diagnostic_test_detected_infections_student,
      'teacher':get_diagnostic_test_detected_infections_teacher,
      'family_member':get_diagnostic_test_detected_infections_family_member,
      'resident':get_diagnostic_test_detected_infections_resident,
      'employee':get_diagnostic_test_detected_infections_employee,
      'unistudent':get_diagnostic_test_detected_infections_unistudent,
      'lecturer':get_diagnostic_test_detected_infections_lecturer
  }
  
  def get_preventive_test_detected_infections_student(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['student']
  def get_preventive_test_detected_infections_teacher(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['teacher']
  def get_preventive_test_detected_infections_family_member(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['family_member']
  def get_preventive_test_detected_infections_resident(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['resident']
  def get_preventive_test_detected_infections_employee(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['employee']
  def get_preventive_test_detected_infections_unistudent(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['unistudent']
  def get_preventive_test_detected_infections_lecturer(model):
      return model.positive_tests[model.Testing.preventive_screening_test_type]['lecturer']
  
  preventive_test_detected_infections_funcs = {
      'student':get_preventive_test_detected_infections_student,
      'teacher':get_preventive_test_detected_infections_teacher,
      'family_member':get_preventive_test_detected_infections_family_member,
      'resident':get_preventive_test_detected_infections_resident,
      'employee':get_preventive_test_detected_infections_employee,
      'unistudent':get_preventive_test_detected_infections_unistudent,
      'lecturer':get_preventive_test_detected_infections_lecturer
  }
  
  
  # parameter sanity check functions
  
  
  def check_positive(var):
  	assert var >= 0, 'negative number'
  	return var
  
  
  def check_bool(var):
  	assert type(var) == bool, 'not a bool'
  	return var
  
  
  def check_positive_int(var):
      if var == None:
          return var
      assert type(var) == int, 'not an integer'
      assert var >= 0, 'negative number'
      return var
  
  
  def check_contact_type_dict(var):
  	assert type(var) == dict, 'not a dictionary'
  	assert set(var.keys()).issubset({'very_far', 'far', 'intermediate', 'close'}), \
  		'does not contain the correct contact types (has to be very_far, far, intermediate or close)'
  	assert all((isinstance(i, int) or isinstance(i, float)) for i in var.values()), \
  		'contact type weights are not numeric'
  
  	return var
  
  
  def check_K1_contact_types(var):
      for area in var:
          assert area in ['very_far', 'far', 'intermediate',
              'close'], 'K1 contact type not recognised'
      return var
  
  
  def check_testing(var):
      assert var in ['diagnostic', 'background', 'preventive',
                     'background+preventive', False], \
          'unknown testing mode: {}'.format(var)
  
      return var
  
  
  
  def check_probability(var):
  	assert (type(var) == float) or (var == 0) or (var == 1), \
  		 '{} not a float'.format(var)
  	assert var >= 0, 'probability negative'
  	assert var <= 1, 'probability larger than 1'
  	return var
  
  
  def check_graph(var):
      assert type(var) in [nx.Graph, nx.MultiGraph], 'not a networkx graph'
      assert len(var.nodes) > 0, 'graph has no nodes'
      assert len(var.edges) > 0, 'graph has no edges'
      areas = [e[2]['contact_type'] for e in var.edges(data=True)]
      areas = set(areas)
      for a in areas:
          assert a in {'very_far', 'far', 'intermediate',
              'close'}, 'contact type {} not recognised'.format(a)
      return var
  
  
  def check_index_case(var, agent_types):
  	allowed_strings = agent_types[:]
  	allowed_strings.extend(['continuous'])
  	assert var in allowed_strings, 'unknown index case mode'
  	return var
  
  
  def check_discount(var):
      if var['slope'] != None:
          assert var['slope'] <= 0, 'slope needs to be <= 0 or None'
          assert np.abs(var['slope']) <= 1, 'absolute value of slope needs to be <= 1'
      assert var['intercept'], 'intercept needs to be positive'
      assert var['intercept'], 'intercept needs to be <= 1'
      return var
  
  
  def get_weibull_shape(k, mu, var):
      '''
      Calculates the shape parameter of a Weibull distribution, given its mean
      mu and its variance var
      '''
      return var / mu**2 - gamma(1 + 2/k) / gamma(1+1/k)**2 + 1
  
  
  
  def get_weibull_scale(mu, k):
      '''
      Calculates the scale parameter of a Weibull distribution, given its mean
      mu and its shape parameter k
      '''
      return mu / gamma(1 + 1/k)
  
  
  def weibull_two_param(shape, scale):
      '''
      A two-parameter Weibull distribution, based on numpy ramdon's single
      parameter distribution. We use this distribution in the simulation to draw
      random epidemiological parameters for agents from the given distribution
      See https://numpy.org/doc/stable/reference/random/generated/numpy.random.weibull.html
      '''
      return scale * np.random.weibull(shape)
  
  
  class SEIRX(Model):
      '''
      A model with a number of different agents that reproduces
      the SEIRX dynamics of pandemic spread in a facility. Note:
      all times are set to correspond to days
  
      G: networkx undirected graph, interaction graph between agents. Edges have
      to have edge the edge attribute 'contact_type' specifying the closeness of
      contacts, which can be ['very far', 'far', 'intermediate' and 'close'].
      Nodes have to have the node attribute 'type' which specifies the agent type
      of the given node (for example 'student' or 'teacher' in a school scenario).
      In addition, nodes can have the attribute 'unit', which assigns them to a
      unit in space (for example a 'class' in a school scenario).
  
      verbosity: integer in [0, 1, 2], controls text output to std out to track
      simulation progress and transmission dynamics. Default = 0.
  
      testing, default = 'diagnostic'
          'diagnostic':   only diagnostic tests for symptomatic agents
          'background':   adds background screens of all agents after a positive
                          diagnostic test
          'preventive':   adds preventive screens of agent groups to diagnostic
                          testing. Screens happen in time intervals specified 
                          separately for each agent group in the variable 
                          'screening_interval'.
          'background+preventive': preventive screens AND background screens on
                          top of diagnostic testing.
  
      infection_duration, default = 11 NOTE: includes the time an agent is exposed
      but not yet infectious at the beginning of an infection
          positive integer:   mean or median of the infection duration in days
          list of two floats: mean and standard deviation of a distribution
                              specifying the infection duration in days. These
                              numbers will be used to construct a Weibull
                              distribution from which the infection duration will
                              be drawn for every agent individually
  
      exposure_duration, default = 4. Sets the time from transmission to becoming
      infectious
          positive integer:   mean or median of the exposure duration in days
          list of two floats: mean and standard deviation of a distribution
                              specifying the exposure duration in days. These
                              numbers will be used to construct a Weibull
                              distributoin from which the exposure duration will
                              be drawn for every agent individually.
  
      time_until_symptoms, default = 6. Sets the time from transmission to
      (potentially) developing symptoms. Symptom probability has to be set for
      each agent group individually using the parameter 'symptom_probability'
          positive integer:   mean or median of the time until symptoms in days
          list of two floats: mean and standard deviation of a distribution
                              specifying the time until symptoms in days. These
                              numbers will be used to construct a Weibull
                              distribution from which the time until symptoms will
                              be drawn for every agent individually.
  
      quarantine_duration, default = 14. Positive integer, sets the time a
      positively tested agent is quarantined in days
  
      infection_risk_contact_type_weights: dictionary of the form
      {'very_far':float, 'far':float, 'intermediate':float, 'close':float}
      that sets transmission risk multipliers for different contact types of
      agents specified in the contact network G. Default: {'very_far': 0.1,
      'far': 0.5, 'intermediate': 1, 'close': 3}
  
      subclinical_modifier: default = 1.0. Float, modifies the infectiousness of
      asymptomatic cases. Example: if subclinical_modifier = 0.5, the
      infectiousness of an asymptomatic case will be reduced to 50%.
  
      K1_contact_types: list of strings from ['very_far', 'far', 'intermediate',
      'close']. Definition of contact types for which agents are considered
      "K1 contact persons" if they had contact to a positively tested person wtith
      a specified contact intensity. Default = ['close'].
  
      diagnostic_test_type, default = 'one_day_PCR'. String, specifies the test
      technology and test result turnover time used for diagnostic testing. For
      example 'same_day_antigen' or 'two_day_PCR'. See module "Testing" for
      different implemented testing techologies.
  
      preventive_screening_test_type:, default = 'one_day_PCR', String, specifies
      the test technology and test result turnover time used for preventive
      sreening. For example 'same_day_antigen' or 'two_day_PCR'. See module
      "Testing" for different implemented testing techologies.
  
      follow_up_testing_interval, default = None. Positive integer, sets the time
      a follow-up screen (background screen) is initiated after an initial screen
      triggered by a positive test result. Only applies if the testing strategy is
      'background' or preventive.
  
      liberating_testing, default = False. Boolean, flag that specifies, whether
      or not an agent is released from quarantine after returning a negative test
      result.
  
  	index_case, default = 'employee' (nursing home scenario) or 'teacher'
      (school scenario). Specifies how infections are introduced into the facility.
          agent_type:     If an agent type (for example 'student' or 'teacher' in
                          the school scenario) is specified, a single randomly
                          chosen agent from this agent group will become the index
                          case and no further index cases will be introduced into
                          the scenario.
          'continuous':   In this case, agents have a continuous risk to become
                          index cases in every simulation step. The risk has to
                          be specified for every agent group individually, using
                          the 'index_probability' parameter. If only a single
                          agent group has a non-zero index probability, then only
                          agents from this group can become index cases.
  
  
      agent_types: dictionary of the structure
          {
          agent type:
              {
              screening interval : integer, number of days between each preventive
              screen in this agent group
  
              index probability : float in the range [0, 1], sets the probability
              to become an index case in each time step
  
              mask : bool
                  whether or not the agent type is wearing a mask
              }
          }
  
      The dictionary's keys are the names of the agent types which have to
      correspond to the node attributes in the contact graph. The screening
      interval sets the time-delay between preventive screens of this agent group,
      the index probability sets the probability of a member of this agent group
      becoming an index case in every time step
  
      seed: positive integer, fixes the seed of the simulation to enable
      repeatable simulation runs. If seed = None, the simulation will be
      initialized at random.
      '''
  
      def __init__(self, G,
          verbosity = 0,
          base_transmission_risk = 0.05,
          testing='diagnostic',
          exposure_duration = [5.0, 1.9],
          time_until_symptoms = [6.4, 0.8],
          infection_duration = [10.91, 3.95],
          quarantine_duration = 10,
          subclinical_modifier = 0.6,
          infection_risk_contact_type_weights = {
              'very_far': 0.1,
              'far': 0.25,
              'intermediate': 0.5,
              'close': 1},
          K1_contact_types = ['close'],
          diagnostic_test_type = 'one_day_PCR',
          preventive_screening_test_type = 'same_day_antigen',
          follow_up_testing_interval = None,
          liberating_testing = False,
          index_case = 'teacher',
          agent_types = {
              'teacher':      {'screening_interval': None,
                               'index_probability': 0,
                               'mask':False,
                               'vaccination_ratio': 0},
              'student':      {'screening_interval': None,
                               'index_probability': 0,
                               'mask':False,
                               'vaccination_ratio': 0},
              'family_member':{'screening_interval': None,
                               'index_probability': 0,
                               'mask':False,
                               'vaccination_ratio': 0}},
          age_transmission_risk_discount = \
               {'slope':-0.02,
                'intercept':1},
          age_symptom_modification = \
               {'slope':-0.02545,
                'intercept':0.854545},
          mask_filter_efficiency = {'exhale':0, 'inhale':0},
          transmission_risk_ventilation_modifier = 0,
          transmission_risk_vaccination_modifier = {
              'reception':1,
              'transmission':0},
          seed = None):
  
          # mesa models already implement fixed seeds through their own random
          # number generations. Sadly, we need to use the Weibull distribution
          # here, which is not implemented in mesa's random number generation
          # module. Therefore, we need to initialize the numpy random number
          # generator with the given seed as well
          if seed != None:
              np.random.seed(seed)
  
          # sets the (daily) transmission risk for a household contact without
          # any precautions. Target infection ratios are taken from literature
          # and the value of the base_transmission_risk is calibrated such that
          # the simulation produces the correct infection ratios in a household
          # setting with the given distributions for epidemiological parameters
          # of agents
          self.base_transmission_risk = base_transmission_risk
      	# sets the level of detail of text output to stdout (0 = no output)
          self.verbosity = check_positive_int(verbosity)
          # flag to turn off the testing & tracing strategy
          self.testing = check_testing(testing)
          self.running = True  # needed for the batch runner implemented by mesa
          # set the interaction mode to simultaneous activation
          self.schedule = SimultaneousActivation(self)
  
  
          # internal step counter used to launch screening tests
          self.Nstep = 0
  
          # since we may have weekday-specific contact networks, we need
          # to keep track of the day of the week. Since the index case
          # per default is introduced at step 0 in index case mode, we
          # need to offset the starting weekday by a random number of weekdays
          # to prevent artifacts from always starting on the same day of the week
  
          self.weekday_offset = self.random.randint(1, 8)
          self.weekday = self.Nstep + self.weekday_offset
  
          ## epidemiological parameters: can be either a single integer or the
          # mean and standard deviation of a distribution
          self.epi_params = {}
          # counter to track the number of pathological parameter combinations
          # that had to be re-rolled (only here for debugging and control reasons)
          self.param_rerolls = 0
  
          for param, param_name in zip([exposure_duration, time_until_symptoms,
                  infection_duration],['exposure_duration', 'time_until_symptoms',
                  'infection_duration']):
  
              if isinstance(param, int):
                  self.epi_params[param_name] = check_positive_int(param)
  
              elif isinstance(param, list) and len(param) == 2:
  
                  mu = check_positive(param[0])
                  var = check_positive(param[1]**2)
                  shape = root_scalar(get_weibull_shape, args=(mu, var),
                              method='toms748', bracket=[0.2, 500]).root
                  scale = get_weibull_scale(mu, shape)
  
                  self.epi_params[param_name] = [shape, scale]
              else:
                  print('{} format not recognized, should be either a single '+\
                    'int or a tuple of two positive numbers'.format(param_name))
  
  
          # duration of quarantine
          self.quarantine_duration = check_positive_int(quarantine_duration)
  
          self.infection_risk_area_weights = check_contact_type_dict(
              infection_risk_contact_type_weights)
  
          # modifier for infectiosness for asymptomatic cases
          self.subclinical_modifier = check_positive(subclinical_modifier)
          # modifiers for the infection risk, depending on contact type
          self.infection_risk_contact_type_weights = infection_risk_contact_type_weights
  
          # modifications for age-dependent transmission and reception risks and
          # symptom probabilities
          self.age_transmission_risk_discount = \
              check_discount(age_transmission_risk_discount)
  
          self.age_symptom_modification = age_symptom_modification
              #check_discount(age_symptom_modification)
  
          self.mask_filter_efficiency = mask_filter_efficiency
          self.transmission_risk_ventilation_modifier = \
              transmission_risk_ventilation_modifier
          self.transmission_risk_vaccination_modifier = \
              transmission_risk_vaccination_modifier
          ## agents and their interactions
          # interaction graph of agents
          self.G = check_graph(G)
          # add weights as edge attributes so they can be visualised easily
          if type(self.G) == nx.MultiGraph:
              for (u, v, key, contact_type) in self.G.edges(keys=True,
                      data='contact_type'):
                  self.G[u][v][key]['weight'] = \
                      self.infection_risk_contact_type_weights[contact_type]
          else:
              for e in G.edges(data=True):
                  G[e[0]][e[1]]['weight'] = self.infection_risk_contact_type_weights\
                  	[G[e[0]][e[1]]['contact_type']]
  
          # extract the different agent types from the contact graph
          self.agent_types = list(agent_types.keys())
          # dictionary of available agent classes with agent types and classes
          self.agent_classes = {}
          if 'resident' in agent_types:
              from scseirx.agent_resident import resident
              self.agent_classes['resident'] = resident
          if 'employee' in agent_types:
              from scseirx.agent_employee import employee
              self.agent_classes['employee'] = employee
          if 'student' in agent_types:
              from scseirx.agent_student import student
              self.agent_classes['student'] = student
          if 'teacher' in agent_types:
              from scseirx.agent_teacher import teacher
              self.agent_classes['teacher'] = teacher
          if 'family_member' in agent_types:
              from scseirx.agent_family_member import family_member
              self.agent_classes['family_member'] = family_member
          if 'lecturer' in agent_types:
              from scseirx.agent_lecturer import lecturer
              self.agent_classes['lecturer'] = lecturer
          if 'unistudent' in agent_types:
              from scseirx.agent_unistudent import unistudent
              self.agent_classes['unistudent'] = unistudent
  
          ## set agent characteristics for all agent groups
          # list of agent characteristics
          params = ['screening_interval','index_probability', 'mask' ,'vaccination_ratio',
                    'voluntary_testing_rate']
  
          # default values that are used in case a characteristic is not specified
          # for an agent group
          defaults = {'screening_interval':None,
                      'index_probability':0,
                      'mask':False,
                      'vaccination_ratio':0,
                      'voluntary_testing_rate':1
                      }
  
          # sanity checks that are applied to parameters passed to the class
          # constructor to make sure they conform to model expectations
          check_funcs = [check_positive_int, check_probability, check_bool,
                         check_probability, check_probability]
  
          # member dicts that store the parameter values for each agent group
          self.screening_intervals = {}
          self.index_probabilities = {}
          self.masks = {}
          self.vaccination_probabilities = {}
          self.voluntary_testing_rates = {}
  
  
          param_dicts = [self.screening_intervals, self.index_probabilities,
                      self.masks, self.vaccination_probabilities, self.voluntary_testing_rates]
  
          # iterate over all possible agent parameters and agent groups: set the
          # respective value to the value passed through the constructor or to
          # the default value if no value has been passed
          for param,param_dict,check_func in zip(params,param_dicts,check_funcs):
              for at in self.agent_types:
                  try:
                      param_dict.update({at:check_func(agent_types[at][param])})
                  except KeyError:
                      param_dict.update({at:defaults[param]})
  
          # pass all parameters relevant for the testing strategy to the testing
          # class. NOTE: this separation is not a strictly necessary design
          # decision but I like to keep the parameters related to testing and
          # tracing in a separate place
          self.Testing = Testing(self, diagnostic_test_type,
               preventive_screening_test_type,
               check_positive_int(follow_up_testing_interval),
               self.screening_intervals,
               check_bool(liberating_testing),
               check_K1_contact_types(K1_contact_types),
               verbosity)
  
  
          # specifies either continuous probability for index cases in agent
          # groups based on the 'index_probability' for each agent group, or a
          # single (randomly chosen) index case in the passed agent group
          self.index_case = check_index_case(index_case, self.agent_types)
  
          self.num_agents = {}
  
          ## add agents
          # extract the agent nodes from the graph and add them to the scheduler
          for agent_type in self.agent_types:
              IDs = [x for x,y in G.nodes(data=True) if y['type'] == agent_type]
              self.num_agents.update({agent_type:len(IDs)})
  
              # get the agent locations (units) from the graph node attributes
              units = [self.G.nodes[ID]['unit'] for ID in IDs]
  
              # determine the agents that will be vaccinated, given the 
              # vaccination ratio of the respective agent group
              vaccination_status = np.asarray([False] * len(IDs))
              if self.vaccination_probabilities[agent_type] > 0:
                  n = round(self.vaccination_probabilities[agent_type] * len(IDs))
                  idx = list(range(len(IDs)))
                  rnd_idx = np.asarray(self.random.sample(idx, n))
                  vaccination_status[rnd_idx] = True
  
  
              for ID, unit, vaccinated in zip(IDs, units, vaccination_status):
  
                  tmp_epi_params = {}
                  # for each of the three epidemiological parameters, check if
                  # the parameter is an integer (if yes, pass it directly to the
                  # agent constructor), or if it is specified by the shape and
                  # scale parameters of a Weibull distribution. In the latter
                  # case, draw a new number for every agent from the distribution
                  # NOTE: parameters drawn from the distribution are rounded to
                  # the nearest integer
                  while True:
                      for param_name, param in self.epi_params.items():
                          if isinstance(param, int):
                              tmp_epi_params[param_name] = param
  
                          else:
                              tmp_epi_params[param_name] = \
                                  round(weibull_two_param(param[0], param[1]))
  
                      if tmp_epi_params['exposure_duration'] > 0 and \
                         tmp_epi_params['time_until_symptoms'] >= \
                         tmp_epi_params['exposure_duration'] and\
                         tmp_epi_params['infection_duration'] > \
                         tmp_epi_params['exposure_duration']:
                         break
                      else:
                          self.param_rerolls += 1
                          if verbosity > 1:
                              print('pathological epi-param case found!')
                              print(tmp_epi_params)
  
                  # check if the agent participates in voluntary testing
                  p = self.voluntary_testing_rates[agent_type]
                  voluntary_testing = np.random.choice([True, False],
                           p=[p, 1-p])
  
                  # construct the agent object
                  a = self.agent_classes[agent_type](ID, unit, self,
                      tmp_epi_params['exposure_duration'],
                      tmp_epi_params['time_until_symptoms'],
                      tmp_epi_params['infection_duration'],
                      vaccinated,
                      voluntary_testing,
                      verbosity)
                  self.schedule.add(a)
  
  
  		# infect the first agent in single index case mode
          if self.index_case != 'continuous':
              infection_targets = [
                  a for a in self.schedule.agents if a.type == index_case]
              # pick a random agent to infect in the selected agent group
              target = self.random.randint(0, len(infection_targets) - 1)
              infection_targets[target].exposed = True
              if self.verbosity > 0:
                  print('{} exposed: {}'.format(index_case,
                      infection_targets[target].ID))
  
  
          # list of agents that were tested positive this turn
          self.newly_positive_agents = []
          # flag that indicates if there were new positive tests this turn
          self.new_positive_tests = False
          # dictionary of flags that indicate whether a given agent group has
          # been creened this turn
          self.screened_agents= {
              'reactive':{agent_type: False for agent_type in self.agent_types},
              'follow_up':{agent_type: False for agent_type in self.agent_types},
              'preventive':{agent_type: False for agent_type in self.agent_types}}
  
  
          # dictionary of counters that count the days since a given agent group
          # was screened. Initialized differently for different index case modes
          if (self.index_case == 'continuous') or \
        	   (not np.any(list(self.Testing.screening_intervals.values()))):
          	self.days_since_last_agent_screen = {agent_type: 0 for agent_type in
          	self.agent_types}
          # NOTE: if we initialize these variables with 0 in the case of a single
          # index case, we introduce a bias since in 'single index case mode' the
          # first index case will always become exposed in step 0. To realize
          # random states of the preventive sceening procedure with respect to the
          # incidence of the index case, we have to randomly pick the days since
          # the last screen for the agent group from which the index case is
          else:
          	self.days_since_last_agent_screen = {}
          	for agent_type in self.agent_types:
          		if self.Testing.screening_intervals[agent_type] != None:
          			self.days_since_last_agent_screen.update({
          				agent_type: self.random.choice(range(0,
          				 self.Testing.screening_intervals[agent_type] + 1))})
          		else:
          			self.days_since_last_agent_screen.update({agent_type: 0})
  
          # dictionary of flags that indicates whether a follow-up screen for a
          # given agent group is scheduled
          self.scheduled_follow_up_screen = {agent_type: False for agent_type in
          	self.agent_types}
  
          # counters
          self.number_of_diagnostic_tests = 0
          self.number_of_preventive_screening_tests = 0
          self.positive_tests = {self.Testing.preventive_screening_test_type:
                              {agent_type:0 for agent_type in self.agent_types},
                                 self.Testing.diagnostic_test_type:
                              {agent_type:0 for agent_type in self.agent_types}}
  
          self.undetected_infections = 0
          self.predetected_infections = 0
          self.pending_test_infections = 0
          self.quarantine_counters = {agent_type:0 for agent_type in agent_types.keys()}
          self.false_negative = 0
  
          # data collectors to save population counts and agent states every
          # time step
  
          model_reporters = {
                  'N_diagnostic_tests':get_N_diagnostic_tests,
                  'N_preventive_screening_tests':get_N_preventive_screening_tests,
                  'undetected_infections':get_undetected_infections,
                  'predetected_infections':get_predetected_infections,
                  'pending_test_infections':get_pending_test_infections
                  }
  
          for agent_type in self.agent_types:
              model_reporters.update({
                  'diagnostic_test_detected_infections_{}'.format(agent_type):\
                  diagnostic_test_detected_infections_funcs[agent_type]
                  })
              model_reporters.update({
                  'preventive_test_detected_infections_{}'.format(agent_type):\
                  preventive_test_detected_infections_funcs[agent_type]
                  })
  
  
          self.datacollector = DataCollector(
              model_reporters=model_reporters,
              agent_reporters=
              	{
              	'infection_state': get_infection_state,
                  'quarantine_state': get_quarantine_state
                  })
  
  
      ## transmission risk modifiers
      def get_transmission_risk_contact_type_modifier(self, source, target):
          # construct the edge key as combination between agent IDs and weekday
          n1 = source.ID
          n2 = target.ID
          tmp = [n1, n2]
          tmp.sort()
          n1, n2 = tmp
          key = '{}{}d{}'.format(n1, n2, self.weekday)
          contact_weight = self.G.get_edge_data(n1, n2, key)['weight']
  
          # the link weight is a multiplicative modifier of the link strength.
          # contacts of type "close" have, by definition, a weight of 1. Contacts
          # of type intermediate, far or very far have a weight < 1 and therefore
          # are less likely to transmit an infection. For example, if the contact
          # type far has a weight of 0.2, a contact of type far has only a 20%
          # chance of transmitting an infection, when compared to a contact of
          # type close. To calculate the probability of success p in the Bernoulli
          # trial, we need to reduce the base risk (or base probability of success)
          # by the modifications introduced by preventive measures. These
          # modifications are formulated in terms of "probability of failure", or
          # "q". A low contact weight has a high probability of failure, therefore
          # we return q = 1 - contact_weight here.
          q1 = 1 - contact_weight
  
          return q1
  
  
      def get_transmission_risk_age_modifier_transmission(self, source):
          '''linear function such that at age 18 the risk is that of an adult (=1).
          The slope of the line needs to be calibrated.
          '''
          age = source.age
          max_age = 18
          if age <= max_age:
              age_weight = self.age_transmission_risk_discount['slope'] * \
                   np.abs(age - max_age) + self.age_transmission_risk_discount['intercept']
  
              # The age weight can be interpreted as multiplicative factor that
              # reduces the chance for transmission with decreasing age. The slope
              # of the age_transmission_discount function is the decrease (in % of
              # the transmission risk for an 18 year old or above) of transmission
              # risk with every year a person is younger than 18 (the intercept is
              # 1 by definition).
              # To calculate the probability of success p in the Bernoulli
              # trial, we need to reduce the base risk (or base probability of 
              # success) by the modifications introduced by preventive measures. 
              # These modifications are formulated in terms of "probability of 
              # failure", or "q". A low age weight has a high probability of 
              # failure, therefore we return q = 1 - age_weight here.
              q2 = 1 - age_weight
          else:
              q2 = 0
  
          return q2
  
  
      def get_transmission_risk_age_modifier_reception(self, target):
          '''linear function such that at age 18 the risk is that of an adult (=1).
          The slope of the line needs to be calibrated.
          '''
          age = target.age
          max_age = 18
          if age <= max_age:
              age_weight = self.age_transmission_risk_discount['slope'] * \
              np.abs(age - max_age) + self.age_transmission_risk_discount['intercept']
              # see description in get_transmission_risk_age_modifier_transmission
              q3 = 1 - age_weight
          else:
              q3 = 0
  
          return q3
  
  
      # infectiousness is constant and high until symptom onset and then
      # decreases monotonically until agents are not infectious anymore
      # at the end of the infection_duration
      def get_transmission_risk_progression_modifier(self, source):
          if source.days_since_exposure < source.exposure_duration:
              progression_weight = 0
          elif source.days_since_exposure <= source.time_until_symptoms:
              progression_weight = 1
          elif source.days_since_exposure > source.time_until_symptoms and \
               source.days_since_exposure <= source.infection_duration:
              # we add 1 in the denominator, such that the source is also
              # (slightly) infectious on the last day of the infection_duration
              progression_weight = \
                   (source.days_since_exposure - source.time_until_symptoms) / \
                   (source.infection_duration - source.time_until_symptoms + 1)
          else:
              progression_weight = 0
          # see description in get_transmission_risk_age_modifier_transmission
          q4 = 1 - progression_weight
  
          return q4
  
      def get_transmission_risk_subclinical_modifier(self, source):
          if source.symptomatic_course == False:
              subclinical_weight = self.subclinical_modifier
          else:
              subclinical_weight = 1
          # see description in get_transmission_risk_age_modifier_transmission
          q5 = 1 - subclinical_weight
          return q5
  
      def get_transmission_risk_exhale_modifier(self, source):
          if source.mask:
              exhale_weight = self.mask_filter_efficiency['exhale']
          else:
              exhale_weight = 1
          # see description in get_transmission_risk_age_modifier_transmission
          q6 = 1 - exhale_weight
          return q6
  
  
      def get_transmission_risk_inhale_modifier(self, target):
          if target.mask:
              inhale_weight = self.mask_filter_efficiency['inhale']
          else:
              inhale_weight = 1
          # see description in get_transmission_risk_age_modifier_transmission
          q7 = 1 - inhale_weight
          return q7
  
  
      def get_transmission_risk_ventilation_modifier(self):
          ventilation_weight = self.transmission_risk_ventilation_modifier
          # see description in get_transmission_risk_age_modifier_transmission
          q8 = 1 - ventilation_weight
          return q8
  
      def get_transmission_risk_vaccination_modifier_reception(self, a):
          if a.vaccinated:
              q9 = self.transmission_risk_vaccination_modifier['reception']
          else:
              q9 = 0
          return q9
  
      def get_transmission_risk_vaccination_modifier_transmission(self, a):
          if a.vaccinated:
              q10 = self.transmission_risk_vaccination_modifier['transmission']
          else:
              q10 = 0
          return q10
  
      def test_agent(self, a, test_type):
          a.tested = True
          a.pending_test = test_type
          if test_type == self.Testing.diagnostic_test_type:
              self.number_of_diagnostic_tests += 1
          else:
              self.number_of_preventive_screening_tests += 1
  
          if a.exposed:
              # tests that happen in the period of time in which the agent is
              # exposed but not yet infectious. 
              # Note: tests[test_type]['time_until_testable'] is negative for
              # tests that can detect an infection before agents become infectious
              if a.days_since_exposure >= a.exposure_duration + \
                      self.Testing.tests[test_type]['time_until_testable']:
                  
                  if self.verbosity > 1:
                      print('{} {} sent positive sample (even though not infectious yet)'
                      .format(a.type, a.ID))
                  a.sample = 'positive'
                  self.predetected_infections += 1
                  self.positive_tests[test_type][a.type] += 1
              else:
                  if self.verbosity > 1: print('{} {} sent negative sample'
                      .format(a.type, a.ID))
                  a.sample = 'negative'
  
          elif a.infectious:
              # tests that happen in the period of time in which the agent is
              # infectious and the infection is detectable by a given test
              # Note: tests[test_type]['time_until_testable'] is negative for 
              # tests that can detect an infection before agents become 
              # infectious. tests[test_type]['time_testable'] is negative for
              # tests that cease to detect an infection before agents stop being
              # infectious
              if a.days_since_exposure >= a.exposure_duration + \
                      self.Testing.tests[test_type]['time_until_testable'] and \
                 a.days_since_exposure <= a.infection_duration + \
                      self.Testing.tests[test_type]['time_testable']:
                  if self.verbosity > 1:
                      print('{} {} sent positive sample'.format(a.type, a.ID))
                  a.sample = 'positive'
                  self.positive_tests[test_type][a.type] += 1
  
              # track the undetected infections to assess how important they are
              # for infection spread
              else:
                  if self.verbosity > 1:
                      print('{} {} sent negative sample (even though infectious)'
                      .format(a.type, a.ID))
                  a.sample = 'negative'
                  self.undetected_infections += 1
  
          else:
              if self.verbosity > 1: print('{} {} sent negative sample'
                  .format(a.type, a.ID))
              a.sample = 'negative'
  
          # for same-day testing, immediately act on the results of the test
          if a.days_since_tested >= self.Testing.tests[test_type]['time_until_test_result']:
              a.act_on_test_result()
  
      def screen_agents(self, agent_group, test_type, screen_type):
          # only test agents that have not been tested already in this simulation
          # step and that are not already known positive cases
  
          if self.verbosity > 0:
              print('initiating {} {} screen'\
                                  .format(screen_type, agent_group))
  
          untested_agents = [a for a in self.schedule.agents if
              (a.tested == False and a.known_positive == False
                  and a.type == agent_group)]
  
          if len(untested_agents) > 0:
              self.screened_agents[screen_type][agent_group] = True
              self.days_since_last_agent_screen[agent_group] = 0
  
              # only test agents if they participate in voluntary testing
              if screen_type == 'preventive':
                  for a in untested_agents:
                      if a.voluntary_testing:
                          self.test_agent(a, test_type)
                      else:
                          if self.verbosity > 1:
                              print('not testing {} {}, not participating in voluntary testing'\
                                  .format(agent_group, a.ID))
              else:
                  for a in untested_agents:
                      self.test_agent(a, test_type)
  
              if self.verbosity > 0:
                  print()
          else:
              if self.verbosity > 0:
                  print('no agents tested because all agents have already been tested')
  
      # the type of the test used in the pending test result is stored in the
      # variable pending_test
  
      def collect_test_results(self):
          agents_with_test_results = [a for a in self.schedule.agents if
              (a.pending_test and
               a.days_since_tested >= self.Testing.tests[a.pending_test]['time_until_test_result'])]
  
          return agents_with_test_results
  
      def trace_contacts(self, a):
          if a.quarantined == False:
              a.quarantined = True
              a.quarantine_start = self.Nstep
  
              if self.verbosity > 0:
                  print('qurantined {} {}'.format(a.type, a.ID))
  
          # find all agents that share edges with the agent
          # that are classified as K1 contact types in the testing
          # strategy
          if a in self.G.nodes():
              K1_contacts = [e[1] for e in self.G.edges(a.ID, data=True) if
                  e[2]['contact_type'] in self.Testing.K1_contact_types]
              K1_contacts = [a for a in self.schedule.agents if a.ID in K1_contacts]
  
              for K1_contact in K1_contacts:
                  if self.verbosity > 0:
                      print('quarantined {} {} (K1 contact of {} {})'
                          .format(K1_contact.type, K1_contact.ID, a.type, a.ID))
                  K1_contact.quarantined = True
                  K1_contact.quarantine_start = self.Nstep
  
      def test_symptomatic_agents(self):
          # find symptomatic agents that have not been tested yet and are not
          # in quarantine and test them
          newly_symptomatic_agents = np.asarray([a for a in self.schedule.agents
              if (a.symptoms == True and a.tested == False and a.quarantined == False)])
  
          for a in newly_symptomatic_agents:
              # all symptomatic agents are quarantined by default
              if self.verbosity > 0:
                  print('quarantined: {} {}'.format(a.type, a.ID))
              a.quarantined = True
              a.quarantine_start = self.Nstep
  
              self.test_agent(a, self.Testing.diagnostic_test_type)
  
      def quarantine_contacts(self):
          # trace and quarantine contacts of newly positive agents
          if len(self.newly_positive_agents) > 0:
              if self.verbosity > 0: print('new positive test(s) from {}'
                  .format([a.ID for a in self.newly_positive_agents]))
  
              # send all K1 contacts of positive agents into quarantine
              for a in self.newly_positive_agents:
                  self.trace_contacts(a)
  
              # indicate that a screen should happen because there are new
              # positive test results
              self.new_positive_tests = True
              self.newly_positive_agents = []
  
          else:
              self.new_positive_tests = False
  
  
      def step(self):
          self.weekday = (self.Nstep + self.weekday_offset) % 7 + 1
          # if the connection graph is time-resloved, set the graph that is
          # used to determine connections in this step to the sub-graph corres-
          # ponding to the current day of the week
          if self.dynamic_connections:
              self.G = self.weekday_connections[self.weekday]
  
          if self.verbosity > 0:
              print('weekday {}'.format(self.weekday))
  
          if self.testing:
              for agent_type in self.agent_types:
                  for screen_type in ['reactive', 'follow_up', 'preventive']:
                      self.screened_agents[screen_type][agent_type] = False
  
              if self.verbosity > 0:
                  print('* testing and tracing *')
  
              self.test_symptomatic_agents()
  
  
              # collect and act on new test results
              agents_with_test_results = self.collect_test_results()
              for a in agents_with_test_results:
                  a.act_on_test_result()
  
              self.quarantine_contacts()
  
              # screening:
              # a screen should take place if
              # (a) there are new positive test results
              # (b) as a follow-up screen for a screen that was initiated because
              # of new positive cases
              # (c) if there is a preventive screening policy and it is time for
              # a preventive screen in a given agent group
  
              # (a)
              if (self.testing == 'background' or self.testing == 'background+preventive')\
                 and self.new_positive_tests == True:
                  for agent_type in self.screening_agents:
                      self.screen_agents(
                          agent_type, self.Testing.diagnostic_test_type, 'reactive')
                      self.scheduled_follow_up_screen[agent_type] = True
  
              # (b)
              elif (self.testing == 'background' or self.testing == 'background+preventive') and \
                  self.Testing.follow_up_testing_interval != None and \
                  sum(list(self.scheduled_follow_up_screen.values())) > 0:
                  for agent_type in self.screening_agents:
                      if self.scheduled_follow_up_screen[agent_type] and\
                         self.days_since_last_agent_screen[agent_type] >=\
                         self.Testing.follow_up_testing_interval:
                          self.screen_agents(
                              agent_type, self.Testing.diagnostic_test_type, 'follow_up')
                      else:
                          if self.verbosity > 0:
                              print('not initiating {} follow-up screen (last screen too close)'\
                                  .format(agent_type))
  
              # (c) 
              elif (self.testing == 'preventive' or self.testing == 'background+preventive')and \
                  np.any(list(self.Testing.screening_intervals.values())):
  
                  for agent_type in self.screening_agents:
                      interval = self.Testing.screening_intervals[agent_type]
                      assert interval in [7, 3, 2, None], \
                          'testing interval {} for agent type {} not supported!'\
                          .format(interval, agent_type)
  
                      # (c.1) testing every 7 days = testing on Mondays
                      if interval == 7 and self.weekday == 1:
                          self.screen_agents(agent_type,
                              self.Testing.preventive_screening_test_type,\
                               'preventive')
                      # (c.2) testing every 3 days = testing on Mo & Turs
                      elif interval == 3 and self.weekday in [1, 4]:
                              self.screen_agents(agent_type,
                              self.Testing.preventive_screening_test_type,\
                               'preventive')
                      # (c.3) testing every 2 days = testing on Mo, Wed & Fri
                      elif interval == 2 and self.weekday in [1, 3, 5]:
                              self.screen_agents(agent_type,
                              self.Testing.preventive_screening_test_type,\
                               'preventive')
                      # No interval specified = no testing, even if testing
                      # mode == preventive
                      elif interval == None:
                          pass
                      else:
                          if self.verbosity > 0:
                              print('not initiating {} preventive screen (wrong weekday)'\
                                      .format(agent_type))
              else:
                  # do nothing
                  pass
  
              for agent_type in self.agent_types:
                  if not (self.screened_agents['reactive'][agent_type] or \
                          self.screened_agents['follow_up'][agent_type] or \
                          self.screened_agents['preventive'][agent_type]):
                          self.days_since_last_agent_screen[agent_type] += 1
  
  
          if self.verbosity > 0: print('* agent interaction *')
          self.datacollector.collect(self)
          self.schedule.step()
          self.Nstep += 1
  
apis: numpy.random.weibull numpy.random.seed numpy.random.choice numpy.abs numpy.asarray scipy.optimize.root_scalar