from soil.agents import FSM, NetworkAgent, state, default_state, prob import logging class DumbViewer(FSM, NetworkAgent): ''' A viewer that gets infected via TV (if it has one) and tries to infect its neighbors once it's infected. ''' defaults = { 'prob_neighbor_spread': 0.5, 'prob_tv_spread': 0.1, } @default_state @state def neutral(self): if self['has_tv']: if self.prob(self.model['prob_tv_spread']): return self.infected @state def infected(self): for neighbor in self.get_neighboring_agents(state_id=self.neutral.id): if self.prob(self.model['prob_neighbor_spread']): neighbor.infect() def infect(self): ''' This is not a state. It is a function that other agents can use to try to infect this agent. DumbViewer always gets infected, but other agents like HerdViewer might not become infected right away ''' self.set_state(self.infected) class HerdViewer(DumbViewer): ''' A viewer whose probability of infection depends on the state of its neighbors. ''' def infect(self): '''Notice again that this is NOT a state. See DumbViewer.infect for reference''' infected = self.count_neighboring_agents(state_id=self.infected.id) total = self.count_neighboring_agents() prob_infect = self.model['prob_neighbor_spread'] * infected/total self.debug('prob_infect', prob_infect) if self.prob(prob_infect): self.set_state(self.infected) class WiseViewer(HerdViewer): ''' A viewer that can change its mind. ''' defaults = { 'prob_neighbor_spread': 0.5, 'prob_neighbor_cure': 0.25, 'prob_tv_spread': 0.1, } @state def cured(self): prob_cure = self.model['prob_neighbor_cure'] for neighbor in self.get_neighboring_agents(state_id=self.infected.id): if self.prob(prob_cure): try: neighbor.cure() except AttributeError: self.debug('Viewer {} cannot be cured'.format(neighbor.id)) def cure(self): self.set_state(self.cured.id) @state def infected(self): cured = max(self.count_neighboring_agents(self.cured.id), 1.0) infected = max(self.count_neighboring_agents(self.infected.id), 1.0) prob_cure = self.model['prob_neighbor_cure'] * (cured/infected) if self.prob(prob_cure): return self.cured return self.set_state(super().infected)