forked from ckaandorp/CompSc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
disease.py
153 lines (140 loc) · 4.84 KB
/
disease.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
from mesa import Model
from mesa.datacollection import DataCollector
from mesa.time import RandomActivation
from mesa.space import SingleGrid
import numpy as np
from math import floor
from diseaseAgent import DiseaseAgent
from wall import wall
from helperFunctions import disease_collector
class DiseaseModel(Model):
"""
A model with some number of agents.
highS: Number of agents with high sociability.
middleS: Number of agents with middle sociability.
lowS: Number of agents with low sociability.
width: Width of the grid.
height: Height of the grid.
edu_setting: Classrooms and set schedule if true, else random free movement.
cureProb: Probability of agent getting better.
cureProbFac: Factor of cureProb getting higher.
mutateProb: Probability of a disease mutating.
diseaseRate: Rate at which the disease spreads.
"""
def __init__(self, highS, middleS, lowS, width, height, edu_setting=True, cureProb=0.1, cureProbFac=2/1440, mutateProb=0.0050, diseaseRate=0.38):
super().__init__()
self.num_agents = highS + middleS + lowS
self.lowS = lowS
self.middleS = middleS
self.highS = highS
self.initialCureProb = cureProb
self.cureProbFac = cureProbFac
self.mutateProb = mutateProb
self.diseaseRate = diseaseRate
self.edu_setting = edu_setting
self.maxDisease = 0# amount of mutations
self.counter = 540 # keeps track of timesteps
self.removed = []
self.exit = (width-1,floor(height/2))
# Check if agents fit within grid
if self.num_agents > width * height:
raise ValueError("Number of agents exceeds grid capacity.")
# Create grid with random activation
self.grid = SingleGrid(width, height, True)
self.schedule = RandomActivation(self)
if edu_setting:
# Create walls
numberRooms = 3
self.add_walls(numberRooms, width, height)
self.midWidthRoom = floor(width / numberRooms / 2)
self.midHeightRoom = floor(height / numberRooms / 2)
# Calculate the centers of the 6 rooms
roomLeftDown = (5 * self.midWidthRoom, self.midHeightRoom)
roomLeftMid = (3 * self.midWidthRoom, self.midHeightRoom)
roomLeftUp = (self.midWidthRoom, self.midHeightRoom)
roomRightDown = (5 * self.midWidthRoom, 5 * self.midHeightRoom, )
roomRightMid = (3 * self.midWidthRoom, 5 * self.midHeightRoom)
roomRightUp = (self.midWidthRoom, 5 * self.midHeightRoom)
# Set 3 goals per roster
self.roster = [[roomLeftDown, roomLeftUp, roomRightMid], [roomRightMid, roomLeftDown, roomRightDown],
[roomRightUp, roomRightDown, roomLeftUp]]
# Create agents
self.addAgents(lowS, 0, 0)
self.addAgents(middleS, lowS, 1)
self.addAgents(highS, lowS + highS, 2)
self.datacollector = DataCollector(
model_reporters={"diseasepercentage": disease_collector},
agent_reporters={"disease": "disease"})
def heuristic(self, start, goal):
"""
Returns manhattan distance.
start: current location (x,y)
goal: goal location (x,y)
"""
dx = abs(start[0] - goal[0])
dy = abs(start[1] - goal[1])
return dx + dy
def get_vertex_neighbors(self, pos):
"""
Returns all neighbors.
pos: current position
"""
n = self.grid.get_neighborhood(pos, moore=False)
neighbors = []
for item in n:
if not abs(item[0]-pos[0]) > 1 and not abs(item[1]-pos[1]) > 1:
neighbors += [item]
return neighbors
def move_cost(self, location):
"""
Return the cost of a location.
"""
if self.grid.is_cell_empty(location):
return 1 # Normal movement cost
else:
return 100
def add_walls(self, n, widthGrid, heightGrid):
"""
Add walls in grid.
n: number of rooms horizontally
widthGrid: width of the grid
heightGrid: height of the grid
"""
widthRooms = floor(widthGrid/n)
heightRooms = floor(heightGrid/n)
widthHall = widthGrid - 2 * widthRooms
heightHall = heightGrid - 2 * heightRooms
# Add horizontal walls
for i in range(n - 1):
for y in range(heightRooms):
brick = wall(self.num_agents, self)
self.grid.place_agent(brick, ((i + 1) * widthRooms, y))
self.grid.place_agent(brick, ((i + 1) * widthRooms, y + heightRooms + heightHall))
doorWidth = 2
# Add vertical walls
for x in range(widthGrid):
if (x % widthRooms) < (widthRooms - doorWidth):
brick = wall(self.num_agents, self)
self.grid.place_agent(brick, (x, heightRooms))
self.grid.place_agent(brick, (x, heightRooms + heightHall - 1))
def addAgents(self, n, startID, sociability):
"""
Add agents with a sociability.
n: number of agents
startID: ID of the first added agent
sociability: sociability of the agents
"""
disease_list = np.random.randint(0,2,n)
for i in range(n):
a = DiseaseAgent(i + startID, sociability,self,disease_list[i])
self.schedule.add(a)
# Add the agent to a random grid cell
location = self.grid.find_empty()
self.grid.place_agent(a, location)
def step(self):
"""
Continue one step in simulation.
"""
self.counter += 1
self.datacollector.collect(self)
self.schedule.step()