You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
soil/examples/newsspread/newsspread.py

80 lines
2.2 KiB
Python

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.
'''
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)