An extension of the SIRS epidemic model with quarantined/immunised state

models.py

class SIRSQModel(SIRSState):
    def __init__(self, g, beta=1, gamma=0.1, mu=0.1, r=0,
                 init_prob = None, qtime=14):
        """SIRSQ compartmental epidemic model. Extends the SIRSState epidemic 
        model from graph-tool to deal with quarantined individuals. We also
        provide a function to perform group tests.
        
        For an SIRQ model, set mu to 0.
        
        Parameters
        ----------
        g       - graph_tool.graph to be used for the dynamics
        beta    - float, transmission probability for each edge (default 1)
        gamma   - float, I to R recovery probability per timestep (default 0.1)
        mu      - float, R to S relapse probability per timestep (default 0.1)
        r       - float, spontaneous infection probability (default 0)
        s       - VertexPropertyMap, initial global state for all individuals
                  (default SUSCEPTIBLE for all individuals)
        qtime   - quarantine time (in number of time steps)
        """
        self.graph = g
        self.pop_size = self.graph.num_vertices()  # population size
        seed = _get_seed(init_prob) if init_prob else None
        # Initialise parent SIRS model
        super().__init__(g, beta=beta, gamma=gamma, mu=mu, r=r, exposed=False,
                         epsilon=0.1, v0=None, s=seed, constant_beta=True)
        # Initialise vertex property maps to track quarantined individuals
        self.quarantined = self.graph.new_vertex_property("bool", False)
        self.quarantine_end = self.graph.new_vertex_property("double", np.inf)
        # Set filter to hide all quarantined individuals in graph
        self.graph.set_vertex_filter(self.quarantined, inverted=True)
        # Set additional attributes
        self.clock = 0  # to keep track of time steps
        self.qtime = qtime  # quarantine time
    
    def _get_seed(self, prob):
        """
        Return VertexPropertyMap in which each individual is True with
        probability 'prob'.
        """
        seed = self.graph.new_vertex_property("bool", False)
        seed.a = np.random.binomial(1, prob, self.pop_size)
        return seed

    def quarantine(self, people, qtime=None):
        """
        Place all individuals in 'people' into quarantine. Does not change
        the status of individuals who are already quarantined.
        
        Note: does not check whether individuals in 'people' are already
        quarantined and will simply extend the length of quarantine time for
        such people.

        people - boolean numpy array of length n (i.e. unfiltered!)
                 OR
                 int numpy array (wrt unfiltered graph!)
        qtime  - time that people will be quarantined
                 (if None, set to self.qtime)
        """
        # Retrieve quarantine time, seting to global default if qtime is None
        if qtime is None: qtime = self.qtime
        # Change quarantined boolean flag to True
        self.quarantined.a[people] = True
        # For each newly quarantined individual, compute end time of infection
        self.quarantine_end.a[people] = self.clock + qtime

    def _release(self):
        """Release all individuals from quarantine who have done their time"""
        mask = self.quarantine_end.a <= self.clock
        self.quarantined.a[mask] = False
        self.quarantine_end.a[mask] = np.inf
    
    def group_test(self, group):
        """
        Test group to see whether anyone is ill. Returns positive if the
        group contains an infected individual.
        
        Note: does not check whether anyone in group is quarantined.
        
        group - boolean numpy array of length n (i.e. unfiltered!)
                OR
                int numpy array (wrt unfiltered graph!)
        """
        # get status of everyone in group
        group_state = self.get_state().a[group]  
        return (group_state == INFECTED).any()  # True if anyone is infected

    def iterate(self, ntimesteps=1):
        """Advance the model by ntimesteps."""
        for _ in range(ntimesteps):
            # Release individuals from quarantine who have finished their time
            self._release()
            # Advance the parent model by one time step
            super().iterate_sync()
            # Increment clock by one time step
            self.clock += 1