11.3. Implementation of Schelling’s Model¶
To implement Schelling’s model, we have yet another class that inherits from Cell2D:
class Schelling(Cell2D):
def __init__(self, n, p):
self.p = p
choices = [0, 1, 2]
probs = [0.1, 0.45, 0.45]
self.array = np.random.choice(choices, (n, n), p=probs)
n is the size of the grid, and p is the threshold on the fraction of similar neighbors. For example, if p=0.3, an agent will be unhappy if fewer than 30% of their neighbors are the same color.
array is a NumPy array where each cell is 0 if empty, 1 if occupied by a red agent, and 2 if occupied by a blue agent. Initially 10% of the cells are empty, 45% red, and 45% blue.
The step function for Schelling’s model is substantially more complicated than previous examples. If you are not interested in the details, you can skip to the next section. But if you stick around, you might pick up some NumPy tips.
First, let’s make boolean arrays that indicate which cells are red, blue, and empty:
a = self.array
red = a==1
blue = a==2
empty = a==0
Then we can use correlate2d to count, for each location, the number of neighboring cells that are red, blue, and non-empty. We saw correlate2d in Section 8.7.
options = dict(mode='same', boundary='wrap')
kernel = np.array([[1, 1, 1],
[1, 0, 1],
[1, 1, 1]], dtype=np.int8)
num_red = correlate2d(red, kernel, **options)
num_blue = correlate2d(blue, kernel, **options)
num_neighbors = num_red + num_blue
options is a dictionary that contains the options we pass to correlate2d. With mode='same', the result is the same size as the input. With boundary='wrap', the top edge is wrapped to meet the bottom, and the left edge is wrapped to meet the right.
kernel indicates that we want to consider the eight neighbors that surround each cell.
After computing num_red and num_blue, we can compute the fraction of neighbors, for each location, that are red and blue.
frac_red = num_red / num_neighbors
frac_blue = num_blue / num_neighbors
Then, we can compute the fraction of neighbors, for each agent, that are the same color as the agent, using np.where, which is like an element-wise if expression. The first parameter is a condition that selects elements from the second or third parameter.
frac_same = np.where(red, frac_red, frac_blue)
frac_same[empty] = np.nan
In this case, wherever red is True, frac_same gets the corresponding element of frac_red. Where red is False, frac_same gets the corresponding element of frac_blue. Finally, where empty indicates that a cell is empty, frac_same is set to np.nan, which is a special value that indicates “Not a Number”.
Now we can identify the locations of the unhappy agents:
unhappy = frac_same < self.p
unhappy_locs = locs_where(unhappy)
locs_where is a wrapper function for np.nonzero:
def locs_where(condition):
return list(zip(*np.nonzero(condition)))
np.nonzero takes an array and returns the coordinates of all non-zero cells; the result is a tuple of arrays, one for each dimension. Then locs_where uses list and zip to convert this result to a list of coordinate pairs.
Similarly, empty_locs is an array that contains the coordinates of the empty cells:
empty_locs = locs_where(empty)
Now we get to the core of the simulation. We loop through the unhappy agents and move them:
num_empty = np.sum(empty)
for source in unhappy_locs:
i = np.random.randint(num_empty)
dest = empty_locs[i]
a[dest] = a[source]
a[source] = 0
empty_locs[i] = source
i is the index of a random empty cell; dest is a tuple containing the coordinates of the empty cell.
In order to move an agent, we copy its value (1 or 2) from source to dest, and then set the value of source to 0 (since it is now empty).
Finally, we replace the entry in empty_locs with source, so the cell that just became empty can be chosen by the next agent.