from soil.agents import FSM, state, default_state, prob import logging class DumbViewer(FSM): ''' 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 prob(self.env['prob_tv_spread']): self.set_state(self.infected) @state def infected(self): for neighbor in self.get_neighboring_agents(state_id=self.neutral.id): if prob(self.env['prob_neighbor_spread']): neighbor.infect() def infect(self): self.set_state(self.infected) class HerdViewer(DumbViewer): ''' A viewer whose probability of infection depends on the state of its neighbors. ''' level = logging.DEBUG def infect(self): infected = self.count_neighboring_agents(state_id=self.infected.id) total = self.count_neighboring_agents() prob_infect = self.env['prob_neighbor_spread'] * infected/total self.debug('prob_infect', prob_infect) if prob(prob_infect): self.set_state(self.infected.id) 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.env['prob_neighbor_cure'] for neighbor in self.get_neighboring_agents(state_id=self.infected.id): if 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.env['prob_neighbor_cure'] * (cured/infected) if prob(prob_cure): return self.cure() return self.set_state(super().infected)