forked from notionparallax/programming---a-new-hope
-
Notifications
You must be signed in to change notification settings - Fork 0
/
classySchelling.py
141 lines (123 loc) · 5.12 KB
/
classySchelling.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
import matplotlib.pyplot as plt
import itertools
import random
import copy
class Schelling:
def __init__(self, width, height, empty_ratio,
similarity_threshold, n_iterations, races=2):
self.width = width
self.height = height
self.races = races
self.empty_ratio = empty_ratio
self.similarity_threshold = similarity_threshold
self.n_iterations = n_iterations
self.empty_houses = []
self.agents = {}
def populate(self):
self.all_houses = list(itertools.product(range(self.width), range(self.height)))
random.shuffle(self.all_houses)
self.n_empty = int( self.empty_ratio * len(self.all_houses) )
self.empty_houses = self.all_houses[: self.n_empty]
self.remaining_houses = self.all_houses[self.n_empty:]
houses_by_race = [self.remaining_houses[i::self.races] for i in range(self.races)]
for i in range(self.races):
# create agents for each race
self.agents = dict(
self.agents.items() +
dict(zip(houses_by_race[i], [i+1]*len(houses_by_race[i]))).items()
)
def is_unsatisfied(self, x, y):
race = self.agents[(x, y)]
count_similar = 0
count_different = 0
if x > 0 and y > 0 and (x-1, y-1) not in self.empty_houses:
if self.agents[(x-1, y-1)] == race:
count_similar += 1
else:
count_different += 1
if y > 0 and (x, y-1) not in self.empty_houses:
if self.agents[(x, y-1)] == race:
count_similar += 1
else:
count_different += 1
if x < (self.width-1) and y > 0 and (x+1, y-1) not in self.empty_houses:
if self.agents[(x+1, y-1)] == race:
count_similar += 1
else:
count_different += 1
if x > 0 and (x-1, y) not in self.empty_houses:
if self.agents[(x-1, y)] == race:
count_similar += 1
else:
count_different += 1
if x < (self.width-1) and (x+1, y) not in self.empty_houses:
if self.agents[(x+1, y)] == race:
count_similar += 1
else:
count_different += 1
if x > 0 and y < (self.height-1) and (x-1, y+1) not in self.empty_houses:
if self.agents[(x-1, y+1)] == race:
count_similar += 1
else:
count_different += 1
if x > 0 and y < (self.height-1) and (x, y+1) not in self.empty_houses:
if self.agents[(x, y+1)] == race:
count_similar += 1
else:
count_different += 1
if x < (self.width-1) and y < (self.height-1) and (x+1, y+1) not in self.empty_houses:
if self.agents[(x+1, y+1)] == race:
count_similar += 1
else:
count_different += 1
if (count_similar+count_different) == 0:
return False
else:
return float(count_similar)/(count_similar+count_different) < self.happy_threshold
def update(self):
for i in range(self.n_iterations):
self.old_agents = copy.deepcopy(self.agents)
n_changes = 0
for agent in self.old_agents:
if self.is_unhappy(agent[0], agent[1]):
agent_race = self.agents[agent]
empty_house = random.choice(self.empty_houses)
self.agents[empty_house] = agent_race
del self.agents[agent]
self.empty_houses.remove(empty_house)
self.empty_houses.append(agent)
n_changes += 1
print n_changes
if n_changes == 0:
break
def move_to_empty(self, x, y):
race = self.agents[(x, y)]
empty_house = random.choice(self.empty_houses)
self.updated_agents[empty_house] = race
del self.updated_agents[(x, y)]
self.empty_houses.remove(empty_house)
self.empty_houses.append((x, y))
def plot(self, title, file_name):
fig, ax = plt.subplots()
# If you want to run the simulation with more than 7 colors,
# you should set agent_colors accordingly
agent_colors = {1: 'b', 2: 'r', 3: 'g', 4: 'c', 5: 'm', 6: 'y', 7: 'k'}
for agent in self.agents:
ax.scatter(agent[0]+0.5,
agent[1]+0.5,
color=agent_colors[self.agents[agent]])
ax.set_title(title, fontsize=10, fontweight='bold')
ax.set_xlim([0, self.width])
ax.set_ylim([0, self.height])
ax.set_xticks([])
ax.set_yticks([])
plt.savefig(file_name)
schelling_1 = Schelling(50, 50, 0.3, 0.3, 500, 2)
schelling_1.populate()
schelling_1.plot("schelling_1", "schelling_1.png")
schelling_2 = Schelling(50, 50, 0.3, 0.5, 500, 2)
schelling_2.populate()
schelling_2.plot("schelling_2", "schelling_2.png")
schelling_3 = Schelling(50, 50, 0.3, 0.9, 1000, 2)
schelling_3.populate()
schelling_3.plot("schelling_3", "schelling_3.png")